unsloth/tests/studio/install/test_pr5940_followups.py
Daniel Han a6dc10dad2
Reduce and tighten comments and docstrings across the test suite (#6429)
* Reduce and tighten comments and docstrings in tests

Shorten verbose comments and docstrings across the test suite without
changing any test logic. Remove narration that restates the next line,
collapse long module and test docstrings to a single line, and drop banner
separators. Keep regression context (issue and PR references, run ids),
skip reasons, mocking and timing rationale, license headers, lint and type
directives, and commented-out code.

Comments and docstrings only: an AST signature check confirms no code,
assertions, or string literals changed, and the suite byte-compiles cleanly.

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
2026-06-18 01:07:09 -07:00

319 lines
12 KiB
Python

# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
"""AMD-Windows installer follow-ups (PR #5940): hf-hub validation-model fetch
+ urllib fallback, amd-smi RunAsInvoker injection, name->arch table parity.
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, so it loads without the studio backend.
_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 name->arch case table into ordered [(substr_tokens, arch), ...]."""
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():
"""install.sh name->arch table must match the PowerShell source on the Strix
Halo/Point split and 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's name->arch table must stay row-for-row identical to
install.sh's (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}"
)
# Guards historical drift: Strix Point -> gfx1150, RX 7700S -> 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 pops a UAC/DiskPart prompt RunAsInvoker
# can't suppress, so _amd_smi_allowed() skips it unless HIP-SDK or opt-in.
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.
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.
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 like _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 naming `amd-smi`
# AND spawning a subprocess must gate it behind _amd_smi_allowed().
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.
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"])