diff --git a/studio/backend/main.py b/studio/backend/main.py index 2313686a69..54946fa992 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -15,6 +15,15 @@ from dataclasses import asdict # Suppress C-level dependency warnings globally os.environ["PYTHONWARNINGS"] = "ignore" +# Pin GPU index ordering to PCI bus id before any torch import creates a CUDA +# context. Without this, torch/CUDA default to FASTEST_FIRST while nvidia-smi +# (and Studio's VRAM probes) use PCI-bus order, so a GPU index chosen from +# nvidia-smi data can resolve to a different physical card via +# CUDA_VISIBLE_DEVICES. setdefault so an explicit user override wins. See +# utils/hardware/hardware.py for the full rationale; set here too so the entry +# process is covered before its heavy ML imports. +os.environ.setdefault("CUDA_DEVICE_ORDER", "PCI_BUS_ID") + # ── Windows AMD ROCm DLL injection ────────────────────────────────────────── # Python 3.8+ ignores PATH for extension modules; register ROCm bin dirs with # os.add_dll_directory() so amdhip64.dll etc. are found before any torch import. diff --git a/studio/backend/tests/test_utils.py b/studio/backend/tests/test_utils.py index c66d56528a..b34618cd4d 100644 --- a/studio/backend/tests/test_utils.py +++ b/studio/backend/tests/test_utils.py @@ -303,6 +303,87 @@ class TestLogGpuMemory: assert "No GPU available" in captured.out +# ========== CUDA_DEVICE_ORDER pinning ========== + + +class TestCudaDeviceOrder: + """Importing the hardware module pins CUDA_DEVICE_ORDER=PCI_BUS_ID when unset, + but setdefault keeps an explicit user override, so nvidia-smi indices, torch + ordinals, and CUDA_VISIBLE_DEVICES agree on a mixed-GPU host.""" + + @staticmethod + def _order_after_fresh_import(preset): + # Fresh interpreter so the module-level setdefault runs against a clean env. + import os, subprocess, sys + from pathlib import Path + + env = os.environ.copy() + backend = str(Path(__file__).resolve().parents[1]) + existing = env.get("PYTHONPATH", "") + # Avoid a trailing os.pathsep (empty entry -> cwd on sys.path) when unset. + env["PYTHONPATH"] = (backend + os.pathsep + existing) if existing else backend + if preset is None: + env.pop("CUDA_DEVICE_ORDER", None) + else: + env["CUDA_DEVICE_ORDER"] = preset + out = subprocess.run( + [ + sys.executable, + "-c", + "import os, utils.hardware.hardware; print(os.environ.get('CUDA_DEVICE_ORDER'))", + ], + env = env, + capture_output = True, + text = True, + check = True, + ) + return out.stdout.strip().splitlines()[-1] + + def test_import_pins_pci_bus_id_when_unset(self): + assert self._order_after_fresh_import(None) == "PCI_BUS_ID" + + def test_import_respects_explicit_user_override(self): + assert self._order_after_fresh_import("FASTEST_FIRST") == "FASTEST_FIRST" + + +# ========== _print_cuda_device_list() ========== + + +class TestPrintCudaDeviceList: + """The startup console lists every CUDA GPU with its index, not just + device 0, so a multi-GPU host shows the full available set.""" + + @needs_torch + def test_lists_all_devices_when_multi_gpu(self, capsys): + props = [ + MagicMock(name = "p0"), + MagicMock(name = "p1"), + ] + props[0].name = "NVIDIA GeForce RTX 5090" + props[1].name = "NVIDIA RTX PRO 6000 Blackwell Workstation Edition" + with ( + patch("torch.cuda.device_count", return_value = 2), + patch("torch.cuda.get_device_properties", side_effect = lambda i: props[i]), + ): + _hw_module._print_cuda_device_list(is_rocm = False) + out = capsys.readouterr().out + assert "[0] NVIDIA GeForce RTX 5090" in out + assert "[1] NVIDIA RTX PRO 6000 Blackwell Workstation Edition" in out + assert "CUDA_DEVICE_ORDER=" in out + + @needs_torch + def test_silent_on_single_gpu(self, capsys): + with patch("torch.cuda.device_count", return_value = 1): + _hw_module._print_cuda_device_list(is_rocm = False) + assert capsys.readouterr().out == "" + + @needs_torch + def test_never_raises_on_probe_failure(self, capsys): + with patch("torch.cuda.device_count", side_effect = RuntimeError("no cuda")): + _hw_module._print_cuda_device_list(is_rocm = False) + assert capsys.readouterr().out == "" + + # ========== format_error_message() ========== diff --git a/studio/backend/utils/hardware/hardware.py b/studio/backend/utils/hardware/hardware.py index 86dfa8a93f..db608c3f53 100644 --- a/studio/backend/utils/hardware/hardware.py +++ b/studio/backend/utils/hardware/hardware.py @@ -35,6 +35,21 @@ from typing import Optional, Dict, Any logger = get_logger(__name__) +# ── GPU index ordering ────────────────────────────────────────────────────── +# CUDA defaults to CUDA_DEVICE_ORDER=FASTEST_FIRST, numbering GPUs by compute +# performance. nvidia-smi -- and every free-VRAM probe in Studio -- numbers GPUs +# by PCI bus id instead. On a mixed-GPU host (e.g. an RTX 5090 alongside an RTX +# PRO 6000) the two orderings disagree, so an index picked from nvidia-smi data +# ("the emptiest card is GPU 1") gets written into CUDA_VISIBLE_DEVICES and then +# reinterpreted by CUDA against FASTEST_FIRST -- landing the model on a different +# physical GPU than the one selected. Pinning PCI_BUS_ID makes torch, nvidia-smi, +# and CUDA_VISIBLE_DEVICES share a single index space, matching what users see in +# `nvidia-smi -L`. Set at import (before any torch.cuda call latches the order +# at context creation) and inherited by child processes, since the llama-server +# and spawn workers copy os.environ. setdefault so an explicit user override wins. +os.environ.setdefault("CUDA_DEVICE_ORDER", "PCI_BUS_ID") + + # ========== Device Enum ========== @@ -91,6 +106,35 @@ def _has_mlx() -> bool: return False +def _print_cuda_device_list(is_rocm: bool) -> None: + """Print every visible CUDA/ROCm GPU with its index. + + The "Hardware detected" banner names only device 0, which hides the other + cards on a multi-GPU host and obscures which physical GPU each index maps to. + Listing all devices in the pinned PCI_BUS_ID order (see CUDA_DEVICE_ORDER at + module top) makes the available set explicit and matches `nvidia-smi -L`. + No-ops on single-GPU hosts and never raises -- it is purely informational. + """ + try: + import torch + + count = torch.cuda.device_count() + if count <= 1: + return + label = "ROCm" if is_rocm else "CUDA" + order = os.environ.get("CUDA_DEVICE_ORDER", "default") + lines = [f"{label} devices ({count}, CUDA_DEVICE_ORDER={order}):"] + for i in range(count): + try: + name = torch.cuda.get_device_properties(i).name + except Exception: + name = "" + lines.append(f" [{i}] {name}") + print("\n".join(lines)) + except Exception: + return # purely informational; never disrupt startup + + def detect_hardware() -> DeviceType: """ Detect the best compute device and set the module-level DEVICE global. @@ -123,6 +167,7 @@ def detect_hardware() -> DeviceType: print(f"Hardware detected: ROCm (HIP {_hip_label}) -- {device_name}") else: print(f"Hardware detected: CUDA -- {device_name}") + _print_cuda_device_list(IS_ROCM) return DEVICE # --- XPU: Intel GPU ---