diff --git a/invokeai/app/services/session_processor/session_processor_default.py b/invokeai/app/services/session_processor/session_processor_default.py index 1a9a7be9cff..4c98e1a66d3 100644 --- a/invokeai/app/services/session_processor/session_processor_default.py +++ b/invokeai/app/services/session_processor/session_processor_default.py @@ -629,8 +629,12 @@ def _process( # which nodes and the model loader consult) resolves to this GPU. CUDA's current device is per-thread. if worker.device is not None: TorchDevice.set_session_device(worker.device) - if worker.device.type == "cuda": - torch.cuda.set_device(worker.device) + + # torch.cuda.set_device() initializes CUDA on the device, which can permanently reserve + # VRAM in an otherwise idle process (#9413). Defer the CUDA-side pin until this worker + # claims its first queue item; the pin is per-thread and this thread persists, so pinning + # once before the first item is equivalent to pinning here. + cuda_pin_needed = worker.device is not None and worker.device.type == "cuda" worker.cancel_event.clear() @@ -669,6 +673,10 @@ def _process( poll_now_event.wait(self._polling_interval) continue + if cuda_pin_needed: + torch.cuda.set_device(worker.device) + cuda_pin_needed = False + # A cancellation can race the claim: it may have marked the row terminal before # this worker recorded `queue_item`, so _on_queue_item_status_changed couldn't set # our cancel_event. A fresh DB status read is the authority — skip running an diff --git a/invokeai/backend/model_manager/load/model_cache/model_cache.py b/invokeai/backend/model_manager/load/model_cache/model_cache.py index 47dd4d45a8d..8f8296b8674 100644 --- a/invokeai/backend/model_manager/load/model_cache/model_cache.py +++ b/invokeai/backend/model_manager/load/model_cache/model_cache.py @@ -1157,9 +1157,13 @@ def _calc_ram_available_to_model_cache(self) -> int: # the default value if desired. # Lookup the total VRAM size for the CUDA execution device. + # This runs at startup (one ModelCache is built per generation device before any request is + # served), and torch.cuda.mem_get_info() would create a CUDA context that permanently holds + # ~100-300 MiB of VRAM in an otherwise idle process. cudaGetDeviceProperties reports the same + # total without creating a context (#9413). total_cuda_vram_bytes: int | None = None if self._execution_device.type == "cuda": - _, total_cuda_vram_bytes = torch.cuda.mem_get_info(self._execution_device) + total_cuda_vram_bytes = torch.cuda.get_device_properties(self._execution_device).total_memory # Apply heuristic 1. # ------------------ diff --git a/tests/app/services/session_processor/test_session_processor_cancel_guard.py b/tests/app/services/session_processor/test_session_processor_cancel_guard.py index 4add666385e..d6f8f1aa830 100644 --- a/tests/app/services/session_processor/test_session_processor_cancel_guard.py +++ b/tests/app/services/session_processor/test_session_processor_cancel_guard.py @@ -193,3 +193,72 @@ def test_shutdown_race_cancels_fresh_claim_instead_of_running_it(): assert run_items == [] assert canceled == [42] + + +def test_cuda_device_pin_is_deferred_until_first_claim_and_runs_once(): + """An idle CUDA worker must not create a context, then must pin before its first run only.""" + from threading import BoundedSemaphore, Event + from unittest.mock import MagicMock, patch + + import torch + + from invokeai.app.services.session_processor.session_processor_default import _SessionWorker + + stop_event = Event() + resume_event = Event() + resume_event.set() + poll_now_event = Event() + events: list[str] = [] + + first = SimpleNamespace(item_id=1, session_id="first", queue_id="default") + second = SimpleNamespace(item_id=2, session_id="second", queue_id="default") + + class _ClaimQueue: + def __init__(self): + self.items = iter([None, first, second]) + + def dequeue(self, device=None): + item = next(self.items) + events.append(f"dequeue:{item.item_id if item else 'empty'}") + return item + + def get_queue_item(self, item_id: int): + return SimpleNamespace(item_id=item_id, status="in_progress") + + runner = MagicMock() + + def run_item(item): + events.append(f"run:{item.item_id}") + if item.item_id == second.item_id: + stop_event.set() + + runner.workflow_call_queue_lifecycle.run_queue_item.side_effect = run_item + worker = _SessionWorker(device=torch.device("cuda:1"), runner=runner) + processor = DefaultSessionProcessor() + processor._invoker = SimpleNamespace( # type: ignore[attr-defined] + services=SimpleNamespace(session_queue=_ClaimQueue(), logger=MagicMock(), image_moves=None) + ) + processor._polling_interval = 0 + processor._thread_semaphore = BoundedSemaphore(1) + + with ( + patch( + "invokeai.app.services.session_processor.session_processor_default.torch.cuda.set_device", + side_effect=lambda device: events.append(f"pin:{device}"), + ) as mock_set_device, + patch( + "invokeai.app.services.session_processor.session_processor_default.GENERATION_DEVICE_POOL.acquire_session" + ), + patch( + "invokeai.app.services.session_processor.session_processor_default.GENERATION_DEVICE_POOL.release_session" + ), + ): + processor._process( + worker=worker, + stop_event=stop_event, + poll_now_event=poll_now_event, + resume_event=resume_event, + ) + + assert events == ["dequeue:empty", "dequeue:1", "pin:cuda:1", "run:1", "dequeue:2", "run:2"] + mock_set_device.assert_called_once_with(torch.device("cuda:1")) diff --git a/tests/backend/model_manager/load/model_cache/test_model_cache_ram_budget.py b/tests/backend/model_manager/load/model_cache/test_model_cache_ram_budget.py index 656a10be534..d57e1970d0a 100644 --- a/tests/backend/model_manager/load/model_cache/test_model_cache_ram_budget.py +++ b/tests/backend/model_manager/load/model_cache/test_model_cache_ram_budget.py @@ -252,6 +252,7 @@ def test_get_vram_in_use_queries_this_caches_execution_device(mock_logger): mc = "invokeai.backend.model_manager.load.model_cache.model_cache" with ( patch(f"{mc}.torch.cuda.mem_get_info", return_value=(10 * GB, 48 * GB)), + patch(f"{mc}.torch.cuda.get_device_properties", return_value=MagicMock(total_memory=48 * GB)), patch(f"{mc}.torch.cuda.memory_allocated", return_value=42) as mock_alloc, ): cache = ModelCache( @@ -269,6 +270,30 @@ def test_get_vram_in_use_queries_this_caches_execution_device(mock_logger): cache.shutdown() +def test_cuda_cache_init_queries_total_vram_without_mem_get_info(mock_logger): + """CUDA cache sizing must not call the VRAM-holding mem_get_info API during idle startup.""" + import torch + + mc = "invokeai.backend.model_manager.load.model_cache.model_cache" + with ( + patch(f"{mc}.torch.cuda.get_device_properties", return_value=MagicMock(total_memory=48 * GB)) as mock_props, + patch(f"{mc}.torch.cuda.mem_get_info", return_value=(10 * GB, 48 * GB)) as mock_mem_get_info, + ): + cache = ModelCache( + execution_device_working_mem_gb=3.0, + enable_partial_loading=True, + keep_ram_copy_of_weights=True, + execution_device="cuda:1", + storage_device="cpu", + logger=mock_logger, + ) + try: + mock_props.assert_called_once_with(torch.device("cuda:1")) + mock_mem_get_info.assert_not_called() + finally: + cache.shutdown() + + def _mock_total_ram(total_bytes: int): """Patch psutil.virtual_memory().total as seen by model_cache.""" vm = MagicMock()