From 9ab258359ee7bb26e7e2dd1a74f18a4d4a757cd2 Mon Sep 17 00:00:00 2001 From: Joseph Magly <1159087+jmagly@users.noreply.github.com> Date: Wed, 26 Aug 2026 22:48:59 -0400 Subject: [PATCH] fix(gpu): admit benchmark loads through lifecycle broker --- app.py | 91 ++++++++++++++++++++------- ci/test-risk-map.json | 1 + tests/test_app_benchmark_lifecycle.py | 67 ++++++++++++++++++++ 3 files changed, 136 insertions(+), 23 deletions(-) create mode 100644 tests/test_app_benchmark_lifecycle.py diff --git a/app.py b/app.py index c205b50..b59808a 100644 --- a/app.py +++ b/app.py @@ -937,6 +937,26 @@ def _clear_gpu(*, release_lifecycle: bool = True): _gpu_lifecycle.release(reason="unload") +def _release_benchmark_pipeline(pipeline_ref, *, reason: str) -> None: + """Free a benchmark-local model before releasing supervisor ownership.""" + pipeline = pipeline_ref[0] + if pipeline is not None: + if getattr(pipeline, "handle", None): + pipeline.handle.model = None + pipeline.handle.tokenizer = None + gc.collect() + if torch.cuda.is_available(): + torch.cuda.synchronize() + dev.empty_cache() + memory = measure_torch_memory(torch) + if memory.allocated_bytes or memory.reserved_bytes: + _gpu_lifecycle.resize(memory) + raise RuntimeError( + "benchmark CUDA allocations remain after cleanup; retaining GPU lease" + ) + _gpu_lifecycle.release(reason=reason) + + def _checkpoint_is_available(checkpoint: str | None) -> bool: """Return whether *checkpoint* is a recoverable model directory.""" if not checkpoint: @@ -1425,6 +1445,10 @@ 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": + memory = measure_torch_memory(torch) + _gpu_lifecycle.resize(memory) + _gpu_lifecycle.ready(memory) quantization = load_settings.quantization @@ -1471,9 +1495,25 @@ 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", + ) + except Exception as cleanup_error: + if run_error is None: + run_error = cleanup_error + try: + _gpu_lifecycle.loading(model_id) + except Exception as error: + run_error = error + _gpu_lifecycle.release(reason="benchmark_admission_failed") 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 +1528,8 @@ def benchmark( ) time.sleep(0.5) - worker.join() + if worker_started: + worker.join() elapsed = time.time() - t_start # Collect results @@ -1584,15 +1625,7 @@ 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") yield ( f"**{method_key} complete** ({mi + 1}/{len(methods_to_test)}) \u2014 {_bench_elapsed()}", @@ -1789,7 +1822,10 @@ 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": + memory = measure_torch_memory(torch) + _gpu_lifecycle.resize(memory) + _gpu_lifecycle.ready(memory) try: load_settings = _resolve_ui_load_settings( @@ -1842,9 +1878,25 @@ 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", + ) + except Exception as cleanup_error: + if run_error is None: + run_error = cleanup_error + try: + _gpu_lifecycle.loading(model_id) + except Exception as error: + run_error = error + _gpu_lifecycle.release(reason="benchmark_admission_failed") 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 +1910,8 @@ def benchmark_multi_model( ) time.sleep(0.5) - worker.join() + if worker_started: + worker.join() elapsed = time.time() - t_start entry = { @@ -1950,15 +2003,7 @@ 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") yield ( f"**{model_id} complete** ({mi + 1}/{len(model_choices)}) \u2014 {_mm_elapsed()}", diff --git a/ci/test-risk-map.json b/ci/test-risk-map.json index 540a5ef..2645a80 100644 --- a/ci/test-risk-map.json +++ b/ci/test-risk-map.json @@ -72,6 +72,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", diff --git a/tests/test_app_benchmark_lifecycle.py b/tests/test_app_benchmark_lifecycle.py new file mode 100644 index 0000000..26b5e45 --- /dev/null +++ b/tests/test_app_benchmark_lifecycle.py @@ -0,0 +1,67 @@ +"""GPU admission and cleanup contracts for application benchmark paths.""" + +from __future__ import annotations + +import subprocess +import sys + + +def test_benchmark_gpu_lifecycle_contracts(): + script = r''' +from types import SimpleNamespace +import inspect + +import app +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: + 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", + ) +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, + ) + assert result.returncode == 0, result.stdout + result.stderr