diff --git a/studio/backend/tests/test_install_resolve_prebuilt.py b/studio/backend/tests/test_install_resolve_prebuilt.py index 090d2932ea..e97ca47717 100644 --- a/studio/backend/tests/test_install_resolve_prebuilt.py +++ b/studio/backend/tests/test_install_resolve_prebuilt.py @@ -483,3 +483,237 @@ def test_resolve_prebuilt_intel_host_routes_to_upstream(monkeypatch, capsys): seen, out = _run_resolve_capture_host(monkeypatch, capsys) assert seen["repo"] == UPSTREAM assert out["repo"] == UPSTREAM + + +# --------------------------------------------------------------------------- +# windows_intel_gpu_in_registry: the in-process Windows Intel probe. A fake +# winreg module stands in for the real registry so the walk runs anywhere. +# --------------------------------------------------------------------------- + + +class _FakeRegKey: + def __init__( + self, + subkeys = None, + values = None, + denied = False, + ): + self.subkeys = subkeys or {} + self.values = values or {} + self.denied = denied + + def __enter__(self): + return self + + def __exit__(self, *exc): + return False + + +class _FakeWinreg: + HKEY_LOCAL_MACHINE = object() + + def __init__(self, root_key): + self._root_key = root_key + + def OpenKey(self, parent, name): + if parent is self.HKEY_LOCAL_MACHINE: + # Pin the production constant: a typo'd class GUID must fail here, + # not silently return the fake tree. + if name != ilp._WINDOWS_DISPLAY_CLASS_KEY: + raise FileNotFoundError(name) + if self._root_key is None: + raise FileNotFoundError(name) + return self._root_key + key = parent.subkeys.get(name) + if key is None: + # Real winreg raises OSError, never KeyError, for a missing key. + raise FileNotFoundError(name) + if key.denied: + raise PermissionError(name) + return key + + def QueryInfoKey(self, key): + return (len(key.subkeys), len(key.values), 0) + + def EnumKey(self, key, index): + return list(key.subkeys)[index] + + def QueryValueEx(self, key, value_name): + if value_name not in key.values: + raise FileNotFoundError(value_name) + return (key.values[value_name], 1) + + +def _probe_with_display_class(monkeypatch, adapters): + # The helper lazily does `import winreg`; plant the fake in sys.modules the + # same way unsloth_cli/tests/test_start.py fakes it for _refresh_windows_path. + monkeypatch.setitem(sys.modules, "winreg", _FakeWinreg(_FakeRegKey(subkeys = adapters))) + return ilp.windows_intel_gpu_in_registry() + + +def test_windows_intel_registry_matches_vendor_id(monkeypatch): + assert ( + _probe_with_display_class( + monkeypatch, + { + "0000": _FakeRegKey( + values = { + "MatchingDeviceId": r"PCI\VEN_8086&DEV_56A0&SUBSYS_12345678", + "DriverDesc": "Intel(R) Arc(TM) A770 Graphics", + } + ), + }, + ) + is True + ) + + +def test_windows_intel_registry_matches_driver_desc_without_device_id(monkeypatch): + assert ( + _probe_with_display_class( + monkeypatch, + { + "0000": _FakeRegKey(values = {"DriverDesc": "Intel(R) UHD Graphics 630"}), + }, + ) + is True + ) + + +def test_windows_intel_registry_ignores_non_intel_adapters(monkeypatch): + assert ( + _probe_with_display_class( + monkeypatch, + { + "0000": _FakeRegKey( + values = { + "MatchingDeviceId": r"PCI\VEN_10DE&DEV_2684", + "DriverDesc": "NVIDIA GeForce RTX 4090", + } + ), + "0001": _FakeRegKey( + values = { + "MatchingDeviceId": r"PCI\VEN_1002&DEV_744C", + "DriverDesc": "AMD Radeon RX 7900 XTX", + } + ), + }, + ) + is False + ) + + +def test_windows_intel_registry_skips_restricted_properties_subkey(monkeypatch): + # The real class key carries an ACL-restricted "Properties" subkey and can + # deny access to individual adapter keys; neither may abort the walk. + assert ( + _probe_with_display_class( + monkeypatch, + { + "Properties": _FakeRegKey(denied = True), + "0000": _FakeRegKey(denied = True), + "0001": _FakeRegKey( + values = { + "MatchingDeviceId": r"PCI\VEN_8086&DEV_56A0", + } + ), + }, + ) + is True + ) + + +def test_windows_intel_registry_missing_class_key_is_false(monkeypatch): + monkeypatch.setitem(sys.modules, "winreg", _FakeWinreg(None)) + assert ilp.windows_intel_gpu_in_registry() is False + + +def _detect_windows_host( + monkeypatch, + winreg_fake, + powershell_stdout = "", +): + """Drive the real detect_host() as a GPU-less Windows host with a fake + registry, recording every run_capture invocation. Pins the wiring the + unit tests above cannot see: registry-first, CIM only on a registry miss.""" + monkeypatch.setitem(sys.modules, "winreg", winreg_fake) + monkeypatch.setattr(ilp.platform, "system", lambda: "Windows") + monkeypatch.setattr(ilp.platform, "machine", lambda: "AMD64") + for _env in ( + "CUDA_VISIBLE_DEVICES", + "HIP_VISIBLE_DEVICES", + "ROCR_VISIBLE_DEVICES", + "HIP_PATH", + "ROCM_PATH", + ): + monkeypatch.delenv(_env, raising = False) + monkeypatch.setattr( + ilp.shutil, + "which", + lambda name: "powershell" if name in ("powershell", "pwsh") else None, + ) + captured = [] + + def _fake_run_capture(command, **kwargs): + captured.append(command[0]) + if command[0] == "powershell": + return SimpleNamespace(returncode = 0, stdout = powershell_stdout, stderr = "") + return SimpleNamespace(returncode = 1, stdout = "", stderr = "") + + monkeypatch.setattr(ilp, "run_capture", _fake_run_capture) + return ilp.detect_host(), captured + + +def test_detect_host_registry_intel_skips_cim_probe(monkeypatch): + winreg = _FakeWinreg( + _FakeRegKey( + subkeys = { + "0000": _FakeRegKey(values = {"MatchingDeviceId": r"PCI\VEN_8086&DEV_56A0"}), + } + ) + ) + host, captured = _detect_windows_host(monkeypatch, winreg) + assert host.has_intel_gpu is True + assert "powershell" not in captured + + +def test_detect_host_cim_fallback_fires_on_registry_miss(monkeypatch): + winreg = _FakeWinreg( + _FakeRegKey( + subkeys = { + "0000": _FakeRegKey(values = {"MatchingDeviceId": r"PCI\VEN_10DE&DEV_2684"}), + } + ) + ) + host, captured = _detect_windows_host( + monkeypatch, winreg, powershell_stdout = "Intel(R) Arc(TM) A770 Graphics" + ) + assert host.has_intel_gpu is True + assert "powershell" in captured + + +def test_windows_intel_registry_unexpected_error_is_false(monkeypatch): + # The probe is advisory: even a non-OSError bug in the walk must return + # False (deferring to the CIM fallback), never crash detect_host. + class _ExplodingWinreg: + HKEY_LOCAL_MACHINE = object() + + def OpenKey(self, parent, name): + raise TypeError(name) + + monkeypatch.setitem(sys.modules, "winreg", _ExplodingWinreg()) + assert ilp.windows_intel_gpu_in_registry() is False + + +def test_detect_host_cim_rescues_exploding_registry(monkeypatch): + class _ExplodingWinreg: + HKEY_LOCAL_MACHINE = object() + + def OpenKey(self, parent, name): + raise TypeError(name) + + host, captured = _detect_windows_host( + monkeypatch, _ExplodingWinreg(), powershell_stdout = "Intel(R) Arc(TM) A770 Graphics" + ) + assert host.has_intel_gpu is True + assert "powershell" in captured diff --git a/studio/install_llama_prebuilt.py b/studio/install_llama_prebuilt.py index 40caebc040..ca1fe79efa 100644 --- a/studio/install_llama_prebuilt.py +++ b/studio/install_llama_prebuilt.py @@ -2762,6 +2762,64 @@ def _pick_rocm_gfx_target(out: str) -> str | None: return _tokens[0] +# Display-adapter device class: one NNNN subkey per installed display driver +# config, each carrying the driver's DriverDesc and PCI MatchingDeviceId. +_WINDOWS_DISPLAY_CLASS_KEY = ( + r"SYSTEM\CurrentControlSet\Control\Class\{4d36e968-e325-11ce-bfc1-08002be10318}" +) + + +def windows_intel_gpu_in_registry() -> bool: + """Whether the Windows registry lists an Intel display adapter. + + In-process Windows counterpart of the Linux DRM vendor-id check (0x8086), + with weaker semantics: the class key lists installed display-driver + configs, which can outlive removed hardware, where sysfs lists present + devices. A stale Intel entry at worst routes to the upstream Vulkan + prebuilt instead of the fork CPU bundle: inference still works (the + Vulkan build runs on CPU when no Vulkan device exists), at the cost of + fork-only extras such as the DiffusionGemma visual server. detect_host's + PowerShell + WMI probe can silently miss a real Intel GPU: a cold + powershell.exe start plus the first CIM query routinely exceeds the 15s + budget on hosts with slow AV scanning or a degraded WMI repository, and + the probe swallows the timeout (#4452, Arc A770 routed to the CPU + prebuilt). Reading the display-adapter class key needs no subprocess and + answers in microseconds. Matches the PCI vendor id in MatchingDeviceId + (ven_8086) or an Intel DriverDesc. + """ + try: + import winreg + except ImportError: + return False + try: + with winreg.OpenKey(winreg.HKEY_LOCAL_MACHINE, _WINDOWS_DISPLAY_CLASS_KEY) as class_key: + for index in range(winreg.QueryInfoKey(class_key)[0]): + try: + name = winreg.EnumKey(class_key, index) + if not name.isdigit(): + # "Properties" is ACL-restricted and not an adapter. + continue + with winreg.OpenKey(class_key, name) as adapter_key: + for value_name, needle in ( + ("MatchingDeviceId", "ven_8086"), + ("DriverDesc", "intel"), + ): + try: + value, _ = winreg.QueryValueEx(adapter_key, value_name) + except OSError: + continue + if needle in str(value).lower(): + return True + except OSError: + continue + except Exception: + # Advisory probe: any unexpected failure must degrade to the CIM + # fallback, never crash the installer (mirrors detect_host's own + # swallow around the CIM probe). + return False + return False + + def detect_host() -> HostInfo: system = platform.system() machine = platform.machine().lower() @@ -2974,9 +3032,10 @@ def detect_host() -> HostInfo: # since the HIP SDK can be installed without an AMD GPU. # Detect an Intel GPU; gates the Vulkan prebuilt. Linux reads the DRM sysfs - # vendor id (0x8086); Windows queries the WMI video controller list. Only - # probed with no usable NVIDIA and no ROCm (matching the Vulkan branches), - # keeping the probe (notably the Windows powershell call) off that path. + # vendor id (0x8086); Windows reads the display-adapter registry class, + # then falls back to the WMI video controller list. Only probed with no + # usable NVIDIA and no ROCm (matching the Vulkan branches), keeping the + # probe (notably the Windows powershell call) off that path. has_intel_gpu = False if not has_usable_nvidia and not has_rocm: if is_linux: @@ -2989,23 +3048,28 @@ def detect_host() -> HostInfo: except OSError: continue elif is_windows: - _ps = shutil.which("powershell") or shutil.which("pwsh") - if _ps: - try: - _result = run_capture( - [ - _ps, - "-NoProfile", - "-Command", - "Get-CimInstance Win32_VideoController | " - "Select-Object -ExpandProperty Name", - ], - timeout = 15, - ) - if _result.returncode == 0 and "intel" in _result.stdout.lower(): - has_intel_gpu = True - except Exception: - pass + # Registry first (in-process; see windows_intel_gpu_in_registry). + # The CIM query stays as the fallback when the registry shows no + # Intel adapter. + has_intel_gpu = windows_intel_gpu_in_registry() + if not has_intel_gpu: + _ps = shutil.which("powershell") or shutil.which("pwsh") + if _ps: + try: + _result = run_capture( + [ + _ps, + "-NoProfile", + "-Command", + "Get-CimInstance Win32_VideoController | " + "Select-Object -ExpandProperty Name", + ], + timeout = 15, + ) + if _result.returncode == 0 and "intel" in _result.stdout.lower(): + has_intel_gpu = True + except Exception: + pass return HostInfo( system = system,