diff --git a/obliteratus/benchmark_lifecycle.py b/obliteratus/benchmark_lifecycle.py index f6c7d09..a51e795 100644 --- a/obliteratus/benchmark_lifecycle.py +++ b/obliteratus/benchmark_lifecycle.py @@ -7,6 +7,12 @@ import gc from obliteratus.gpu_lifecycle import measure_torch_memory +# PyTorch keeps a small allocator/driver context alive in a long-running UI +# process after every model tensor has been freed. The root supervisor remains +# authoritative and independently refuses release above 2 GiB of process VRAM. +MAX_RELEASABLE_ALLOCATOR_RESIDUE_BYTES = 128 * 1024 * 1024 + + def admit_benchmark(lifecycle, model_id: str) -> Exception | None: """Request admission before a benchmark worker may allocate CUDA memory.""" try: @@ -42,7 +48,10 @@ def release_benchmark_pipeline( torch_module.cuda.synchronize() device_module.empty_cache() memory = measure_torch_memory(torch_module) - if memory.allocated_bytes or memory.reserved_bytes: + if ( + memory.allocated_bytes > MAX_RELEASABLE_ALLOCATOR_RESIDUE_BYTES + or memory.reserved_bytes > MAX_RELEASABLE_ALLOCATOR_RESIDUE_BYTES + ): lifecycle.resize(memory) raise RuntimeError( "benchmark CUDA allocations remain after cleanup; retaining GPU lease" diff --git a/tests/test_app_benchmark_lifecycle.py b/tests/test_app_benchmark_lifecycle.py index 7ceb590..30e8c04 100644 --- a/tests/test_app_benchmark_lifecycle.py +++ b/tests/test_app_benchmark_lifecycle.py @@ -6,6 +6,7 @@ from types import SimpleNamespace import pytest from obliteratus.benchmark_lifecycle import ( + MAX_RELEASABLE_ALLOCATOR_RESIDUE_BYTES, admit_benchmark, mark_benchmark_ready, release_benchmark_pipeline, @@ -99,9 +100,36 @@ def test_cleanup_releases_only_after_cuda_is_gone(monkeypatch): assert lifecycle.events == [("release", "benchmark_complete")] +def test_cleanup_releases_with_bounded_cuda_context_residue(monkeypatch): + lifecycle = _LifecycleRecorder() + memory = MemoryUsage( + allocated_bytes=25_559_040, + reserved_bytes=62_914_560, + device_count=3, + ) + assert memory.reserved_bytes < MAX_RELEASABLE_ALLOCATOR_RESIDUE_BYTES + monkeypatch.setattr( + "obliteratus.benchmark_lifecycle.measure_torch_memory", lambda _torch: memory + ) + + release_benchmark_pipeline( + [SimpleNamespace(handle=SimpleNamespace(model=object(), tokenizer=object()))], + reason="benchmark_complete", + lifecycle=lifecycle, + torch_module=SimpleNamespace(cuda=SimpleNamespace(is_available=lambda: False)), + device_module=SimpleNamespace(empty_cache=lambda: None), + ) + + 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) + memory = MemoryUsage( + allocated_bytes=MAX_RELEASABLE_ALLOCATOR_RESIDUE_BYTES + 1, + reserved_bytes=MAX_RELEASABLE_ALLOCATOR_RESIDUE_BYTES + 1, + device_count=1, + ) monkeypatch.setattr( "obliteratus.benchmark_lifecycle.measure_torch_memory", lambda _torch: memory )