unsloth/tests/python/test_tokenizers_and_torch_constraint.py
Daniel Han a9be36830e
Installer: allow torch 2.11.x on the CUDA install path (fresh install + studio) (#6959)
* Studio: allow torch 2.11.x on the CUDA install path

The CUDA torch repair path (_ensure_cuda_torch) installs torch/torchvision/
torchaudio from an exclusive --index-url, so _CUDA_TORCH_PKG_SPEC decides
exactly which torch the Studio venv gets. It was capped at torch<2.11.0, so on
a cu128/cu130 host the venv resolved torch 2.10.x even though the CUDA indexes
now publish torch 2.11.0. That left the Studio venv a torch minor behind the
torch 2.11.0 Docker base image, so the CUDA dedup step would relink base libs
under a mismatched torch.

Raise the upper bound to <2.12.0 (torchvision <0.27.0, torchaudio <2.12.0) so
the CUDA install path lands on torch 2.11.x, matching the rocm7.2 spec and the
base image. The torchao selector already maps torch 2.11 -> torchao 0.17.0, and
_ensure_flash_attn degrades gracefully when no prebuilt wheel matches (Blackwell
skips it outright; non-Blackwell prints a warning and continues), so no other
pin needs to move.

Add test_cuda_torch_spec.py to lock the bound (torch 2.11.x in, 2.12.x out) and
assert the CUDA and rocm7.2 upper bounds stay in lockstep.

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

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

* test: use zip(strict=True) so a spec length mismatch fails loudly

* install.sh: widen the CUDA torch ceiling to <2.12.0 so a fresh install matches the base

Raising _CUDA_TORCH_PKG_SPEC alone was not enough: that spec only feeds
_ensure_cuda_torch(), the ROCm-poisoning repair path that early-returns on a
normal NVIDIA host. A fresh CUDA install (including the studio Docker build,
which runs `bash install.sh --local`) takes its torch from install.sh's
TORCH_CONSTRAINT, which was still capped at torch>=2.4,<2.11.0, so cu12x/cu13x
resolved torch 2.10.x and the venv landed a minor behind the torch 2.11.0 base
image.

Extend the existing `case "$TORCH_INDEX_URL"` block (which already relaxes
rocm7.2) with a `*/cu[0-9]*` branch that widens the ceiling to <2.12.0, keeping
the >=2.4 floor so an older CUDA index (e.g. cu118) that tops out below 2.11
still resolves. The CPU wheel and older ROCm tags stay on <2.11.0 (the glob
does not match /cpu). torchvision/torchaudio are bare on this install line and
resolve their compatible companions via wheel metadata, matching the rocm7.2
pattern.

Add behavioral tests (Python + shell) exercising the case block: cu118/124/126/
128/130 widen to <2.12.0, rocm7.2 stays 2.11.x, and /cpu plus older ROCm keep
the default <2.11.0.

* install.sh: key the CUDA torch widening off the index leaf, not the full URL

The `*/cu[0-9]*` glob matched a `cu<digit>` segment anywhere in TORCH_INDEX_URL,
so a custom UNSLOTH_PYTORCH_MIRROR whose base path contains e.g. cu128 but whose
final leaf is cpu or an older ROCm tag would still widen TORCH_CONSTRAINT to
<2.12.0, contradicting the block's own comment and letting a CPU / older-ROCm
mirror resolve torch 2.11.x. Match on _torch_index_leaf (the final path segment
the backend classification just above already computes) so only a real cu*/
rocm7.2 leaf is affected; cpu and older ROCm keep the default <2.11.0. Update
the Python + shell tests to mirror the leaf-anchored case and add regression
cases for a mirror base that contains cu128 but resolves to a cpu / rocm7.1 leaf.

* install: freeze the torch trio during the with-deps unsloth installs

Released unsloth wheels can pin an older torch than Step 1 installed
(unsloth 2026.7.2 declares torch<2.11.0), so the with-deps resolve from
PyPI silently downgrades the pinned +cuXXX torch trio to PyPI's default
wheel. The flavor guard cannot catch every such swap: PyPI's torch 2.10
default is itself cu128-flavored, so the cuXXX tag comparison still
matches while the version silently drops. Freeze the just-installed trio
with uv --overrides (overrides replace dependency requirements during
resolution), keeping torch 2.11.0+cuXXX in place while unsloth's other
dependencies resolve normally. Verified on the cu128 path: without the
override torch drops 2.11.0+cu128 -> 2.10.0; with it the trio survives
and unsloth 2026.7.2 + unsloth-zoo install cleanly.

* install: fold UV_OVERRIDE env files into the torch-trio overrides file

The CLI --overrides flag is the command-line form of UV_OVERRIDE, so
passing it replaced any overrides file already exported for the process;
macOS arm64 exports UV_OVERRIDE=overrides-darwin-arm64.txt for the same
generic install path and would have lost those pins. Concatenate any
UV_OVERRIDE files into the temp trio file so both keep applying.

* install: extend the torch-trio overrides guard to migrated installs

Four follow-ups to the Step-2 --overrides guard, all empirically verified:

1. The migrated-environment with-deps unsloth install resolved
   unsloth>=2026.7.2 (which pins torch<2.11.0) without the overrides file,
   so a migrated CUDA venv on torch 2.11 was silently downgraded -- the
   exact bug this branch fixes on the fresh path. The overrides build is
   now a function (_build_unsloth_torch_overrides, reading the trio
   installed at call time) invoked by both with-deps paths; the migrated
   no-torch path installs --no-deps and stays unguarded.

2. The overrides temp file is now cleaned by the EXIT trap (same pattern
   as _UV_OVERRIDE_TMPDIR, pre-initialized empty so an inherited value can
   never reach the trap's rm); previously any Step-2 failure leaked it.

3. Folding UV_OVERRIDE files used cat, which joins the last requirement of
   a file lacking a trailing newline onto the next file's first requirement
   (reproduced: idna==3.10certifi==2025.1.31 makes uv fail parsing).

4. Inherited torch/torchvision/torchaudio override lines are now filtered
   out when folding: uv intersects duplicate overrides rather than
   last-wins (verified on uv 0.10.12: direct conflict is unsatisfiable,
   transitive conflict silently backtracks), so a conflicting inherited
   trio pin would break the resolve the generated exact pins protect.
   Both 3 and 4 are handled by a single newline-terminating awk filter
   that preserves non-trio overrides (torchmetrics, torchao, ...).

test_unsloth_torch_override.sh extended: migrated-path coverage, trap
assertion, and a functional fold test (14 checks).

* installer: tighten comments

* install: keep the existing torch release when re-running the installer

Re-running `curl -fsSL https://unsloth.ai/install.sh | sh` over an existing
install rebuilds the venv for clean state, which silently moved users to the
newest torch in range (2.10 -> 2.11 once the constraint widened). A torch the
user already validated must survive an unsloth update.

Before the old venv is moved aside for rollback, its torch version is probed
(last stdout line only, so sitecustomize noise cannot corrupt it). After the
index leaf is chosen, _previous_torch_pin turns that version into a
torch==X.Y.Z pin, but only when it cannot do harm:

- cu*/cpu leaves only; rocm leaves keep their floors (rocm7.2 must land 2.11
  for the Strix _grouped_mm fix) and the Radeon wheel-matching path is
  untouched.
- The wheel's flavor tag must match the freshly chosen leaf, so a flavor
  change (cpu -> cuda, cu126 -> cu130) still installs the correct new build.
- The base must look like a release, so probe noise never becomes a pin.
- UNSLOTH_TORCH_UPGRADE=1 opts out and restores the old always-newest
  behavior; the substep line advertises it.

The supported range is kept in _PREV_FALLBACK_CONSTRAINT: if the exact
release is not resolvable from the chosen index (custom mirrors prune old
wheels), the install warns and falls back to the newest supported release
instead of failing the whole run. The later flavor-mismatch repair reuses
TORCH_CONSTRAINT, so a mid-install clobber is repaired back to the kept
release rather than the newest one.

Verified end to end: a venv seeded with torch 2.10.0+cu130 re-run through the
full installer finishes with torch 2.10.0+cu130 (previously 2.11.0+cu130).

Tests: tests/sh/test_previous_torch_pin.sh covers keep/flavor-change/rocm/
noise/opt-out plus wiring (probe ordering before venv replacement, fallback
present, SKIP_TORCH gate).

* install: constrain kept torch pins to the supported window

Review caught that _previous_torch_pin pinned the previous venv's torch on
flavor match alone, so a release outside the installer's active range (a
2.3.x manual install below the >=2.4 floor, or a 2.12.x manual upgrade above
the ceiling) replaced the bounds computed just above it and a rerun kept a
torch the installer otherwise deliberately excludes.

New _torch_release_in_window checks the probed base against the active
TORCH_CONSTRAINT ("torch>=A.B[,<C.D.F]") at major.minor granularity, which
is exact for the windows this script uses (ceilings are always X.Y.0; a
non-.0 ceiling would only make it conservative). Anything unparseable
answers no, so probe noise or a malformed window fails toward the supported
range instead of becoming a pin. _previous_torch_pin takes the active
constraint as a third argument and refuses out-of-window releases; the
in-window keep behavior is unchanged.

Tests: out-of-window rows (2.3.x floor, 2.12.x ceiling, boundary keeps, cpu
and macOS windows, malformed/empty windows) plus direct
_torch_release_in_window coverage.

---------

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

649 lines
26 KiB
Python

"""Install fixes: tokenizers in no-torch-runtime.txt, and TORCH_CONSTRAINT in install.sh."""
from __future__ import annotations
import pathlib
import re
import subprocess
import textwrap
import pytest
# Locate source files relative to this test.
_TESTS_DIR = pathlib.Path(__file__).resolve().parent.parent # tests/
_REPO_ROOT = _TESTS_DIR.parent # unsloth/
_INSTALL_SH = _REPO_ROOT / "install.sh"
_INSTALL_PS1 = _REPO_ROOT / "install.ps1"
_SETUP_PS1 = _REPO_ROOT / "studio" / "setup.ps1"
_NO_TORCH_RT = _REPO_ROOT / "studio" / "backend" / "requirements" / "no-torch-runtime.txt"
def _read(path: pathlib.Path) -> str:
return path.read_text(encoding = "utf-8")
def _lines(path: pathlib.Path) -> list[str]:
"""Return non-comment, non-blank lines stripped."""
return [
ln.strip()
for ln in _read(path).splitlines()
if ln.strip() and not ln.strip().startswith("#")
]
# Group 1 -- Structural checks (no network, instant)
class TestStructuralTokenizers:
"""Verify tokenizers presence and ordering in no-torch-runtime.txt."""
def test_tokenizers_present(self):
"""tokenizers must be a standalone package line."""
pkgs = _lines(_NO_TORCH_RT)
bare_names = [p.split(">")[0].split("<")[0].split("!")[0].split("=")[0] for p in pkgs]
assert "tokenizers" in bare_names
def test_tokenizers_before_transformers(self):
"""tokenizers should appear before transformers (install order intent)."""
pkgs = _lines(_NO_TORCH_RT)
bare_names = [p.split(">")[0].split("<")[0].split("!")[0].split("=")[0] for p in pkgs]
idx_tok = bare_names.index("tokenizers")
idx_tf = bare_names.index("transformers")
assert idx_tok < idx_tf, (
f"tokenizers at index {idx_tok} should appear before " f"transformers at index {idx_tf}"
)
def test_torch_not_in_no_torch_file(self):
"""torch itself must NOT be listed in the no-torch requirements."""
pkgs = _lines(_NO_TORCH_RT)
bare_names = [p.split(">")[0].split("<")[0].split("!")[0].split("=")[0] for p in pkgs]
assert "torch" not in bare_names
class TestStructuralTorchConstraint:
"""Verify TORCH_CONSTRAINT wiring in install.sh."""
_sh = _read(_INSTALL_SH)
def test_default_assignment_exists(self):
assert 'TORCH_CONSTRAINT="torch>=2.4,<2.11.0"' in self._sh
def test_tightened_assignment_exists(self):
assert 'TORCH_CONSTRAINT="torch>=2.6,<2.11.0"' in self._sh
def test_cuda_constraint_widened_to_2_12(self):
"""A fresh CUDA install widens the ceiling to <2.12.0 so cu12x/cu13x
land torch 2.11.x (matches the base image and _CUDA_TORCH_PKG_SPEC);
without it cu128/cu130 resolves torch 2.10.x."""
assert 'TORCH_CONSTRAINT="torch>=2.4,<2.12.0"' in self._sh
def test_cuda_case_widens_via_index_leaf(self):
"""The cu* branch of the _torch_index_leaf case sets the widened
constraint (parallel to rocm7.2), anchored on the leaf."""
m = re.search(
r'cu\[0-9\]\*\)\s*TORCH_CONSTRAINT="torch>=2\.4,<2\.12\.0"',
self._sh,
)
assert m is not None, "CUDA (cu*) TORCH_CONSTRAINT widening case not found"
def test_variable_used_in_pip_install(self):
"""$TORCH_CONSTRAINT must appear in a uv pip install line."""
assert '"$TORCH_CONSTRAINT"' in self._sh
def test_hardcoded_torch_constraint_only_once(self):
"""The hard-coded torch>=2.4,<2.11.0 string should appear exactly once
in install.sh (the default assignment), not in pip install lines."""
count = self._sh.count('"torch>=2.4,<2.11.0"')
assert count == 1, f"Expected 1, found {count}"
def test_tightening_guarded_by_skip_torch(self):
"""The block must check SKIP_TORCH=false."""
# Find the tightening if-block
m = re.search(
r"if\s.*SKIP_TORCH.*=\s*false.*&&.*OS.*=.*macos.*&&.*_ARCH.*=.*arm64",
self._sh,
)
assert m is not None, "Guard not found: SKIP_TORCH + macos + arm64"
def test_tightening_guarded_by_arch(self):
m = re.search(r"_ARCH.*=.*arm64", self._sh)
assert m is not None
def test_tightening_guarded_by_os(self):
m = re.search(r"OS.*=.*macos", self._sh)
assert m is not None
class TestStructuralInstallPs1Unchanged:
"""install.ps1 should NOT have TORCH_CONSTRAINT variable."""
_ps1 = _read(_INSTALL_PS1)
def test_no_torch_constraint_variable(self):
assert "TORCH_CONSTRAINT" not in self._ps1
assert "$TorchConstraint" not in self._ps1
def test_hardcoded_torch_constraint_present(self):
assert '"torch>=2.4,<2.11.0"' in self._ps1
class TestInstallPs1UvDefaultIndex:
"""Installer-managed torch indexes must override inherited uv defaults."""
_ps1 = _read(_INSTALL_PS1)
def test_torch_installs_use_default_index(self):
assert "--default-index $TorchIndexUrl" in self._ps1
assert "--default-index $ROCmIndexUrl" in self._ps1
def test_torch_installs_do_not_use_deprecated_index_url(self):
assert "--index-url $TorchIndexUrl" not in self._ps1
assert "--index-url $ROCmIndexUrl" not in self._ps1
def test_torch_installs_neutralize_all_uv_index_env_vars(self):
# Extra-index vars outrank --default-index, so pinned installs must clear them.
for var in ("UV_DEFAULT_INDEX", "UV_INDEX_URL", "UV_INDEX", "UV_EXTRA_INDEX_URL"):
assert var in self._ps1
assert 'Remove-Item "Env:$n"' in self._ps1
class TestSetupPs1FastInstallIndex:
"""setup.ps1 Fast-Install must neutralize inherited uv indexes when pinning."""
_ps1 = _read(_SETUP_PS1)
def test_fast_install_clears_all_uv_index_env_vars(self):
for var in ("UV_DEFAULT_INDEX", "UV_INDEX_URL", "UV_INDEX", "UV_EXTRA_INDEX_URL"):
assert var in self._ps1
# Must truly remove the vars (child sees no value), not set them empty.
assert 'Remove-Item "Env:$n"' in self._ps1
class TestInstallShUvDefaultIndex:
"""Linux/Mac installer torch indexes must override inherited uv defaults."""
_sh = _read(_INSTALL_SH)
def test_torch_installs_use_default_index(self):
assert '--default-index "$TORCH_INDEX_URL"' in self._sh
def test_torch_installs_do_not_use_deprecated_index_url(self):
assert '--index-url "$TORCH_INDEX_URL"' not in self._sh
def test_torch_installs_neutralize_all_uv_index_env_vars(self):
# --default-index installs run with all uv index env vars unset via `env -u`.
assert (
"env -u UV_DEFAULT_INDEX -u UV_INDEX_URL -u UV_INDEX -u UV_EXTRA_INDEX_URL" in self._sh
)
# Group 2 -- Shell snippet tests (bash subprocess, mocked python)
class TestTorchConstraintShell:
"""Test the TORCH_CONSTRAINT block via bash with mocked python minor versions."""
# Snippet tested in isolation: override OS/_ARCH/SKIP_TORCH and a mock python.
_SNIPPET_TEMPLATE = textwrap.dedent(r"""
#!/bin/bash
set -e
SKIP_TORCH={skip_torch}
OS="{os}"
_ARCH="{arch}"
VENV_DIR="{venv_dir}"
TORCH_CONSTRAINT="torch>=2.4,<2.11.0"
if [ "$SKIP_TORCH" = false ] && [ "$OS" = "macos" ] && [ "$_ARCH" = "arm64" ]; then
_PY_MINOR=$("$VENV_DIR/bin/python" -c \
"import sys; print(sys.version_info.minor)" 2>/dev/null || echo "0")
if [ "$_PY_MINOR" -ge 13 ] 2>/dev/null; then
TORCH_CONSTRAINT="torch>=2.6,<2.11.0"
fi
fi
echo "$TORCH_CONSTRAINT"
""").strip()
@staticmethod
def _make_mock_python(tmp_path: pathlib.Path, minor: int) -> pathlib.Path:
"""Create a mock python that prints a controlled minor version."""
venv = tmp_path / "venv"
bin_dir = venv / "bin"
bin_dir.mkdir(parents = True, exist_ok = True)
mock_py = bin_dir / "python"
mock_py.write_text(
textwrap.dedent(f"""\
#!/bin/bash
# Mock python: always report minor={minor}
if echo "$@" | grep -q "sys.version_info.minor"; then
echo "{minor}"
else
echo "0"
fi
""")
)
mock_py.chmod(0o755)
return venv
def _run(
self,
tmp_path: pathlib.Path,
*,
py_minor: int = 12,
os_val: str = "macos",
arch: str = "arm64",
skip_torch: str = "false",
) -> str:
venv = self._make_mock_python(tmp_path, py_minor)
script = self._SNIPPET_TEMPLATE.format(
skip_torch = skip_torch,
os = os_val,
arch = arch,
venv_dir = str(venv),
)
script_file = tmp_path / "test_snippet.sh"
script_file.write_text(script)
script_file.chmod(0o755)
result = subprocess.run(
["bash", str(script_file)],
capture_output = True,
text = True,
timeout = 10,
)
assert result.returncode == 0, f"Script failed: {result.stderr}"
return result.stdout.strip()
def test_arm64_macos_py313_tightened(self, tmp_path):
out = self._run(tmp_path, py_minor = 13, os_val = "macos", arch = "arm64")
assert out == "torch>=2.6,<2.11.0"
def test_arm64_macos_py314_tightened(self, tmp_path):
out = self._run(tmp_path, py_minor = 14, os_val = "macos", arch = "arm64")
assert out == "torch>=2.6,<2.11.0"
def test_arm64_macos_py312_default(self, tmp_path):
out = self._run(tmp_path, py_minor = 12, os_val = "macos", arch = "arm64")
assert out == "torch>=2.4,<2.11.0"
def test_arm64_macos_py311_default(self, tmp_path):
out = self._run(tmp_path, py_minor = 11, os_val = "macos", arch = "arm64")
assert out == "torch>=2.4,<2.11.0"
# Linux is unaffected by the tightening.
def test_linux_x86_py313_default(self, tmp_path):
out = self._run(tmp_path, py_minor = 13, os_val = "linux", arch = "x86_64")
assert out == "torch>=2.4,<2.11.0"
def test_linux_aarch64_py313_default(self, tmp_path):
out = self._run(tmp_path, py_minor = 13, os_val = "linux", arch = "aarch64")
assert out == "torch>=2.4,<2.11.0"
# Intel Mac: arch mismatch, no tightening.
def test_intel_mac_x86_py313_default(self, tmp_path):
out = self._run(tmp_path, py_minor = 13, os_val = "macos", arch = "x86_64")
assert out == "torch>=2.4,<2.11.0"
# SKIP_TORCH bypasses the tightening.
def test_skip_torch_arm64_macos_py313_default(self, tmp_path):
out = self._run(
tmp_path,
py_minor = 13,
os_val = "macos",
arch = "arm64",
skip_torch = "true",
)
assert out == "torch>=2.4,<2.11.0"
def test_wsl_py313_default(self, tmp_path):
out = self._run(tmp_path, py_minor = 13, os_val = "wsl", arch = "x86_64")
assert out == "torch>=2.4,<2.11.0"
def test_py_minor_0_fallback_default(self, tmp_path):
"""Failed python query (returns 0) keeps the default constraint."""
out = self._run(tmp_path, py_minor = 0, os_val = "macos", arch = "arm64")
assert out == "torch>=2.4,<2.11.0"
def test_boundary_py_minor_12_not_tightened(self, tmp_path):
out = self._run(tmp_path, py_minor = 12, os_val = "macos", arch = "arm64")
assert out == "torch>=2.4,<2.11.0"
def test_boundary_py_minor_13_tightened(self, tmp_path):
out = self._run(tmp_path, py_minor = 13, os_val = "macos", arch = "arm64")
assert out == "torch>=2.6,<2.11.0"
def test_mock_uv_receives_correct_constraint(self, tmp_path):
"""A mock uv receives the tightened constraint on py3.13 arm64 macOS."""
venv = self._make_mock_python(tmp_path, minor = 13)
# Mock uv logs its arguments.
mock_uv = tmp_path / "mock_uv"
log_file = tmp_path / "uv_log.txt"
mock_uv.write_text(
textwrap.dedent(f"""\
#!/bin/bash
echo "$@" >> {log_file}
""")
)
mock_uv.chmod(0o755)
script = textwrap.dedent(f"""\
#!/bin/bash
set -e
SKIP_TORCH=false
OS="macos"
_ARCH="arm64"
VENV_DIR="{venv}"
TORCH_CONSTRAINT="torch>=2.4,<2.11.0"
if [ "$SKIP_TORCH" = false ] && [ "$OS" = "macos" ] && [ "$_ARCH" = "arm64" ]; then
_PY_MINOR=$("$VENV_DIR/bin/python" -c \\
"import sys; print(sys.version_info.minor)" 2>/dev/null || echo "0")
if [ "$_PY_MINOR" -ge 13 ] 2>/dev/null; then
TORCH_CONSTRAINT="torch>=2.6,<2.11.0"
fi
fi
# Simulate the uv pip install line
{mock_uv} pip install --python "$VENV_DIR/bin/python" "$TORCH_CONSTRAINT" torchvision torchaudio
""")
script_file = tmp_path / "test_uv.sh"
script_file.write_text(script)
script_file.chmod(0o755)
result = subprocess.run(
["bash", str(script_file)],
capture_output = True,
text = True,
timeout = 10,
)
assert result.returncode == 0, f"Script failed: {result.stderr}"
logged = log_file.read_text()
assert "torch>=2.6,<2.11.0" in logged, f"uv log: {logged}"
def test_mock_uv_receives_default_constraint(self, tmp_path):
"""On py3.12 arm64 macOS, uv should receive the default constraint."""
venv = self._make_mock_python(tmp_path, minor = 12)
mock_uv = tmp_path / "mock_uv"
log_file = tmp_path / "uv_log.txt"
mock_uv.write_text(
textwrap.dedent(f"""\
#!/bin/bash
echo "$@" >> {log_file}
""")
)
mock_uv.chmod(0o755)
script = textwrap.dedent(f"""\
#!/bin/bash
set -e
SKIP_TORCH=false
OS="macos"
_ARCH="arm64"
VENV_DIR="{venv}"
TORCH_CONSTRAINT="torch>=2.4,<2.11.0"
if [ "$SKIP_TORCH" = false ] && [ "$OS" = "macos" ] && [ "$_ARCH" = "arm64" ]; then
_PY_MINOR=$("$VENV_DIR/bin/python" -c \\
"import sys; print(sys.version_info.minor)" 2>/dev/null || echo "0")
if [ "$_PY_MINOR" -ge 13 ] 2>/dev/null; then
TORCH_CONSTRAINT="torch>=2.6,<2.11.0"
fi
fi
{mock_uv} pip install --python "$VENV_DIR/bin/python" "$TORCH_CONSTRAINT" torchvision torchaudio
""")
script_file = tmp_path / "test_uv.sh"
script_file.write_text(script)
script_file.chmod(0o755)
result = subprocess.run(
["bash", str(script_file)],
capture_output = True,
text = True,
timeout = 10,
)
assert result.returncode == 0, f"Script failed: {result.stderr}"
logged = log_file.read_text()
assert "torch>=2.4,<2.11.0" in logged, f"uv log: {logged}"
# Mirrors the _torch_index_leaf case in install.sh: rocm7.2 -> 2.11.x floor,
# CUDA -> widened <2.12.0 ceiling, else (CPU/older ROCm) -> default. Anchored
# on the final path segment, so a mirror base path containing cu*/rocm7.2 but
# ending in a cpu/older-rocm leaf keeps the default.
_INDEX_SNIPPET = textwrap.dedent(r"""
#!/bin/bash
set -e
TORCH_INDEX_URL="{index_url}"
TORCH_CONSTRAINT="torch>=2.4,<2.11.0"
_torch_index_leaf="${TORCH_INDEX_URL%/}"
_torch_index_leaf="${_torch_index_leaf##*/}"
case "$_torch_index_leaf" in
rocm7.2) TORCH_CONSTRAINT="torch>=2.11.0,<2.12.0" ;;
cu[0-9]*) TORCH_CONSTRAINT="torch>=2.4,<2.12.0" ;;
esac
echo "$TORCH_CONSTRAINT"
""").strip()
def _resolve_index(self, tmp_path: pathlib.Path, index_url: str) -> str:
script_file = tmp_path / "index_snippet.sh"
script_file.write_text(self._INDEX_SNIPPET.replace("{index_url}", index_url))
script_file.chmod(0o755)
result = subprocess.run(
["bash", str(script_file)],
capture_output = True,
text = True,
timeout = 10,
)
assert result.returncode == 0, f"Script failed: {result.stderr}"
return result.stdout.strip()
@pytest.mark.parametrize("leaf", ["cu118", "cu124", "cu126", "cu128", "cu130"])
def test_cuda_index_widens_to_2_12(self, tmp_path, leaf):
url = f"https://download.pytorch.org/whl/{leaf}"
assert self._resolve_index(tmp_path, url) == "torch>=2.4,<2.12.0"
def test_rocm72_index_uses_211_floor(self, tmp_path):
url = "https://download.pytorch.org/whl/rocm7.2"
assert self._resolve_index(tmp_path, url) == "torch>=2.11.0,<2.12.0"
def test_cpu_index_keeps_default(self, tmp_path):
# /cpu must NOT match the */cu[0-9]* branch.
url = "https://download.pytorch.org/whl/cpu"
assert self._resolve_index(tmp_path, url) == "torch>=2.4,<2.11.0"
def test_older_rocm_index_keeps_default(self, tmp_path):
url = "https://download.pytorch.org/whl/rocm7.1"
assert self._resolve_index(tmp_path, url) == "torch>=2.4,<2.11.0"
def test_cuda_index_custom_mirror_widens(self, tmp_path):
url = "https://internal.example.com/pytorch/cu128"
assert self._resolve_index(tmp_path, url) == "torch>=2.4,<2.12.0"
@pytest.mark.parametrize(
"url",
[
"https://internal.example.com/pytorch/cu128/cpu",
"https://internal.example.com/cu128/whl/rocm7.1",
],
)
def test_cuda_in_mirror_path_but_noncuda_leaf_keeps_default(self, tmp_path, url):
# A cu128 in the mirror base path must not widen when the leaf is cpu /
# older ROCm: the case anchors on _torch_index_leaf, not the whole URL.
assert self._resolve_index(tmp_path, url) == "torch>=2.4,<2.11.0"
# Group 3 -- E2E tokenizers fix (requires network, ~2-5 min)
@pytest.mark.e2e
class TestE2ETokenizersFix:
"""Real uv venvs verify tokenizers + transformers work without torch installed."""
@staticmethod
def _create_venv(tmp_path: pathlib.Path, name: str, py: str) -> pathlib.Path:
venv = tmp_path / name
result = subprocess.run(
["uv", "venv", str(venv), "--python", py],
capture_output = True,
text = True,
timeout = 120,
)
if result.returncode != 0:
pytest.skip(f"uv venv creation failed for {py}: {result.stderr}")
return venv
@staticmethod
def _pip_install(venv: pathlib.Path, *args: str) -> subprocess.CompletedProcess:
py = str(venv / "bin" / "python")
cmd = ["uv", "pip", "install", "--python", py, *args]
return subprocess.run(cmd, capture_output = True, text = True, timeout = 300)
@staticmethod
def _run_python(venv: pathlib.Path, code: str) -> subprocess.CompletedProcess:
py = str(venv / "bin" / "python")
return subprocess.run(
[py, "-c", code],
capture_output = True,
text = True,
timeout = 60,
)
@pytest.mark.parametrize("py_version", ["3.12", "3.13"])
def test_autoconfig_works_with_no_torch_runtime(self, tmp_path, py_version):
"""Install no-torch-runtime.txt with --no-deps, then AutoConfig must import."""
venv = self._create_venv(tmp_path, f"tok-{py_version}", py_version)
r = self._pip_install(venv, "--no-deps", "-r", str(_NO_TORCH_RT))
assert r.returncode == 0, f"Install failed: {r.stderr}"
result = self._run_python(venv, "from transformers import AutoConfig; print('OK')")
assert (
result.returncode == 0
), f"AutoConfig import failed:\nstdout: {result.stdout}\nstderr: {result.stderr}"
assert "OK" in result.stdout
@pytest.mark.parametrize("py_version", ["3.12", "3.13"])
def test_tokenizers_directly_importable(self, tmp_path, py_version):
venv = self._create_venv(tmp_path, f"tok-imp-{py_version}", py_version)
r = self._pip_install(venv, "--no-deps", "-r", str(_NO_TORCH_RT))
assert r.returncode == 0, f"Install failed: {r.stderr}"
result = self._run_python(venv, "import tokenizers; print('OK')")
assert result.returncode == 0, f"Failed: {result.stderr}"
@pytest.mark.parametrize("py_version", ["3.12", "3.13"])
def test_torch_not_importable(self, tmp_path, py_version):
"""In the no-torch scenario, torch should not be available."""
venv = self._create_venv(tmp_path, f"no-torch-{py_version}", py_version)
r = self._pip_install(venv, "--no-deps", "-r", str(_NO_TORCH_RT))
assert r.returncode == 0, f"Install failed: {r.stderr}"
result = self._run_python(venv, "import torch")
assert result.returncode != 0, "torch should NOT be importable"
def test_negative_control_no_tokenizers(self, tmp_path):
"""Without the tokenizers line, AutoConfig must fail (negative control)."""
venv = self._create_venv(tmp_path, "neg-ctrl", "3.12")
req_no_tokenizers = tmp_path / "no-tokenizers.txt"
req_no_tokenizers.write_text(
"\n".join(
line for line in _read(_NO_TORCH_RT).splitlines() if line.strip() != "tokenizers"
),
encoding = "utf-8",
)
r = self._pip_install(venv, "--no-deps", "-r", str(req_no_tokenizers))
assert r.returncode == 0, f"Install failed: {r.stderr}"
result = self._run_python(venv, "from transformers import AutoConfig")
assert result.returncode != 0, "AutoConfig should fail without tokenizers installed"
assert "tokenizers" in result.stderr.lower() or "ModuleNotFoundError" in result.stderr
# Group 4 -- Integration: install.sh reads no-torch-runtime.txt correctly
class TestInstallShNoTorchIntegration:
"""Verify install.sh has the correct no-torch-runtime.txt wiring."""
_sh = _read(_INSTALL_SH)
def test_find_no_torch_runtime_exists(self):
assert "_find_no_torch_runtime()" in self._sh
def test_no_deps_invocation_for_migrated(self):
"""Migrated path should use --no-deps -r."""
assert '--no-deps -r "$_NO_TORCH_RT"' in self._sh
def test_no_deps_invocation_for_fresh(self):
"""Fresh install path should also use --no-deps -r."""
count = self._sh.count('--no-deps -r "$_NO_TORCH_RT"')
assert count >= 2, f"Expected >=2 no-deps -r invocations, found {count}"
def test_mock_uv_skip_torch_reads_requirements(self, tmp_path):
"""SKIP_TORCH=true blocks must call _find_no_torch_runtime."""
skip_blocks = re.findall(
r'if \[ "\$SKIP_TORCH" = true \].*?(?=\n (?:else|elif|fi))',
self._sh,
re.DOTALL,
)
found = any("_find_no_torch_runtime" in block for block in skip_blocks)
assert found, "SKIP_TORCH=true block should call _find_no_torch_runtime"
# Group 5 -- Full no-torch sandbox (requires network, ~5 min)
@pytest.mark.e2e
class TestE2EFullNoTorchSandbox:
"""Creates venvs and installs the actual no-torch-runtime.txt."""
@staticmethod
def _create_venv(tmp_path: pathlib.Path, name: str) -> pathlib.Path:
venv = tmp_path / name
result = subprocess.run(
["uv", "venv", str(venv), "--python", "3.12"],
capture_output = True,
text = True,
timeout = 120,
)
if result.returncode != 0:
pytest.skip(f"uv venv creation failed: {result.stderr}")
return venv
@staticmethod
def _pip_install(venv: pathlib.Path, *args: str) -> subprocess.CompletedProcess:
py = str(venv / "bin" / "python")
cmd = ["uv", "pip", "install", "--python", py, *args]
return subprocess.run(cmd, capture_output = True, text = True, timeout = 600)
@staticmethod
def _run_python(venv: pathlib.Path, code: str) -> subprocess.CompletedProcess:
py = str(venv / "bin" / "python")
return subprocess.run(
[py, "-c", code],
capture_output = True,
text = True,
timeout = 60,
)
def test_autoconfig_succeeds(self, tmp_path):
"""Install with --no-deps and verify AutoConfig imports (the bug fix)."""
venv = self._create_venv(tmp_path, "full-no-torch")
r = self._pip_install(venv, "--no-deps", "-r", str(_NO_TORCH_RT))
assert r.returncode == 0, f"Install failed: {r.stderr}"
result = self._run_python(venv, "from transformers import AutoConfig; print('OK')")
assert (
result.returncode == 0
), f"AutoConfig failed:\nstdout: {result.stdout}\nstderr: {result.stderr}"
def test_torch_not_importable(self, tmp_path):
"""With --no-deps, torch must not be pulled in."""
venv = self._create_venv(tmp_path, "no-torch-check")
r = self._pip_install(venv, "--no-deps", "-r", str(_NO_TORCH_RT))
assert r.returncode == 0, f"Install failed: {r.stderr}"
result = self._run_python(venv, "import torch")
assert result.returncode != 0, "torch should NOT be importable"
def test_tokenizers_importable(self, tmp_path):
venv = self._create_venv(tmp_path, "tok-check")
r = self._pip_install(venv, "--no-deps", "-r", str(_NO_TORCH_RT))
assert r.returncode == 0, f"Install failed: {r.stderr}"
result = self._run_python(venv, "import tokenizers; print('OK')")
assert result.returncode == 0, f"tokenizers import failed: {result.stderr}"
def test_safetensors_importable(self, tmp_path):
venv = self._create_venv(tmp_path, "st-check")
r = self._pip_install(venv, "--no-deps", "-r", str(_NO_TORCH_RT))
assert r.returncode == 0, f"Install failed: {r.stderr}"
result = self._run_python(venv, "import safetensors; print('OK')")
assert result.returncode == 0, f"safetensors import failed: {result.stderr}"
def test_huggingface_hub_importable(self, tmp_path):
venv = self._create_venv(tmp_path, "hfhub-check")
r = self._pip_install(venv, "--no-deps", "-r", str(_NO_TORCH_RT))
assert r.returncode == 0, f"Install failed: {r.stderr}"
result = self._run_python(venv, "import huggingface_hub; print('OK')")
assert result.returncode == 0, f"huggingface_hub import failed: {result.stderr}"