From d90fd8b563ac18c896bd6603524fdd7793db1efb Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Mon, 20 Jul 2026 12:07:08 +0000
Subject: [PATCH] install: harden the preservation probe and the Windows
 override handoff

Three review follow-ups:

- install.sh's existing-venv torch probe now reads the version from dist
  metadata via the bounded runner instead of importing torch: a wedged
  CUDA/ROCm driver can hang the import indefinitely, while the metadata
  read never touches the driver. A failed or timed-out probe simply
  yields no preservable version.

- install.ps1's New-UnslothTorchOverridesFile folds caller-supplied
  UV_OVERRIDE files into the temporary overrides file (minus their
  torch-trio lines), matching install.sh: --overrides replaces the env
  file wholesale, so without the merge a caller's own dependency
  overrides were silently dropped during the unsloth resolution.

- install.ps1 clears an inherited UNSLOTH_KEPT_TORCH before the pin
  decision and sets it only on a fresh Get-PreviousTorchPin result: an
  interrupted earlier run could leak a stale exact pin into setup.ps1
  even when the current run found nothing to preserve or the upgrade
  opt-in was set.

Verified: the metadata probe returns the installed release through the
bounded runner; the override filter keeps unrelated pins (numpy,
torchao) while dropping torch/torchvision/torchaudio lines in all spec
forms. Full sh, ps1 and pytest installer batteries pass (host-defaults
and the tokenizers negative-control are the known pre-existing
failures).
---
 install.ps1 | 14 ++++++++++++++
 install.sh  |  8 +++++---
 2 files changed, 19 insertions(+), 3 deletions(-)

diff --git a/install.ps1 b/install.ps1
index ab60b06e0a..63b88efa35 100644
--- a/install.ps1
+++ b/install.ps1
@@ -2112,6 +2112,17 @@ exit 0
         $pins = & $PythonExe -c "from importlib.metadata import version, PackageNotFoundError`nfor _p in ('torch', 'torchvision', 'torchaudio'):`n    try:`n        print(_p + '==' + version(_p))`n    except PackageNotFoundError:`n        pass" 2>$null
         $lines = @($pins | Where-Object { $_ -match '^torch' })
         if ($lines.Count -eq 0 -or $lines[0] -notmatch '^torch==') { return $null }
+        # --overrides replaces any UV_OVERRIDE env file, so fold caller-supplied
+        # override files in (minus their torch-trio lines) like install.sh does.
+        if ($env:UV_OVERRIDE) {
+            foreach ($ovFile in ($env:UV_OVERRIDE -split '\s+' | Where-Object { $_ })) {
+                if (Test-Path -LiteralPath $ovFile -PathType Leaf) {
+                    $lines += @(Get-Content -LiteralPath $ovFile | Where-Object {
+                        $_ -notmatch '^\s*torch(vision|audio)?([\s<>=!~;@[]|$)'
+                    })
+                }
+            }
+        }
         $f = [System.IO.Path]::GetTempFileName()
         Set-Content -LiteralPath $f -Value ($lines -join "`n") -Encoding ascii
         return $f
@@ -2168,6 +2179,9 @@ exit 0
         # index/floor choice, incl. the ROCm reroute, so a raised floor rejects an older release. The
         # kept release is exported for setup.ps1 (UNSLOTH_KEPT_TORCH) and cleared after setup runs.
         $script:PrevTorchPin = $null
+        # Internal handoff variable: always clear an inherited value first (an interrupted
+        # earlier run can leak a stale pin into setup.ps1) and set it only on a fresh decision.
+        Remove-Item Env:UNSLOTH_KEPT_TORCH -ErrorAction SilentlyContinue
         if (-not $SkipTorch -and $script:PrevTorchVer) {
             $_routeWindow = $_pinTorchSpec
             if ($ROCmIndexUrl -and $ROCmTorchFloor) { $_routeWindow = $ROCmTorchFloor }
diff --git a/install.sh b/install.sh
index 3919d4b53e..2ebb9be571 100755
--- a/install.sh
+++ b/install.sh
@@ -1770,9 +1770,11 @@ 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 replacement (see _previous_torch_pin); last line only so stdout noise can't corrupt it.
-    _PREV_TORCH_VER=$("$VENV_DIR/bin/python" -c \
-        "import torch; print(torch.__version__)" 2>/dev/null | tail -n 1 || true)
+    # Record the existing venv's torch BEFORE replacement (see _previous_torch_pin); last line
+    # only so stdout noise can't corrupt it. Reads dist metadata instead of importing torch (a
+    # wedged CUDA/ROCm driver can hang the import indefinitely) and bounds the interpreter run.
+    _PREV_TORCH_VER=$(_run_bounded "$VENV_DIR/bin/python" -c \
+        "import importlib.metadata as m; print(m.version('torch'))" 2>/dev/null | tail -n 1 || true)
     substep "preserving existing environment for rollback..."
     _start_studio_venv_replacement "$VENV_DIR"
 elif [ "$_STUDIO_HOME_REDIRECT" != "env" ] && [ -x "$STUDIO_HOME/.venv/bin/python" ]; then