From a9be36830eb4ec731dcd10008d2f6dc3bb102d40 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 19 Jul 2026 06:19:29 -0700 Subject: [PATCH] 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` 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[, --- install.sh | 168 +++++++++++++++++- studio/backend/tests/test_cuda_torch_spec.py | 73 ++++++++ .../test_tokenizers_and_torch_constraint.py | 80 +++++++++ tests/sh/test_previous_torch_pin.sh | 101 +++++++++++ tests/sh/test_torch_constraint.sh | 14 ++ tests/sh/test_unsloth_torch_override.sh | 131 ++++++++++++++ 6 files changed, 558 insertions(+), 9 deletions(-) create mode 100644 studio/backend/tests/test_cuda_torch_spec.py create mode 100644 tests/sh/test_previous_torch_pin.sh create mode 100644 tests/sh/test_unsloth_torch_override.sh 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