Support long code traces in Frontier replay

2026-07-24 00:57:08 +08:00
parent b80d3f03de
commit 7aed90f9e6
4 changed files with 646 additions and 0 deletions
--- a/runs/frontier-code-trace-v0/test_frontier_code_replay.py
+++ b/runs/frontier-code-trace-v0/test_frontier_code_replay.py
@@ -0,0 +1,114 @@
+#!/usr/bin/env python3
+
+from __future__ import annotations
+
+import csv
+import importlib.util
+import json
+import sys
+import tempfile
+import unittest
+from pathlib import Path
+
+
+ROOT = Path(__file__).resolve().parent
+S3_REAL = ROOT.parent / "frontier-s3-real-v0"
+
+
+def load_replay_module():
+    sys.path.insert(0, str(S3_REAL))
+    spec = importlib.util.spec_from_file_location(
+        "frontier_prefix_replay", S3_REAL / "run_frontier_prefix_replay.py"
+    )
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module
+
+
+class FrontierCodeReplayTest(unittest.TestCase):
+    def test_summary_excludes_undefined_single_token_tpot(self) -> None:
+        module = load_replay_module()
+        with tempfile.TemporaryDirectory() as tmp:
+            root = Path(tmp)
+            trace = root / "frontier.csv"
+            with trace.open("w", newline="") as stream:
+                writer = csv.DictWriter(
+                    stream,
+                    fieldnames=(
+                        "arrived_at",
+                        "num_prefill_tokens",
+                        "num_decode_tokens",
+                        "session_id",
+                        "block_hash_ids",
+                    ),
+                )
+                writer.writeheader()
+                writer.writerow(
+                    {
+                        "arrived_at": 0,
+                        "num_prefill_tokens": 7,
+                        "num_decode_tokens": 1,
+                        "session_id": 1,
+                        "block_hash_ids": "[]",
+                    }
+                )
+                writer.writerow(
+                    {
+                        "arrived_at": 1,
+                        "num_prefill_tokens": 16,
+                        "num_decode_tokens": 2,
+                        "session_id": 2,
+                        "block_hash_ids": 123,
+                    }
+                )
+            metrics = root / "metrics" / "cell"
+            metrics.mkdir(parents=True)
+            with (metrics / "request_metrics.csv").open("w", newline="") as stream:
+                writer = csv.DictWriter(
+                    stream,
+                    fieldnames=(
+                        "request_waiting_time_total",
+                        "ttft",
+                        "tpot",
+                        "request_e2e_time",
+                        "request_cached_prefill_tokens",
+                        "request_prefix_cache_query_blocks",
+                        "request_prefix_cache_hit_blocks",
+                    ),
+                )
+                writer.writeheader()
+                writer.writerow(
+                    {
+                        "request_waiting_time_total": 0,
+                        "ttft": 10,
+                        "tpot": "",
+                        "request_e2e_time": 10,
+                        "request_cached_prefill_tokens": 0,
+                        "request_prefix_cache_query_blocks": 0,
+                        "request_prefix_cache_hit_blocks": 0,
+                    }
+                )
+                writer.writerow(
+                    {
+                        "request_waiting_time_total": 0,
+                        "ttft": 20,
+                        "tpot": 5,
+                        "request_e2e_time": 25,
+                        "request_cached_prefill_tokens": 0,
+                        "request_prefix_cache_query_blocks": 1,
+                        "request_prefix_cache_hit_blocks": 0,
+                    }
+                )
+            (metrics / "system_metrics.json").write_text("{}")
+            (metrics / "frontier_stage_batch_ledger.jsonl").write_text(
+                json.dumps({"request_num_tokens": [1]}) + "\n"
+            )
+            summary = module.summarize(trace, root / "metrics", 2)
+
+        self.assertEqual(summary["requests"], 2)
+        self.assertEqual(summary["latency_ms"]["tpot"]["count"], 1)
+        self.assertEqual(summary["latency_ms"]["tpot"]["mean"], 5)
+
+
+if __name__ == "__main__":
+    unittest.main()