Merge pull request #172 from jmagly/fix/171-benchmark-gpu-lifecycle

fix(gpu): broker benchmark model loads

app.py

@@ -59,6 +59,11 @@ import gradio as gr
 import torch
 from obliteratus import device as dev
 from obliteratus.gpu_lifecycle import from_environment, measure_torch_memory
+from obliteratus.benchmark_lifecycle import (
+    admit_benchmark,
+    mark_benchmark_ready,
+    release_benchmark_pipeline,
+)
 from obliteratus.credential_sources import resolve_first, resolve_secret, secret_available
 from obliteratus.ui_vram import (
     DEFAULT_VRAM_REFRESH_INTERVAL,
@@ -1425,6 +1430,8 @@ def benchmark(
             stage_key = result.stage
             if result.status == "running":
                 run_logs.append(f"{stage_key.upper()} — {result.message}")
+            elif stage_key == "summon" and result.status == "done":
+                mark_benchmark_ready(_gpu_lifecycle, torch)
 
         quantization = load_settings.quantization
 
@@ -1471,9 +1478,24 @@ def benchmark(
             except Exception as e:
                 nonlocal run_error
                 run_error = e
+            finally:
+                try:
+                    release_benchmark_pipeline(
+                        pipeline_ref,
+                        reason="benchmark_complete" if run_error is None else "benchmark_failed",
+                        lifecycle=_gpu_lifecycle,
+                        torch_module=torch,
+                        device_module=dev,
+                    )
+                except Exception as cleanup_error:
+                    if run_error is None:
+                        run_error = cleanup_error
 
+        run_error = admit_benchmark(_gpu_lifecycle, model_id)
         worker = threading.Thread(target=run_pipeline, daemon=True)
-        worker.start()
+        worker_started = run_error is None
+        if worker_started:
+            worker.start()
 
         # Stream log updates while pipeline runs
         last_count = len(all_logs)
@@ -1488,7 +1510,8 @@ def benchmark(
                 )
             time.sleep(0.5)
 
-        worker.join()
+        if worker_started:
+            worker.join()
         elapsed = time.time() - t_start
 
         # Collect results
@@ -1584,15 +1607,13 @@ def benchmark(
         # before the next benchmark iteration. _clear_gpu() only clears
         # _state["model"], not the benchmark-local pipeline object.
         if pipeline_ref[0] is not None:
-            try:
-                if hasattr(pipeline_ref[0], "handle") and pipeline_ref[0].handle:
-                    pipeline_ref[0].handle.model = None
-                    pipeline_ref[0].handle.tokenizer = None
-            except Exception:
-                pass
-            pipeline_ref[0] = None
-        gc.collect()
-        dev.empty_cache()
+            release_benchmark_pipeline(
+                pipeline_ref,
+                reason="benchmark_complete",
+                lifecycle=_gpu_lifecycle,
+                torch_module=torch,
+                device_module=dev,
+            )
 
         yield (
             f"**{method_key} complete** ({mi + 1}/{len(methods_to_test)}) \u2014 {_bench_elapsed()}",
@@ -1789,7 +1810,8 @@ def benchmark_multi_model(
             all_logs.append(f"  [{_mid.split('/')[-1]}] {msg}")
 
         def on_stage(result):
-            pass
+            if result.stage == "summon" and result.status == "done":
+                mark_benchmark_ready(_gpu_lifecycle, torch)
 
         try:
             load_settings = _resolve_ui_load_settings(
@@ -1842,9 +1864,24 @@ def benchmark_multi_model(
             except Exception as e:
                 nonlocal run_error
                 run_error = e
+            finally:
+                try:
+                    release_benchmark_pipeline(
+                        pipeline_ref,
+                        reason="benchmark_complete" if run_error is None else "benchmark_failed",
+                        lifecycle=_gpu_lifecycle,
+                        torch_module=torch,
+                        device_module=dev,
+                    )
+                except Exception as cleanup_error:
+                    if run_error is None:
+                        run_error = cleanup_error
 
+        run_error = admit_benchmark(_gpu_lifecycle, model_id)
         worker = threading.Thread(target=run_pipeline, daemon=True)
-        worker.start()
+        worker_started = run_error is None
+        if worker_started:
+            worker.start()
 
         last_count = len(all_logs)
         while worker.is_alive():
@@ -1858,7 +1895,8 @@ def benchmark_multi_model(
                 )
             time.sleep(0.5)
 
-        worker.join()
+        if worker_started:
+            worker.join()
         elapsed = time.time() - t_start
 
         entry = {
@@ -1950,15 +1988,13 @@ def benchmark_multi_model(
 
         # Explicitly free pipeline and model before next iteration
         if pipeline_ref[0] is not None:
-            try:
-                if hasattr(pipeline_ref[0], "handle") and pipeline_ref[0].handle:
-                    pipeline_ref[0].handle.model = None
-                    pipeline_ref[0].handle.tokenizer = None
-            except Exception:
-                pass
-            pipeline_ref[0] = None
-        gc.collect()
-        dev.empty_cache()
+            release_benchmark_pipeline(
+                pipeline_ref,
+                reason="benchmark_complete",
+                lifecycle=_gpu_lifecycle,
+                torch_module=torch,
+                device_module=dev,
+            )
 
         yield (
             f"**{model_id} complete** ({mi + 1}/{len(model_choices)}) \u2014 {_mm_elapsed()}",

ci/test-risk-map.json

@@ -55,6 +55,7 @@
       ],
       "paths": [
         "obliteratus/abliterate.py",
+        "obliteratus/benchmark_lifecycle.py",
         "obliteratus/models/qwen35_contracts.py",
         "obliteratus/models/offload_surgery.py",
         "obliteratus/persistence_contracts.py",
@@ -72,6 +73,7 @@
       "required_tests": [
         "tests/test_abliterate.py",
         "tests/test_abliterate_extended.py",
+        "tests/test_app_benchmark_lifecycle.py",
         "tests/test_qwen35_contracts.py",
         "tests/test_auto_obliterate.py",
         "tests/test_bayesian_optimizer_contracts.py",
@@ -756,6 +758,17 @@
       ],
       "conditional_gates": []
     },
+    {
+      "path": "obliteratus/benchmark_lifecycle.py",
+      "risk_class": "cpu-contract",
+      "risk": "benchmark GPU admission, readiness publication, and fail-closed release ordering",
+      "contract_owner": "operator interface maintainers",
+      "required_tests": [
+        "tests/test_app_benchmark_lifecycle.py",
+        "tests/test_gpu_lifecycle.py"
+      ],
+      "conditional_gates": []
+    },
     {
       "path": "obliteratus/service_contracts.py",
       "risk_class": "cpu-contract",

obliteratus/benchmark_lifecycle.py

@@ -0,0 +1,50 @@
+"""GPU admission helpers shared by the UI benchmark entry points."""
+
+from __future__ import annotations
+
+import gc
+
+from obliteratus.gpu_lifecycle import measure_torch_memory
+
+
+def admit_benchmark(lifecycle, model_id: str) -> Exception | None:
+    """Request admission before a benchmark worker may allocate CUDA memory."""
+    try:
+        lifecycle.loading(model_id)
+    except Exception as error:
+        lifecycle.release(reason="benchmark_admission_failed")
+        return error
+    return None
+
+
+def mark_benchmark_ready(lifecycle, torch_module) -> None:
+    """Publish measured residency after the pipeline finishes model loading."""
+    memory = measure_torch_memory(torch_module)
+    lifecycle.resize(memory)
+    lifecycle.ready(memory)
+
+
+def release_benchmark_pipeline(
+    pipeline_ref,
+    *,
+    reason: str,
+    lifecycle,
+    torch_module,
+    device_module,
+) -> None:
+    """Free a benchmark-local model before releasing supervisor ownership."""
+    pipeline = pipeline_ref[0]
+    if pipeline is not None and getattr(pipeline, "handle", None):
+        pipeline.handle.model = None
+        pipeline.handle.tokenizer = None
+    gc.collect()
+    if torch_module.cuda.is_available():
+        torch_module.cuda.synchronize()
+    device_module.empty_cache()
+    memory = measure_torch_memory(torch_module)
+    if memory.allocated_bytes or memory.reserved_bytes:
+        lifecycle.resize(memory)
+        raise RuntimeError(
+            "benchmark CUDA allocations remain after cleanup; retaining GPU lease"
+        )
+    lifecycle.release(reason=reason)

tests/test_app_benchmark_lifecycle.py

@@ -0,0 +1,120 @@
+"""GPU admission and cleanup contracts for application benchmark paths."""
+
+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
+
+
+class _LifecycleRecorder:
+    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_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_admission_success_allows_worker_start():
+    lifecycle = _LifecycleRecorder()
+    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())
+
+    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 == ["synchronize", "empty_cache"]
+    assert lifecycle.events == [("release", "benchmark_complete")]
+
+
+def test_cleanup_retains_lease_when_cuda_remains(monkeypatch):
+    lifecycle = _LifecycleRecorder()
+    memory = MemoryUsage(allocated_bytes=1, reserved_bytes=2, device_count=1)
+    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"):
+        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)]