Studio: pin CUDA_DEVICE_ORDER=PCI_BUS_ID and list GPUs at startup (#6353)

* Studio: pin CUDA_DEVICE_ORDER=PCI_BUS_ID and list GPUs at startup

On a mixed-GPU host, Studio could load a model onto a different physical
GPU than the one it selected. The free-VRAM probe numbers GPUs via
nvidia-smi (PCI-bus order), but CUDA defaults to FASTEST_FIRST ordering,
so a selected index written into CUDA_VISIBLE_DEVICES resolved to the
wrong card. Example: 5090 + RTX PRO 6000, the picker chose the emptier
RTX PRO 6000 (nvidia-smi index 1) but CUDA read index 1 as the 5090.

Pin CUDA_DEVICE_ORDER=PCI_BUS_ID at import (before any CUDA context is
created) in both the Studio entrypoint and the hardware module, so torch,
nvidia-smi, and CUDA_VISIBLE_DEVICES share one index space. setdefault
keeps an explicit user override intact. Child processes inherit it via
os.environ.

Also list every detected CUDA GPU with its index at startup instead of
naming only device 0, matching nvidia-smi -L and making the selected
index unambiguous on multi-GPU hosts.

* Studio: make CUDA_DEVICE_ORDER tests exercise module import and respect user override

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: guard full _print_cuda_device_list body and fix test PYTHONPATH trailing separator

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: Daniel Han <danielhanchen@gmail.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

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.

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() ==========
 
 

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 = "<unavailable>"
+            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 ---