diff --git a/install.ps1 b/install.ps1 index a27af9dd3b..5e3d4b6a50 100644 --- a/install.ps1 +++ b/install.ps1 @@ -92,6 +92,7 @@ function Install-UnslothStudio { $RepoRoot = "" $TauriMode = $false $SkipTorch = $false + $ShortcutsOnly = $false $argList = $args for ($i = 0; $i -lt $argList.Count; $i++) { switch ($argList[$i]) { @@ -100,6 +101,7 @@ function Install-UnslothStudio { "--no-torch" { $SkipTorch = $true } "--verbose" { $script:UnslothVerbose = $true } "-v" { $script:UnslothVerbose = $true } + "--shortcuts-only" { $ShortcutsOnly = $true } "--package" { $i++ if ($i -ge $argList.Count) { @@ -871,6 +873,19 @@ shell.Run cmd, 0, False } } + # Regen .lnk + launcher only; used by `unsloth studio update`. + if ($ShortcutsOnly) { + if ($TauriMode) { return } + $UnslothExe = Join-Path $VenvDir "Scripts\unsloth.exe" + if (-not (Test-Path -LiteralPath $UnslothExe)) { + Write-Host "[ERROR] unsloth.exe missing at $UnslothExe; run install.ps1 first." -ForegroundColor Red + # throw (not Exit-InstallFailure) so non-Tauri callers see rc != 0. + throw "unsloth.exe missing" + } + New-StudioShortcuts -UnslothExePath $UnslothExe + return + } + # ── Check winget ── Write-TauriLog "STEP" "Checking system dependencies" if (-not (Get-Command winget -ErrorAction SilentlyContinue)) { diff --git a/install.sh b/install.sh index dd4f83fab6..cfd76fa945 100755 --- a/install.sh +++ b/install.sh @@ -45,6 +45,7 @@ TAURI_MODE=false _USER_PYTHON="" _NO_TORCH_FLAG=false _VERBOSE=false +_SHORTCUTS_ONLY=false _next_is_package=false _next_is_python=false for arg in "$@"; do @@ -65,6 +66,7 @@ for arg in "$@"; do --python) _next_is_python=true ;; --no-torch) _NO_TORCH_FLAG=true ;; --verbose|-v) _VERBOSE=true ;; + --shortcuts-only) _SHORTCUTS_ONLY=true ;; esac done @@ -1233,6 +1235,20 @@ elif grep -qi microsoft /proc/version 2>/dev/null; then fi step "platform" "$OS" +# Regen launcher/shortcuts only; used by `unsloth studio update`. +if [ "$_SHORTCUTS_ONLY" = true ]; then + # Tauri owns its own shortcuts. + if [ "$TAURI_MODE" != true ]; then + VENV_ABS_BIN="$VENV_DIR/bin" + if [ ! -x "$VENV_ABS_BIN/unsloth" ]; then + echo "ERROR: unsloth binary missing at '$VENV_ABS_BIN/unsloth'; run install.sh first." >&2 + exit 1 + fi + create_studio_shortcuts "$VENV_ABS_BIN/unsloth" "$OS" + fi + exit 0 +fi + # ── Architecture detection & Python version ── _ARCH=$(uname -m) MAC_INTEL=false diff --git a/studio/setup.ps1 b/studio/setup.ps1 index 40788a0ecb..16df87bbd5 100644 --- a/studio/setup.ps1 +++ b/studio/setup.ps1 @@ -1842,6 +1842,21 @@ if ($CuTag -eq "cpu") { } } +# Rename running unsloth.exe so pip can replace it (Windows refuses to delete a mapped .exe). +$VenvScriptsDir = Join-Path $VenvDir "Scripts" +$RunningUnslothExe = Join-Path $VenvScriptsDir "unsloth.exe" +if (Test-Path -LiteralPath $RunningUnslothExe -PathType Leaf) { + $StaleUnslothExe = "$RunningUnslothExe.deleteme" + if (Test-Path -LiteralPath $StaleUnslothExe) { + Remove-Item -LiteralPath $StaleUnslothExe -Force -ErrorAction SilentlyContinue + } + try { + Rename-Item -LiteralPath $RunningUnslothExe -NewName "unsloth.exe.deleteme" -Force -ErrorAction Stop + } catch { + substep "could not rename unsloth.exe ($($_.Exception.Message)); pip may fail with WinError 32" "Yellow" + } +} + # Ordered heavy dependency installation -- shared cross-platform script substep "running ordered dependency installation..." python "$PSScriptRoot\install_python_stack.py" @@ -1851,6 +1866,28 @@ $ErrorActionPreference = $prevEAP if ($stackExit -ne 0) { Write-Host "[FAILED] Python dependency installation failed (exit code $stackExit)" -ForegroundColor Red Write-Host " Re-run the installer or check the error above for details." -ForegroundColor Red + # Restore the pre-rename unsloth.exe so the user keeps a working CLI. + # Treat a zero-byte exe as "pip half-wrote a broken binary" -- prefer the + # stale-but-working copy in .deleteme. + if (Test-Path -LiteralPath "$RunningUnslothExe.deleteme") { + $needRestore = -not (Test-Path -LiteralPath $RunningUnslothExe) + if (-not $needRestore) { + try { + $needRestore = (Get-Item -LiteralPath $RunningUnslothExe -ErrorAction Stop).Length -eq 0 + } catch { $needRestore = $true } + } + if ($needRestore) { + try { + if (Test-Path -LiteralPath $RunningUnslothExe) { + Remove-Item -LiteralPath $RunningUnslothExe -Force -ErrorAction SilentlyContinue + } + Rename-Item -LiteralPath "$RunningUnslothExe.deleteme" -NewName "unsloth.exe" -Force -ErrorAction Stop + substep "restored unsloth.exe after failed install" + } catch { + substep "could not restore unsloth.exe ($($_.Exception.Message))" "Yellow" + } + } + } exit 1 } diff --git a/studio/src-tauri/src/update.rs b/studio/src-tauri/src/update.rs index 390ea5878a..6136c5ea81 100644 --- a/studio/src-tauri/src/update.rs +++ b/studio/src-tauri/src/update.rs @@ -65,6 +65,10 @@ fn spawn_update( // the same install the desktop app uses, not an inherited custom root. cmd.env_remove("UNSLOTH_STUDIO_HOME"); cmd.env_remove("STUDIO_HOME"); + // Signal to unsloth_cli that this update was initiated by the Tauri + // desktop bundle so it skips re-creating CLI launchers/.app/.desktop + // shortcuts (Tauri owns its own bundle entries). + cmd.env("UNSLOTH_TAURI_UPDATE", "1"); #[cfg(windows)] let mut child: Box = { diff --git a/unsloth_cli/commands/studio.py b/unsloth_cli/commands/studio.py index 714ace4533..67395a8378 100644 --- a/unsloth_cli/commands/studio.py +++ b/unsloth_cli/commands/studio.py @@ -6,6 +6,7 @@ import hashlib import json import os import platform +import re import secrets import sqlite3 import subprocess @@ -13,6 +14,8 @@ import sys import tempfile import time import types +import urllib.error +import urllib.request from datetime import datetime, timezone from pathlib import Path from typing import List, Optional @@ -1064,6 +1067,170 @@ def _run_setup_script(*, verbose: bool = False) -> None: raise typer.Exit(result.returncode) +_INSTALLER_URL_BASH = "https://unsloth.ai/install.sh" +_INSTALLER_URL_PWSH = "https://unsloth.ai/install.ps1" + + +def _refresh_desktop_shortcuts(*, verbose: bool = False) -> None: + """Re-run installer with --shortcuts-only to refresh launchers post-update.""" + env = {**os.environ} + if verbose: + env["UNSLOTH_VERBOSE"] = "1" + + is_windows = platform.system() == "Windows" + installer_name = "install.ps1" if is_windows else "install.sh" + installer_url = _INSTALLER_URL_PWSH if is_windows else _INSTALLER_URL_BASH + + # Prefer local checkout, fall back to package dir, then network fetch. + local_repo = (os.environ.get("STUDIO_LOCAL_REPO") or "").strip() + candidates: list[Path] = [] + if local_repo: + candidates.append(Path(local_repo) / installer_name) + candidates.append(_PACKAGE_ROOT / installer_name) + + args = ["--shortcuts-only"] + if verbose: + args.append("--verbose") + + if is_windows: + ps_argv: list[str] = ["powershell.exe"] + if _should_hide_windows_subprocesses(): + ps_argv.extend( + ["-NoLogo", "-NoProfile", "-NonInteractive", "-WindowStyle", "Hidden"] + ) + + for script in candidates: + try: + if script.is_file(): + quoted = str(script).replace("'", "''") + argv = list(ps_argv) + argv.extend( + [ + "-ExecutionPolicy", + "Bypass", + "-Command", + f"& '{quoted}' {' '.join(args)} *>&1", + ] + ) + result = subprocess.run( + argv, + env = env, + check = False, + **_windows_hidden_subprocess_kwargs(), + ) + if result.returncode != 0: + typer.echo( + f" refresh-launcher install.ps1 exited {result.returncode}" + ) + return + except OSError: + continue + + # PyPI installs lack install.ps1: fetch + pipe to powershell stdin. + try: + request = urllib.request.Request( + installer_url, headers = {"User-Agent": "unsloth-studio-update"} + ) + with urllib.request.urlopen(request, timeout = 30) as response: + installer = response.read().decode("utf-8", errors = "replace") + except (urllib.error.URLError, TimeoutError, OSError) as exc: + typer.echo( + f" refresh-launcher skipped: could not fetch {installer_url} ({exc})" + ) + return + + # install.ps1 auto-invokes `Install-UnslothStudio @args` at EOF; over + # stdin `$args` is empty so that triggers the full installer flow + # (deps, venv, prompts) before our shortcuts-only call. Strip it. + installer = re.sub( + r"(?m)^[ \t]*Install-UnslothStudio[ \t]+@args[ \t]*\r?\n?", + "", + installer, + ) + # stdin-piped scripts have empty $args, so call Install-UnslothStudio explicitly. + marker_args = " ".join(args) + wrapper = installer + f"\nInstall-UnslothStudio {marker_args}\n" + + # Write to a UTF-8 BOM tempfile and use -File rather than -Command -. + # `powershell.exe -Command -` reads stdin via [Console]::InputEncoding + # (CP1252/OEM on most Windows boxes), which mangles box-drawing chars + # in install.ps1. -File reads the BOM and decodes correctly. The + # prefix gives AV/EDR engines (and grep'ing users) a clear identity. + ps1_fd, ps1_path = tempfile.mkstemp( + prefix = "unsloth-studio-refresh-", + suffix = ".ps1", + ) + try: + with os.fdopen(ps1_fd, "wb") as fh: + fh.write(b"\xef\xbb\xbf" + wrapper.encode("utf-8")) + argv = list(ps_argv) + argv.extend(["-ExecutionPolicy", "Bypass", "-File", ps1_path]) + try: + result = subprocess.run( + argv, + env = env, + check = False, + **_windows_hidden_subprocess_kwargs(), + ) + if result.returncode != 0: + typer.echo( + f" refresh-launcher fetched install.ps1 exited {result.returncode}" + ) + except OSError as exc: + typer.echo( + f" refresh-launcher skipped: powershell exec failed ({exc})" + ) + finally: + try: + os.unlink(ps1_path) + except OSError: + pass + return + + for script in candidates: + try: + if script.is_file(): + result = subprocess.run( + ["bash", str(script), *args], + env = env, + check = False, + ) + if result.returncode != 0: + typer.echo( + f" refresh-launcher install.sh exited {result.returncode}" + ) + return + except OSError: + continue + + # PyPI installs lack install.sh: fetch upstream. + try: + request = urllib.request.Request( + installer_url, headers = {"User-Agent": "unsloth-studio-update"} + ) + with urllib.request.urlopen(request, timeout = 30) as response: + installer = response.read() + except (urllib.error.URLError, TimeoutError, OSError) as exc: + typer.echo( + f" refresh-launcher skipped: could not fetch {installer_url} ({exc})" + ) + return + + try: + result = subprocess.run( + ["bash", "-s", "--", *args], + input = installer, + env = env, + check = False, + ) + if result.returncode != 0: + typer.echo( + f" refresh-launcher fetched install.sh exited {result.returncode}" + ) + except OSError as exc: + typer.echo(f" refresh-launcher skipped: bash exec failed ({exc})") + + @studio_app.command(hidden = True) def setup( verbose: bool = typer.Option( @@ -1093,6 +1260,9 @@ def update( ), ): """Update Unsloth Studio dependencies and rebuild.""" + # Re-export UNSLOTH_STUDIO_HOME for env-mode installs so the refresh + # subprocess resolves the same install root the user originally chose. + _ensure_studio_env_exported() # Ensure SKIP_STUDIO_BASE is not inherited from a parent install.ps1 session os.environ.pop("SKIP_STUDIO_BASE", None) os.environ["STUDIO_PACKAGE_NAME"] = package @@ -1105,7 +1275,88 @@ def update( else: os.environ["STUDIO_LOCAL_INSTALL"] = "0" os.environ.pop("STUDIO_LOCAL_REPO", None) - _run_setup_script(verbose = verbose) + _release_self_exe_lock_windows() + try: + _run_setup_script(verbose = verbose) + except BaseException: + # Restore unsloth.exe from .deleteme if setup failed before pip + # produced a replacement; otherwise the user has no CLI for recovery. + _restore_self_exe_lock_windows() + raise + # On Windows clear the .deleteme orphan now that pip wrote a fresh + # unsloth.exe; on next update os.replace would overwrite it anyway, + # but leaving a stale binary around invites cross-version restore + # confusion from _restore_self_exe_lock_windows. + _cleanup_self_exe_lock_windows() + # Tauri desktop owns its own bundle entries; skip CLI launcher refresh + # so a Tauri-initiated update doesn't create duplicate shortcuts. + if os.environ.get("UNSLOTH_TAURI_UPDATE") == "1": + if verbose: + typer.echo(" refresh-launcher skipped (Tauri update)") + return + _refresh_desktop_shortcuts(verbose = verbose) + + +def _release_self_exe_lock_windows() -> None: + """Rename running unsloth.exe so pip can replace it. setup.ps1 also retries.""" + if platform.system() != "Windows": + return + try: + venv_scripts = Path(sys.executable).resolve().parent + except OSError: + return + exe = venv_scripts / "unsloth.exe" + if not exe.exists(): + return + stale = exe.with_suffix(".exe.deleteme") + try: + # os.replace is atomic-overwrite on Windows; os.rename would raise + # FileExistsError if a prior aborted update left a .deleteme behind. + os.replace(exe, stale) + except OSError as e: + # Not fatal; setup.ps1 retries from a sibling process. + print(f"[update] could not rename {exe.name} -> {stale.name}: {e}") + + +def _restore_self_exe_lock_windows() -> None: + """If setup failed before pip wrote a working unsloth.exe, restore .deleteme.""" + if platform.system() != "Windows": + return + try: + venv_scripts = Path(sys.executable).resolve().parent + except OSError: + return + exe = venv_scripts / "unsloth.exe" + stale = exe.with_suffix(".exe.deleteme") + if not stale.exists(): + return + # Treat a missing or zero-byte exe as "pip didn't produce a usable + # replacement"; otherwise leave the new binary alone. + if exe.exists(): + try: + if exe.stat().st_size > 0: + return + except OSError: + return + try: + os.replace(stale, exe) + except OSError as e: + print(f"[update] could not restore {stale.name} -> {exe.name}: {e}") + + +def _cleanup_self_exe_lock_windows() -> None: + """Remove the .deleteme orphan after a successful update on Windows.""" + if platform.system() != "Windows": + return + try: + venv_scripts = Path(sys.executable).resolve().parent + except OSError: + return + stale = (venv_scripts / "unsloth.exe").with_suffix(".exe.deleteme") + try: + stale.unlink(missing_ok = True) + except OSError: + pass # ── unsloth studio reset-password ────────────────────────────────────