diff --git a/runs/action-aware-v0/analyze_pilot.py b/runs/action-aware-v0/analyze_pilot.py index ba5511a..18e295d 100644 --- a/runs/action-aware-v0/analyze_pilot.py +++ b/runs/action-aware-v0/analyze_pilot.py @@ -55,6 +55,18 @@ def numeric(values: Iterable[float]) -> dict[str, Any]: } +def distribution(values: Iterable[float]) -> dict[str, Any]: + finite = [float(value) for value in values] + summary = numeric(finite) + return { + **summary, + "mean": statistics.fmean(finite), + "p50": quantile(finite, 0.50), + "p95": quantile(finite, 0.95), + "p99": quantile(finite, 0.99), + } + + def quantile(values: Iterable[float], probability: float) -> float: ordered = sorted(float(value) for value in values) if not ordered: @@ -169,6 +181,77 @@ def telemetry_coverage( return coverage, covered +def mechanism_summary(records: list[Mapping[str, Any]]) -> dict[str, Any]: + executed = [record for record in records if bool(record["model_executed"])] + if not executed: + raise ValueError("mechanism summary requires executed steps") + prefill = [record for record in executed if int(record["prefill_tokens"]) > 0] + decode_only = [ + record for record in executed if int(record["prefill_tokens"]) == 0 + ] + if not prefill or not decode_only: + raise ValueError("mechanism summary requires prefill and decode-only steps") + + def durations_ms(selected: list[Mapping[str, Any]]) -> list[float]: + values = [ + (int(record["complete_mono_ns"]) - int(record["submit_mono_ns"])) + / 1e6 + for record in selected + ] + if any(value < 0.0 for value in values): + raise ValueError("engine step duration must be non-negative") + return values + + chunk_keys = ("first", "middle", "final", "unsplit", "tokens") + chunks = { + key: sum(int(record["chunked_prefill"][key]) for record in executed) + for key in chunk_keys + } + prefill_tokens = [int(record["prefill_tokens"]) for record in prefill] + prefill_requests = sum(int(record["prefill_requests"]) for record in prefill) + prefix_queries = sum( + int(record["prefix"]["local"]["queries"]) for record in executed + ) + prefix_hits = sum( + int(record["prefix"]["local"]["hits"]) for record in executed + ) + invariants = { + "nonnegative_counts": all( + value >= 0 + for value in ( + *chunks.values(), + prefill_requests, + prefix_queries, + prefix_hits, + ) + ), + "chunk_tokens_match_prefill_tokens": chunks["tokens"] + == sum(prefill_tokens), + "prefix_hits_bounded": 0 <= prefix_hits <= prefix_queries, + } + return { + "executed_steps": len(executed), + "step_duration_ms": distribution(durations_ms(executed)), + "prefill_steps": len(prefill), + "prefill_step_duration_ms": distribution(durations_ms(prefill)), + "decode_only_steps": len(decode_only), + "decode_only_step_duration_ms": distribution(durations_ms(decode_only)), + "prefill": { + "requests": prefill_requests, + "requests_per_step": prefill_requests / len(prefill), + "tokens": sum(prefill_tokens), + "tokens_per_step": distribution(prefill_tokens), + "chunks": chunks, + }, + "prefix": { + "queries": prefix_queries, + "hits": prefix_hits, + "hit_rate": prefix_hits / prefix_queries if prefix_queries else 0.0, + }, + "sanity": {"invariants": invariants}, + } + + def request_summary(path: Path, expected_count: int) -> dict[str, Any]: rows = load_jsonl(path) if len(rows) != expected_count: @@ -259,6 +342,10 @@ def analyze_run( binding = binding_summary( full_records, mns=int(config["mns"]), mbbt=int(config["mbbt"]) ) + mechanism = mechanism_summary(full_records) + invariants["mechanism_summary"] = all( + mechanism["sanity"]["invariants"].values() + ) phases = {} for fraction in phase_fractions: phase_end = start_ns + round(duration_s * fraction * 1e9) @@ -301,6 +388,7 @@ def analyze_run( **latency, }, "binding": binding, + "mechanism": mechanism, "phases": phases, "state": state, "coverage": coverage, diff --git a/runs/action-aware-v0/test_pilot.py b/runs/action-aware-v0/test_pilot.py index afb1e29..3b25456 100644 --- a/runs/action-aware-v0/test_pilot.py +++ b/runs/action-aware-v0/test_pilot.py @@ -100,6 +100,45 @@ def main() -> None: missing, start_ns=0, end_ns=1_100_000_000 )[1] + mechanism = analysis.mechanism_summary( + [ + { + "model_executed": True, + "submit_mono_ns": 0, + "complete_mono_ns": 2_000_000, + "prefill_tokens": 8, + "prefill_requests": 2, + "chunked_prefill": { + "first": 1, + "middle": 0, + "final": 0, + "unsplit": 1, + "tokens": 8, + }, + "prefix": {"local": {"queries": 10, "hits": 2}}, + }, + { + "model_executed": True, + "submit_mono_ns": 2_000_000, + "complete_mono_ns": 3_000_000, + "prefill_tokens": 0, + "prefill_requests": 0, + "chunked_prefill": { + "first": 0, + "middle": 0, + "final": 0, + "unsplit": 0, + "tokens": 0, + }, + "prefix": {"local": {"queries": 0, "hits": 0}}, + }, + ] + ) + assert mechanism["prefill"]["requests_per_step"] == 2.0 + assert mechanism["prefill"]["chunks"]["first"] == 1 + assert mechanism["prefix"]["hit_rate"] == 0.2 + assert all(mechanism["sanity"]["invariants"].values()) + manifest = { "repetitions": {str(index): {} for index in (1, 2, 3)}, "regimes": {