Files
unslothai--unsloth/studio/backend/tests/test_install_resolve_prebuilt.py
T
wehub-resource-sync e93507a09c
Lockfile supply-chain audit / lockfile supply-chain audit (push) Has been cancelled
Windows Studio GGUF CI / GPU prebuilt resolves without Visual Studio (push) Has been cancelled
Windows Studio GGUF CI / setup.ps1 unit tests (VS 2026 / CMake guard) (push) Has been cancelled
Windows Studio GGUF CI / real-VS detection (VS 2022) (push) Has been cancelled
Windows Studio GGUF CI / real-VS detection (VS 2026) (push) Has been cancelled
Windows Studio GGUF CI / VC++ runtime detect + install round-trip (windows-2025-vs2026) (push) Has been cancelled
Windows Studio GGUF CI / VC++ runtime detect + install round-trip (windows-latest) (push) Has been cancelled
Windows Studio Update CI / Studio Updating Tests (push) Has been cancelled
Wheel CI / Wheel build + content sanity + import smoke (push) Has been cancelled
Lint CI / Source lint (Python + shell + YAML + JSON + safety nets) (push) Has been cancelled
MLX CI on Mac M1 / dispatch (push) Has been cancelled
Security audit / advisory audit (pip + npm + cargo) (push) Has been cancelled
Security audit / pip scan-packages :: extras (push) Has been cancelled
Security audit / pip scan-packages :: studio (push) Has been cancelled
Security audit / pip scan-packages :: hf-stack (push) Has been cancelled
Security audit / npm scan-packages (Studio frontend tarballs) (push) Has been cancelled
Security audit / workflow-trigger lint (pull_request_target / cache-poisoning) (push) Has been cancelled
Security audit / pytest tests/security (push) Has been cancelled
Security audit / npm provenance + new install-script diff (push) Has been cancelled
Studio API CI / Studio API & Auth Tests (push) Has been cancelled
Backend CI / (Python 3.10) (push) Has been cancelled
Backend CI / (Python 3.11) (push) Has been cancelled
Backend CI / (Python 3.12) (push) Has been cancelled
Backend CI / (Python 3.13) (push) Has been cancelled
Backend CI / Repo tests (CPU) (push) Has been cancelled
Frontend CI / Frontend build + bundle sanity (push) Has been cancelled
Studio GGUF CI / OpenAI, Anthropic API tests (push) Has been cancelled
Studio GGUF CI / Tool calling Tests (push) Has been cancelled
Studio GGUF CI / JSON, images (push) Has been cancelled
Mac Studio GGUF CI / OpenAI, Anthropic API tests (push) Has been cancelled
Mac Studio GGUF CI / Tool calling Tests (push) Has been cancelled
Mac Studio GGUF CI / JSON, images (push) Has been cancelled
Mac Studio Install Matrix CI / Install + load (macos-14) (push) Has been cancelled
Mac Studio Install Matrix CI / Install + load (macos-15) (push) Has been cancelled
Mac Studio Install Matrix CI / Install + load (macos-26) (push) Has been cancelled
Mac Studio Install Matrix CI / Install + load (macos-15-intel) (push) Has been cancelled
Mac Studio API CI / Studio API & Auth Tests (push) Has been cancelled
Mac Studio Install Matrix CI / Install + load (macos-26-intel) (push) Has been cancelled
Mac Studio UI CI / Chat UI Tests (push) Has been cancelled
Studio Tauri CI / Tauri Linux debug build (no codesign) (push) Has been cancelled
Mac Studio Update CI / Studio Updating Tests (push) Has been cancelled
Studio UI CI / Chat UI Tests (push) Has been cancelled
Windows Studio API CI / Studio API & Auth Tests (push) Has been cancelled
Windows Studio UI CI / Chat UI Tests (push) Has been cancelled
Studio Update CI / Studio Updating Tests (push) Has been cancelled
Core / Core (HF=default + TRL=default) (push) Has been cancelled
Core / Core (HF=4.57.6 + TRL<1) (push) Has been cancelled
Core / Core (HF=latest + TRL=latest) (push) Has been cancelled
Core / llama.cpp build + smoke (push) Has been cancelled
Windows Studio GGUF CI / OpenAI, Anthropic API tests (push) Has been cancelled
Windows Studio GGUF CI / Tool calling Tests (push) Has been cancelled
Windows Studio GGUF CI / JSON, images (push) Has been cancelled
Windows Studio GGUF CI / Studio install + inference without Visual Studio (push) Has been cancelled
Studio export capability / capability (macos-latest) (push) Has been cancelled
Studio export capability / capability (ubuntu-latest) (push) Has been cancelled
Studio export capability / capability (windows-latest) (push) Has been cancelled
Cross-platform parity / parity (macos-latest) (push) Has been cancelled
Cross-platform parity / parity (windows-latest) (push) Has been cancelled
Scorecard supply-chain security / Scorecard analysis (push) Has been cancelled
Studio load-orchestrator CI / test (push) Has been cancelled
chore: import upstream snapshot with attribution
2026-07-13 12:59:56 +08:00

720 lines
25 KiB
Python

# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""install_llama_prebuilt.py: the --resolve-prebuilt probe (plans against the fork
by default; --published-repo overrides).
These back the in-app update for source-build (markerless) installs: the backend
asks the installer whether an official prebuilt exists for this host without
downloading. Network and host detection are stubbed; no GPU or internet needed.
"""
from __future__ import annotations
import importlib
import json
import sys
from pathlib import Path
from types import SimpleNamespace
import pytest
_studio = Path(__file__).resolve().parent.parent.parent
if str(_studio) not in sys.path:
sys.path.insert(0, str(_studio))
ilp = importlib.import_module("install_llama_prebuilt")
if not hasattr(ilp, "resolve_simple_install_release_plans"):
pytest.skip("PR symbols not present - check branch", allow_module_level = True)
FORK = ilp.DEFAULT_PUBLISHED_REPO # unslothai/llama.cpp
UPSTREAM = ilp.UPSTREAM_REPO # ggml-org/llama.cpp
def _host(**kw):
base = dict(
system = "Linux",
machine = "x86_64",
is_windows = False,
is_linux = False,
is_macos = False,
is_x86_64 = False,
is_arm64 = False,
nvidia_smi = None,
driver_cuda_version = None,
compute_caps = [],
visible_cuda_devices = None,
has_physical_nvidia = False,
has_usable_nvidia = False,
has_rocm = False,
rocm_gfx_target = None,
macos_version = None,
)
base.update(kw)
return ilp.HostInfo(**base)
def test_force_cpu_clears_all_gpu_attributes_including_intel():
# --cpu-fallback is the "select the CPU prebuilt even when a GPU is present"
# escape hatch. It must drop EVERY GPU attribute, including has_intel_gpu, or
# the planner still prepends the Vulkan asset on an Intel-GPU host.
host = _host(
is_linux = True,
is_x86_64 = True,
has_usable_nvidia = True,
has_physical_nvidia = True,
has_rocm = True,
rocm_gfx_target = "gfx1100",
has_intel_gpu = True,
)
forced = ilp._apply_host_overrides(host, force_cpu = True)
assert forced.has_usable_nvidia is False
assert forced.has_physical_nvidia is False
assert forced.has_rocm is False
assert forced.rocm_gfx_target is None
assert forced.has_intel_gpu is False
def test_macos_upstream_pin_only_for_explicit_pre26_upstream():
pre26 = _host(
system = "Darwin",
is_macos = True,
is_arm64 = True,
machine = "arm64",
macos_version = (15, 5),
)
assert ilp.pinned_macos_release_tag(pre26, UPSTREAM) == "b9415"
assert ilp.pinned_macos_release_tag(pre26, FORK) is None
tahoe = _host(
system = "Darwin",
is_macos = True,
is_arm64 = True,
machine = "arm64",
macos_version = (26, 0),
)
assert ilp.pinned_macos_release_tag(tahoe, UPSTREAM) is None
def _run_resolve(monkeypatch, capsys, plans_or_exc):
monkeypatch.setattr(
ilp,
"detect_host",
lambda: _host(system = "Darwin", is_macos = True, is_arm64 = True, machine = "arm64"),
)
def _resolver(tag, host, repo, published_release_tag):
if isinstance(plans_or_exc, Exception):
raise plans_or_exc
return ("b9585", plans_or_exc)
monkeypatch.setattr(ilp, "resolve_simple_install_release_plans", _resolver)
monkeypatch.setattr(
sys,
"argv",
["install_llama_prebuilt.py", "--resolve-prebuilt", "latest", "--output-format", "json"],
)
rc = ilp.main()
assert rc == ilp.EXIT_SUCCESS
return json.loads(capsys.readouterr().out.strip().splitlines()[-1])
def test_resolve_prebuilt_available(monkeypatch, capsys):
plan = SimpleNamespace(
release_tag = "b9585",
llama_tag = "b9585",
attempts = [
SimpleNamespace(name = "llama-b9585-bin-macos-arm64.tar.gz", install_kind = "macos-arm64")
],
)
out = _run_resolve(monkeypatch, capsys, [plan])
assert out["prebuilt_available"] is True
assert out["repo"] == FORK
assert out["release_tag"] == "b9585"
assert out["asset"] == "llama-b9585-bin-macos-arm64.tar.gz"
assert out["install_kind"] == "macos-arm64"
def test_resolve_prebuilt_unavailable(monkeypatch, capsys):
out = _run_resolve(monkeypatch, capsys, ilp.PrebuiltFallback("no macOS asset"))
assert out["prebuilt_available"] is False
assert out["repo"] == FORK
def _run_resolve_capture_host(monkeypatch, capsys):
"""Drive --resolve-prebuilt and return the host the resolver was handed."""
seen = {}
def _resolver(tag, host, repo, published_release_tag):
seen["repo"] = repo
seen["host"] = host
raise ilp.PrebuiltFallback("no asset")
monkeypatch.setattr(ilp, "resolve_simple_install_release_plans", _resolver)
monkeypatch.setattr(
sys,
"argv",
["install_llama_prebuilt.py", "--resolve-prebuilt", "latest", "--output-format", "json"],
)
assert ilp.main() == ilp.EXIT_SUCCESS
out = json.loads(capsys.readouterr().out.strip().splitlines()[-1])
return seen, out
def test_resolve_prebuilt_cpu_linux_routes_to_fork(monkeypatch, capsys):
# CPU-only Linux host (no GPU): the dispatch routes to the fork, which now
# ships the CPU prebuilt -- it no longer falls back to ggml-org upstream.
monkeypatch.setattr(ilp, "detect_host", lambda: _host(is_linux = True, is_x86_64 = True))
seen, out = _run_resolve_capture_host(monkeypatch, capsys)
assert seen["repo"] == FORK
assert out["repo"] == FORK
def test_resolve_prebuilt_rocm_sdk_only_host_still_offered_cpu(monkeypatch, capsys):
# A CPU-only host that merely has ROCm/HIP SDK tools on PATH (no AMD GPU, so
# detect_host leaves has_rocm False) is a valid CPU-prebuilt target. The probe
# must NOT reclassify it as ROCm from tool presence alone and suppress the CPU
# bundle -- that would deny the fork CPU prebuilt to a legitimate CPU source
# build. The host is left CPU-only and resolves against the fork.
monkeypatch.setattr(ilp, "detect_host", lambda: _host(is_linux = True, is_x86_64 = True))
monkeypatch.setattr(
ilp.shutil, "which", lambda tool: "/opt/rocm/bin/hipconfig" if tool == "hipconfig" else None
)
seen, out = _run_resolve_capture_host(monkeypatch, capsys)
assert seen["repo"] == FORK
assert seen["host"].has_rocm is False
# Blackwell floor is sm_100 (data-center B100/B200, B300/GB300), below consumer
# sm_120 -- 120 wrongly excluded data-center hosts from the prebuilt selection.
def _gpu_linux_host(caps):
return _host(
is_linux = True,
is_x86_64 = True,
has_physical_nvidia = True,
has_usable_nvidia = True,
driver_cuda_version = (13, 1),
compute_caps = caps,
)
def test_host_is_blackwell_includes_datacenter_parts():
assert ilp._host_is_blackwell(_gpu_linux_host(["10.0"])) is True # B200 sm_100
assert ilp._host_is_blackwell(_gpu_linux_host(["10.3"])) is True # B300 sm_103
assert ilp._host_is_blackwell(_gpu_linux_host(["12.0"])) is True # RTX 50 sm_120
assert ilp._host_is_blackwell(_gpu_linux_host(["12.1"])) is True # DGX Spark sm_121
assert ilp._host_is_blackwell(_gpu_linux_host(["9.0"])) is False # Hopper
assert ilp._host_is_blackwell(_gpu_linux_host(["8.0"])) is False # Ampere
assert ilp._host_is_blackwell(_gpu_linux_host(["9.0", "10.0"])) is True # highest cap wins
def _linux_cuda_artifact(runtime_line, supported_sms, min_sm, max_sm, profile):
return ilp.PublishedLlamaArtifact(
asset_name = f"app-b9739-linux-x64-{profile}.tar.gz",
install_kind = "linux-cuda",
runtime_line = runtime_line,
coverage_class = "newer",
supported_sms = supported_sms,
min_sm = min_sm,
max_sm = max_sm,
bundle_profile = profile,
rank = 50,
)
def test_linux_blackwell_override_prefers_cuda13_for_datacenter(monkeypatch):
# Both bundles cover sm_100 and torch reports cuda12, so coverage alone can't
# decide -- only the sm_100 Blackwell floor lifts cuda13 to the front.
cuda12 = _linux_cuda_artifact(
"cuda12", ["86", "89", "90", "100", "120"], 86, 120, "cuda12-newer"
)
cuda13 = _linux_cuda_artifact(
"cuda13", ["86", "89", "90", "100", "103", "120"], 86, 120, "cuda13-newer"
)
release = ilp.PublishedReleaseBundle(
repo = FORK,
release_tag = "b9739-mix",
upstream_tag = "b9739",
assets = {cuda12.asset_name: "https://x/cuda12", cuda13.asset_name: "https://x/cuda13"},
artifacts = [cuda12, cuda13],
)
monkeypatch.setattr(
ilp,
"detected_linux_runtime_lines",
lambda: (["cuda13", "cuda12"], {"cuda13": ["/usr/lib"], "cuda12": ["/usr/lib"]}),
)
selection = ilp.linux_cuda_choice_from_release(
_gpu_linux_host(["10.0"]), release, preferred_runtime_line = "cuda12"
)
assert selection is not None
assert selection.primary.runtime_line == "cuda13"
assert selection.primary.bundle_profile == "cuda13-newer"
def test_drop_blackwell_incapable_windows_cuda_applies_to_datacenter():
# B200 (sm_100) on Windows must drop the cuda-12.4 build and keep cuda13.
host = _host(
system = "Windows",
is_windows = True,
is_x86_64 = True,
has_physical_nvidia = True,
has_usable_nvidia = True,
compute_caps = ["10.0"],
)
cuda124 = ilp.AssetChoice(
repo = FORK,
tag = "b9739",
name = "llama-b9739-bin-win-cuda-12.4-x64.zip",
url = "https://x/124",
source_label = "published",
install_kind = "windows-cuda",
)
cuda13 = ilp.AssetChoice(
repo = FORK,
tag = "b9739",
name = "app-b9739-windows-x64-cuda13-newer.zip",
url = "https://x/13",
source_label = "published",
install_kind = "windows-cuda",
max_sm = 120,
)
kept = ilp._drop_blackwell_incapable_windows_cuda(host, [cuda124, cuda13])
assert [a.name for a in kept] == [cuda13.name]
def test_blackwell_min_toolkit_is_sm_aware():
# Family floor is 12.8; sm_103/sm_121 (no native target before 12.9) lift it.
f = ilp._blackwell_min_toolkit_for_host
assert f(_gpu_linux_host(["10.0"])) == (12, 8) # B200
assert f(_gpu_linux_host(["12.0"])) == (12, 8) # RTX 50
assert f(_gpu_linux_host(["10.3"])) == (12, 9) # B300
assert f(_gpu_linux_host(["12.1"])) == (12, 9) # DGX Spark
assert f(_gpu_linux_host(["10.0", "10.3"])) == (12, 9) # max across SMs wins
def test_sm103_host_drops_cuda128_windows_build():
# B300 (sm_103) needs cuda-12.9: a legacy win-cuda-12.8 build must be dropped.
host = _host(
system = "Windows",
is_windows = True,
is_x86_64 = True,
has_physical_nvidia = True,
has_usable_nvidia = True,
compute_caps = ["10.3"],
)
cuda128 = ilp.AssetChoice(
repo = FORK,
tag = "b9739",
name = "llama-b9739-bin-win-cuda-12.8-x64.zip",
url = "https://x/128",
source_label = "published",
install_kind = "windows-cuda",
)
cuda129 = ilp.AssetChoice(
repo = FORK,
tag = "b9739",
name = "llama-b9739-bin-win-cuda-12.9-x64.zip",
url = "https://x/129",
source_label = "published",
install_kind = "windows-cuda",
)
kept = ilp._drop_blackwell_incapable_windows_cuda(host, [cuda128, cuda129])
assert [a.name for a in kept] == [cuda129.name]
# sm_100 stays on the 12.8 family floor and keeps the same 12.8 build.
b200 = _host(
system = "Windows",
is_windows = True,
is_x86_64 = True,
has_physical_nvidia = True,
has_usable_nvidia = True,
compute_caps = ["10.0"],
)
kept_b200 = ilp._drop_blackwell_incapable_windows_cuda(b200, [cuda128, cuda129])
assert [a.name for a in kept_b200] == [cuda128.name, cuda129.name]
def _upstream_release(tag, asset_names):
return {
"tag_name": tag,
"assets": [
{"name": n, "browser_download_url": f"https://example/{n}"} for n in asset_names
],
}
def test_direct_upstream_arm64_intel_prefers_vulkan():
# Auto-detected Intel GPU on Linux arm64 -> Vulkan prebuilt first, CPU
# second (mirrors the x86_64 branch; ggml-org ships the arm64 Vulkan asset).
host = _host(is_linux = True, is_arm64 = True, machine = "aarch64", has_intel_gpu = True)
rel = _upstream_release(
"b9925",
["llama-b9925-bin-ubuntu-vulkan-arm64.tar.gz", "llama-b9925-bin-ubuntu-arm64.tar.gz"],
)
plan = ilp.direct_upstream_release_plan(rel, host, UPSTREAM, "latest")
kinds = [a.install_kind for a in plan.attempts]
assert kinds[0] == "linux-vulkan", kinds
assert "linux-arm64" in kinds
assert plan.attempts[0].name == "llama-b9925-bin-ubuntu-vulkan-arm64.tar.gz"
def test_direct_upstream_intel_with_hidden_nvidia_is_cpu_only():
# A host with a physical NVIDIA hidden via CUDA_VISIBLE_DEVICES (physical
# True, usable False) + an Intel iGPU must NOT get the Vulkan archive even
# when planning directly against upstream: Vulkan ignores CUDA_VISIBLE_DEVICES
# and could grab the reserved card. It falls through to the CPU asset.
host = _host(
is_linux = True,
is_x86_64 = True,
has_intel_gpu = True,
has_physical_nvidia = True,
has_usable_nvidia = False,
)
rel = _upstream_release(
"b9925",
["llama-b9925-bin-ubuntu-vulkan-x64.tar.gz", "llama-b9925-bin-ubuntu-x64.tar.gz"],
)
plan = ilp.direct_upstream_release_plan(rel, host, UPSTREAM, "latest")
assert [a.install_kind for a in plan.attempts] == ["linux-cpu"]
def test_direct_upstream_arm64_without_intel_is_cpu_only():
host = _host(is_linux = True, is_arm64 = True, machine = "aarch64")
rel = _upstream_release(
"b9925",
["llama-b9925-bin-ubuntu-vulkan-arm64.tar.gz", "llama-b9925-bin-ubuntu-arm64.tar.gz"],
)
plan = ilp.direct_upstream_release_plan(rel, host, UPSTREAM, "latest")
assert [a.install_kind for a in plan.attempts] == ["linux-arm64"]
def test_direct_upstream_x86_intel_prefers_vulkan():
host = _host(is_linux = True, is_x86_64 = True, has_intel_gpu = True)
rel = _upstream_release(
"b9925",
["llama-b9925-bin-ubuntu-vulkan-x64.tar.gz", "llama-b9925-bin-ubuntu-x64.tar.gz"],
)
plan = ilp.direct_upstream_release_plan(rel, host, UPSTREAM, "latest")
kinds = [a.install_kind for a in plan.attempts]
assert kinds[0] == "linux-vulkan", kinds
assert "linux-cpu" in kinds
def test_linux_vulkan_health_glob_matches_bare_cpu_lib():
# The widened glob must cover both arch-suffixed (x64) and bare (arm64) CPU
# libs so a valid Vulkan install is not re-flagged unhealthy every check.
choice = ilp.AssetChoice(
repo = UPSTREAM,
tag = "b9925",
name = "llama-b9925-bin-ubuntu-vulkan-arm64.tar.gz",
url = "https://example/x",
source_label = "upstream",
install_kind = "linux-vulkan",
)
groups = ilp.runtime_payload_health_groups(choice)
assert ["libggml-cpu*.so*"] in groups
assert ["libggml-cpu-*.so*"] not in groups
def test_route_to_vulkan_prebuilt_auto_intel_goes_upstream_and_drops_fork_pin():
# Routing fork -> upstream also drops the fork release pin, which is in a
# different tag namespace and would make the upstream resolver miss.
host = _host(is_linux = True, is_x86_64 = True, has_intel_gpu = True)
routed, repo, tag = ilp._route_to_vulkan_prebuilt(host, FORK, "b9596-mix-abc", force_cpu = False)
assert repo == UPSTREAM
assert tag == ""
assert routed.has_intel_gpu is True
def test_route_to_vulkan_prebuilt_preserves_explicit_upstream_pin():
# A pin set WITH an explicit upstream repo is already on upstream -> kept.
host = _host(is_linux = True, is_x86_64 = True, has_intel_gpu = True)
_routed, repo, tag = ilp._route_to_vulkan_prebuilt(host, UPSTREAM, "b9596", force_cpu = False)
assert repo == UPSTREAM
assert tag == "b9596"
def test_route_to_vulkan_prebuilt_cpu_fallback_wins():
# --cpu-fallback suppresses Vulkan routing even for an Intel host.
host = _host(is_linux = True, is_x86_64 = True, has_intel_gpu = True)
routed, repo, tag = ilp._route_to_vulkan_prebuilt(host, FORK, "b9596-mix-abc", force_cpu = True)
assert repo == FORK
assert tag == "b9596-mix-abc"
assert routed is host
def test_route_to_vulkan_prebuilt_hidden_nvidia_not_rerouted():
# A mixed NVIDIA+Intel host that hid NVIDIA (CUDA_VISIBLE_DEVICES=""/-1):
# physical NVIDIA present but not usable. Must NOT auto-route to Vulkan, or
# Vulkan (which ignores CUDA_VISIBLE_DEVICES) could grab the reserved GPU.
host = _host(
is_linux = True,
is_x86_64 = True,
has_intel_gpu = True,
has_physical_nvidia = True,
has_usable_nvidia = False,
)
_routed, repo, _tag = ilp._route_to_vulkan_prebuilt(host, FORK, "", force_cpu = False)
assert repo == FORK
def test_route_to_vulkan_prebuilt_rocm_host_not_rerouted():
# An Intel iGPU alongside a usable ROCm GPU stays on its ROCm/fork path.
host = _host(is_linux = True, is_x86_64 = True, has_intel_gpu = True, has_rocm = True)
_routed, repo, _tag = ilp._route_to_vulkan_prebuilt(host, FORK, "", force_cpu = False)
assert repo == FORK
def test_route_to_vulkan_prebuilt_non_intel_unchanged():
host = _host(is_linux = True, is_x86_64 = True)
routed, repo, _tag = ilp._route_to_vulkan_prebuilt(host, FORK, "", force_cpu = False)
assert repo == FORK
assert routed is host
def test_resolve_prebuilt_intel_host_routes_to_upstream(monkeypatch, capsys):
# The --resolve-prebuilt probe must agree with the install path: an
# auto-detected Intel host resolves against upstream (Vulkan), not the fork.
monkeypatch.setattr(
ilp, "detect_host", lambda: _host(is_linux = True, is_x86_64 = True, has_intel_gpu = True)
)
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