From f31d1824562709a720d402bcc98f54829f6df7f5 Mon Sep 17 00:00:00 2001 From: Joseph Magly <1159087+jmagly@users.noreply.github.com> Date: Wed, 26 Aug 2026 22:57:37 -0400 Subject: [PATCH] test(gpu): cover benchmark lifecycle orchestration --- app.py | 71 ++++++-------- ci/test-risk-map.json | 12 +++ obliteratus/benchmark_lifecycle.py | 50 ++++++++++ tests/test_app_benchmark_lifecycle.py | 130 ++++++++++++++++---------- 4 files changed, 172 insertions(+), 91 deletions(-) create mode 100644 obliteratus/benchmark_lifecycle.py diff --git a/app.py b/app.py index b59808a..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, @@ -937,26 +942,6 @@ 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: @@ -1446,9 +1431,7 @@ def benchmark( 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) + mark_benchmark_ready(_gpu_lifecycle, torch) quantization = load_settings.quantization @@ -1497,19 +1480,18 @@ def benchmark( run_error = e finally: try: - _release_benchmark_pipeline( + 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 - try: - _gpu_lifecycle.loading(model_id) - except Exception as error: - run_error = error - _gpu_lifecycle.release(reason="benchmark_admission_failed") + run_error = admit_benchmark(_gpu_lifecycle, model_id) worker = threading.Thread(target=run_pipeline, daemon=True) worker_started = run_error is None if worker_started: @@ -1625,7 +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: - _release_benchmark_pipeline(pipeline_ref, reason="benchmark_complete") + 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()}", @@ -1823,9 +1811,7 @@ def benchmark_multi_model( def on_stage(result): if result.stage == "summon" and result.status == "done": - memory = measure_torch_memory(torch) - _gpu_lifecycle.resize(memory) - _gpu_lifecycle.ready(memory) + mark_benchmark_ready(_gpu_lifecycle, torch) try: load_settings = _resolve_ui_load_settings( @@ -1880,19 +1866,18 @@ def benchmark_multi_model( run_error = e finally: try: - _release_benchmark_pipeline( + 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 - try: - _gpu_lifecycle.loading(model_id) - except Exception as error: - run_error = error - _gpu_lifecycle.release(reason="benchmark_admission_failed") + run_error = admit_benchmark(_gpu_lifecycle, model_id) worker = threading.Thread(target=run_pipeline, daemon=True) worker_started = run_error is None if worker_started: @@ -2003,7 +1988,13 @@ def benchmark_multi_model( # Explicitly free pipeline and model before next iteration if pipeline_ref[0] is not None: - _release_benchmark_pipeline(pipeline_ref, reason="benchmark_complete") + 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 2645a80..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", @@ -757,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 index 75c3ca4..7ceb590 100644 --- a/tests/test_app_benchmark_lifecycle.py +++ b/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)]