#!/usr/bin/env python3 from __future__ import annotations import asyncio import importlib.util import json import sys import tempfile import unittest from unittest import mock from pathlib import Path from types import SimpleNamespace from aiohttp import web HERE = Path(__file__).parent def load_module(name: str, filename: str): spec = importlib.util.spec_from_file_location(name, HERE / filename) module = importlib.util.module_from_spec(spec) sys.modules[name] = module assert spec.loader is not None spec.loader.exec_module(module) return module client = load_module("phase3_client", "opprof_phase3_client.py") controller = load_module("phase3_controller", "opprof_phase3_controller.py") sys.modules["opprof_phase3_controller"] = controller matrix = load_module("phase3_matrix", "opprof_phase3_matrix.py") def rows(n: int, arrival: str = "steady") -> list[dict]: return [ { "schema": 1, "request_id": f"T-{i:05d}", "pattern_id": "T", "kind": "synthetic", "input_tokens": 8, "output_tokens": 2, "arrival": arrival, "token_seed": i + 1, } for i in range(n) ] class MockServer: def __init__(self, delay: float = 0.01, trace_dir: Path | None = None) -> None: self.active = 0 self.max_active = 0 self.payloads = [] self.runner = None self.port = None self.delay = delay self.trace_dir = trace_dir self.profile_count = 0 async def completion(self, request): payload = await request.json() self.payloads.append(payload) self.active += 1 self.max_active = max(self.max_active, self.active) await asyncio.sleep(self.delay) response = web.StreamResponse( status=200, headers={"Content-Type": "text/event-stream"} ) await response.prepare(request) await response.write( b'data: {"choices":[{"text":"x"}],"usage":null}\n\n' ) usage = json.dumps( { "choices": [], "usage": { "prompt_tokens": len(payload["prompt"]), "completion_tokens": payload["max_tokens"], }, } ).encode() await response.write(b"data: " + usage + b"\n\n") await response.write(b"data: [DONE]\n\n") await response.write_eof() self.active -= 1 return response async def start_profile(self, request): self.profile_count += 1 path = self.trace_dir / f"window-{self.profile_count}.pt.trace.json" path.write_text('{"traceEvents": []}') return web.Response(status=200) async def stop_profile(self, request): return web.Response(status=200) async def start(self): app = web.Application() app.router.add_post("/v1/completions", self.completion) if self.trace_dir is not None: app.router.add_post("/start_profile", self.start_profile) app.router.add_post("/stop_profile", self.stop_profile) self.runner = web.AppRunner(app) await self.runner.setup() site = web.TCPSite(self.runner, "127.0.0.1", 0) await site.start() self.port = site._server.sockets[0].getsockname()[1] async def stop(self): await self.runner.cleanup() def run_args( manifest: Path, result_dir: Path, port: int, load_point: str = "saturation", saturation_result: Path | None = None, ): return SimpleNamespace( manifest=str(manifest), base_url=f"http://127.0.0.1:{port}", model="mock", load_point=load_point, request_rate="inf" if load_point == "saturation" else None, saturation_result=None if saturation_result is None else str(saturation_result), rate_fraction=0.60, max_concurrency=3, ignore_eos=True, temperature=0, warmup_seconds=0.04, clean_segment_seconds=0.04, num_clean_segments=3, profile_after_clean=False, num_profile_windows=0, profile_warmup_iterations=2, profile_active_iterations=8, profile_trace_dir=None, profile_timeout_seconds=1, recovery_seconds=0, post_clean_seconds=0, drain_timeout_seconds=1, workload_seed=20260712, server_seed=20260712, result_dir=str(result_dir), ) class Phase3ToolTests(unittest.TestCase): def write_warmup_stream(self, root: Path, tokens: list[int]) -> None: stream_dir = root / "opprof" stream_dir.mkdir() records = [] step = 0 for bin_index, token_count in enumerate(tokens): per_step, remainder = divmod(token_count, 16) for in_bin in range(16): records.append( { "step_index": step, "model_executed": True, "submit_mono_ns": int( 1e9 + (45 + bin_index * 5 + (in_bin + 0.5) / 16 * 5) * 1e9 ), "prefill_tokens": per_step + (in_bin < remainder), "decode_tokens": 0, } ) step += 1 (stream_dir / "test.jsonl").write_text( "".join(json.dumps(record) + "\n" for record in records) ) def test_p10_warmup_stability_accepts_flat_trailing_quartile(self): with tempfile.TemporaryDirectory() as tmp: root = Path(tmp) self.write_warmup_stream(root, [1600, 1600, 1600]) result = matrix.p10_warmup_stability(root, int(1e9)) self.assertTrue(result["passed"]) self.assertEqual(result["step_counts"], [16, 16, 16]) self.assertEqual(result["scheduled_token_throughput"], [320, 320, 320]) self.assertEqual(result["normalized_drift"], 0) def test_p10_warmup_stability_rejects_large_drift(self): with tempfile.TemporaryDirectory() as tmp: root = Path(tmp) self.write_warmup_stream(root, [1600, 3200, 4800]) result = matrix.p10_warmup_stability(root, int(1e9)) self.assertFalse(result["passed"]) self.assertGreater(result["normalized_drift"], 0.10) def test_only_clean_window_request_failures_are_hard_failures(self): requests = [ {"success": False, "completed_s": 59.9, "error_kind": "warmup"}, {"success": False, "completed_s": 60.0, "error_kind": "clean"}, {"success": False, "completed_s": 299.9, "error_kind": "clean"}, {"success": False, "completed_s": 300.0, "error_kind": "recovery"}, {"success": True, "completed_s": 100.0, "error_kind": None}, ] summary = matrix.summarize_request_failures(requests, 60.0, 300.0) self.assertEqual(summary["failed"], 4) self.assertEqual(summary["clean_failed"], 2) self.assertEqual(summary["excluded"], 2) self.assertEqual(summary["excluded_kinds"], {"warmup": 1, "recovery": 1}) def test_synthetic_prompt_exact_and_unique_early_prefix(self): generated = [client.synthetic_prompt(row) for row in rows(200)] self.assertTrue(all(len(value) == 8 for value in generated)) self.assertEqual(len({tuple(value[:3]) for value in generated}), 200) def test_p05_manifest_has_exact_half_modes(self): with tempfile.TemporaryDirectory() as tmp: out = Path(tmp) / "P05.jsonl" args = SimpleNamespace( id="P05", kind="synthetic", input_uniform=None, input_fixed=None, input_mixture='{"uniform:128:512":0.5,"uniform:4096:8192":0.5}', output_fixed=64, prefix="none", arrival="steady", num_requests=100, workload_seed=20260712, num_prefixes=0, prefix_len=0, suffix_fixed=0, out=str(out), ) client.materialize(args) manifest = client.load_manifest(out) self.assertEqual(sum(r["input_tokens"] <= 512 for r in manifest), 50) self.assertEqual(sum(r["input_tokens"] >= 4096 for r in manifest), 50) def test_prefix_pool_balanced(self): with tempfile.TemporaryDirectory() as tmp: out = Path(tmp) / "P08.jsonl" args = SimpleNamespace( id="P08", kind="prefix-pool", input_uniform=None, input_fixed=None, input_mixture=None, output_fixed=512, prefix="none", arrival="burst:8", num_requests=80, workload_seed=20260712, num_prefixes=8, prefix_len=1024, suffix_fixed=256, out=str(out), ) client.materialize(args) manifest = client.load_manifest(out) counts = {i: 0 for i in range(8)} for row in manifest: counts[row["prefix_id"]] += 1 self.assertEqual(row["input_tokens"], 1280) self.assertEqual(set(counts.values()), {10}) def test_fixed_duration_saturation_and_redaction(self): async def case(): with tempfile.TemporaryDirectory() as tmp: root = Path(tmp) manifest = root / "m.jsonl" client.atomic_jsonl(manifest, rows(1000)) mock = MockServer() await mock.start() try: result = await client.run_load( run_args(manifest, root / "result", mock.port) ) finally: await mock.stop() self.assertEqual(result["max_in_flight"], 3) self.assertEqual(mock.max_active, 3) self.assertAlmostEqual(result["clean"]["duration_s"], 0.12, places=9) self.assertEqual(len(result["segments"]), 3) self.assertLessEqual(result["drain_seconds"], 1) self.assertTrue( all(p["max_tokens"] == 2 and p["ignore_eos"] for p in mock.payloads) ) text = (root / "result/requests.jsonl").read_text() self.assertNotIn('"prompt":', text) self.assertNotIn('"text":', text) asyncio.run(case()) def test_profile_control_plane_is_not_starved_by_data_connector(self): async def case(): with tempfile.TemporaryDirectory() as tmp: root = Path(tmp) manifest = root / "m.jsonl" trace_dir = root / "traces" trace_dir.mkdir() client.atomic_jsonl(manifest, rows(1000)) server = MockServer(delay=0.5, trace_dir=trace_dir) await server.start() args = run_args(manifest, root / "result", server.port) args.max_concurrency = 1 args.warmup_seconds = 0.01 args.clean_segment_seconds = 0.01 args.num_clean_segments = 1 args.profile_after_clean = True args.num_profile_windows = 1 args.profile_trace_dir = str(trace_dir) args.recovery_seconds = 0 try: result = await client.run_load(args) finally: await server.stop() self.assertEqual(server.max_active, 1) self.assertEqual(server.profile_count, 1) self.assertLess(result["profiles"][0]["start_return_s"], 0.25) asyncio.run(case()) def test_drain_timeout_is_a_hard_sanity_failure(self): async def case(): with tempfile.TemporaryDirectory() as tmp: root = Path(tmp) manifest = root / "m.jsonl" client.atomic_jsonl(manifest, rows(1000)) mock = MockServer() await mock.start() try: args = run_args(manifest, root / "result", mock.port) args.drain_timeout_seconds = 0 with self.assertRaisesRegex(RuntimeError, "drain_within_timeout"): await client.run_load(args) finally: await mock.stop() sanity = json.loads((root / "result/sanity.json").read_text()) self.assertFalse(sanity["invariants"]["drain_within_timeout"]) asyncio.run(case()) def test_burst_schedule_is_eight_at_once(self): async def case(): with tempfile.TemporaryDirectory() as tmp: root = Path(tmp) manifest = root / "m.jsonl" client.atomic_jsonl(manifest, rows(1000, "burst:8")) sat = root / "sat.json" client.atomic_json( sat, {"clean": {"completed_throughput_rps": 100.0}} ) mock = MockServer() await mock.start() try: args = run_args( manifest, root / "result", mock.port, "moderate", sat ) args.warmup_seconds = 0 args.clean_segment_seconds = 0.4 / 3 result = await client.run_load(args) finally: await mock.stop() records = [ json.loads(line) for line in (root / "result/requests.jsonl").read_text().splitlines() ] groups = {} for record in records: groups.setdefault(round(record["scheduled_s"], 6), 0) groups[round(record["scheduled_s"], 6)] += 1 self.assertTrue(all(value == 8 for value in groups.values())) starts = sorted(groups) if len(starts) > 1: self.assertAlmostEqual(starts[1] - starts[0], 8 / 60, places=5) self.assertAlmostEqual(result["request_rate"], 60.0) asyncio.run(case()) def test_manifest_exhaustion_stops_instead_of_wrapping(self): async def case(): with tempfile.TemporaryDirectory() as tmp: root = Path(tmp) manifest = root / "m.jsonl" client.atomic_jsonl(manifest, rows(2)) mock = MockServer() await mock.start() try: with self.assertRaises(client.ManifestExhausted): await client.run_load( run_args(manifest, root / "result", mock.port) ) finally: await mock.stop() asyncio.run(case()) def test_cpu_affinity_sets_are_disjoint_and_cover_host(self): values = [] for spec in controller.CPU_MAP.values(): lo, hi = [int(x) for x in spec.split("-")] values.extend(range(lo, hi + 1)) self.assertEqual(sorted(values), list(range(160))) self.assertEqual(len(values), len(set(values))) def test_preflight_waits_for_three_stable_zero_samples(self): def sample(memory): return [ { "index": 0, "uuid": "GPU-test", "memory_used_mib": memory, "utilization_pct": 0, } ] with tempfile.TemporaryDirectory() as tmp: with ( mock.patch.object( controller, "gpu_query", side_effect=[sample(4), sample(0), sample(0), sample(0)], ), mock.patch.object(controller, "compute_apps", return_value=[]), mock.patch.object(controller, "run_text", return_value=""), mock.patch.object(controller.time, "sleep"), ): root = Path(tmp) controller.preflight([0], root) samples = json.loads((root / "gpu-before-samples.json").read_text()) self.assertEqual(len(samples), 4) self.assertEqual(samples[0]["gpus"][0]["memory_used_mib"], 4) self.assertEqual(samples[-1]["gpus"][0]["memory_used_mib"], 0) def test_kernel_mapping_priority(self): cases = { "void vllm::moe::topkGating": "moe_router", "ncclDevKernel_AllReduce": "collective", "flash_fwd_kernel": "attention", "nvjet_sm90_tst": "moe_gemm", "argmax_kernel": "sampler", "cutlass_gemm": "dense_gemm", "triton_red_fused_add_rms_norm": "norm_elementwise", "cache_swap_kernel": "kv_memory", "unknown": "other", } self.assertEqual( {name: controller.classify_kernel(name) for name in cases}, cases ) def test_atomic_state_replacement(self): with tempfile.TemporaryDirectory() as tmp: path = Path(tmp) / "state.json" controller.atomic_json(path, {"schema": 1, "value": 1}) controller.atomic_json(path, {"schema": 1, "value": 2}) self.assertEqual(json.loads(path.read_text())["value"], 2) self.assertFalse(list(path.parent.glob("*.tmp.*"))) def test_fingerprint_survives_json_roundtrip(self): original_run_text = controller.run_text original_hash = controller.sha256_file try: controller.run_text = lambda *args, **kwargs: "deadbeef\n" controller.sha256_file = lambda path: "a" * 64 fingerprint = controller.make_fingerprint() finally: controller.run_text = original_run_text controller.sha256_file = original_hash self.assertEqual(json.loads(json.dumps(fingerprint)), fingerprint) self.assertEqual(set(fingerprint["cpu_map"]), {str(i) for i in range(8)}) def test_server_shutdown_signals_parent_before_group(self): process = SimpleNamespace(pid=12345, poll=lambda: None) with ( mock.patch.object(controller.os, "kill") as kill, mock.patch.object(controller.os, "killpg") as killpg, mock.patch.object(controller, "_process_group_alive", return_value=False), ): controller.stop_servers([process]) kill.assert_called_once_with(12345, controller.signal.SIGINT) killpg.assert_not_called() if __name__ == "__main__": unittest.main(verbosity=2)