Audit telemetry residual tuning premise

runs/telemetry-residual/test_common_state.py

@@ -0,0 +1,183 @@
+#!/usr/bin/env python3
+from __future__ import annotations
+
+import csv
+import json
+import math
+import tempfile
+from pathlib import Path
+
+from common_state import numeric, residual, summarize_engine, summarize_frontier
+from run_frontier_state import enable_state_outputs
+
+
+def engine_record(
+    step: int,
+    timestamp_ns: int,
+    *,
+    batch: int,
+    prefill: int,
+    decode: int,
+    waiting: int,
+    running: int,
+    kv: float,
+) -> dict[str, object]:
+    return {
+        "step_index": step,
+        "submit_mono_ns": timestamp_ns,
+        "model_executed": True,
+        "scheduled_requests": batch,
+        "decode_batch_size": decode,
+        "prefill_tokens": prefill,
+        "decode_tokens": decode,
+        "preemptions": 0,
+        "queues": {"waiting": waiting, "running": running},
+        "kv": {"usage": kv},
+        "cudagraph": {
+            "runtime_mode": "NONE" if step == 0 else "FULL",
+            "bucket_tokens": prefill + decode + 1,
+            "padding_tokens": 1,
+        },
+        "dropped_records_before": 0,
+    }
+
+
+def main() -> None:
+    assert numeric((-2.0, -1.0))["cv"] >= 0.0
+    command = [
+        "python3",
+        "-m",
+        "frontier.main",
+        "--no-metrics_config_store_frontier_stage_batch_ledger",
+    ]
+    state_command = enable_state_outputs(command)
+    assert "--metrics_config_store_frontier_stage_batch_ledger" in state_command
+    assert "--metrics_config_keep_individual_batch_metrics" in state_command
+    assert "--no-metrics_config_store_frontier_stage_batch_ledger" not in state_command
+
+    engine = summarize_engine(
+        [
+            engine_record(
+                0,
+                0,
+                batch=2,
+                prefill=6,
+                decode=0,
+                waiting=1,
+                running=2,
+                kv=0.1,
+            ),
+            engine_record(
+                1,
+                1_000_000_000,
+                batch=4,
+                prefill=0,
+                decode=4,
+                waiting=3,
+                running=4,
+                kv=0.3,
+            ),
+        ],
+        start_ns=0,
+        end_ns=2_000_000_000,
+        request_count=2,
+    )
+    assert math.isclose(engine["common"]["queue_waiting_mean"], 2.0)
+    assert math.isclose(engine["common"]["queue_running_mean"], 3.0)
+    assert math.isclose(
+        engine["common"]["queue_waiting_time_per_request_ms"], 2000.0
+    )
+    assert engine["common"]["batch_size"]["distinct_n"] == 2
+    assert math.isclose(engine["common"]["prefill_token_fraction"], 0.6)
+
+    with tempfile.TemporaryDirectory() as temporary:
+        root = Path(temporary)
+        system = root / "system.json"
+        requests = root / "requests.csv"
+        batches = root / "batches.csv"
+        ledger = root / "ledger.jsonl"
+        system.write_text(
+            json.dumps(
+                {
+                    "throughput_metrics": {
+                        "total_duration_seconds": 2.0,
+                    }
+                }
+            ),
+            encoding="utf-8",
+        )
+        with requests.open("w", encoding="utf-8", newline="") as output:
+            writer = csv.DictWriter(
+                output,
+                fieldnames=[
+                    "request_e2e_time",
+                    "request_waiting_time_total",
+                    "request_total_preemption_count",
+                ],
+            )
+            writer.writeheader()
+            writer.writerow(
+                {
+                    "request_e2e_time": 2000,
+                    "request_waiting_time_total": 1000,
+                    "request_total_preemption_count": 0,
+                }
+            )
+            writer.writerow(
+                {
+                    "request_e2e_time": 3000,
+                    "request_waiting_time_total": 3000,
+                    "request_total_preemption_count": 0,
+                }
+            )
+        with batches.open("w", encoding="utf-8", newline="") as output:
+            writer = csv.DictWriter(
+                output,
+                fieldnames=[
+                    "batch_size",
+                    "batch_num_tokens",
+                    "batch_num_prefill_tokens",
+                    "batch_num_decode_tokens",
+                ],
+            )
+            writer.writeheader()
+            writer.writerow(
+                {
+                    "batch_size": 2,
+                    "batch_num_tokens": 6,
+                    "batch_num_prefill_tokens": 6,
+                    "batch_num_decode_tokens": 0,
+                }
+            )
+            writer.writerow(
+                {
+                    "batch_size": 4,
+                    "batch_num_tokens": 4,
+                    "batch_num_prefill_tokens": 0,
+                    "batch_num_decode_tokens": 4,
+                }
+            )
+        ledger.write_text(
+            json.dumps({"stage_start_ts": 0.0, "stage_end_ts": 1.0}) + "\n"
+            + json.dumps({"stage_start_ts": 1.0, "stage_end_ts": 2.0})
+            + "\n",
+            encoding="utf-8",
+        )
+        simulator = summarize_frontier(
+            system_metrics_path=system,
+            request_metrics_path=requests,
+            batch_metrics_path=batches,
+            ledger_path=ledger,
+        )
+    assert math.isclose(simulator["common"]["queue_waiting_mean"], 2.0)
+    assert math.isclose(simulator["common"]["queue_running_mean"], 0.5)
+    assert simulator["common"]["batch_size"]["distinct_n"] == 2
+    difference = residual(engine, simulator)
+    assert difference["coverage"]["missing"] == 0
+    assert difference["coverage"]["available"] == 16
+    assert math.isclose(difference["values"]["queue_waiting_mean"], 0.0)
+    print("telemetry common state: PASS")
+
+
+if __name__ == "__main__":
+    main()