diff --git a/app.py b/app.py index c205b50..aad4e32 100644 --- a/app.py +++ b/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()}", diff --git a/ci/test-risk-map.json b/ci/test-risk-map.json index 540a5ef..5f0651b 100644 --- a/ci/test-risk-map.json +++ b/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", diff --git a/obliteratus/benchmark_lifecycle.py b/obliteratus/benchmark_lifecycle.py new file mode 100644 index 0000000..f6c7d09 --- /dev/null +++ b/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) diff --git a/tests/test_app_benchmark_lifecycle.py b/tests/test_app_benchmark_lifecycle.py new file mode 100644 index 0000000..7ceb590 --- /dev/null +++ b/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)]