Studio: harden GPU startup detection and clarify multi-GPU listing (#6427)

* Studio: harden GPU startup detection and clarify multi-GPU listing

* Log swallowed CUDA device property probe failures at debug level
This commit is contained in:
Daniel Han 2026-06-18 02:35:13 -07:00 committed by GitHub
commit 5f822c2e90
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 47 additions and 9 deletions

View file

@ -88,6 +88,19 @@ class TestGetDevice:
):
assert _reset_and_detect() == DeviceType.CUDA
@needs_torch
def test_detect_survives_device0_probe_failure(self, capsys):
# is_available() True but the device-0 name probe raises: startup must
# still resolve CUDA rather than crash.
with (
patch("utils.hardware.hardware._has_torch", return_value = True),
patch("torch.cuda.is_available", return_value = True),
patch("torch.cuda.device_count", return_value = 1),
patch("torch.cuda.get_device_properties", side_effect = RuntimeError("probe")),
):
assert _reset_and_detect() == DeviceType.CUDA
assert "<unavailable>" in capsys.readouterr().out
@needs_mlx
def test_returns_mlx_when_on_apple_silicon_with_mlx(self):
with (
@ -383,6 +396,22 @@ class TestPrintCudaDeviceList:
_hw_module._print_cuda_device_list(is_rocm = False)
assert capsys.readouterr().out == ""
@needs_torch
def test_rocm_label_omits_cuda_device_order(self, capsys):
# CUDA_DEVICE_ORDER governs CUDA only, so the ROCm listing must not claim it.
props = [MagicMock(), MagicMock()]
props[0].name = "AMD Instinct MI300X"
props[1].name = "AMD Instinct MI300X"
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 = True)
out = capsys.readouterr().out
assert "ROCm devices (2):" in out
assert "CUDA_DEVICE_ORDER" not in out
assert "[0] AMD Instinct MI300X" in out
# ========== format_error_message() ==========

View file

@ -107,12 +107,13 @@ def _has_mlx() -> bool:
def _print_cuda_device_list(is_rocm: bool) -> None:
"""Print every visible CUDA/ROCm GPU with its index.
"""List every visible CUDA/ROCm GPU with its index at startup.
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`.
cards on a multi-GPU host. This lists the full visible set in CUDA-ordinal
order, matching `nvidia-smi -L` when no CUDA_VISIBLE_DEVICES mask is set
(under a mask the indices are visible ordinals, not physical PCI ids).
CUDA_DEVICE_ORDER governs only CUDA, so it is shown for CUDA but not ROCm.
No-ops on single-GPU hosts and never raises -- it is purely informational.
"""
try:
@ -121,13 +122,17 @@ def _print_cuda_device_list(is_rocm: bool) -> None:
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}):"]
if is_rocm:
header = f"ROCm devices ({count}):"
else:
order = os.environ.get("CUDA_DEVICE_ORDER", "default")
header = f"CUDA devices ({count}, CUDA_DEVICE_ORDER={order}):"
lines = [header]
for i in range(count):
try:
name = torch.cuda.get_device_properties(i).name
except Exception:
except Exception as e:
logger.debug("CUDA device %d property probe failed: %s", i, e)
name = "<unavailable>"
lines.append(f" [{i}] {name}")
print("\n".join(lines))
@ -156,7 +161,11 @@ def detect_hardware() -> DeviceType:
if torch.cuda.is_available():
DEVICE = DeviceType.CUDA
CHAT_ONLY = False
device_name = torch.cuda.get_device_properties(0).name
try:
device_name = torch.cuda.get_device_properties(0).name
except Exception as e:
logger.debug("CUDA device 0 property probe failed: %s", e)
device_name = "<unavailable>"
# Distinguish ROCm from CUDA for display only (DeviceType stays CUDA).
# AMD SDK wheels don't set torch.version.hip, so fall back to __version__.