diff --git a/install.sh b/install.sh index 5972379d26..6076721540 100755 --- a/install.sh +++ b/install.sh @@ -472,11 +472,13 @@ _on_install_exit() { _restore_studio_venv_replacement fi [ -n "${_UV_OVERRIDE_TMPDIR:-}" ] && rm -rf "$_UV_OVERRIDE_TMPDIR" 2>/dev/null || true + [ -n "${_UNSLOTH_TORCH_OVERRIDES:-}" ] && rm -f "$_UNSLOTH_TORCH_OVERRIDES" 2>/dev/null || true exit "$_status" } -# Empty so an inherited value can never reach the trap's rm; only a temp dir -# this script creates below (Apple Silicon, spaced path) is ever removed. +# Empty so an inherited value never reaches the trap's rm; only temp paths this +# script creates below (spaced-path dir, torch-trio overrides) are removed. _UV_OVERRIDE_TMPDIR="" +_UNSLOTH_TORCH_OVERRIDES="" trap _on_install_exit EXIT # ── Helper: download a URL to a file (supports curl and wget) ── @@ -1821,6 +1823,8 @@ tauri_log "STEP" "Creating virtual environment" mkdir -p "$STUDIO_HOME" _MIGRATED=false +# Empty so an inherited value can never masquerade as a probed torch version. +_PREV_TORCH_VER="" if [ -x "$VENV_DIR/bin/python" ]; then # why: matching guard to the .venv branch below -- in env-mode @@ -1838,6 +1842,12 @@ if [ -x "$VENV_DIR/bin/python" ]; then echo " Move it aside or choose an empty UNSLOTH_STUDIO_HOME." >&2 exit 1 fi + # Record the existing venv's torch BEFORE the replacement moves it aside: a re-run + # rebuilds the venv for clean state, but must keep the torch release the user + # already has (see _previous_torch_pin below). Last line only: sitecustomize or + # import-hook noise on stdout must not corrupt the version. + _PREV_TORCH_VER=$("$VENV_DIR/bin/python" -c \ + "import torch; print(torch.__version__)" 2>/dev/null | tail -n 1 || true) # New layout already exists — replace only after preserving rollback copy. substep "preserving existing environment for rollback..." _start_studio_venv_replacement "$VENV_DIR" @@ -2187,6 +2197,68 @@ _torch_flavor_tag() { esac } +# Whether release base $1 (X.Y[.Z...]) falls inside constraint window $2 +# ("torch>=A.B[.C],="*",<"*) ;; + *) echo "no"; return ;; + esac + _trw_floor="${_trw_con#torch>=}"; _trw_floor="${_trw_floor%%,*}" + _trw_ceil="${_trw_con##*,<}" + _v_maj="${1%%.*}"; _v_rest="${1#*.}"; _v_min="${_v_rest%%.*}" + _f_maj="${_trw_floor%%.*}"; _f_rest="${_trw_floor#*.}"; _f_min="${_f_rest%%.*}" + _c_maj="${_trw_ceil%%.*}"; _c_rest="${_trw_ceil#*.}"; _c_min="${_c_rest%%.*}" + for _trw_n in "$_v_maj" "$_v_min" "$_f_maj" "$_f_min" "$_c_maj" "$_c_min"; do + case "$_trw_n" in ''|*[!0-9]*) echo "no"; return ;; esac + done + if [ "$_v_maj" -gt "$_f_maj" ] || { [ "$_v_maj" -eq "$_f_maj" ] && [ "$_v_min" -ge "$_f_min" ]; }; then + if [ "$_v_maj" -lt "$_c_maj" ] || { [ "$_v_maj" -eq "$_c_maj" ] && [ "$_v_min" -lt "$_c_min" ]; }; then + echo "yes" + return + fi + fi + echo "no" +} + +# Whether a re-run should keep the previous venv's torch: echo "torch==X.Y.Z" when the +# probed previous version ($1) has a flavor tag matching the freshly chosen cu*/cpu index +# leaf ($2) AND sits inside the active constraint window ($3), else "". Re-running +# `curl | sh` rebuilds the venv for clean state, but a healthy torch the user already +# validated must not be silently moved to a newer release (2.10 -> 2.11); a flavor +# change (cpu <-> cuda, cu126 -> cu130) still installs the correct new build, rocm +# leaves keep their floors (rocm7.2 must land 2.11 for the Strix _grouped_mm fix), and +# a release outside the window (2.3.x manual install, 2.12.x manual upgrade) is never +# kept: the installer's own bounds win. Opt out with UNSLOTH_TORCH_UPGRADE=1 to get +# the newest release. +_previous_torch_pin() { + _ptp_ver="$1" + _ptp_leaf="$2" + _ptp_con="$3" + [ -n "$_ptp_ver" ] || { echo ""; return; } + [ "${UNSLOTH_TORCH_UPGRADE:-0}" = "1" ] && { echo ""; return; } + case "$_ptp_leaf" in + cu[0-9]*|cpu) ;; + *) echo ""; return ;; + esac + _ptp_base="${_ptp_ver%%+*}" + # The base must look like a release (probe noise / garbage must never become a pin). + case "$_ptp_base" in + [0-9]*.[0-9]*) ;; + *) echo ""; return ;; + esac + [ "$(_torch_release_in_window "$_ptp_base" "$_ptp_con")" = "yes" ] || { echo ""; return; } + if [ "$(_torch_flavor_tag "$_ptp_ver")" = "$_ptp_leaf" ]; then + echo "torch==$_ptp_base" + else + echo "" + fi +} + # Expected tag from the index leaf ($1): cuXXX / cpu / rocm (rocmX.Y and gfx* -> # rocm). Empty on an unknown leaf (odd mirror) so the repair safely no-ops. _expected_torch_flavor_tag() { @@ -2478,12 +2550,32 @@ case "$_torch_index_leaf" in *) export UNSLOTH_TORCH_BACKEND="cuda" ;; esac -# rocm7.2 ships torch 2.11.0 -- adjust the constraint to allow it. -# All other ROCm tags and CUDA stay within <2.11.0. -case "$TORCH_INDEX_URL" in - */rocm7.2) TORCH_CONSTRAINT="torch>=2.11.0,<2.12.0" ;; +# rocm7.2 and the CUDA cu12x/cu13x indexes now ship torch 2.11.x, so widen the +# ceiling to <2.12.0 (matches the base image and _CUDA_TORCH_PKG_SPEC in +# studio/install_python_stack.py). Keep the >=2.4 floor so an older CUDA index +# (e.g. cu118) still resolves. Match on _torch_index_leaf, not the full URL, so +# a mirror whose base path contains cu*/rocm7.2 but resolves to a cpu/older-rocm +# leaf keeps the default <2.11.0. +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 +# Re-run over an existing install: keep the previous venv's torch release instead of +# resolving the newest in range. The range stays in _PREV_FALLBACK_CONSTRAINT so the +# install can fall back when the exact release is not on the chosen index (custom +# mirrors may prune old wheels). Skipped for --no-torch (no previous probe runs). +_PREV_TORCH_PIN="" +_PREV_FALLBACK_CONSTRAINT="$TORCH_CONSTRAINT" +if [ "$SKIP_TORCH" = false ]; then + _prev_pin=$(_previous_torch_pin "$_PREV_TORCH_VER" "$_torch_index_leaf" "$TORCH_CONSTRAINT") + if [ -n "$_prev_pin" ]; then + _PREV_TORCH_PIN="$_prev_pin" + TORCH_CONSTRAINT="$_prev_pin" + substep "existing install has torch $_PREV_TORCH_VER -- keeping it (set UNSLOTH_TORCH_UPGRADE=1 to get the newest release)" + fi +fi + # Auto-detect GPU for AMD ROCm based # get_torch_index_url must have chosen */rocm* # (gfx in rocminfo or amd-smi list). Then require rocminfo "Marketing Name:.*Radeon". @@ -2705,6 +2797,43 @@ esac # ── Install unsloth directly into the venv (no activation needed) ── tauri_log "STEP" "Installing PyTorch" _VENV_PY="$VENV_DIR/bin/python" + +# A released unsloth wheel can pin an older torch (unsloth 2026.7.2 declares +# torch<2.11.0); a with-deps PyPI resolve then downgrades the whole trio, +# swapping the pinned +cuXXX/+rocm build for PyPI's default. The flavor guard +# below misses this (PyPI's torch 2.10 default is itself cu128-flavored), so +# freeze the trio via uv --overrides (overrides replace dependency requirements +# during resolution) while unsloth's other deps resolve normally. Sets +# _UNSLOTH_TORCH_OVERRIDES from the trio in the venv; every with-deps unsloth +# install (migrated and fresh) must call this before resolving and rm it after. +_build_unsloth_torch_overrides() { + _UNSLOTH_TORCH_OVERRIDES="" + [ "$SKIP_TORCH" = false ] || return 0 + _torch_trio_pins=$("$_VENV_PY" -c " +from importlib.metadata import version, PackageNotFoundError +for _p in ('torch', 'torchvision', 'torchaudio'): + try: + print(_p + '==' + version(_p)) + except PackageNotFoundError: + pass +" 2>/dev/null) || _torch_trio_pins="" + case "$_torch_trio_pins" in + torch==*) + _UNSLOTH_TORCH_OVERRIDES=$(mktemp) + printf '%s\n' "$_torch_trio_pins" > "$_UNSLOTH_TORCH_OVERRIDES" + # The CLI --overrides flag replaces any UV_OVERRIDE env file (same + # uv setting; macOS arm64 exports one here), so fold its pins in. + # awk, not cat: it drops inherited torch-trio lines (uv intersects + # duplicate overrides, so a conflicting pin would make resolution + # unsatisfiable) and newline-terminates the last line so an + # unterminated file cannot join two requirements into one. + for _ov_file in ${UV_OVERRIDE:-}; do + [ -f "$_ov_file" ] && awk '!/^[[:space:]]*torch(vision|audio)?([[:space:]<>=!~;@[]|$)/' "$_ov_file" >> "$_UNSLOTH_TORCH_OVERRIDES" + done + ;; + esac +} + if [ "$_MIGRATED" = true ]; then # Migrated env: force-reinstall unsloth+unsloth-zoo to ensure clean state # in the new venv location, while preserving existing torch/CUDA @@ -2729,9 +2858,13 @@ if [ "$_MIGRATED" = true ]; then else # Pin mlx-lm away from 0.31.3 here too: a curl-piped migration has no # overrides file, so UV_OVERRIDE is unset and this positional is the only cover. + _build_unsloth_torch_overrides run_install_cmd_retry "install unsloth (migrated)" uv pip install --python "$_VENV_PY" \ + ${_UNSLOTH_TORCH_OVERRIDES:+--overrides "$_UNSLOTH_TORCH_OVERRIDES"} \ --reinstall-package unsloth --reinstall-package unsloth-zoo \ "unsloth>=2026.7.3" "unsloth-zoo>=2026.7.3" ${_MLX_LM_EXCLUDE_ARG:-} + [ -n "$_UNSLOTH_TORCH_OVERRIDES" ] && rm -f "$_UNSLOTH_TORCH_OVERRIDES" + _UNSLOTH_TORCH_OVERRIDES="" fi if [ "$STUDIO_LOCAL_INSTALL" = true ]; then substep "overlaying local repo (editable)..." @@ -2913,8 +3046,20 @@ elif [ -n "$TORCH_INDEX_URL" ]; then fi else substep "installing PyTorch ($TORCH_INDEX_URL)..." - run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" "$TORCH_CONSTRAINT" torchvision torchaudio \ - --default-index "$TORCH_INDEX_URL" + if [ -n "$_PREV_TORCH_PIN" ]; then + # Kept previous release: fall back to the supported range if the exact + # release is not resolvable from the chosen index (pruned mirror). + if ! run_install_cmd_retry "install PyTorch (kept release)" uv pip install --python "$_VENV_PY" "$TORCH_CONSTRAINT" torchvision torchaudio \ + --default-index "$TORCH_INDEX_URL"; then + substep "[WARN] $_PREV_TORCH_PIN is not installable from $TORCH_INDEX_URL -- installing the newest supported release instead" "$C_WARN" + TORCH_CONSTRAINT="$_PREV_FALLBACK_CONSTRAINT" + run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" "$TORCH_CONSTRAINT" torchvision torchaudio \ + --default-index "$TORCH_INDEX_URL" + fi + else + run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" "$TORCH_CONSTRAINT" torchvision torchaudio \ + --default-index "$TORCH_INDEX_URL" + fi fi # AMD ROCm: install bitsandbytes (once, after torch, for all ROCm paths). # Gate on SKIP_TORCH=false so a user running with --no-torch on a ROCm @@ -2927,9 +3072,10 @@ elif [ -n "$TORCH_INDEX_URL" ]; then ;; esac fi - # Fresh: Step 2 - install unsloth, preserving pre-installed torch + # Fresh: Step 2 - install unsloth, preserving the torch Step 1 installed tauri_log "STEP" "Installing Unsloth" substep "installing unsloth (this may take a few minutes)..." + _build_unsloth_torch_overrides if [ "$SKIP_TORCH" = true ]; then # No-torch: install unsloth + unsloth-zoo with --no-deps, then # runtime deps (typer, safetensors, transformers, etc.) with --no-deps. @@ -2953,6 +3099,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then fi elif [ "$STUDIO_LOCAL_INSTALL" = true ]; then run_install_cmd_retry "install unsloth (local)" uv pip install --python "$_VENV_PY" \ + ${_UNSLOTH_TORCH_OVERRIDES:+--overrides "$_UNSLOTH_TORCH_OVERRIDES"} \ --upgrade-package unsloth "unsloth>=2026.7.3" "unsloth-zoo>=2026.7.3" substep "overlaying local repo (editable)..." run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps @@ -2962,8 +3109,11 @@ elif [ -n "$TORCH_INDEX_URL" ]; then "unsloth-zoo @ git+https://github.com/unslothai/unsloth-zoo" else run_install_cmd_retry "install unsloth" uv pip install --python "$_VENV_PY" \ + ${_UNSLOTH_TORCH_OVERRIDES:+--overrides "$_UNSLOTH_TORCH_OVERRIDES"} \ --upgrade-package unsloth -- "$PACKAGE_NAME" ${_MLX_LM_EXCLUDE_ARG:-} fi + [ -n "$_UNSLOTH_TORCH_OVERRIDES" ] && rm -f "$_UNSLOTH_TORCH_OVERRIDES" + _UNSLOTH_TORCH_OVERRIDES="" # AMD ROCm: repair torch if the unsloth/unsloth-zoo install pulled in # CUDA torch from PyPI, overwriting the ROCm wheels installed in Step 1. if [ "$SKIP_TORCH" = false ]; then diff --git a/studio/backend/tests/test_cuda_torch_spec.py b/studio/backend/tests/test_cuda_torch_spec.py new file mode 100644 index 0000000000..928cef787e --- /dev/null +++ b/studio/backend/tests/test_cuda_torch_spec.py @@ -0,0 +1,73 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Tests for _CUDA_TORCH_PKG_SPEC in install_python_stack.py. + +The CUDA repair path installs the torch trio from an exclusive --index-url (no +PyPI fallback), so these pinned ranges decide which torch the venv gets. The +upper bound is locked to the 2.11.x family to match the base image and rocm7.2 +spec and to keep the companions off a torch-2.12 wheel that would ABI-mismatch. +""" + +from __future__ import annotations + +import sys +from pathlib import Path + +import pytest +from packaging.requirements import Requirement + +# install_python_stack.py lives at repo_root/studio/install_python_stack.py +_INSTALL_SCRIPT = Path(__file__).resolve().parents[2] / "install_python_stack.py" + + +def _load_module(monkeypatch): + """(Re-)import and return install_python_stack (mirrors test_torchao_select).""" + sys.modules.pop("install_python_stack", None) + monkeypatch.syspath_prepend(str(_INSTALL_SCRIPT.parent)) + import install_python_stack + + return install_python_stack + + +def _spec_of(pkg_spec: str): + """Parse 'torch>=2.4,<2.12.0' into a packaging SpecifierSet.""" + return Requirement(pkg_spec).specifier + + +@pytest.mark.parametrize( + "index, allowed, rejected", + [ + # torch: 2.11.x allowed (matches base image); 2.12.x excluded. + (0, ["2.11.0", "2.11.2", "2.10.0", "2.4.0"], ["2.12.0", "2.3.0", "1.13.1"]), + # torchvision: 0.26.x (torch 2.11 companion) allowed; 0.27.x (torch 2.12) out. + (1, ["0.26.0", "0.26.1", "0.19.0"], ["0.27.0", "0.18.0"]), + # torchaudio: same 2.11.x window as torch. + (2, ["2.11.0", "2.10.0", "2.4.0"], ["2.12.0", "2.3.0"]), + ], +) +def test_cuda_spec_bounds(monkeypatch, index, allowed, rejected): + mod = _load_module(monkeypatch) + spec = _spec_of(mod._CUDA_TORCH_PKG_SPEC[index]) + for v in allowed: + assert spec.contains(v, prereleases = True), f"{v} should satisfy {spec}" + for v in rejected: + assert not spec.contains(v, prereleases = True), f"{v} should not satisfy {spec}" + + +def test_cuda_spec_matches_rocm72_upper_bound(monkeypatch): + """CUDA and rocm7.2 target the same torch 2.11.x family, so their upper + bounds must stay in lockstep (bump both together at 2.12.x).""" + mod = _load_module(monkeypatch) + rocm72 = mod._ROCM_TORCH_PKG_SPECS["rocm7.2"] + + def _upper(pkg_spec: str) -> str: + for clause in _spec_of(pkg_spec): + if clause.operator == "<": + return clause.version + raise AssertionError(f"no upper bound in {pkg_spec!r}") + + for cuda_pkg, rocm_pkg in zip(mod._CUDA_TORCH_PKG_SPEC, rocm72, strict = True): + assert _upper(cuda_pkg) == _upper( + rocm_pkg + ), f"CUDA {cuda_pkg!r} upper bound must match rocm7.2 {rocm_pkg!r}" diff --git a/tests/python/test_tokenizers_and_torch_constraint.py b/tests/python/test_tokenizers_and_torch_constraint.py index 4322f0c7d6..c58808689b 100644 --- a/tests/python/test_tokenizers_and_torch_constraint.py +++ b/tests/python/test_tokenizers_and_torch_constraint.py @@ -69,6 +69,21 @@ class TestStructuralTorchConstraint: 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 @@ -384,6 +399,71 @@ class TestTorchConstraintShell: 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 diff --git a/tests/sh/test_previous_torch_pin.sh b/tests/sh/test_previous_torch_pin.sh new file mode 100644 index 0000000000..253ede8a27 --- /dev/null +++ b/tests/sh/test_previous_torch_pin.sh @@ -0,0 +1,101 @@ +#!/bin/bash +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 +# Unit tests for install.sh's _previous_torch_pin, which keeps the previous +# venv's torch release on a re-run (curl | sh over an existing install) instead +# of silently moving the user to a newer release. Helpers are extracted from +# install.sh and sourced. +set -e + +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" +INSTALL_SH="$SCRIPT_DIR/../../install.sh" +PASS=0 +FAIL=0 + +# Extract _previous_torch_pin and its dependencies _torch_flavor_tag and +# _torch_release_in_window. +_FUNC_FILE=$(mktemp) +{ + sed -n '/^_torch_flavor_tag()/,/^}/p' "$INSTALL_SH" + echo "" + sed -n '/^_torch_release_in_window()/,/^}/p' "$INSTALL_SH" + echo "" + sed -n '/^_previous_torch_pin()/,/^}/p' "$INSTALL_SH" +} > "$_FUNC_FILE" +# shellcheck disable=SC1090 +. "$_FUNC_FILE" +rm -f "$_FUNC_FILE" + +assert_eq() { + _label="$1"; _expected="$2"; _actual="$3" + if [ "$_actual" = "$_expected" ]; then + echo " PASS: $_label"; PASS=$((PASS + 1)) + else + echo " FAIL: $_label (expected '$_expected', got '$_actual')"; FAIL=$((FAIL + 1)) + fi +} + +unset UNSLOTH_TORCH_UPGRADE + +echo "=== _previous_torch_pin: matching flavor keeps the release ===" +assert_eq "cu126 wheel on cu126 leaf" "torch==2.10.0" "$(_previous_torch_pin '2.10.0+cu126' 'cu126' 'torch>=2.4,<2.12.0')" +assert_eq "cu130 wheel on cu130 leaf" "torch==2.10.0" "$(_previous_torch_pin '2.10.0+cu130' 'cu130' 'torch>=2.4,<2.12.0')" +assert_eq "cpu wheel on cpu leaf" "torch==2.10.0" "$(_previous_torch_pin '2.10.0+cpu' 'cpu' 'torch>=2.4,<2.12.0')" +assert_eq "untagged wheel on cpu leaf" "torch==2.10.0" "$(_previous_torch_pin '2.10.0' 'cpu' 'torch>=2.4,<2.12.0')" +assert_eq "local suffix stripped" "torch==2.9.1" "$(_previous_torch_pin '2.9.1+cu128' 'cu128' 'torch>=2.4,<2.12.0')" + +echo "=== _previous_torch_pin: flavor change installs the new build ===" +assert_eq "cu126 wheel on cu130 leaf" "" "$(_previous_torch_pin '2.10.0+cu126' 'cu130' 'torch>=2.4,<2.12.0')" +assert_eq "cpu wheel on cu126 leaf" "" "$(_previous_torch_pin '2.10.0+cpu' 'cu126' 'torch>=2.4,<2.12.0')" +assert_eq "cu126 wheel on cpu leaf" "" "$(_previous_torch_pin '2.10.0+cu126' 'cpu' 'torch>=2.4,<2.12.0')" + +echo "=== _previous_torch_pin: rocm and unknown leaves never pin ===" +assert_eq "rocm7.2 leaf keeps its floor" "" "$(_previous_torch_pin '2.11.0+rocm7.2' 'rocm7.2' 'torch>=2.4,<2.12.0')" +assert_eq "gfx leaf keeps its floor" "" "$(_previous_torch_pin '2.11.0+rocm7.2' 'gfx120X-all' 'torch>=2.4,<2.12.0')" +assert_eq "unknown mirror leaf" "" "$(_previous_torch_pin '2.10.0+cu126' 'simple' 'torch>=2.4,<2.12.0')" + +echo "=== _previous_torch_pin: probe noise never becomes a pin ===" +assert_eq "empty version" "" "$(_previous_torch_pin '' 'cu126' 'torch>=2.4,<2.12.0')" +assert_eq "garbage version" "" "$(_previous_torch_pin 'not-a-version' 'cpu' 'torch>=2.4,<2.12.0')" +assert_eq "traceback fragment" "" "$(_previous_torch_pin "ModuleNotFoundError: No module named 'torch'" 'cpu' 'torch>=2.4,<2.12.0')" + +echo "=== _previous_torch_pin: out-of-window releases never pin ===" +assert_eq "2.3.x below the cu floor" "" "$(_previous_torch_pin '2.3.1+cu118' 'cu118' 'torch>=2.4,<2.12.0')" +assert_eq "2.12.x above the cu ceiling" "" "$(_previous_torch_pin '2.12.0+cu130' 'cu130' 'torch>=2.4,<2.12.0')" +assert_eq "floor boundary 2.4.0 kept" "torch==2.4.0" "$(_previous_torch_pin '2.4.0+cu126' 'cu126' 'torch>=2.4,<2.12.0')" +assert_eq "ceiling-adjacent 2.11.x kept" "torch==2.11.1" "$(_previous_torch_pin '2.11.1+cu130' 'cu130' 'torch>=2.4,<2.12.0')" +assert_eq "cpu window excludes 2.11.x" "" "$(_previous_torch_pin '2.11.0+cpu' 'cpu' 'torch>=2.4,<2.11.0')" +assert_eq "mac floor excludes 2.5.x" "" "$(_previous_torch_pin '2.5.1' 'cpu' 'torch>=2.6,<2.11.0')" +assert_eq "malformed window never pins" "" "$(_previous_torch_pin '2.10.0+cu126' 'cu126' 'torch')" +assert_eq "empty window never pins" "" "$(_previous_torch_pin '2.10.0+cu126' 'cu126' '')" + +echo "=== _torch_release_in_window ===" +assert_eq "in window" "yes" "$(_torch_release_in_window '2.10.0' 'torch>=2.4,<2.12.0')" +assert_eq "at floor" "yes" "$(_torch_release_in_window '2.4.0' 'torch>=2.4,<2.12.0')" +assert_eq "below floor" "no" "$(_torch_release_in_window '2.3.1' 'torch>=2.4,<2.12.0')" +assert_eq "at ceiling" "no" "$(_torch_release_in_window '2.12.0' 'torch>=2.4,<2.12.0')" +assert_eq "next major" "no" "$(_torch_release_in_window '3.0.0' 'torch>=2.4,<2.12.0')" +assert_eq "patch-level floor" "yes" "$(_torch_release_in_window '2.11.5' 'torch>=2.11.0,<2.12.0')" +assert_eq "no ceiling -> no" "no" "$(_torch_release_in_window '2.10.0' 'torch>=2.4')" +assert_eq "garbage minor -> no" "no" "$(_torch_release_in_window '2.x' 'torch>=2.4,<2.12.0')" + +echo "=== _previous_torch_pin: UNSLOTH_TORCH_UPGRADE=1 opts out ===" +assert_eq "upgrade env set" "" "$(UNSLOTH_TORCH_UPGRADE=1 _previous_torch_pin '2.10.0+cu126' 'cu126' 'torch>=2.4,<2.12.0')" +assert_eq "upgrade env 0" "torch==2.10.0" "$(UNSLOTH_TORCH_UPGRADE=0 _previous_torch_pin '2.10.0+cu126' 'cu126' 'torch>=2.4,<2.12.0')" + +echo "=== install.sh wiring ===" +# The probe must run against the OLD venv, before it is moved aside for rollback. +_probe_line=$(grep -n '_PREV_TORCH_VER=\$(' "$INSTALL_SH" | head -1 | cut -d: -f1) +_move_line=$(grep -n '_start_studio_venv_replacement "\$VENV_DIR"' "$INSTALL_SH" | head -1 | cut -d: -f1) +assert_eq "probe exists" "yes" "$([ -n "$_probe_line" ] && echo yes)" +assert_eq "probe before venv replacement" "yes" "$([ -n "$_probe_line" ] && [ -n "$_move_line" ] && [ "$_probe_line" -lt "$_move_line" ] && echo yes)" +# A kept release that vanished from the index must fall back to the supported range. +assert_eq "resolve-failure fallback wired" "yes" "$(grep -q 'TORCH_CONSTRAINT="\$_PREV_FALLBACK_CONSTRAINT"' "$INSTALL_SH" && echo yes)" +assert_eq "pin gated on SKIP_TORCH" "yes" "$(grep -q 'if \[ "\$SKIP_TORCH" = false \]; then' "$INSTALL_SH" && echo yes)" + +echo "" +if [ "$FAIL" -gt 0 ]; then + echo "$FAIL check(s) FAILED" + exit 1 +fi +echo "All $PASS checks passed" diff --git a/tests/sh/test_torch_constraint.sh b/tests/sh/test_torch_constraint.sh index 293a709360..d60dfc9f90 100644 --- a/tests/sh/test_torch_constraint.sh +++ b/tests/sh/test_torch_constraint.sh @@ -108,6 +108,20 @@ assert_eq "\$TORCH_CONSTRAINT used in pip install" "yes" "$_has_var" _hardcoded=$(grep -c '"torch>=2.4,<2.11.0"' "$INSTALL_SH" || true) assert_eq "hardcoded torch>=2.4 appears exactly once" "1" "$_hardcoded" +# 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). +_cuda_widen=$(grep -c 'TORCH_CONSTRAINT="torch>=2.4,<2.12.0"' "$INSTALL_SH" || true) +assert_eq "CUDA TORCH_CONSTRAINT widened to <2.12.0" "1" "$_cuda_widen" + +# Widening keys off the final leaf (_torch_index_leaf), not the full URL, so a +# mirror base path with cu*/rocm7.2 but a cpu/older-rocm leaf is not mis-widened. +_cuda_case=$(grep -c 'cu\[0-9\]\*)' "$INSTALL_SH" || true) +_has_cuda_case=$([ "$_cuda_case" -ge 1 ] && echo "yes" || echo "no") +assert_eq "cu* index case adjusts TORCH_CONSTRAINT" "yes" "$_has_cuda_case" +_leaf_case=$(grep -c 'case "\$_torch_index_leaf" in' "$INSTALL_SH" || true) +_has_leaf_constraint=$([ "$_leaf_case" -ge 2 ] && echo "yes" || echo "no") +assert_eq "constraint case anchors on _torch_index_leaf" "yes" "$_has_leaf_constraint" + echo "" echo "=== Structural: tokenizers in no-torch-runtime.txt ===" diff --git a/tests/sh/test_unsloth_torch_override.sh b/tests/sh/test_unsloth_torch_override.sh new file mode 100644 index 0000000000..7e8e3f5b5b --- /dev/null +++ b/tests/sh/test_unsloth_torch_override.sh @@ -0,0 +1,131 @@ +#!/bin/bash +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 +# Tests for the torch-trio --overrides guard on the Step-2 unsloth installs in +# install.sh. A released unsloth wheel can pin an older torch (2026.7.2 declares +# torch<2.11.0); without the overrides file a with-deps PyPI resolve downgrades +# the trio Step 1 installed, and the flavor guard misses it (PyPI's torch 2.10 +# default is itself cu128-flavored). Same assertion pattern as test_torch_constraint.sh. +set -e + +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" +INSTALL_SH="$SCRIPT_DIR/../../install.sh" +PASS=0 +FAIL=0 + +assert_true() { + _label="$1"; _ok="$2" + if [ "$_ok" = "0" ]; then + echo " PASS: $_label" + PASS=$((PASS + 1)) + else + echo " FAIL: $_label" + FAIL=$((FAIL + 1)) + fi +} + +echo "=== test_unsloth_torch_override ===" + +# 1. Every with-deps unsloth install carries the overrides expansion (local, +# generic, migrated); the --no-deps no-torch paths need no guard. +_local_block=$(grep -A2 '"install unsloth (local)"' "$INSTALL_SH") +printf '%s' "$_local_block" | grep -q -- '--overrides "\$_UNSLOTH_TORCH_OVERRIDES"' +assert_true "local (with-deps) unsloth install passes --overrides" "$?" + +_generic_block=$(grep -A2 '"install unsloth" uv pip install' "$INSTALL_SH") +printf '%s' "$_generic_block" | grep -q -- '--overrides "\$_UNSLOTH_TORCH_OVERRIDES"' +assert_true "generic (with-deps) unsloth install passes --overrides" "$?" + +_migrated_block=$(grep -A3 '"install unsloth (migrated)"' "$INSTALL_SH") +printf '%s' "$_migrated_block" | grep -q -- '--overrides "\$_UNSLOTH_TORCH_OVERRIDES"' +assert_true "migrated (with-deps) unsloth install passes --overrides" "$?" + +_no_torch_block=$(grep -A2 '"install unsloth (no-torch)"' "$INSTALL_SH") +if printf '%s' "$_no_torch_block" | grep -q -- '--overrides'; then _rc=1; else _rc=0; fi +assert_true "no-torch (--no-deps) unsloth install has no overrides" "$_rc" + +_migrated_nt_block=$(grep -A2 '"install unsloth (migrated no-torch)"' "$INSTALL_SH") +if printf '%s' "$_migrated_nt_block" | grep -q -- '--overrides'; then _rc=1; else _rc=0; fi +assert_true "migrated no-torch (--no-deps) unsloth install has no overrides" "$_rc" + +# 2. The overrides file is only built when SKIP_TORCH=false. +grep -B2 '_torch_trio_pins=\$(' "$INSTALL_SH" | grep -q 'SKIP_TORCH" = false' +assert_true "overrides file build is gated on SKIP_TORCH=false" "$?" + +# 3. The pin-collection snippet emits exact ==pins for the installed trio (run +# the embedded python against this test's interpreter). +_snippet=$(sed -n '/_torch_trio_pins=\$("\$_VENV_PY" -c "/,/^" 2>\/dev\/null)/p' "$INSTALL_SH" \ + | sed '1s/.*-c "//' | sed '$d') +_out=$(python3 -c "$_snippet" 2>&1) || true +# torch may or may not be importable on the test host; the snippet must not +# crash and every line it does emit must be an exact pkg==version pin. +if [ -n "$_out" ]; then + printf '%s\n' "$_out" | grep -vqE '^(torch|torchvision|torchaudio)==.+$' && _rc=1 || _rc=0 +else + _rc=0 +fi +assert_true "pin snippet emits only exact trio ==pins (or nothing)" "$_rc" + +# 4. The temp overrides file is cleaned up after Step 2. +grep -q 'rm -f "\$_UNSLOTH_TORCH_OVERRIDES"' "$INSTALL_SH" +assert_true "overrides temp file is removed after the unsloth installs" "$?" + +# 5. Any UV_OVERRIDE env file is folded in (the CLI --overrides flag would +# otherwise replace it, dropping e.g. the macOS arm64 darwin overrides). +grep -q 'for _ov_file in \${UV_OVERRIDE:-}' "$INSTALL_SH" +assert_true "UV_OVERRIDE env files are merged into the overrides file" "$?" + +# 6. The EXIT trap also removes the overrides file, so a failed Step 2 (set -e +# fires before the normal-path rm) cannot leak it. +sed -n '/_on_install_exit() {/,/^}/p' "$INSTALL_SH" \ + | grep -q 'rm -f "\$_UNSLOTH_TORCH_OVERRIDES"' +assert_true "EXIT trap removes the overrides temp file on failure" "$?" + +# 7. The UV_OVERRIDE fold filters inherited files instead of cat-ing them (run +# the extracted awk program on sample files): (a) inherited torch-trio lines +# are dropped so the generated exact pins win (uv intersects duplicates); +# (b) every line is newline-terminated so an unterminated file cannot join +# two requirements into one. +_awk_prog=$(sed -n "s/.*awk '\(.*\)' \"\$_ov_file\".*/\1/p" "$INSTALL_SH") +[ -n "$_awk_prog" ] +assert_true "UV_OVERRIDE fold uses the trio-filtering awk program" "$?" + +_ov_dir=$(mktemp -d) +printf '%s' 'transformers>=4.57.6' > "$_ov_dir/ov1.txt" # no trailing newline +cat > "$_ov_dir/ov2.txt" <<'EOF' +# comment survives +torch<2.11.0 +torchvision==0.25.0 +torchaudio!=2.11.0 +torchmetrics==1.0 +anyio<4.14.0 +EOF +_merged="$_ov_dir/merged.txt" +printf '%s\n' 'torch==2.11.0+cu128' > "$_merged" +for _f in "$_ov_dir/ov1.txt" "$_ov_dir/ov2.txt"; do + awk "$_awk_prog" "$_f" >> "$_merged" +done + +grep -qx 'transformers>=4.57.6' "$_merged" +assert_true "no-trailing-newline override stays a separate requirement line" "$?" + +if grep -qx 'torchmetrics==1.0' "$_merged" && grep -qx 'anyio<4.14.0' "$_merged"; then + _rc=0 +else + _rc=1 +fi +assert_true "unrelated inherited overrides are preserved" "$_rc" + +if grep -qE '^(torch|torchvision|torchaudio)([[:space:]<>=!~;@[]|$)' "$_merged" \ + && [ "$(grep -cE '^(torch|torchvision|torchaudio)([[:space:]<>=!~;@[]|$)' "$_merged")" != "1" ]; then + _rc=1 +else + _rc=0 +fi +grep -qx 'torch==2.11.0+cu128' "$_merged" || _rc=1 +assert_true "inherited torch-trio lines are dropped; generated pin wins" "$_rc" +rm -rf "$_ov_dir" + +echo "" +echo "Results: $PASS passed, $FAIL failed" +[ "$FAIL" -eq 0 ] || exit 1