| import json |
| from types import SimpleNamespace |
|
|
| import mlx.core as mx |
|
|
| from h3_mlx import benchmark, runtime |
|
|
|
|
| def test_summary_excludes_warmup_and_retains_every_resident_sample( |
| tmp_path, monkeypatch |
| ): |
| calls = [] |
| cleared = [] |
|
|
| class FakePipeline: |
| def __init__(self, *args, **kwargs): |
| assert kwargs["kernel"] == "metal-qk-exact" |
| assert kwargs["tune_matmul"] is True |
|
|
| def generate(self, prompt, **kwargs): |
| calls.append(kwargs) |
| return None, { |
| "png_sha256": "a" * 64, |
| "conditioning_cache_hit": len(calls) > 1, |
| "dit_was_resident": len(calls) > 1, |
| } |
|
|
| def clear(self): |
| cleared.append(True) |
|
|
| clock = iter([0.0, 50.0, 50.0, 55.0, 55.0, 62.0]) |
| monkeypatch.setattr( |
| benchmark, "time", SimpleNamespace(perf_counter=lambda: next(clock)) |
| ) |
| monkeypatch.setattr(runtime, "Pipeline", FakePipeline) |
| monkeypatch.setattr(runtime, "environment", lambda: {"machine": "test"}) |
| monkeypatch.setattr(mx, "synchronize", lambda: None) |
| output = tmp_path / "run" |
| benchmark.main( |
| [ |
| "--model", |
| "model", |
| "--text-encoder", |
| "encoder", |
| "--prompt", |
| "test", |
| "--warmup", |
| "1", |
| "--repeats", |
| "2", |
| "--output", |
| str(output), |
| ] |
| ) |
| summary = json.loads((output / "summary.json").read_text()) |
| assert summary["median_seconds"] == 6.0 |
| assert [s["wall_seconds"] for s in summary["samples"]] == [50.0, 5.0, 7.0] |
| assert [s["phase"] for s in summary["samples"]] == [ |
| "warmup", |
| "measured", |
| "measured", |
| ] |
| assert all(s["dit_was_resident"] for s in summary["samples"][1:]) |
| assert all(c["save_trace"] for c in calls) |
| assert cleared == [True] |
|
|
