Spaces:
Running on Zero
Running on Zero
File size: 2,152 Bytes
993f563 d63e724 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 | """Accelerator cleanup behavior for regular and ZeroGPU processes."""
from __future__ import annotations
from types import SimpleNamespace
from core import runtime
class FakeCuda:
"""Record cleanup calls without importing or initializing CUDA."""
def __init__(self, initialized: bool) -> None:
self.initialized = initialized
self.empty_cache_calls = 0
self.ipc_collect_calls = 0
def is_initialized(self) -> bool:
return self.initialized
def empty_cache(self) -> None:
self.empty_cache_calls += 1
def ipc_collect(self) -> None:
self.ipc_collect_calls += 1
def test_cleanup_does_not_initialize_cuda(monkeypatch) -> None:
cuda = FakeCuda(initialized=False)
monkeypatch.setattr(runtime, "import_torch", lambda: SimpleNamespace(cuda=cuda))
runtime.cleanup_memory()
assert cuda.empty_cache_calls == 0
assert cuda.ipc_collect_calls == 0
def test_cleanup_releases_an_existing_cuda_context(monkeypatch) -> None:
cuda = FakeCuda(initialized=True)
monkeypatch.setattr(runtime, "import_torch", lambda: SimpleNamespace(cuda=cuda))
runtime.cleanup_memory()
assert cuda.empty_cache_calls == 1
assert cuda.ipc_collect_calls == 1
def test_cleanup_skips_an_unavailable_mps_backend(monkeypatch) -> None:
cuda = FakeCuda(initialized=False)
calls = 0
def empty_cache() -> None:
nonlocal calls
calls += 1
mps = SimpleNamespace(empty_cache=empty_cache)
mps_backend = SimpleNamespace(is_available=lambda: False)
torch = SimpleNamespace(
cuda=cuda,
mps=mps,
backends=SimpleNamespace(mps=mps_backend),
)
monkeypatch.setattr(runtime, "import_torch", lambda: torch)
runtime.cleanup_memory()
assert calls == 0
def test_cleanup_errors_never_escape(monkeypatch) -> None:
cuda = FakeCuda(initialized=True)
def fail_cleanup() -> None:
raise RuntimeError("cleanup failed")
cuda.empty_cache = fail_cleanup # type: ignore[method-assign]
monkeypatch.setattr(runtime, "import_torch", lambda: SimpleNamespace(cuda=cuda))
runtime.cleanup_memory()
|