209 lines
8.2 KiB
Python
209 lines
8.2 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__).parent
|
|
REPO_ROOT = ROOT.parents[1]
|
|
|
|
|
|
def load(name: str):
|
|
path = ROOT / name
|
|
spec = importlib.util.spec_from_file_location(path.stem, path)
|
|
assert spec and spec.loader
|
|
module = importlib.util.module_from_spec(spec)
|
|
sys.modules[spec.name] = module
|
|
spec.loader.exec_module(module)
|
|
return module
|
|
|
|
|
|
class FidelityEnvelopeTest(unittest.TestCase):
|
|
def test_block_identities_are_parent_sensitive_and_prefix_stable(self) -> None:
|
|
module = load("prepare_exact_trace.py")
|
|
prefix = list(range(32))
|
|
left = module.block_identities(prefix + [100, 101], 16)
|
|
right = module.block_identities(prefix + [200, 201], 16)
|
|
self.assertEqual(left[:2], right[:2])
|
|
self.assertNotEqual(left[2], right[2])
|
|
changed_parent = module.block_identities([999] + prefix[1:] + [100, 101], 16)
|
|
self.assertNotEqual(left[0], changed_parent[0])
|
|
self.assertNotEqual(left[1], changed_parent[1])
|
|
|
|
def test_root_sessions_follow_parent_chain(self) -> None:
|
|
module = load("prepare_exact_trace.py")
|
|
rows = [
|
|
{"chat_id": 10, "parent_chat_id": -1},
|
|
{"chat_id": 11, "parent_chat_id": 10},
|
|
{"chat_id": 12, "parent_chat_id": 11},
|
|
{"chat_id": 20, "parent_chat_id": -1},
|
|
]
|
|
self.assertEqual(module.root_sessions(rows), {10: 10, 11: 10, 12: 10, 20: 20})
|
|
|
|
def test_materialize_allreduce(self) -> None:
|
|
module = load("materialize_frontier_allreduce.py")
|
|
rows = []
|
|
for tp in (2, 4):
|
|
for tokens in (1, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192):
|
|
rows.append(
|
|
{
|
|
"tensor_parallel_size": tp,
|
|
"num_tokens": tokens,
|
|
"hidden_dim": 2048,
|
|
"payload_bytes": tokens * 2048 * 2,
|
|
"critical_path_median_ms": tp + tokens / 1000,
|
|
}
|
|
)
|
|
with tempfile.TemporaryDirectory() as temporary:
|
|
root = Path(temporary)
|
|
source = root / "allreduce.json"
|
|
source.write_text(
|
|
json.dumps(
|
|
{
|
|
"schema_version": "qwen30_vllm020_allreduce_frozen.v1",
|
|
"rows": rows,
|
|
}
|
|
)
|
|
)
|
|
output = root / "all_reduce.csv"
|
|
manifest = module.convert(source, output)
|
|
self.assertEqual(manifest["rows"], 24)
|
|
with output.open(newline="") as handle:
|
|
converted = list(csv.DictReader(handle))
|
|
self.assertEqual(converted[0]["num_workers"], "2")
|
|
self.assertEqual(converted[0]["size"], "4096")
|
|
self.assertEqual(converted[-1]["num_workers"], "4")
|
|
self.assertEqual(converted[-1]["size"], str(8192 * 2048 * 2))
|
|
self.assertEqual(
|
|
converted[-1]["time_stats.all_reduce.median"],
|
|
str(4 + 8192 / 1000),
|
|
)
|
|
|
|
def test_batch_profile_retains_one_single_request_anchor_per_tp(self) -> None:
|
|
runner = REPO_ROOT / "runs/frontier-phase-factorial-v0/run_frontier_qwen30_prefill_surface.py"
|
|
spec = importlib.util.spec_from_file_location("qwen30_surface_runner", runner)
|
|
assert spec and spec.loader
|
|
module = importlib.util.module_from_spec(spec)
|
|
sys.modules[spec.name] = module
|
|
spec.loader.exec_module(module)
|
|
profile = ROOT / "profiles/profile-v3-batch-final"
|
|
coverage = module.validate_profile(module.profile_paths(profile))
|
|
self.assertEqual(
|
|
coverage["attention"],
|
|
{
|
|
"1": {"exact_prefill_2048_rows": 1, "profile_batch_size": 1},
|
|
"2": {"exact_prefill_2048_rows": 1, "profile_batch_size": 1},
|
|
"4": {"exact_prefill_2048_rows": 1, "profile_batch_size": 1},
|
|
},
|
|
)
|
|
|
|
def test_exact_trace_parser_and_joint_slo_score(self) -> None:
|
|
module = load("run_frontier_qwen30_exact_trace_surface.py")
|
|
with tempfile.TemporaryDirectory() as temporary:
|
|
root = Path(temporary)
|
|
trace = root / "trace.csv"
|
|
with trace.open("w", newline="") as handle:
|
|
writer = csv.DictWriter(
|
|
handle,
|
|
fieldnames=[
|
|
"arrived_at",
|
|
"num_prefill_tokens",
|
|
"num_decode_tokens",
|
|
"session_id",
|
|
"block_hash_ids",
|
|
],
|
|
lineterminator="\n",
|
|
)
|
|
writer.writeheader()
|
|
writer.writerows(
|
|
[
|
|
{
|
|
"arrived_at": 1.0,
|
|
"num_prefill_tokens": 512,
|
|
"num_decode_tokens": 1,
|
|
"session_id": 1,
|
|
"block_hash_ids": "10|11",
|
|
},
|
|
{
|
|
"arrived_at": 2.0,
|
|
"num_prefill_tokens": 2048,
|
|
"num_decode_tokens": 2,
|
|
"session_id": 2,
|
|
"block_hash_ids": "20|21",
|
|
},
|
|
]
|
|
)
|
|
parsed = module.parse_trace(f"u0p01={trace}")
|
|
self.assertEqual(parsed["requests"], 2)
|
|
self.assertEqual(parsed["shapes"], [(512, 1), (2048, 2)])
|
|
uniform = module.parse_trace(
|
|
f"r1={trace}", rate_contract="uniform-spacing"
|
|
)
|
|
self.assertEqual(uniform["offered_request_rate"], 1.0)
|
|
|
|
metrics = root / "request_metrics.csv"
|
|
with metrics.open("w", newline="") as handle:
|
|
writer = csv.DictWriter(
|
|
handle,
|
|
fieldnames=[
|
|
"Request Id",
|
|
"request_num_prefill_tokens",
|
|
"request_num_decode_tokens",
|
|
"ttft",
|
|
"request_e2e_time",
|
|
],
|
|
lineterminator="\n",
|
|
)
|
|
writer.writeheader()
|
|
writer.writerows(
|
|
[
|
|
{
|
|
"Request Id": 1,
|
|
"request_num_prefill_tokens": 2048,
|
|
"request_num_decode_tokens": 2,
|
|
"ttft": 1200,
|
|
"request_e2e_time": 1300,
|
|
},
|
|
{
|
|
"Request Id": 0,
|
|
"request_num_prefill_tokens": 512,
|
|
"request_num_decode_tokens": 1,
|
|
"ttft": 1000,
|
|
"request_e2e_time": 1000,
|
|
},
|
|
]
|
|
)
|
|
scored = module.score(metrics, parsed["shapes"])
|
|
self.assertEqual(scored["slos"]["tpot_150ms"]["passed"], 2)
|
|
self.assertTrue(scored["slos"]["tpot_150ms"]["feasible"])
|
|
|
|
def test_fixed_trace_materializer_uses_disjoint_prefixes(self) -> None:
|
|
module = load("prepare_fixed_frontier_traces.py")
|
|
with tempfile.TemporaryDirectory() as temporary:
|
|
root = Path(temporary)
|
|
manifest = module.materialize(
|
|
root,
|
|
input_tokens=32,
|
|
output_tokens=4,
|
|
requests=2,
|
|
rates=[2.0],
|
|
block_size=16,
|
|
)
|
|
with (root / "r2.csv").open(newline="") as handle:
|
|
rows = list(csv.DictReader(handle))
|
|
self.assertEqual(rows[0]["block_hash_ids"], "1|2")
|
|
self.assertEqual(rows[1]["block_hash_ids"], "3|4")
|
|
self.assertEqual(rows[1]["arrived_at"], "0.500000000000")
|
|
self.assertEqual(manifest["contract"]["prefix_caching"], False)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|