From a36b95be9913728ea230e4ac6188b5c092c1b241 Mon Sep 17 00:00:00 2001 From: silveroxides Date: Sun, 9 Aug 2026 11:27:11 +0200 Subject: [PATCH] fix: limit default GPU management Avoid querying every visible CUDA device unless multi-GPU operation is explicitly requested. Refs #15255 --- comfy/cli_args.py | 2 +- comfy/model_management.py | 27 ++++++++-- comfy/multigpu.py | 2 +- main.py | 4 +- server.py | 4 +- tests/test_model_management_devices.py | 69 ++++++++++++++++++++++++++ 6 files changed, 97 insertions(+), 11 deletions(-) create mode 100644 tests/test_model_management_devices.py diff --git a/comfy/cli_args.py b/comfy/cli_args.py index ee9e1ce9f..b15a6d746 100644 --- a/comfy/cli_args.py +++ b/comfy/cli_args.py @@ -74,7 +74,7 @@ parser.add_argument("--temp-directory", type=str, default=None, help="Set the Co parser.add_argument("--input-directory", type=str, default=None, help="Set the ComfyUI input directory. Overrides --base-directory.") parser.add_argument("--auto-launch", action="store_true", help="Automatically launch ComfyUI in the default browser.") parser.add_argument("--disable-auto-launch", action="store_true", help="Disable auto launching the browser.") -parser.add_argument("--cuda-device", type=str, default=None, metavar="DEVICE_ID", help="Set the ids of cuda devices this instance will use, as a comma-separated list (e.g. '0' or '0,1'). All other devices will not be visible.") +parser.add_argument("--cuda-device", type=str, default=None, metavar="DEVICE_ID", help="Set the ids of cuda devices this instance will use, as a comma-separated list (e.g. '0' or '0,1'). Passing multiple ids enables core multi-GPU support. All other devices will not be visible.") parser.add_argument("--default-device", type=int, default=None, metavar="DEFAULT_DEVICE_ID", help="Set the id of the default device, all other devices will stay visible.") cm_group = parser.add_mutually_exclusive_group() cm_group.add_argument("--cuda-malloc", action="store_true", help="Enable cudaMallocAsync (enabled by default for torch 2.0 and up).") diff --git a/comfy/model_management.py b/comfy/model_management.py index 9f8e7f07b..91cd99aac 100644 --- a/comfy/model_management.py +++ b/comfy/model_management.py @@ -243,6 +243,23 @@ def get_all_torch_devices(exclude_current=False): devices.remove(current) return devices +def get_core_torch_devices(exclude_current=False): + """Return devices ComfyUI core is allowed to manage. + + Other visible devices remain available to custom nodes. Core multi-GPU + management is enabled only by explicitly passing multiple devices through + --cuda-device. + """ + if args.cuda_device is not None and "," in args.cuda_device: + devices = get_all_torch_devices() + else: + devices = [get_torch_device()] + if exclude_current: + current = get_torch_device() + if current in devices: + devices.remove(current) + return devices + def get_gpu_device_options(): """Return list of device option strings for node widgets. @@ -250,7 +267,7 @@ def get_gpu_device_options(): adds "gpu:0", "gpu:1", etc. (vendor-agnostic labels). """ options = ["default", "cpu"] - devices = get_all_torch_devices() + devices = get_core_torch_devices() if len(devices) > 1: for i in range(len(devices)): options.append(f"gpu:{i}") @@ -283,7 +300,7 @@ def resolve_gpu_device_option(option: str): idx = int(option[4:]) except ValueError: return None - devices = get_all_torch_devices() + devices = get_core_torch_devices() if 0 <= idx < len(devices): return devices[idx] return None @@ -603,7 +620,7 @@ try: except: logging.warning("Could not pick default device.") try: - for device in get_all_torch_devices(exclude_current=True): + for device in get_core_torch_devices(exclude_current=True): logging.info("Device: {}".format(get_torch_device_name(device))) except: pass @@ -2041,7 +2058,7 @@ def soft_empty_cache(force=False): torch.cuda.ipc_collect() def unload_all_models(): - for device in get_all_torch_devices(): + for device in get_core_torch_devices(): free_memory(1e30, device) def unload_model_and_clones(model: ModelPatcher, unload_additional_models=True, all_devices=False): @@ -2067,7 +2084,7 @@ def unload_model_and_clones(model: ModelPatcher, unload_additional_models=True, if not all_devices: free_memory(1e30, get_torch_device(), keep_loaded) else: - for device in get_all_torch_devices(): + for device in get_core_torch_devices(): free_memory(1e30, device, keep_loaded) def debug_memory_summary(): diff --git a/comfy/multigpu.py b/comfy/multigpu.py index 2b6d8260d..da635a8a5 100644 --- a/comfy/multigpu.py +++ b/comfy/multigpu.py @@ -136,7 +136,7 @@ def create_multigpu_deepclones(model: ModelPatcher, max_gpus: int, gpu_options: # Exclude the primary model's actual device, not the global current device: # after SelectModelDevice(gpu:N) the primary may not live on the process's # current CUDA device, and excluding the wrong device picks bad extras. - all_devices = comfy.model_management.get_all_torch_devices(exclude_current=False) + all_devices = comfy.model_management.get_core_torch_devices(exclude_current=False) full_extra_devices = [d for d in all_devices if d != model.load_device] limit_extra_devices = full_extra_devices[:max_gpus-1] extra_devices = limit_extra_devices.copy() diff --git a/main.py b/main.py index 361b1fc89..67169119e 100644 --- a/main.py +++ b/main.py @@ -253,10 +253,10 @@ if args.enable_dynamic_vram or (enables_dynamic_vram() and comfy.model_managemen logging.warning("Unsupported Pytorch detected. DynamicVRAM support requires Pytorch version 2.8 or later. Falling back to legacy ModelPatcher. VRAM estimates may be unreliable especially on Windows") else: try: - aimdo_initialized = comfy_aimdo.control.init_devices((d.index, int(args.vram_headroom * 1024 ** 3)) for d in comfy.model_management.get_all_torch_devices()) + aimdo_initialized = comfy_aimdo.control.init_devices((d.index, int(args.vram_headroom * 1024 ** 3)) for d in comfy.model_management.get_core_torch_devices()) except TypeError: # comfy-aimdo 0.4.9 protocol. - aimdo_initialized = comfy_aimdo.control.init_devices(d.index for d in comfy.model_management.get_all_torch_devices()) + aimdo_initialized = comfy_aimdo.control.init_devices(d.index for d in comfy.model_management.get_core_torch_devices()) if aimdo_initialized: if console_log_level == 'DEBUG': diff --git a/server.py b/server.py index c9ffcaa0d..99e47cf93 100644 --- a/server.py +++ b/server.py @@ -694,9 +694,9 @@ class PromptServer(): required_templates_version = FrontendManager.get_required_templates_version() comfy_package_versions = FrontendManager.get_comfy_package_versions() - # Report every torch device visible to multigpu, with the primary + # Report every device managed by ComfyUI core, with the primary # device first so existing clients that read devices[0] keep working. - torch_devices = comfy.model_management.get_all_torch_devices() + torch_devices = comfy.model_management.get_core_torch_devices() if primary_device in torch_devices: torch_devices = [primary_device] + [d for d in torch_devices if d != primary_device] else: diff --git a/tests/test_model_management_devices.py b/tests/test_model_management_devices.py new file mode 100644 index 000000000..b09a1e622 --- /dev/null +++ b/tests/test_model_management_devices.py @@ -0,0 +1,69 @@ +import pytest +import torch + +from comfy.cli_args import args +import comfy.model_management as model_management + + +@pytest.mark.parametrize("cuda_device", [None, "0", "GPU-example"]) +def test_core_devices_default_to_current_without_enumerating(monkeypatch, cuda_device): + current = torch.device("cuda", 0) + monkeypatch.setattr(args, "cuda_device", cuda_device) + monkeypatch.setattr(args, "default_device", 1) + monkeypatch.setattr(model_management, "get_torch_device", lambda: current) + monkeypatch.setattr( + model_management, + "get_all_torch_devices", + lambda: pytest.fail("secondary devices must not be enumerated"), + ) + + assert model_management.get_core_torch_devices() == [current] + + +@pytest.mark.parametrize("cuda_device", ["0,1", "GPU-primary,GPU-secondary"]) +def test_core_devices_use_explicit_multi_device_selection(monkeypatch, cuda_device): + current = torch.device("cuda", 0) + devices = [current, torch.device("cuda", 1)] + monkeypatch.setattr(args, "cuda_device", cuda_device) + monkeypatch.setattr(model_management, "get_torch_device", lambda: current) + monkeypatch.setattr(model_management, "get_all_torch_devices", lambda: devices.copy()) + + assert model_management.get_core_torch_devices() == devices + assert model_management.get_core_torch_devices(exclude_current=True) == devices[1:] + + +def test_all_devices_remain_discoverable(monkeypatch): + monkeypatch.setattr(model_management, "cpu_state", model_management.CPUState.GPU) + monkeypatch.setattr(model_management, "is_nvidia", lambda: True) + monkeypatch.setattr(model_management.torch.cuda, "device_count", lambda: 2) + + assert model_management.get_all_torch_devices() == [ + torch.device("cuda", 0), + torch.device("cuda", 1), + ] + + +def test_device_options_only_offer_core_devices(monkeypatch): + current = torch.device("cuda", 0) + secondary = torch.device("cuda", 1) + monkeypatch.setattr(args, "cuda_device", None) + monkeypatch.setattr(model_management, "get_torch_device", lambda: current) + monkeypatch.setattr(model_management, "get_all_torch_devices", lambda: [current, secondary]) + + assert model_management.get_gpu_device_options() == ["default", "cpu"] + assert model_management.resolve_gpu_device_option("gpu:1") is None + + monkeypatch.setattr(args, "cuda_device", "0,1") + assert model_management.get_gpu_device_options() == ["default", "cpu", "gpu:0", "gpu:1"] + assert model_management.resolve_gpu_device_option("gpu:1") == secondary + + +def test_unload_all_models_only_uses_core_devices(monkeypatch): + current = torch.device("cuda", 0) + freed = [] + monkeypatch.setattr(model_management, "get_core_torch_devices", lambda: [current]) + monkeypatch.setattr(model_management, "free_memory", lambda amount, device: freed.append((amount, device))) + + model_management.unload_all_models() + + assert freed == [(1e30, current)]