test(gpu): cover benchmark lifecycle orchestration

tests/test_app_benchmark_lifecycle.py

@@ -1,92 +1,120 @@
 """GPU admission and cleanup contracts for application benchmark paths."""
 
-from __future__ import annotations
-
-import ast
 from pathlib import Path
 from types import SimpleNamespace
 
 import pytest
 
+from obliteratus.benchmark_lifecycle import (
+    admit_benchmark,
+    mark_benchmark_ready,
+    release_benchmark_pipeline,
+)
 from obliteratus.gpu_lifecycle import MemoryUsage
 
 
-APP_SOURCE = Path("app.py").read_text(encoding="utf-8")
-APP_TREE = ast.parse(APP_SOURCE)
-
-
-def _function_node(name: str) -> ast.FunctionDef:
-    return next(
-        node
-        for node in APP_TREE.body
-        if isinstance(node, ast.FunctionDef) and node.name == name
-    )
-
-
-def _load_cleanup_function(namespace: dict):
-    node = _function_node("_release_benchmark_pipeline")
-    module = ast.fix_missing_locations(ast.Module(body=[node], type_ignores=[]))
-    exec(compile(module, "app.py", "exec"), namespace)
-    return namespace["_release_benchmark_pipeline"]
-
-
 class _LifecycleRecorder:
-    def __init__(self):
+    def __init__(self, *, loading_error=None):
         self.events = []
+        self.loading_error = loading_error
+
+    def loading(self, model_id):
+        self.events.append(("loading", model_id))
+        if self.loading_error:
+            raise self.loading_error
 
     def resize(self, memory):
         self.events.append(("resize", memory))
 
+    def ready(self, memory):
+        self.events.append(("ready", memory))
+
     def release(self, *, reason):
         self.events.append(("release", reason))
 
 
-def test_benchmark_entrypoints_admit_before_worker_start():
-    for name in ("benchmark", "benchmark_multi_model"):
-        source = ast.get_source_segment(APP_SOURCE, _function_node(name))
-        assert source.index("_gpu_lifecycle.loading(model_id)") < source.index("worker.start()")
-        assert "result.status == \"done\"" in source
-        assert "_release_benchmark_pipeline(" in source
+def test_benchmark_entrypoints_use_lifecycle_helpers_before_worker_start():
+    source = Path("app.py").read_text(encoding="utf-8")
+    for marker in ("def benchmark(", "def benchmark_multi_model("):
+        body = source[source.index(marker) :]
+        assert body.index("admit_benchmark(_gpu_lifecycle, model_id)") < body.index(
+            "worker.start()"
+        )
+        assert "mark_benchmark_ready(_gpu_lifecycle, torch)" in body
+        assert "release_benchmark_pipeline(" in body
 
 
-def test_benchmark_cleanup_releases_only_after_cuda_is_gone():
+def test_admission_success_allows_worker_start():
     lifecycle = _LifecycleRecorder()
-    calls = []
-    cleanup = _load_cleanup_function(
-        {
-            "gc": SimpleNamespace(collect=lambda: calls.append("gc")),
-            "torch": SimpleNamespace(cuda=SimpleNamespace(is_available=lambda: False)),
-            "dev": SimpleNamespace(empty_cache=lambda: calls.append("empty_cache")),
-            "measure_torch_memory": lambda _torch: MemoryUsage(),
-            "_gpu_lifecycle": lifecycle,
-        }
+    assert admit_benchmark(lifecycle, "org/model") is None
+    assert lifecycle.events == [("loading", "org/model")]
+
+
+def test_admission_failure_releases_and_returns_error():
+    error = RuntimeError("denied")
+    lifecycle = _LifecycleRecorder(loading_error=error)
+    assert admit_benchmark(lifecycle, "org/model") is error
+    assert lifecycle.events == [
+        ("loading", "org/model"),
+        ("release", "benchmark_admission_failed"),
+    ]
+
+
+def test_ready_publishes_measured_memory(monkeypatch):
+    memory = MemoryUsage(allocated_bytes=3, reserved_bytes=4, device_count=1)
+    monkeypatch.setattr(
+        "obliteratus.benchmark_lifecycle.measure_torch_memory", lambda _torch: memory
     )
+    lifecycle = _LifecycleRecorder()
+    mark_benchmark_ready(lifecycle, object())
+    assert lifecycle.events == [("resize", memory), ("ready", memory)]
+
+
+def test_cleanup_releases_only_after_cuda_is_gone(monkeypatch):
+    lifecycle = _LifecycleRecorder()
+    monkeypatch.setattr(
+        "obliteratus.benchmark_lifecycle.measure_torch_memory",
+        lambda _torch: MemoryUsage(),
+    )
+    calls = []
+    torch_module = SimpleNamespace(
+        cuda=SimpleNamespace(
+            is_available=lambda: True,
+            synchronize=lambda: calls.append("synchronize"),
+        )
+    )
+    device_module = SimpleNamespace(empty_cache=lambda: calls.append("empty_cache"))
     handle = SimpleNamespace(model=object(), tokenizer=object())
 
-    cleanup([SimpleNamespace(handle=handle)], reason="benchmark_complete")
+    release_benchmark_pipeline(
+        [SimpleNamespace(handle=handle)],
+        reason="benchmark_complete",
+        lifecycle=lifecycle,
+        torch_module=torch_module,
+        device_module=device_module,
+    )
 
     assert handle.model is None and handle.tokenizer is None
-    assert calls == ["gc", "empty_cache"]
+    assert calls == ["synchronize", "empty_cache"]
     assert lifecycle.events == [("release", "benchmark_complete")]
 
 
-def test_benchmark_cleanup_retains_lease_when_cuda_remains():
+def test_cleanup_retains_lease_when_cuda_remains(monkeypatch):
     lifecycle = _LifecycleRecorder()
     memory = MemoryUsage(allocated_bytes=1, reserved_bytes=2, device_count=1)
-    cleanup = _load_cleanup_function(
-        {
-            "gc": SimpleNamespace(collect=lambda: None),
-            "torch": SimpleNamespace(cuda=SimpleNamespace(is_available=lambda: False)),
-            "dev": SimpleNamespace(empty_cache=lambda: None),
-            "measure_torch_memory": lambda _torch: memory,
-            "_gpu_lifecycle": lifecycle,
-        }
+    monkeypatch.setattr(
+        "obliteratus.benchmark_lifecycle.measure_torch_memory", lambda _torch: memory
     )
+    torch_module = SimpleNamespace(cuda=SimpleNamespace(is_available=lambda: False))
+    device_module = SimpleNamespace(empty_cache=lambda: None)
 
     with pytest.raises(RuntimeError, match="retaining GPU lease"):
-        cleanup(
+        release_benchmark_pipeline(
             [SimpleNamespace(handle=SimpleNamespace(model=object(), tokenizer=object()))],
             reason="benchmark_complete",
+            lifecycle=lifecycle,
+            torch_module=torch_module,
+            device_module=device_module,
         )
 
     assert lifecycle.events == [("resize", memory)]