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 1/3] 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 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 2/3] 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)] 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 3/3] 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)]