5736d5da96
The pinned-index work landed for install.sh and install_python_stack.py, but the Windows installers still picked the wheel index from GPU probing. Extend the same UNSLOTH_TORCH_INDEX_URL / _FAMILY contract so a pinned index wins on every platform: - install.ps1: Get-TorchIndexUrl returns the pinned URL/family before nvidia-smi probing; the AMD ROCm reroute is skipped when the index is pinned, so an explicit cpu/cu* pin on an AMD host is not overwritten. - studio/setup.ps1: add shared Get-PinnedTorchIndexUrl / Get-TorchIndexLeaf helpers; the stale-venv check, the install selection and the AMD reroute all honor the pin, and the CPU/CUDA install pulls from the resolved index URL. - tests: parity test that all four installers read both override vars and the two Windows installers gate the AMD reroute on the pinned flag.
212 lines
8.5 KiB
Python
212 lines
8.5 KiB
Python
"""Cross-platform parity tests between install.sh and install.ps1."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import re
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
|
|
REPO_ROOT = Path(__file__).resolve().parents[2]
|
|
INSTALL_SH = REPO_ROOT / "install.sh"
|
|
INSTALL_PS1 = REPO_ROOT / "install.ps1"
|
|
SETUP_PS1 = REPO_ROOT / "studio" / "setup.ps1"
|
|
STACK_PY = REPO_ROOT / "studio" / "install_python_stack.py"
|
|
|
|
|
|
class TestNoTorchBackendAutoInInstallSh:
|
|
"""install.sh primary paths must not use --torch-backend=auto (only the fallback else-branch may)."""
|
|
|
|
def test_no_torch_backend_auto_outside_fallback(self):
|
|
lines = INSTALL_SH.read_text(encoding = "utf-8").splitlines()
|
|
# Fallback block: from "GPU detection failed" to the next "fi".
|
|
fallback_start = None
|
|
fallback_end = None
|
|
for i, line in enumerate(lines):
|
|
if fallback_start is None and "GPU detection failed" in line:
|
|
fallback_start = i
|
|
elif fallback_start is not None and fallback_end is None and line.strip() == "fi":
|
|
fallback_end = i
|
|
break
|
|
fallback_range = (
|
|
range(fallback_start or 0, (fallback_end or 0) + 1) if fallback_start else range(0)
|
|
)
|
|
|
|
matches = [
|
|
(i + 1, line)
|
|
for i, line in enumerate(lines)
|
|
if "--torch-backend=auto" in line
|
|
and not line.lstrip().startswith("#")
|
|
and i not in fallback_range
|
|
]
|
|
assert matches == [], (
|
|
f"install.sh contains --torch-backend=auto outside the fallback block at lines: "
|
|
f"{[m[0] for m in matches]}"
|
|
)
|
|
|
|
def test_fallback_uses_torch_backend_auto(self):
|
|
"""The fallback branch should use --torch-backend=auto as recovery."""
|
|
text = INSTALL_SH.read_text(encoding = "utf-8")
|
|
assert (
|
|
"GPU detection failed" in text
|
|
), "install.sh should have a fallback branch for when GPU detection fails"
|
|
|
|
|
|
class TestInstallShHasGpuDetection:
|
|
"""install.sh must contain the get_torch_index_url function."""
|
|
|
|
def test_function_exists(self):
|
|
text = INSTALL_SH.read_text(encoding = "utf-8")
|
|
assert (
|
|
"get_torch_index_url()" in text
|
|
), "install.sh is missing the get_torch_index_url() function"
|
|
|
|
def test_torch_index_url_assigned(self):
|
|
text = INSTALL_SH.read_text(encoding = "utf-8")
|
|
assert (
|
|
"TORCH_INDEX_URL=$(get_torch_index_url)" in text
|
|
), "install.sh should assign TORCH_INDEX_URL from get_torch_index_url()"
|
|
|
|
|
|
class TestCudaMappingParity:
|
|
"""CUDA version thresholds must match between install.sh and install.ps1."""
|
|
|
|
@staticmethod
|
|
def _extract_cuda_thresholds_sh(text: str) -> list[str]:
|
|
"""Extract cu* suffixes from the major/minor comparison chain in install.sh."""
|
|
# Only match lines in the if/elif chain that compare _major/_minor
|
|
in_func = False
|
|
results = []
|
|
for line in text.splitlines():
|
|
if "get_torch_index_url()" in line:
|
|
in_func = True
|
|
continue
|
|
if in_func and line.startswith("}"):
|
|
break
|
|
if in_func and ("_major" in line or "_minor" in line):
|
|
m = re.search(r"/(cu\d+|cpu)", line)
|
|
if m:
|
|
results.append(m.group(1))
|
|
return results
|
|
|
|
@staticmethod
|
|
def _extract_cuda_thresholds_ps1(text: str) -> list[str]:
|
|
"""Extract cu* suffixes from the major/minor comparison chain in install.ps1."""
|
|
in_func = False
|
|
depth = 0
|
|
results = []
|
|
for line in text.splitlines():
|
|
if "function Get-TorchIndexUrl" in line:
|
|
in_func = True
|
|
depth = 1
|
|
continue
|
|
if in_func:
|
|
depth += line.count("{") - line.count("}")
|
|
if depth <= 0:
|
|
break
|
|
# Only match the if-chain lines that compare $major/$minor
|
|
if "$major" in line or "$minor" in line:
|
|
m = re.search(r"/(cu\d+|cpu)", line)
|
|
if m:
|
|
results.append(m.group(1))
|
|
return results
|
|
|
|
def test_same_cuda_suffixes(self):
|
|
"""Both scripts should produce the same ordered list of CUDA index suffixes."""
|
|
sh_text = INSTALL_SH.read_text(encoding = "utf-8")
|
|
ps1_text = INSTALL_PS1.read_text(encoding = "utf-8")
|
|
|
|
sh_thresholds = self._extract_cuda_thresholds_sh(sh_text)
|
|
ps1_thresholds = self._extract_cuda_thresholds_ps1(ps1_text)
|
|
|
|
assert len(sh_thresholds) > 0, "Could not extract thresholds from install.sh"
|
|
assert len(ps1_thresholds) > 0, "Could not extract thresholds from install.ps1"
|
|
assert sh_thresholds == ps1_thresholds, (
|
|
f"CUDA mapping mismatch:\n"
|
|
f" install.sh: {sh_thresholds}\n"
|
|
f" install.ps1: {ps1_thresholds}"
|
|
)
|
|
|
|
|
|
class TestPyTorchMirrorEnvVar:
|
|
"""Both install scripts must support the UNSLOTH_PYTORCH_MIRROR env var."""
|
|
|
|
def test_install_sh_has_mirror_var(self):
|
|
text = INSTALL_SH.read_text(encoding = "utf-8")
|
|
assert (
|
|
"UNSLOTH_PYTORCH_MIRROR" in text
|
|
), "install.sh should reference UNSLOTH_PYTORCH_MIRROR"
|
|
|
|
def test_install_ps1_has_mirror_var(self):
|
|
text = INSTALL_PS1.read_text(encoding = "utf-8")
|
|
assert (
|
|
"UNSLOTH_PYTORCH_MIRROR" in text
|
|
), "install.ps1 should reference UNSLOTH_PYTORCH_MIRROR"
|
|
|
|
|
|
class TestUvBytecodeCompileTimeout:
|
|
"""Installers should relax uv bytecode compilation timeout by default."""
|
|
|
|
@staticmethod
|
|
def _version_tuple(version: str) -> tuple[int, ...]:
|
|
return tuple(int(part) for part in version.split("."))
|
|
|
|
def test_install_sh_uses_uv_version_with_timeout_env(self):
|
|
text = INSTALL_SH.read_text(encoding = "utf-8")
|
|
match = re.search(r'^UV_MIN_VERSION="([^"]+)"$', text, re.MULTILINE)
|
|
assert match, "install.sh should declare UV_MIN_VERSION"
|
|
assert self._version_tuple(match.group(1)) >= self._version_tuple("0.7.22")
|
|
|
|
def test_install_ps1_uses_uv_version_with_timeout_env(self):
|
|
text = INSTALL_PS1.read_text(encoding = "utf-8")
|
|
match = re.search(r'^\s*\$UvMinVersion = "([^"]+)"$', text, re.MULTILINE)
|
|
assert match, "install.ps1 should declare $UvMinVersion"
|
|
assert self._version_tuple(match.group(1)) >= self._version_tuple("0.7.22")
|
|
assert "function Test-UvVersionOk" in text
|
|
assert "if (-not (Test-UvVersionOk))" in text
|
|
|
|
def test_install_sh_preserves_timeout_override(self):
|
|
text = INSTALL_SH.read_text(encoding = "utf-8")
|
|
assert (
|
|
': "${UV_COMPILE_BYTECODE_TIMEOUT:=180}"' in text
|
|
), "install.sh should default UV_COMPILE_BYTECODE_TIMEOUT without overwriting callers"
|
|
assert (
|
|
"export UV_COMPILE_BYTECODE_TIMEOUT" in text
|
|
), "install.sh should export UV_COMPILE_BYTECODE_TIMEOUT for uv subprocesses"
|
|
|
|
def test_install_ps1_preserves_timeout_override(self):
|
|
text = INSTALL_PS1.read_text(encoding = "utf-8")
|
|
assert (
|
|
"if (-not $env:UV_COMPILE_BYTECODE_TIMEOUT)" in text
|
|
), "install.ps1 should preserve caller UV_COMPILE_BYTECODE_TIMEOUT overrides"
|
|
assert (
|
|
'$env:UV_COMPILE_BYTECODE_TIMEOUT = "180"' in text
|
|
), "install.ps1 should default UV_COMPILE_BYTECODE_TIMEOUT"
|
|
|
|
|
|
class TestTorchIndexOverrideParity:
|
|
"""Every installer must honor UNSLOTH_TORCH_INDEX_URL / _FAMILY so a pinned wheel
|
|
index wins over GPU probing on all platforms (no asymmetric, per-OS coverage)."""
|
|
|
|
@pytest.mark.parametrize(
|
|
"path",
|
|
[INSTALL_SH, INSTALL_PS1, SETUP_PS1, STACK_PY],
|
|
ids = ["install.sh", "install.ps1", "setup.ps1", "install_python_stack.py"],
|
|
)
|
|
def test_installer_reads_override_env(self, path):
|
|
text = path.read_text(encoding = "utf-8")
|
|
for var in ("UNSLOTH_TORCH_INDEX_URL", "UNSLOTH_TORCH_INDEX_FAMILY"):
|
|
assert var in text, f"{path.name} does not honor {var}"
|
|
|
|
@pytest.mark.parametrize(
|
|
"path",
|
|
[INSTALL_PS1, SETUP_PS1],
|
|
ids = ["install.ps1", "setup.ps1"],
|
|
)
|
|
def test_amd_reroute_guarded_when_pinned(self, path):
|
|
# The AMD ROCm reroute must be skipped when the index is explicitly pinned,
|
|
# so an explicit cpu / cu* / rocm pin on an AMD host is not overwritten.
|
|
text = path.read_text(encoding = "utf-8")
|
|
assert "TorchIndexPinned" in text, (
|
|
f"{path.name} should gate the AMD ROCm reroute on a pinned-index flag"
|
|
)
|