diff --git a/src/keep_gpu/single_gpu_controller/cuda_gpu_controller.py b/src/keep_gpu/single_gpu_controller/cuda_gpu_controller.py index 81ef874..8ad814b 100644 --- a/src/keep_gpu/single_gpu_controller/cuda_gpu_controller.py +++ b/src/keep_gpu/single_gpu_controller/cuda_gpu_controller.py @@ -124,7 +124,12 @@ def keep(self) -> None: name=f"gpu-keeper-{self.rank}", daemon=True, # daemon so program can exit cleanly ) - self._thread.start() + try: + self._thread.start() + except Exception: # noqa: BLE001 - preserve original thread-start failure + self._thread = None + self._stop_evt = None + raise startup_timeout = 5.0 if not startup_evt.wait(startup_timeout): stop_evt = self._stop_evt diff --git a/src/keep_gpu/single_gpu_controller/macm_gpu_controller.py b/src/keep_gpu/single_gpu_controller/macm_gpu_controller.py index c7f78cc..4a47c39 100644 --- a/src/keep_gpu/single_gpu_controller/macm_gpu_controller.py +++ b/src/keep_gpu/single_gpu_controller/macm_gpu_controller.py @@ -81,7 +81,12 @@ def keep(self) -> None: name=f"gpu-keeper-macm-{self.rank}", daemon=True, ) - self._thread.start() + try: + self._thread.start() + except Exception: # noqa: BLE001 - preserve original thread-start failure + self._thread = None + self._stop_evt = None + raise startup_timeout = 5.0 if not startup_evt.wait(startup_timeout): stop_evt = self._stop_evt diff --git a/tests/cuda_controller/test_keep_and_release.py b/tests/cuda_controller/test_keep_and_release.py index 56e743a..4e1a4a3 100644 --- a/tests/cuda_controller/test_keep_and_release.py +++ b/tests/cuda_controller/test_keep_and_release.py @@ -58,6 +58,34 @@ def fail_allocation(*_args, **_kwargs): assert ctrl._stop_evt is None +def test_cuda_keep_clears_state_when_thread_start_fails(monkeypatch): + import keep_gpu.single_gpu_controller.cuda_gpu_controller as cuda_module + + monkeypatch.setattr(cuda_module.torch.cuda, "device_count", lambda: 1) + + class FailingThread: + def __init__(self, *args, **kwargs): + pass + + def start(self): + raise RuntimeError("thread start failed") + + monkeypatch.setattr(cuda_module.threading, "Thread", FailingThread) + + ctrl = CudaGPUController( + rank=0, + interval=0.01, + vram_to_keep=4, + busy_threshold=-1, + ) + + with pytest.raises(RuntimeError, match="thread start failed"): + ctrl.keep() + + assert ctrl._thread is None + assert ctrl._stop_evt is None + + def test_cuda_keep_returns_when_startup_defers_for_unknown_utilization(monkeypatch): import keep_gpu.single_gpu_controller.cuda_gpu_controller as cuda_module diff --git a/tests/macm_controller/test_macm_backoff.py b/tests/macm_controller/test_macm_backoff.py index c16714f..9a9f778 100644 --- a/tests/macm_controller/test_macm_backoff.py +++ b/tests/macm_controller/test_macm_backoff.py @@ -184,6 +184,39 @@ def fail_allocation(*_args, **_kwargs): assert ctrl._stop_evt is None +def test_macm_keep_clears_state_when_thread_start_fails(monkeypatch): + import keep_gpu.single_gpu_controller.macm_gpu_controller as macm_module + + monkeypatch.setattr( + macm_module.torch.backends.mps, + "is_available", + lambda: True, + ) + + class FailingThread: + def __init__(self, *args, **kwargs): + pass + + def start(self): + raise RuntimeError("thread start failed") + + monkeypatch.setattr(macm_module.threading, "Thread", FailingThread) + + ctrl = MacMGPUController( + rank=0, + interval=0.01, + vram_to_keep=4, + busy_threshold=-1, + iterations=1, + ) + + with pytest.raises(RuntimeError, match="thread start failed"): + ctrl.keep() + + assert ctrl._thread is None + assert ctrl._stop_evt is None + + def test_macm_unknown_utilization_backs_off_when_threshold_enabled(): assert MacMGPUController._should_run_batch(None, 10) is False