From 6df26d07e7d98d21ade9153f2acee6e0570cbf51 Mon Sep 17 00:00:00 2001
From: Joseph Magly <1159087+jmagly@users.noreply.github.com>
Date: Wed, 26 Aug 2026 22:51:13 -0400
Subject: [PATCH] test(gpu): keep benchmark contracts CPU-only

---
 tests/test_app_benchmark_lifecycle.py | 121 ++++++++++++++++----------
 1 file changed, 73 insertions(+), 48 deletions(-)

diff --git a/tests/test_app_benchmark_lifecycle.py b/tests/test_app_benchmark_lifecycle.py
index 26b5e45..75c3ca4 100644
--- a/tests/test_app_benchmark_lifecycle.py
+++ b/tests/test_app_benchmark_lifecycle.py
@@ -2,66 +2,91 @@
 
 from __future__ import annotations
 
-import subprocess
-import sys
-
-
-def test_benchmark_gpu_lifecycle_contracts():
-    script = r'''
+import ast
+from pathlib import Path
 from types import SimpleNamespace
-import inspect
 
-import app
+import pytest
+
 from obliteratus.gpu_lifecycle import MemoryUsage
 
-for entrypoint in (app.benchmark, app.benchmark_multi_model):
-    source = inspect.getsource(entrypoint)
-    assert source.index("_gpu_lifecycle.loading(model_id)") < source.index("worker.start()")
-    assert "result.status == \"done\"" in source
-    assert "_release_benchmark_pipeline(" in source
 
-class LifecycleRecorder:
+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):
         self.events = []
+
     def resize(self, memory):
         self.events.append(("resize", memory))
+
     def release(self, *, reason):
         self.events.append(("release", reason))
 
-lifecycle = LifecycleRecorder()
-handle = SimpleNamespace(model=object(), tokenizer=object())
-pipeline_ref = [SimpleNamespace(handle=handle)]
-calls = []
-app._gpu_lifecycle = lifecycle
-app.gc.collect = lambda: calls.append("gc")
-app.torch.cuda.is_available = lambda: False
-app.dev.empty_cache = lambda: calls.append("empty_cache")
-app.measure_torch_memory = lambda _torch: MemoryUsage()
-app._release_benchmark_pipeline(pipeline_ref, reason="benchmark_complete")
-assert handle.model is None and handle.tokenizer is None
-assert calls == ["gc", "empty_cache"]
-assert lifecycle.events == [("release", "benchmark_complete")]
 
-lifecycle = LifecycleRecorder()
-app._gpu_lifecycle = lifecycle
-memory = MemoryUsage(allocated_bytes=1, reserved_bytes=2, device_count=1)
-app.measure_torch_memory = lambda _torch: memory
-try:
-    app._release_benchmark_pipeline(
-        [SimpleNamespace(handle=SimpleNamespace(model=object(), tokenizer=object()))],
-        reason="benchmark_complete",
+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_cleanup_releases_only_after_cuda_is_gone():
+    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,
+        }
     )
-except RuntimeError as error:
-    assert "retaining GPU lease" in str(error)
-else:
-    raise AssertionError("cleanup must fail closed while CUDA allocations remain")
-assert lifecycle.events == [("resize", memory)]
-'''
-    result = subprocess.run(
-        [sys.executable, "-c", script],
-        capture_output=True,
-        text=True,
-        timeout=120,
-        check=False,
+    handle = SimpleNamespace(model=object(), tokenizer=object())
+
+    cleanup([SimpleNamespace(handle=handle)], reason="benchmark_complete")
+
+    assert handle.model is None and handle.tokenizer is None
+    assert calls == ["gc", "empty_cache"]
+    assert lifecycle.events == [("release", "benchmark_complete")]
+
+
+def test_benchmark_cleanup_retains_lease_when_cuda_remains():
+    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,
+        }
     )
-    assert result.returncode == 0, result.stdout + result.stderr
+
+    with pytest.raises(RuntimeError, match="retaining GPU lease"):
+        cleanup(
+            [SimpleNamespace(handle=SimpleNamespace(model=object(), tokenizer=object()))],
+            reason="benchmark_complete",
+        )
+
+    assert lifecycle.events == [("resize", memory)]