diff --git a/.github/workflows/cross-platform-parity-ci.yml b/.github/workflows/cross-platform-parity-ci.yml index bb7dcbf8e4..45ce231743 100644 --- a/.github/workflows/cross-platform-parity-ci.yml +++ b/.github/workflows/cross-platform-parity-ci.yml @@ -1,18 +1,16 @@ # SPDX-License-Identifier: AGPL-3.0-only # Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. -# Runs installer parity and autostart opt-out tests on Windows and macOS. +# Runs installer parity and autostart opt-out tests across all three platforms. # -# Why: that test is the guard that install.sh and install.ps1 stay in -# sync, but today it only runs on ubuntu-latest (auto-discovered by -# studio-backend-ci.yml's "Repo tests (CPU)" job). The test reads both -# installer scripts, and on Windows Path.read_text() defaults to the -# cp1252 locale encoding, so a non-cp1252 byte in install.sh (it already -# contains a U+274C) raises UnicodeDecodeError there even though Linux and -# macOS default to UTF-8. The reads were pinned to encoding="utf-8" in -# #6166; this job keeps that from silently regressing by exercising the -# test on the platforms it claims parity for. Pure pytest, no GPU, -# sub-second, so the matrix is cheap. +# Why: the parity test guards that install.sh and install.ps1 stay in sync. +# It originally ran only on ubuntu-latest through studio-backend-ci.yml. +# On Windows, Path.read_text() defaults to the cp1252 locale encoding, so a +# non-cp1252 byte in install.sh raises UnicodeDecodeError even though Linux +# and macOS default to UTF-8. The reads were pinned to encoding="utf-8" in +# #6166; this matrix keeps that from silently regressing. Pure pytest, no GPU, +# sub-second, so the matrix is cheap. Linux also runs the POSIX rollback test +# under dash, matching the supported curl-to-sh installer path. name: Cross-platform parity @@ -23,6 +21,8 @@ on: - 'install.ps1' - 'tests/test_installer_skip_autostart.py' - 'tests/python/test_cross_platform_parity.py' + - 'tests/sh/test_install_rollback_lifecycle.sh' + - 'tests/studio/test_install_rollback_lifecycle.ps1' - '.github/workflows/cross-platform-parity-ci.yml' push: branches: [main] @@ -31,6 +31,8 @@ on: - 'install.ps1' - 'tests/test_installer_skip_autostart.py' - 'tests/python/test_cross_platform_parity.py' + - 'tests/sh/test_install_rollback_lifecycle.sh' + - 'tests/studio/test_install_rollback_lifecycle.ps1' - '.github/workflows/cross-platform-parity-ci.yml' workflow_dispatch: @@ -47,7 +49,7 @@ jobs: strategy: fail-fast: false matrix: - os: [windows-latest, macos-latest] + os: [ubuntu-latest, windows-latest, macos-latest] runs-on: ${{ matrix.os }} timeout-minutes: 10 steps: @@ -67,3 +69,10 @@ jobs: tests/python/test_cross_platform_parity.py tests/test_installer_skip_autostart.py -q + - name: PowerShell rollback lifecycle tests + if: runner.os == 'Windows' + shell: pwsh + run: pwsh -NoProfile -File tests/studio/test_install_rollback_lifecycle.ps1 + - name: POSIX rollback lifecycle tests + if: runner.os == 'Linux' + run: sh tests/sh/test_install_rollback_lifecycle.sh diff --git a/install.ps1 b/install.ps1 index 36e03ca51d..c06f3a1120 100644 --- a/install.ps1 +++ b/install.ps1 @@ -1416,13 +1416,82 @@ exit 0 $suffix++ $candidate = Join-Path $StudioHome "unsloth_studio.rollback.$stamp.$PID.$suffix" } - Move-Item -LiteralPath $ExistingDir -Destination $candidate -ErrorAction Stop $script:StudioVenvRollbackDir = $candidate $script:StudioVenvRollbackTarget = $ExistingDir $script:StudioVenvRollbackActive = $true + # Publish the rollback state before the atomic rename so interruption + # cannot land after Move-Item but before cleanup knows where the old venv went. + try { + Move-Item -LiteralPath $ExistingDir -Destination $candidate -ErrorAction Stop + } catch { + # A collision or ordinary rename failure leaves the original in place. + # Keep state active only when the rename happened before interruption. + if (Test-Path -LiteralPath $ExistingDir) { + $script:StudioVenvRollbackActive = $false + $script:StudioVenvRollbackDir = $null + } + throw + } substep "previous environment preserved for rollback" } + function Remove-StudioVenvTreeWithRetry { + param( + [Parameter(Mandatory = $true)][string]$Path, + [Parameter(Mandatory = $true)][string]$Label + ) + $lastError = $null + for ($attempt = 1; $attempt -le 3; $attempt++) { + try { + Remove-Item -LiteralPath $Path -Recurse -Force -ErrorAction Stop + } catch { + $lastError = $_.Exception.Message + } + if (-not (Test-Path -LiteralPath $Path)) { return $true } + if ($attempt -lt 3) { Start-Sleep -Milliseconds (250 * $attempt) } + } + Write-Host "[WARN] Could not remove $Label at $Path" -ForegroundColor Yellow + if ($lastError) { Write-Host " $lastError" -ForegroundColor Yellow } + return $false + } + + function Test-StudioVenvRollbackMustBePreserved { + param([Parameter(Mandatory = $true)][System.IO.FileSystemInfo]$Rollback) + # Preserve anything outside the installer's timestamp.PID[.suffix] format. + if ($Rollback.Name -notmatch '^unsloth_studio\.rollback\.[0-9]{14}\.([0-9]+)(?:\.[0-9]+)?$') { + return $true + } + $ownerPid = 0 + if (-not [int]::TryParse($Matches[1], [ref]$ownerPid)) { return $true } + if ($ownerPid -eq $PID) { return $true } + return $null -ne (Get-Process -Id $ownerPid -ErrorAction SilentlyContinue) + } + + function Remove-StaleStudioVenvRollbacks { + try { + $rollbacks = @( + Get-ChildItem -LiteralPath $StudioHome -Directory -Force -ErrorAction Stop | + Where-Object { $_.Name -like 'unsloth_studio.rollback.*' } + ) + } catch { + Write-Host "[WARN] Could not inspect stale environment rollbacks in $StudioHome" -ForegroundColor Yellow + Write-Host " $($_.Exception.Message)" -ForegroundColor Yellow + return + } + foreach ($rollback in $rollbacks) { + if (($rollback.Attributes -band [System.IO.FileAttributes]::ReparsePoint) -ne 0) { + Write-Host "[WARN] Refusing to remove rollback reparse point $($rollback.FullName)" -ForegroundColor Yellow + continue + } + # A concurrent installer may have moved its live venv aside. The PID + # in the generated name keeps this run from deleting its rescue copy. + if (Test-StudioVenvRollbackMustBePreserved -Rollback $rollback) { continue } + if (Remove-StudioVenvTreeWithRetry -Path $rollback.FullName -Label "stale environment rollback") { + substep "removed stale environment rollback $($rollback.Name)" + } + } + } + function Restore-StudioVenvRollback { if (-not $script:StudioVenvRollbackActive) { return } $backup = $script:StudioVenvRollbackDir @@ -1434,7 +1503,9 @@ exit 0 substep "restoring previous environment after failed install..." "Yellow" try { if (Test-Path -LiteralPath $target) { - Remove-Item -LiteralPath $target -Recurse -Force -ErrorAction SilentlyContinue + if (-not (Remove-StudioVenvTreeWithRetry -Path $target -Label "incomplete environment")) { + throw "Could not remove incomplete environment at $target" + } } Move-Item -LiteralPath $backup -Destination $target -Force -ErrorAction Stop substep "restored previous environment" @@ -1449,13 +1520,17 @@ exit 0 function Complete-StudioVenvRollback { if (-not $script:StudioVenvRollbackActive) { return } $backup = $script:StudioVenvRollbackDir - if ($backup -and (Test-Path -LiteralPath $backup)) { - Remove-Item -LiteralPath $backup -Recurse -Force -ErrorAction SilentlyContinue - } + # The replacement is committed. Disable restoration before deleting the + # backup so interruption cannot restore a partially deleted environment. $script:StudioVenvRollbackActive = $false $script:StudioVenvRollbackDir = $null + if ($backup -and (Test-Path -LiteralPath $backup)) { + Remove-StudioVenvTreeWithRetry -Path $backup -Label "environment rollback" | Out-Null + } } + $studioVenvReplacementCommitted = $false + try { if (Test-Path -LiteralPath $VenvPython) { # why: matching guard to the .venv branch below -- in env-mode # $StudioHome is a user-chosen workspace, so refuse to nuke an @@ -2688,6 +2763,13 @@ exit 0 } Refresh-SessionPath # sync current session with registry Complete-StudioVenvRollback + $studioVenvReplacementCommitted = $true + Remove-StaleStudioVenvRollbacks + } finally { + if (-not $studioVenvReplacementCommitted) { + Restore-StudioVenvRollback + } + } # Env-mode session export AFTER Refresh-SessionPath; otherwise a legacy # User PATH entry (Machine > User > current $env:Path) would win. diff --git a/install.sh b/install.sh index d06fff07c9..44d51490f7 100755 --- a/install.sh +++ b/install.sh @@ -475,14 +475,20 @@ _start_studio_venv_replacement() { _stamp=$(date +%Y%m%d%H%M%S 2>/dev/null || echo "time") _candidate="$STUDIO_HOME/unsloth_studio.rollback.$_stamp.$$" _suffix=0 - while [ -e "$_candidate" ]; do + while [ -e "$_candidate" ] || [ -L "$_candidate" ]; do _suffix=$((_suffix + 1)) _candidate="$STUDIO_HOME/unsloth_studio.rollback.$_stamp.$$.$_suffix" done - mv "$_existing_dir" "$_candidate" _VENV_ROLLBACK_DIR="$_candidate" _VENV_ROLLBACK_TARGET="$_existing_dir" _VENV_ROLLBACK_ACTIVE=true + # Publish the rollback state before the atomic rename so a signal cannot + # land after mv but before the exit handlers know where the old venv went. + if ! mv "$_existing_dir" "$_candidate"; then + _VENV_ROLLBACK_ACTIVE=false + _VENV_ROLLBACK_DIR="" + return 1 + fi substep "previous environment preserved for rollback" } @@ -503,13 +509,68 @@ _restore_studio_venv_replacement() { fi } -_commit_studio_venv_replacement() { - [ "$_VENV_ROLLBACK_ACTIVE" = true ] || return 0 - if [ -n "$_VENV_ROLLBACK_DIR" ] && [ -d "$_VENV_ROLLBACK_DIR" ]; then - rm -rf "$_VENV_ROLLBACK_DIR" || true +_studio_venv_rollback_must_be_preserved() { + _rollback_name=${1##*/} + _rollback_metadata=${_rollback_name#unsloth_studio.rollback.} + _rollback_stamp=${_rollback_metadata%%.*} + _rollback_process=${_rollback_metadata#*.} + # Preserve anything outside the installer's timestamp.PID[.suffix] format. + [ "$_rollback_process" != "$_rollback_metadata" ] || return 0 + case "$_rollback_stamp" in + time) ;; + ''|*[!0-9]*) return 0 ;; + *) [ "${#_rollback_stamp}" -eq 14 ] || return 0 ;; + esac + _rollback_pid=${_rollback_process%%.*} + case "$_rollback_pid" in + ''|*[!0-9]*) return 0 ;; + esac + _rollback_suffix=${_rollback_process#*.} + if [ "$_rollback_suffix" != "$_rollback_process" ]; then + case "$_rollback_suffix" in ''|*[!0-9]*) return 0 ;; esac fi - _VENV_ROLLBACK_ACTIVE=false - _VENV_ROLLBACK_DIR="" + kill -0 "$_rollback_pid" 2>/dev/null +} + +_prune_stale_studio_venv_rollbacks() { + for _stale_rollback in "$STUDIO_HOME"/unsloth_studio.rollback.*; do + [ -d "$_stale_rollback" ] || continue + if [ -L "$_stale_rollback" ]; then + echo "⚠️ Refusing to remove rollback symlink $_stale_rollback" >&2 + continue + fi + # A concurrent installer may have moved its live venv aside. The PID in + # the generated name keeps this successful run from deleting its rescue copy. + _studio_venv_rollback_must_be_preserved "$_stale_rollback" && continue + if rm -rf "$_stale_rollback"; then + substep "removed stale environment rollback ${_stale_rollback##*/}" + else + echo "⚠️ Could not remove stale environment rollback $_stale_rollback" >&2 + fi + done +} + +_commit_studio_venv_replacement() { + if [ "$_VENV_ROLLBACK_ACTIVE" = true ]; then + _rollback_to_remove="$_VENV_ROLLBACK_DIR" + # The new environment is already committed. Clear the restore state + # before deletion so an interrupt cannot replace it with a half-deleted backup. + _VENV_ROLLBACK_ACTIVE=false + _VENV_ROLLBACK_DIR="" + if [ -n "$_rollback_to_remove" ] && [ -d "$_rollback_to_remove" ]; then + if ! rm -rf "$_rollback_to_remove"; then + echo "⚠️ Could not remove environment rollback $_rollback_to_remove" >&2 + fi + fi + fi + # Only prune older orphaned copies after the replacement has succeeded, so + # an interrupted install never discards the last known-good environment. + _prune_stale_studio_venv_rollbacks +} + +_cleanup_install_temporaries() { + [ -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 } _on_install_exit() { @@ -517,15 +578,28 @@ _on_install_exit() { if [ "$_status" -ne 0 ]; then _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 + _cleanup_install_temporaries exit "$_status" } + +_on_install_signal() { + _signal_status="$1" + # EXIT is disabled to avoid a second cleanup pass. Ignore further termination + # signals until the old environment is back in place. + trap - EXIT + trap '' HUP INT TERM + _restore_studio_venv_replacement + _cleanup_install_temporaries + exit "$_signal_status" +} # 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 +trap '_on_install_signal 129' HUP +trap '_on_install_signal 130' INT +trap '_on_install_signal 143' TERM # ── Helper: download a URL to a file (supports curl and wget) ── download() { diff --git a/tests/run_all.sh b/tests/run_all.sh index a31103a85b..eaa726f73c 100755 --- a/tests/run_all.sh +++ b/tests/run_all.sh @@ -17,6 +17,7 @@ sh "$TESTS_DIR/sh/test_uninstall_shared_icon.sh" sh "$TESTS_DIR/sh/test_torch_flavor.sh" sh "$TESTS_DIR/sh/test_redact_install_output.sh" sh "$TESTS_DIR/sh/test_install_uv_override_space.sh" +sh "$TESTS_DIR/sh/test_install_rollback_lifecycle.sh" echo "" echo "=== Python tests ===" diff --git a/tests/sh/test_install_rollback_lifecycle.sh b/tests/sh/test_install_rollback_lifecycle.sh new file mode 100644 index 0000000000..d1ccae8e19 --- /dev/null +++ b/tests/sh/test_install_rollback_lifecycle.sh @@ -0,0 +1,202 @@ +#!/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 +# Exercises install.sh's real rollback helpers without downloading the Studio stack. +set -e + +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" +INSTALL_SH="$SCRIPT_DIR/../../install.sh" +INSTALL_PS1="$SCRIPT_DIR/../../install.ps1" +PASS=0 +FAIL=0 + +ok() { echo " PASS: $1"; PASS=$((PASS + 1)); } +bad() { echo " FAIL: $1"; FAIL=$((FAIL + 1)); } + +ROLLBACK_BLOCK=$(sed -n '/^_VENV_ROLLBACK_DIR=""/,/^trap '\''_on_install_signal 143'\'' TERM$/p' "$INSTALL_SH") +if ! printf '%s\n' "$ROLLBACK_BLOCK" | grep -q '^_on_install_signal() {'; then + echo " FAIL: could not extract rollback lifecycle block from install.sh" + exit 1 +fi + +WORK=$(mktemp -d) +trap 'rm -rf "$WORK"' EXIT + +run_signal_case() { + _signal="$1" + _expected_status="$2" + _case_dir="$WORK/signal-$_signal" + mkdir -p "$_case_dir/unsloth_studio" + printf 'old\n' > "$_case_dir/unsloth_studio/generation" + _harness="$_case_dir/harness.sh" + { + printf '%s\n' 'set -e' + printf '%s\n' 'substep() { :; }' + printf '%s\n' 'C_WARN=""' + printf "STUDIO_HOME='%s'\n" "$_case_dir" + printf "VENV_DIR='%s/unsloth_studio'\n" "$_case_dir" + printf '%s\n' "$ROLLBACK_BLOCK" + printf '%s\n' '_start_studio_venv_replacement "$VENV_DIR"' + printf '%s\n' 'mkdir -p "$VENV_DIR"' + printf '%s\n' 'printf "partial\n" > "$VENV_DIR/generation"' + printf 'kill -%s $$\n' "$_signal" + printf '%s\n' 'exit 99' + } > "$_harness" + + set +e + dash "$_harness" >/dev/null 2>&1 + _status=$? + set -e + if [ "$_status" = "$_expected_status" ]; then + ok "dash $_signal exits with $_expected_status" + else + bad "dash $_signal exits with $_expected_status (got $_status)" + fi + if [ "$(cat "$_case_dir/unsloth_studio/generation" 2>/dev/null)" = "old" ]; then + ok "dash $_signal restores the previous environment" + else + bad "dash $_signal did not restore the previous environment" + fi + if ! find "$_case_dir" -maxdepth 1 -name 'unsloth_studio.rollback.*' -print -quit | grep -q .; then + ok "dash $_signal leaves no rollback copy" + else + bad "dash $_signal left a rollback copy" + fi +} + +echo "=== install.sh signal rollback ===" +run_signal_case INT 130 +run_signal_case TERM 143 +run_signal_case HUP 129 + +echo "=== install.sh transition boundaries ===" +START_BOUNDARY_DIR="$WORK/start-boundary" +mkdir -p "$START_BOUNDARY_DIR/unsloth_studio" +printf 'old\n' > "$START_BOUNDARY_DIR/unsloth_studio/generation" +START_BOUNDARY_HARNESS="$START_BOUNDARY_DIR/harness.sh" +{ + printf '%s\n' 'set -e' + printf '%s\n' 'substep() { :; }' + printf '%s\n' 'C_WARN=""' + printf "STUDIO_HOME='%s'\n" "$START_BOUNDARY_DIR" + printf "VENV_DIR='%s/unsloth_studio'\n" "$START_BOUNDARY_DIR" + printf '%s\n' "$ROLLBACK_BLOCK" + printf '%s\n' 'mv() { command mv "$@"; kill -TERM $$; }' + printf '%s\n' '_start_studio_venv_replacement "$VENV_DIR"' +} > "$START_BOUNDARY_HARNESS" +set +e +dash "$START_BOUNDARY_HARNESS" >/dev/null 2>&1 +_start_boundary_status=$? +set -e +if [ "$_start_boundary_status" -eq 143 ] \ + && [ "$(cat "$START_BOUNDARY_DIR/unsloth_studio/generation" 2>/dev/null)" = "old" ]; then + ok "signal immediately after rollback rename restores the old environment" +else + bad "rollback state was not published before rename" +fi + +COMMIT_BOUNDARY_DIR="$WORK/commit-boundary" +mkdir -p "$COMMIT_BOUNDARY_DIR/unsloth_studio" +printf 'old\n' > "$COMMIT_BOUNDARY_DIR/unsloth_studio/generation" +COMMIT_BOUNDARY_HARNESS="$COMMIT_BOUNDARY_DIR/harness.sh" +{ + printf '%s\n' 'set -e' + printf '%s\n' 'substep() { :; }' + printf '%s\n' 'C_WARN=""' + printf "STUDIO_HOME='%s'\n" "$COMMIT_BOUNDARY_DIR" + printf "VENV_DIR='%s/unsloth_studio'\n" "$COMMIT_BOUNDARY_DIR" + printf '%s\n' "$ROLLBACK_BLOCK" + printf '%s\n' '_start_studio_venv_replacement "$VENV_DIR"' + printf '%s\n' 'mkdir -p "$VENV_DIR"' + printf '%s\n' 'printf "new\n" > "$VENV_DIR/generation"' + printf '%s\n' 'rm() { kill -TERM $$; }' + printf '%s\n' '_commit_studio_venv_replacement' +} > "$COMMIT_BOUNDARY_HARNESS" +set +e +dash "$COMMIT_BOUNDARY_HARNESS" >/dev/null 2>&1 +_commit_boundary_status=$? +set -e +if [ "$_commit_boundary_status" -eq 143 ] \ + && [ "$(cat "$COMMIT_BOUNDARY_DIR/unsloth_studio/generation" 2>/dev/null)" = "new" ]; then + ok "signal during committed-backup deletion keeps the new environment" +else + bad "signal during committed-backup deletion restored a partial backup" +fi + +echo "=== install.sh successful cleanup ===" +PRUNE_DIR="$WORK/prune" +mkdir -p "$PRUNE_DIR/unsloth_studio" +printf 'old\n' > "$PRUNE_DIR/unsloth_studio/generation" +PRUNE_HARNESS="$PRUNE_DIR/harness.sh" +{ + printf '%s\n' 'set -e' + printf '%s\n' 'substep() { :; }' + printf '%s\n' 'C_WARN=""' + printf "STUDIO_HOME='%s'\n" "$PRUNE_DIR" + printf "VENV_DIR='%s/unsloth_studio'\n" "$PRUNE_DIR" + printf '%s\n' "$ROLLBACK_BLOCK" + printf '%s\n' '_start_studio_venv_replacement "$VENV_DIR"' + printf '%s\n' 'mkdir -p "$VENV_DIR"' + printf '%s\n' 'printf "new\n" > "$VENV_DIR/generation"' + printf '%s\n' 'mkdir "$STUDIO_HOME/unsloth_studio.rollback.20000101000000.999999999"' + printf '%s\n' 'mkdir "$STUDIO_HOME/unsloth_studio.rollback.20000101000001.$$"' + printf '%s\n' 'mkdir "$STUDIO_HOME/unsloth_studio.rollback.user-data"' + printf '%s\n' 'mkdir "$STUDIO_HOME/outside"' + printf '%s\n' 'ln -s "$STUDIO_HOME/outside" "$STUDIO_HOME/unsloth_studio.rollback.20000101000002.999999998"' + printf '%s\n' '_commit_studio_venv_replacement' +} > "$PRUNE_HARNESS" + +sh "$PRUNE_HARNESS" >/dev/null 2>&1 +if [ "$(cat "$PRUNE_DIR/unsloth_studio/generation" 2>/dev/null)" = "new" ]; then + ok "successful replacement keeps the new environment" +else + bad "successful replacement lost the new environment" +fi +if [ ! -d "$PRUNE_DIR/unsloth_studio.rollback.20000101000000.999999999" ]; then + ok "successful install removes an orphan from a dead PID" +else + bad "successful install left an orphan from a dead PID" +fi +_active_count=$(find "$PRUNE_DIR" -maxdepth 1 -type d -name 'unsloth_studio.rollback.20000101000001.*' | wc -l) +if [ "$_active_count" -eq 1 ]; then + ok "successful install preserves a concurrent installer's rollback" +else + bad "successful install removed a concurrent installer's rollback" +fi +if [ -d "$PRUNE_DIR/unsloth_studio.rollback.user-data" ]; then + ok "stale cleanup preserves names outside the generated format" +else + bad "stale cleanup removed a non-generated rollback name" +fi +if [ -L "$PRUNE_DIR/unsloth_studio.rollback.20000101000002.999999998" ] \ + && [ -d "$PRUNE_DIR/outside" ]; then + ok "stale cleanup does not follow rollback symlinks" +else + bad "stale cleanup mutated a rollback symlink target" +fi + +echo "=== install.ps1 rollback wiring ===" +if grep -q '^ function Remove-StaleStudioVenvRollbacks {' "$INSTALL_PS1" \ + && grep -q '^ Remove-StaleStudioVenvRollbacks$' "$INSTALL_PS1"; then + ok "Windows installer prunes stale rollbacks after success" +else + bad "Windows installer does not wire stale rollback cleanup" +fi +if grep -q '^ } finally {$' "$INSTALL_PS1" \ + && grep -A3 '^ } finally {$' "$INSTALL_PS1" | grep -q 'Restore-StudioVenvRollback'; then + ok "Windows replacement is protected by finally" +else + bad "Windows replacement lacks finally rollback" +fi +if grep -A18 '^ function Remove-StudioVenvTreeWithRetry {' "$INSTALL_PS1" \ + | grep -q 'ErrorAction Stop'; then + ok "Windows rollback deletion failures are observable and retried" +else + bad "Windows rollback deletion still hides failures" +fi + +echo "" +echo " PASS: $PASS" +echo " FAIL: $FAIL" +[ "$FAIL" -eq 0 ] || exit 1 +echo "ALL PASSED" diff --git a/tests/sh/test_unsloth_torch_override.sh b/tests/sh/test_unsloth_torch_override.sh index 7e8e3f5b5b..84e52ca287 100644 --- a/tests/sh/test_unsloth_torch_override.sh +++ b/tests/sh/test_unsloth_torch_override.sh @@ -75,11 +75,22 @@ assert_true "overrides temp file is removed after the unsloth installs" "$?" 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" \ +# 6. Exit and signal traps share cleanup, so a failed or interrupted Step 2 +# cannot leak the overrides file. +sed -n '/_on_install_exit() {/,/^}/p' "$INSTALL_SH" | grep -q '_cleanup_install_temporaries' +_exit_cleanup_rc=$? +sed -n '/_on_install_signal() {/,/^}/p' "$INSTALL_SH" | grep -q '_cleanup_install_temporaries' +_signal_cleanup_rc=$? +sed -n '/_cleanup_install_temporaries() {/,/^}/p' "$INSTALL_SH" \ | grep -q 'rm -f "\$_UNSLOTH_TORCH_OVERRIDES"' -assert_true "EXIT trap removes the overrides temp file on failure" "$?" +_cleanup_body_rc=$? +if [ "$_exit_cleanup_rc" -eq 0 ] && [ "$_signal_cleanup_rc" -eq 0 ] \ + && [ "$_cleanup_body_rc" -eq 0 ]; then + _rc=0 +else + _rc=1 +fi +assert_true "exit and signal traps remove the overrides temp file" "$_rc" # 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 diff --git a/tests/studio/test_install_rollback_lifecycle.ps1 b/tests/studio/test_install_rollback_lifecycle.ps1 new file mode 100644 index 0000000000..d70f384ca7 --- /dev/null +++ b/tests/studio/test_install_rollback_lifecycle.ps1 @@ -0,0 +1,126 @@ +#!/usr/bin/env pwsh +# 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.ps1's venv rollback helpers. The functions are AST-extracted +# so the top-level installer is never executed. + +$ErrorActionPreference = "Stop" +$installPath = [System.IO.Path]::Combine($PSScriptRoot, "..", "..", "install.ps1") +$installPath = (Resolve-Path $installPath).Path + +$tokens = $null; $errors = $null +$ast = [System.Management.Automation.Language.Parser]::ParseFile($installPath, [ref]$tokens, [ref]$errors) +if ($errors) { $errors | ForEach-Object { $_.ToString() }; throw "install.ps1 has parse errors" } + +$helperNames = @( + "Start-StudioVenvRollback", + "Remove-StudioVenvTreeWithRetry", + "Test-StudioVenvRollbackMustBePreserved", + "Remove-StaleStudioVenvRollbacks", + "Restore-StudioVenvRollback", + "Complete-StudioVenvRollback" +) +foreach ($name in $helperNames) { + $fn = $ast.FindAll({ param($node) + $node -is [System.Management.Automation.Language.FunctionDefinitionAst] -and $node.Name -eq $name + }, $true) + if ($fn.Count -ne 1) { throw "expected exactly one $name in install.ps1, found $($fn.Count)" } + Invoke-Expression $fn[0].Extent.Text +} + +function substep { param([string]$Message, [string]$Color) } + +$failures = 0 +function Check($name, $condition) { + if ($condition) { Write-Host " PASS $name" } + else { Write-Host " FAIL $name" -ForegroundColor Red; $script:failures++ } +} + +function Reset-RollbackState($target) { + $script:StudioVenvRollbackDir = $null + $script:StudioVenvRollbackTarget = $target + $script:StudioVenvRollbackActive = $false +} + +$StudioHome = Join-Path ([System.IO.Path]::GetTempPath()) "unsloth-rollback-$([guid]::NewGuid().ToString('N'))" +$VenvDir = Join-Path $StudioHome "unsloth_studio" +[System.IO.Directory]::CreateDirectory($VenvDir) | Out-Null + +try { + Write-Host "Successful replacement" + [System.IO.File]::WriteAllText((Join-Path $VenvDir "generation"), "old") + Reset-RollbackState $VenvDir + Start-StudioVenvRollback -ExistingDir $VenvDir + [System.IO.Directory]::CreateDirectory($VenvDir) | Out-Null + [System.IO.File]::WriteAllText((Join-Path $VenvDir "generation"), "new") + Complete-StudioVenvRollback + Check "new environment remains" ((Get-Content -LiteralPath (Join-Path $VenvDir "generation") -Raw) -eq "new") + Check "current rollback is removed" (-not @(Get-ChildItem -LiteralPath $StudioHome -Directory | + Where-Object { $_.Name -like "unsloth_studio.rollback.*" })) + + Write-Host "Stale cleanup" + $stale = Join-Path $StudioHome "unsloth_studio.rollback.20000101000000.2147483647" + $active = Join-Path $StudioHome "unsloth_studio.rollback.20000101000001.$PID" + $unrecognized = Join-Path $StudioHome "unsloth_studio.rollback.user-data" + [System.IO.Directory]::CreateDirectory($stale) | Out-Null + [System.IO.Directory]::CreateDirectory($active) | Out-Null + [System.IO.Directory]::CreateDirectory($unrecognized) | Out-Null + Remove-StaleStudioVenvRollbacks + Check "dead-owner rollback is removed" (-not (Test-Path -LiteralPath $stale)) + Check "live-owner rollback is preserved" (Test-Path -LiteralPath $active) + Check "unrecognized rollback name is preserved" (Test-Path -LiteralPath $unrecognized) + Microsoft.PowerShell.Management\Remove-Item -LiteralPath $active -Recurse -Force + Microsoft.PowerShell.Management\Remove-Item -LiteralPath $unrecognized -Recurse -Force + + Write-Host "Failure restoration" + [System.IO.File]::WriteAllText((Join-Path $VenvDir "generation"), "old-again") + Reset-RollbackState $VenvDir + $committed = $false + try { + try { + Start-StudioVenvRollback -ExistingDir $VenvDir + [System.IO.Directory]::CreateDirectory($VenvDir) | Out-Null + [System.IO.File]::WriteAllText((Join-Path $VenvDir "generation"), "partial") + throw "simulated install failure" + } finally { + if (-not $committed) { Restore-StudioVenvRollback } + } + } catch { + if ($_.Exception.Message -ne "simulated install failure") { throw } + } + Check "finally restores the previous environment" ( + (Get-Content -LiteralPath (Join-Path $VenvDir "generation") -Raw) -eq "old-again" + ) + Check "failure restoration consumes the rollback" (-not @(Get-ChildItem -LiteralPath $StudioHome -Directory | + Where-Object { $_.Name -like "unsloth_studio.rollback.*" })) + + Write-Host "Locked-file retry" + $retryDir = Join-Path $StudioHome "retry" + [System.IO.Directory]::CreateDirectory($retryDir) | Out-Null + $script:removeAttempts = 0 + function Remove-Item { + param( + [string]$LiteralPath, + [switch]$Recurse, + [switch]$Force, + [object]$ErrorAction + ) + $script:removeAttempts++ + if ($script:removeAttempts -lt 3) { throw "simulated lock" } + Microsoft.PowerShell.Management\Remove-Item -LiteralPath $LiteralPath -Recurse:$Recurse -Force:$Force + } + try { + $removed = Remove-StudioVenvTreeWithRetry -Path $retryDir -Label "test rollback" + } finally { + Microsoft.PowerShell.Management\Remove-Item -LiteralPath Function:\Remove-Item -Force + } + Check "locked rollback deletion retries" ($removed -and $script:removeAttempts -eq 3) +} finally { + if (Test-Path -LiteralPath $StudioHome) { + Microsoft.PowerShell.Management\Remove-Item -LiteralPath $StudioHome -Recurse -Force + } +} + +Write-Host "" +if ($failures -gt 0) { Write-Host "$failures check(s) FAILED" -ForegroundColor Red; exit 1 } +Write-Host "All checks passed" -ForegroundColor Green