115 lines
4.0 KiB
Python
115 lines
4.0 KiB
Python
#!/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()
|