353 lines
13 KiB
Python
353 lines
13 KiB
Python
# SPDX-License-Identifier: AGPL-3.0-only
|
|
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
|
|
|
|
"""Tests for the AMD-Windows installer follow-ups (PR #5940):
|
|
|
|
* the huggingface_hub validation-model fetch + its urllib fallback,
|
|
* run_capture's Windows-only amd-smi __COMPAT_LAYER=RunAsInvoker injection,
|
|
* parity of the name->arch table between install.ps1 and setup.ps1.
|
|
|
|
Mock-only; no AMD hardware or network required.
|
|
"""
|
|
|
|
import importlib.util
|
|
import re
|
|
import subprocess
|
|
import sys
|
|
from pathlib import Path
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import pytest
|
|
|
|
PACKAGE_ROOT = Path(__file__).resolve().parents[3]
|
|
|
|
# install_llama_prebuilt.py is self-contained (stdlib + optional filelock), so it
|
|
# loads without the studio backend on sys.path.
|
|
_PREBUILT_PATH = PACKAGE_ROOT / "studio" / "install_llama_prebuilt.py"
|
|
_SPEC = importlib.util.spec_from_file_location(
|
|
"studio_install_llama_prebuilt_pr5940", _PREBUILT_PATH
|
|
)
|
|
assert _SPEC is not None and _SPEC.loader is not None
|
|
prebuilt = importlib.util.module_from_spec(_SPEC)
|
|
sys.modules[_SPEC.name] = prebuilt
|
|
_SPEC.loader.exec_module(prebuilt)
|
|
|
|
_INSTALL_PS1 = PACKAGE_ROOT / "install.ps1"
|
|
_SETUP_PS1 = PACKAGE_ROOT / "studio" / "setup.ps1"
|
|
_INSTALL_SH = PACKAGE_ROOT / "install.sh"
|
|
|
|
|
|
# ── _hf_resolve_url_parts ────────────────────────────────────────────────────
|
|
|
|
|
|
def test_hf_resolve_url_parts_valid():
|
|
assert prebuilt._hf_resolve_url_parts(
|
|
"https://huggingface.co/ggml-org/models/resolve/main/tinyllamas/stories260K.gguf"
|
|
) == ("ggml-org/models", "main", "tinyllamas/stories260K.gguf")
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"url",
|
|
[
|
|
"https://github.com/owner/repo/releases/download/x.gguf", # not huggingface
|
|
"https://huggingface.co/owner/repo", # no /resolve/<rev>/
|
|
"https://huggingface.co/owner/repo/blob/main/x.gguf", # /blob/ not /resolve/
|
|
"not even a url",
|
|
],
|
|
)
|
|
def test_hf_resolve_url_parts_non_hf_returns_none(url):
|
|
assert prebuilt._hf_resolve_url_parts(url) is None
|
|
|
|
|
|
# ── _fetch_validation_model_bytes ────────────────────────────────────────────
|
|
|
|
|
|
def test_fetch_validation_model_prefers_huggingface_hub(tmp_path):
|
|
model = tmp_path / "stories260K.gguf"
|
|
model.write_bytes(b"GGUF-via-hf")
|
|
fake_hf = MagicMock(return_value = str(model))
|
|
with (
|
|
patch.object(
|
|
prebuilt, "validated_validation_model_bytes", side_effect = lambda b: b
|
|
),
|
|
patch.dict(
|
|
sys.modules, {"huggingface_hub": MagicMock(hf_hub_download = fake_hf)}
|
|
),
|
|
):
|
|
assert prebuilt._fetch_validation_model_bytes() == b"GGUF-via-hf"
|
|
assert fake_hf.called # hf path was taken, urllib not needed
|
|
|
|
|
|
def test_fetch_validation_model_falls_back_to_urllib_on_hf_failure():
|
|
fake_hf = MagicMock(side_effect = RuntimeError("hf unreachable"))
|
|
with (
|
|
patch.object(
|
|
prebuilt, "validated_validation_model_bytes", side_effect = lambda b: b
|
|
),
|
|
patch.dict(
|
|
sys.modules, {"huggingface_hub": MagicMock(hf_hub_download = fake_hf)}
|
|
),
|
|
patch.object(prebuilt, "download_bytes", return_value = b"GGUF-via-urllib") as dl,
|
|
):
|
|
assert prebuilt._fetch_validation_model_bytes() == b"GGUF-via-urllib"
|
|
assert dl.called # fell back to the direct URL download
|
|
|
|
|
|
# ── run_capture amd-smi RunAsInvoker injection ───────────────────────────────
|
|
|
|
|
|
def _capture_env(command, system):
|
|
captured = {"env": "sentinel"}
|
|
|
|
def fake_run(cmd, **kwargs):
|
|
captured["env"] = kwargs.get("env")
|
|
return subprocess.CompletedProcess(cmd, 0, "", "")
|
|
|
|
with (
|
|
patch.object(prebuilt.subprocess, "run", side_effect = fake_run),
|
|
patch.object(prebuilt.platform, "system", return_value = system),
|
|
):
|
|
prebuilt.run_capture(command)
|
|
return captured["env"]
|
|
|
|
|
|
def test_run_capture_injects_runasinvoker_for_amd_smi_on_windows():
|
|
env = _capture_env(["amd-smi", "list"], "Windows")
|
|
assert env is not None and env.get("__COMPAT_LAYER") == "RunAsInvoker"
|
|
|
|
|
|
def test_run_capture_injects_for_full_path_amd_smi_exe_on_windows():
|
|
env = _capture_env(["amd-smi.exe", "version"], "Windows")
|
|
assert env is not None and env.get("__COMPAT_LAYER") == "RunAsInvoker"
|
|
|
|
|
|
def test_run_capture_no_injection_for_non_amd_smi_on_windows():
|
|
assert _capture_env(["rocminfo"], "Windows") is None
|
|
|
|
|
|
def test_run_capture_no_injection_on_linux():
|
|
# amd-smi does not auto-elevate on Linux, so no env override is applied.
|
|
assert _capture_env(["amd-smi", "list"], "Linux") is None
|
|
|
|
|
|
# ── name->arch table parity (install.ps1 vs setup.ps1) ───────────────────────
|
|
|
|
|
|
def _ps_name_arch_rows(text):
|
|
return re.findall(r'@\{\s*P\s*=\s*"([^"]*)"\s*;\s*A\s*=\s*"(gfx[0-9a-z]+)"', text)
|
|
|
|
|
|
def test_ps_name_arch_tables_in_sync():
|
|
t1 = _ps_name_arch_rows(_INSTALL_PS1.read_text(encoding = "utf-8"))
|
|
t2 = _ps_name_arch_rows(_SETUP_PS1.read_text(encoding = "utf-8"))
|
|
assert t1, "no nameArchTable found in install.ps1"
|
|
assert t1 == t2, f"name->arch tables drifted:\ninstall.ps1={t1}\nsetup.ps1={t2}"
|
|
|
|
|
|
def test_rx_7700s_resolves_to_gfx1102_not_gfx1100():
|
|
rows = _ps_name_arch_rows(_INSTALL_PS1.read_text(encoding = "utf-8"))
|
|
name = "AMD Radeon RX 7700S"
|
|
matched = next((arch for pattern, arch in rows if re.search(pattern, name)), None)
|
|
assert matched == "gfx1102", f"RX 7700S matched {matched!r}, expected gfx1102"
|
|
|
|
|
|
def test_radeon_8060s_resolves_to_gfx1151():
|
|
rows = _ps_name_arch_rows(_INSTALL_PS1.read_text(encoding = "utf-8"))
|
|
name = "AMD Radeon(TM) 8060S Graphics"
|
|
matched = next((arch for pattern, arch in rows if re.search(pattern, name)), None)
|
|
assert matched == "gfx1151"
|
|
|
|
|
|
def _sh_name_arch_rows(text, var = "_gpu_disp_gfx"):
|
|
"""Parse a bash `case "$..._mkt" in ... ) <var>="gfxNNNN"` name->arch
|
|
table into [(substr_tokens, arch), ...] preserving order."""
|
|
rows = []
|
|
for line in text.splitlines():
|
|
m = re.search(var + r'="(gfx[0-9a-z]+)"', line)
|
|
if not m or '*"' not in line:
|
|
continue
|
|
tokens = re.findall(r'\*"([^"]+)"\*', line)
|
|
if tokens:
|
|
rows.append((tokens, m.group(1)))
|
|
return rows
|
|
|
|
|
|
def _sh_resolve(rows, name):
|
|
for tokens, arch in rows:
|
|
if any(tok in name for tok in tokens): # bash *"X"* == substring
|
|
return arch
|
|
return None
|
|
|
|
|
|
def test_install_sh_name_arch_agrees_with_ps_for_strix_and_non_amd():
|
|
"""The bash install.sh name->arch table must agree with the PowerShell
|
|
source-of-truth for the Strix Halo (gfx1151) vs Strix Point (gfx1150)
|
|
split, and must never misclassify NVIDIA/Intel as an AMD gfx."""
|
|
sh_rows = _sh_name_arch_rows(_INSTALL_SH.read_text(encoding = "utf-8"))
|
|
ps_rows = _ps_name_arch_rows(_INSTALL_PS1.read_text(encoding = "utf-8"))
|
|
assert sh_rows, "no name->arch case table found in install.sh"
|
|
cases = {
|
|
"AMD Radeon(TM) 8060S Graphics": "gfx1151", # Strix Halo
|
|
"AMD Ryzen AI Max+ PRO 395 w/ Radeon 8060S": "gfx1151",
|
|
"AMD Radeon 890M Graphics": "gfx1150", # Strix Point (NOT gfx1151)
|
|
"AMD Ryzen AI 9 HX 370 w/ Radeon 890M": "gfx1150",
|
|
"AMD Radeon RX 7700S": "gfx1102",
|
|
"NVIDIA GeForce RTX 4090": None,
|
|
"Intel(R) Arc A770 Graphics": None,
|
|
}
|
|
for name, expect in cases.items():
|
|
sh = _sh_resolve(sh_rows, name)
|
|
assert sh == expect, f"install.sh: {name!r} -> {sh!r}, expected {expect!r}"
|
|
if expect is not None: # cross-check bash agrees with the PowerShell table
|
|
ps = next((a for p, a in ps_rows if re.search(p, name)), None)
|
|
assert (
|
|
sh == ps
|
|
), f"install.sh/install.ps1 drift for {name!r}: {sh!r} vs {ps!r}"
|
|
|
|
|
|
def test_setup_sh_name_arch_table_in_sync_with_install_sh():
|
|
"""studio/setup.sh keeps its own copy of the bash name->arch table (over
|
|
`_setup_gfx`); it must stay row-for-row identical to install.sh's, both in
|
|
tokens and in match order (order carries the RX 7700S -> gfx1102 rule)."""
|
|
install_rows = _sh_name_arch_rows(_INSTALL_SH.read_text(encoding = "utf-8"))
|
|
setup_rows = _sh_name_arch_rows(
|
|
(PACKAGE_ROOT / "studio" / "setup.sh").read_text(encoding = "utf-8"),
|
|
var = "_setup_gfx",
|
|
)
|
|
assert setup_rows, "no name->arch case table found in studio/setup.sh"
|
|
assert install_rows == setup_rows, (
|
|
"bash name->arch tables drifted:\n"
|
|
f"install.sh={install_rows}\nstudio/setup.sh={setup_rows}"
|
|
)
|
|
# The historical drift this guards against: Strix Point SKUs must be
|
|
# gfx1150, and the spaceless RX 7700S must match gfx1102 before gfx1100.
|
|
for name, expect in {
|
|
"AMD Radeon 890M Graphics": "gfx1150",
|
|
"AMD Ryzen AI 9 HX 370 w/ Radeon 890M": "gfx1150",
|
|
"AMD Radeon(TM) 8060S Graphics": "gfx1151",
|
|
"AMD Radeon RX 7700S": "gfx1102",
|
|
}.items():
|
|
got = _sh_resolve(setup_rows, name)
|
|
assert got == expect, f"setup.sh: {name!r} -> {got!r}, expected {expect!r}"
|
|
|
|
|
|
# ── amd-smi gating (DiskPart UAC-prompt avoidance) ───────────────────────────
|
|
# On Windows w/o a HIP SDK, amd-smi elevates and pops a UAC/DiskPart prompt
|
|
# RunAsInvoker can't suppress, so _amd_smi_allowed() skips it by default;
|
|
# HIP-SDK hosts and an explicit opt-in keep it.
|
|
|
|
|
|
def _amd_smi_allowed_under(system, hipinfo_present, env):
|
|
which = (
|
|
(lambda name: r"C:\hip\bin\hipinfo.exe" if name == "hipinfo" else None)
|
|
if hipinfo_present
|
|
else (lambda name: None)
|
|
)
|
|
with (
|
|
patch.object(prebuilt.platform, "system", return_value = system),
|
|
patch.object(prebuilt.shutil, "which", side_effect = which),
|
|
patch.dict(prebuilt.os.environ, env, clear = True),
|
|
):
|
|
return prebuilt._amd_smi_allowed()
|
|
|
|
|
|
def test_amd_smi_allowed_on_linux_regardless():
|
|
# Linux amd-smi does not elevate -> always allowed (no regression on Linux).
|
|
assert _amd_smi_allowed_under("Linux", hipinfo_present = False, env = {}) is True
|
|
|
|
|
|
def test_amd_smi_skipped_on_windows_without_hip_sdk():
|
|
# The DiskPart fix: no HIP SDK + no opt-in -> do not spawn amd-smi.
|
|
assert _amd_smi_allowed_under("Windows", hipinfo_present = False, env = {}) is False
|
|
|
|
|
|
def test_amd_smi_allowed_on_windows_with_hip_sdk():
|
|
# hipinfo present => amd-smi runs un-elevated, so it is allowed (no regression
|
|
# for HIP-SDK Windows users, who never saw the prompt).
|
|
assert _amd_smi_allowed_under("Windows", hipinfo_present = True, env = {}) is True
|
|
|
|
|
|
def test_amd_smi_opt_in_forces_on_windows_no_sdk():
|
|
assert (
|
|
_amd_smi_allowed_under(
|
|
"Windows", hipinfo_present = False, env = {"UNSLOTH_ENABLE_AMD_SMI": "1"}
|
|
)
|
|
is True
|
|
)
|
|
|
|
|
|
def test_amd_smi_opt_out_overrides_hip_sdk():
|
|
assert (
|
|
_amd_smi_allowed_under(
|
|
"Windows", hipinfo_present = True, env = {"UNSLOTH_ENABLE_AMD_SMI": "0"}
|
|
)
|
|
is False
|
|
)
|
|
|
|
|
|
def test_ps_installers_gate_amd_smi_on_windows():
|
|
# Both PowerShell installers must gate amd-smi behind HIP-SDK presence + the
|
|
# UNSLOTH_ENABLE_AMD_SMI opt-in, mirroring _amd_smi_allowed().
|
|
for ps in (_INSTALL_PS1, _SETUP_PS1):
|
|
text = ps.read_text(encoding = "utf-8")
|
|
assert (
|
|
"UNSLOTH_ENABLE_AMD_SMI" in text
|
|
), f"{ps.name} missing amd-smi opt-in gate"
|
|
assert "amdSmiAllowed" in text, f"{ps.name} missing amd-smi gate variable"
|
|
|
|
|
|
def test_install_python_stack_gates_every_amd_smi_spawn():
|
|
# Regression for the DiskPart UAC prompt: every function that both names the
|
|
# `amd-smi` command AND spawns a subprocess must gate it behind
|
|
# _amd_smi_allowed(). The "ROCm torch missing" probe once spawned `amd-smi
|
|
# list` ungated on Adrenalin-only hosts; not-spawning is the only fix.
|
|
import ast
|
|
|
|
src = (PACKAGE_ROOT / "studio" / "install_python_stack.py").read_text(
|
|
encoding = "utf-8"
|
|
)
|
|
tree = ast.parse(src)
|
|
|
|
def _names_amd_smi_command(node):
|
|
# Exact "amd-smi"/"amd-smi.exe" constant, not a substring in a log message.
|
|
return any(
|
|
isinstance(n, ast.Constant)
|
|
and isinstance(n.value, str)
|
|
and n.value.lower() in ("amd-smi", "amd-smi.exe")
|
|
for n in ast.walk(node)
|
|
)
|
|
|
|
def _spawns_subprocess(node):
|
|
for n in ast.walk(node):
|
|
if (
|
|
isinstance(n, ast.Call)
|
|
and isinstance(n.func, ast.Attribute)
|
|
and isinstance(n.func.value, ast.Name)
|
|
and n.func.value.id == "subprocess"
|
|
):
|
|
return True
|
|
return False
|
|
|
|
def _references_gate(node):
|
|
return any(
|
|
isinstance(n, ast.Name) and n.id == "_amd_smi_allowed"
|
|
for n in ast.walk(node)
|
|
)
|
|
|
|
offenders = [
|
|
node.name
|
|
for node in ast.walk(tree)
|
|
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef))
|
|
and _names_amd_smi_command(node)
|
|
and _spawns_subprocess(node)
|
|
and not _references_gate(node)
|
|
]
|
|
assert not offenders, (
|
|
"install_python_stack.py spawns amd-smi without an _amd_smi_allowed() "
|
|
f"gate in: {offenders} -- this pops the Windows UAC/DiskPart prompt on "
|
|
"Adrenalin-only (no HIP SDK) hosts."
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
pytest.main([__file__, "-v"])
|