Studio: avoid slow cold MIOpen searches on gfx1151

studio/backend/tests/test_rocm_miopen_cutoff.py

@@ -0,0 +1,83 @@
+# SPDX-License-Identifier: AGPL-3.0-only
+import os
+import sys
+from types import FunctionType, SimpleNamespace
+from unittest.mock import Mock
+
+import pytest
+
+from utils.hardware import hardware as hw
+
+
+@pytest.fixture(autouse = True)
+def isolated_environment(monkeypatch):
+    monkeypatch.setattr(os, "environ", os.environ.copy())
+
+
+@pytest.mark.parametrize(
+    "arches, expected",
+    [
+        (["gfx1151"], "1"),
+        (["gfx1151:xnack-", "gfx1151"], "1"),
+        (["gfx1100"], None),
+        (["gfx1151", "gfx1100"], None),
+        (["gfx1151", ""], None),
+        ([], None),
+    ],
+)
+def test_visible_device_scope(monkeypatch, arches, expected):
+    monkeypatch.delenv("MIOPEN_SEARCH_CUTOFF", raising = False)
+    cuda = SimpleNamespace(
+        device_count = lambda: len(arches),
+        get_device_properties = lambda i: SimpleNamespace(gcnArchName = arches[i]),
+    )
+    hw._configure_rocm_miopen(SimpleNamespace(cuda = cuda))
+    assert os.environ.get("MIOPEN_SEARCH_CUTOFF") == expected
+
+
+@pytest.mark.parametrize("value", ["0", "1", "", "custom"])
+def test_explicit_override_is_preserved_without_probing(monkeypatch, value):
+    monkeypatch.setenv("MIOPEN_SEARCH_CUTOFF", value)
+    torch = Mock()
+    hw._configure_rocm_miopen(torch)
+    torch.cuda.device_count.assert_not_called()
+    assert os.environ["MIOPEN_SEARCH_CUTOFF"] == value
+
+
+def test_probe_failure_does_not_break_startup(monkeypatch):
+    monkeypatch.delenv("MIOPEN_SEARCH_CUTOFF", raising = False)
+    torch = Mock()
+    torch.cuda.device_count.side_effect = RuntimeError("driver unavailable")
+    hw._configure_rocm_miopen(torch)
+    assert "MIOPEN_SEARCH_CUTOFF" not in os.environ
+
+
+@pytest.mark.parametrize(
+    "hip, version, enabled",
+    [("7.13", "2.11", True), (None, "2.11+rocm7.13", True), (None, "2.11+cu130", False)],
+)
+def test_detection_configures_only_rocm(monkeypatch, hip, version, enabled):
+    torch = SimpleNamespace(
+        version = SimpleNamespace(hip = hip),
+        __version__ = version,
+        cuda = SimpleNamespace(
+            is_available = lambda: True, get_device_properties = lambda i: SimpleNamespace(name = "GPU")
+        ),
+    )
+    monkeypatch.setitem(sys.modules, "torch", torch)
+    monkeypatch.setattr(hw, "_has_torch", lambda: True)
+    monkeypatch.setattr(hw, "_print_cuda_device_list", lambda _: None)
+    configure = Mock()
+    monkeypatch.setattr(hw, "_configure_rocm_miopen", configure)
+    monkeypatch.delenv("UNSLOTH_FORCE_XPU", raising = False)
+    monkeypatch.delenv("ZE_AFFINITY_MASK", raising = False)
+    monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising = False)
+    # Run the real body with isolated globals, including any future state fields.
+    detect = FunctionType(hw._detect_hardware_locked.__code__, vars(hw).copy())
+    assert detect() == hw.DeviceType.CUDA
+    assert configure.call_count == int(enabled)
+
+
+def test_desktop_launch_imports_a_shell_override():
+    from utils import desktop_shell_env as dse
+    assert "MIOPEN_SEARCH_CUTOFF" in dse.ROCM_SHELL_ENV_ALLOWLIST

studio/backend/utils/desktop_shell_env.py

@@ -48,6 +48,8 @@
     "MIOPEN_USER_DB_PATH",
     "MIOPEN_CUSTOM_CACHE_DIR",
     "MIOPEN_FIND_MODE",
+    # A shell opt-out (0) or opt-in must beat hardware.py's gfx1151 default.
+    "MIOPEN_SEARCH_CUTOFF",
     "HIP_VISIBLE_DEVICES",
     "ROCR_VISIBLE_DEVICES",
     "GPU_DEVICE_ORDINAL",

studio/backend/utils/hardware/hardware.py

@@ -1407,6 +1407,23 @@ def _xpu_device_name_or_placeholder(torch) -> str:
         return "<unavailable>"
 
 
+def _configure_rocm_miopen(torch) -> None:
+    if "MIOPEN_SEARCH_CUTOFF" in os.environ:
+        return
+    try:
+        count = torch.cuda.device_count()
+        # Process-wide in MIOpen: any untested arch on a mixed host keeps the default.
+        if not count or any(
+            _props_gfx_arch(torch.cuda.get_device_properties(i)) != "gfx1151" for i in range(count)
+        ):
+            return
+    except Exception as exc:
+        logger.debug("MIOpen search cutoff device probe failed: %s", exc)
+        return
+    os.environ.setdefault("MIOPEN_SEARCH_CUTOFF", "1")
+    logger.info("ROCm gfx1151: enabled MIOpen search cutoff (MIOPEN_SEARCH_CUTOFF=1)")
+
+
 def _detect_hardware_locked() -> DeviceType:
     """detect_hardware() body. Call only with _DETECT_LOCK held."""
     global DEVICE, CHAT_ONLY, CHAT_ONLY_REASON, CHAT_ONLY_DETAIL, IS_ROCM
@@ -1471,6 +1488,7 @@ def _detect_hardware_locked() -> DeviceType:
             _hip_ver = getattr(torch.version, "hip", None)
             if _hip_ver is not None or "rocm" in torch.__version__.lower():
                 IS_ROCM = True
+                _configure_rocm_miopen(torch)
                 _hip_label = _hip_ver or torch.__version__
                 print(f"Hardware detected: ROCm (HIP {_hip_label}) -- {device_name}")
             else: