67 lines
2.7 KiB
Python
67 lines
2.7 KiB
Python
from __future__ import annotations
|
|
|
|
import csv
|
|
import importlib.util
|
|
import json
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
|
|
SCRIPT = Path(__file__).with_name("frontier_prefill_grid.py")
|
|
SPEC = importlib.util.spec_from_file_location("frontier_prefill_grid", SCRIPT)
|
|
MODULE = importlib.util.module_from_spec(SPEC)
|
|
assert SPEC.loader is not None
|
|
sys.modules[SPEC.name] = MODULE
|
|
SPEC.loader.exec_module(MODULE)
|
|
|
|
|
|
def test_binary_search_lattice_covers_six_probe_tree() -> None:
|
|
anchors = MODULE.binary_search_lattice()
|
|
assert len(anchors) == 63
|
|
assert anchors[0] == 0.001953125
|
|
assert anchors[-1] == 0.123046875
|
|
assert 0.0625 in anchors
|
|
|
|
|
|
def test_prepare_preserves_stable_arrival_order(tmp_path: Path) -> None:
|
|
trace = tmp_path / "trace.jsonl"
|
|
rows = [
|
|
{"timestamp": 1.0, "sampling_u": 0.02, "input_length": 9000, "id": "a"},
|
|
{"timestamp": 0.5, "sampling_u": 0.03, "input_length": 100, "id": "b"},
|
|
{"timestamp": 1.0, "sampling_u": 0.01, "input_length": 200, "id": "c"},
|
|
{"timestamp": 0.2, "sampling_u": 0.01, "input_length": 40000, "id": "x"},
|
|
]
|
|
trace.write_text("".join(json.dumps(row) + "\n" for row in rows))
|
|
manifest_path = MODULE.prepare_traces(trace, tmp_path / "out", "")
|
|
manifest = json.loads(manifest_path.read_text())
|
|
record = manifest["anchors"][MODULE.anchor_key(0.0625)]
|
|
with Path(record["path"]).open(newline="") as source:
|
|
selected = list(csv.DictReader(source))
|
|
assert [row["source_request_id"] for row in selected] == ["b", "a", "c"]
|
|
assert [float(row["slo_ttft_ms"]) for row in selected] == [1000, 2000, 1000]
|
|
|
|
|
|
def test_score_uses_trace_slo_and_request_id_mapping(tmp_path: Path) -> None:
|
|
trace = tmp_path / "trace.csv"
|
|
with trace.open("w", newline="") as output:
|
|
writer = csv.DictWriter(
|
|
output,
|
|
fieldnames=["arrived_at", "num_prefill_tokens", "num_decode_tokens", "slo_ttft_ms"],
|
|
)
|
|
writer.writeheader()
|
|
writer.writerows(
|
|
[
|
|
{"arrived_at": 0, "num_prefill_tokens": 10, "num_decode_tokens": 1, "slo_ttft_ms": 1000},
|
|
{"arrived_at": 1, "num_prefill_tokens": 9000, "num_decode_tokens": 1, "slo_ttft_ms": 2000},
|
|
]
|
|
)
|
|
metrics = tmp_path / "request_metrics.csv"
|
|
with metrics.open("w", newline="") as output:
|
|
writer = csv.DictWriter(output, fieldnames=["Request Id", "ttft"])
|
|
writer.writeheader()
|
|
writer.writerows([{"Request Id": 1, "ttft": 1900}, {"Request Id": 0, "ttft": 1100}])
|
|
score = MODULE.score_request_metrics(trace, metrics)
|
|
assert score["passed_request_count"] == 1
|
|
assert score["slo_pass_rate"] == 0.5
|
|
assert score["feasible"] is False
|