Studio: detect Windows Intel GPUs via the registry before WMI (#7064)
This commit is contained in:
oobabooga
GitHub
parent
7bfa209623
commit
d105bd7b42
2 changed files with 318 additions and 20 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue