diff --git a/.github/workflows/studio-update-smoke.yml b/.github/workflows/studio-update-smoke.yml index 625c2c7811..047840e41c 100644 --- a/.github/workflows/studio-update-smoke.yml +++ b/.github/workflows/studio-update-smoke.yml @@ -146,6 +146,46 @@ jobs: kill "$PID" 2>/dev/null || true echo "post-update Unsloth /api/health OK" + - name: A complete install reports itself complete + run: | + set -o pipefail + unsloth studio verify-install + unsloth studio desktop-capabilities --json | tee /tmp/caps.json + jq -e '.studio_install_ok == true' /tmp/caps.json + jq -e '.desktop_manageability_version >= 2' /tmp/caps.json + + - name: An incomplete install must not report itself ready + # An installer killed part-way leaves a working CLI but no studio.txt + # deps, which the old preflight called ManagedReady. The manifest is + # written last, so removing it reproduces that state. + run: | + set -o pipefail + # install.sh's default root, resolved explicitly: `python` on PATH + # here is setup-python's, not the managed venv. + MANIFEST="$HOME/.unsloth/studio/unsloth_studio/unsloth_install_manifest.json" + test -f "$MANIFEST" || { echo "::error::installer never wrote $MANIFEST"; exit 1; } + rm -f "$MANIFEST" + unsloth studio desktop-capabilities --json | tee /tmp/caps_bad.json + jq -e '.studio_install_ok == false' /tmp/caps_bad.json + if unsloth studio verify-install; then + echo "::error::verify-install passed on an install with no manifest" + exit 1 + fi + echo "incomplete install correctly reported not-ready" + + - name: Update repairs an incomplete install + # `--local` bypasses setup.sh's PyPI version compare, so this asserts + # the repair OUTCOME. The non-local fast path the desktop Repair button + # uses is covered by tests/studio/install/test_setup_fast_path_guard.py. + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + run: | + set -o pipefail + unsloth studio update --local 2>&1 | tee logs/update_repair.log + unsloth studio verify-install + unsloth studio desktop-capabilities --json | jq -e '.studio_install_ok == true' + echo "update repaired the incomplete install" + - name: Uninstall and verify clean # Round-trip the installer through scripts/uninstall.sh: confirms the # uninstaller actually finds and removes everything install.sh + diff --git a/.github/workflows/wheel-smoke.yml b/.github/workflows/wheel-smoke.yml index cdad617027..f7a7511616 100644 --- a/.github/workflows/wheel-smoke.yml +++ b/.github/workflows/wheel-smoke.yml @@ -127,6 +127,31 @@ jobs: cd /tmp /tmp/v/bin/python -c "from studio.backend.main import app; print('Unsloth backend OK:', app.title)" + - name: CLI without the Studio stack guides instead of tracebacking + # The smoke above installs studio.txt first, so it cannot catch a wheel + # that ships studio/ without declaring what it imports (#4701, #5260, + # #7147). Drop only structlog to reuse that venv without a re-download. + run: | + set -eu + /tmp/v/bin/pip uninstall -y structlog >/dev/null + cd /tmp + status=0 + for args in "export ./nope ./out" "list-checkpoints"; do + echo "--- unsloth $args" + out=$(/tmp/v/bin/unsloth $args 2>&1 || true) + printf '%s\n' "$out" + case "$out" in + *Traceback*) + echo "FAIL: raw traceback instead of guidance"; status=1 ;; + esac + case "$out" in + *'unsloth studio update'*) ;; + *) echo "FAIL: no remediation in the message"; status=1 ;; + esac + done + /tmp/v/bin/pip install -q structlog >/dev/null + exit "$status" + - name: Upload wheel on failure if: failure() uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1 diff --git a/install.ps1 b/install.ps1 index a2aff0b69a..0b06cb3ea1 100644 --- a/install.ps1 +++ b/install.ps1 @@ -28,6 +28,14 @@ function Install-UnslothStudio { } } + function Clear-TauriInstallError { + param([string]$Message) + if ($TauriMode) { + Write-TauriLog "ERROR_CLEAR" $Message + [Console]::Error.WriteLine("[TAURI:ERROR_CLEAR] $Message") + } + } + function Format-TauriDiagBool { param([bool]$Value) if ($Value) { return "true" } @@ -86,7 +94,7 @@ function Install-UnslothStudio { [int]$Code = 1 ) if ($Code -eq 0) { $Code = 1 } - Write-TauriLog "ERROR" $Message + Write-TauriLog "ERROR_DEFAULT" $Message if (Get-Command Restore-StudioVenvRollback -CommandType Function -ErrorAction SilentlyContinue) { Restore-StudioVenvRollback } @@ -485,7 +493,8 @@ function Install-UnslothStudio { # Full command output is shown only when --verbose / UNSLOTH_VERBOSE=1. function Invoke-InstallCommand { param( - [Parameter(Mandatory = $true)][ScriptBlock]$Command + [Parameter(Mandatory = $true)][ScriptBlock]$Command, + [string]$Label = "install command" ) # Installer-pinned index installs (torch) must beat an inherited uv mirror (#6898): # for --default-index, clear the uv index env vars (restore in finally) and set @@ -504,6 +513,7 @@ function Install-UnslothStudio { try { # Reset to avoid stale values from prior native commands. $global:LASTEXITCODE = 0 + Write-TauriLog "OUTPUT_CLEAR" $Label if ($script:UnslothVerbose) { # Merge stderr into stdout so progress/warning output stays visible # without flipping $? on successful native commands (PS 5.1 treats @@ -518,7 +528,13 @@ function Install-UnslothStudio { Write-Host (Redact-InstallOutput $output) -ForegroundColor Red } } - return [int]$LASTEXITCODE + $exitCode = [int]$LASTEXITCODE + if ($exitCode -eq 0) { + Clear-TauriInstallError "$Label recovered" + } else { + Write-TauriLog "ERROR_OUTPUT" "$Label failed (exit code $exitCode)" + } + return $exitCode } finally { $ErrorActionPreference = $prevEap if ($savedUvIndex) { @@ -549,7 +565,7 @@ function Install-UnslothStudio { } $attempt = 1 while ($true) { - $code = Invoke-InstallCommand $Command + $code = Invoke-InstallCommand -Command $Command -Label $Label if ($code -eq 0) { return 0 } if ($attempt -ge $maxAttempts) { return $code } substep ("retrying ""$Label"" after transient failure (attempt $($attempt + 1)/$maxAttempts, waiting ${delay}s)...") "Yellow" @@ -1603,7 +1619,7 @@ exit 0 if (-not (Test-Path -LiteralPath $VenvPython)) { step "venv" "creating Python $($DetectedPython.Version) virtual environment" substep "$VenvDir" - $venvExit = Invoke-InstallCommand { uv venv $VenvDir --python "$($DetectedPython.Path)" } + $venvExit = Invoke-InstallCommand -Label "create virtual environment" { uv venv $VenvDir --python "$($DetectedPython.Path)" } if ($venvExit -ne 0) { Write-Host "[ERROR] Failed to create virtual environment (exit code $venvExit)" -ForegroundColor Red return (Exit-InstallFailure "Failed to create virtual environment (exit code $venvExit)" $venvExit) @@ -2375,7 +2391,7 @@ exit 0 } if ($StudioLocalInstall) { substep "overlaying local repo (editable)..." - $overlayExit = Invoke-InstallCommand { uv pip install --python $VenvPython -e $RepoRoot --no-deps } + $overlayExit = Invoke-InstallCommand -Label "overlay local repo" { uv pip install --python $VenvPython -e $RepoRoot --no-deps } if ($overlayExit -ne 0) { Write-Host "[ERROR] Failed to overlay local repo (exit code $overlayExit)" -ForegroundColor Red return (Exit-InstallFailure "Failed to overlay local repo (exit code $overlayExit)" $overlayExit) @@ -2464,7 +2480,7 @@ exit 0 if ($StudioLocalInstall) { substep "overlaying local repo (editable)..." - $overlayExit = Invoke-InstallCommand { uv pip install --python $VenvPython -e $RepoRoot --no-deps } + $overlayExit = Invoke-InstallCommand -Label "overlay local repo" { uv pip install --python $VenvPython -e $RepoRoot --no-deps } if ($overlayExit -ne 0) { Write-Host "[ERROR] Failed to overlay local repo (exit code $overlayExit)" -ForegroundColor Red return (Exit-InstallFailure "Failed to overlay local repo (exit code $overlayExit)" $overlayExit) @@ -2487,7 +2503,7 @@ exit 0 return (Exit-InstallFailure "Failed to install unsloth (exit code $baseInstallExit)" $baseInstallExit) } substep "overlaying local repo (editable)..." - $overlayExit = Invoke-InstallCommand { uv pip install --python $VenvPython -e $RepoRoot --no-deps } + $overlayExit = Invoke-InstallCommand -Label "overlay local repo" { uv pip install --python $VenvPython -e $RepoRoot --no-deps } if ($overlayExit -ne 0) { Write-Host "[ERROR] Failed to overlay local repo (exit code $overlayExit)" -ForegroundColor Red return (Exit-InstallFailure "Failed to overlay local repo (exit code $overlayExit)" $overlayExit) @@ -2535,7 +2551,7 @@ exit 0 $visionSpec = if ($PinnedRocmVisionSpec) { $PinnedRocmVisionSpec } elseif ($ROCmGfxArch -and $torchvisionFloorMap -and $torchvisionFloorMap.ContainsKey($ROCmGfxArch)) { $torchvisionFloorMap[$ROCmGfxArch] } else { "torchvision" } $audioSpec = if ($PinnedRocmAudioSpec) { $PinnedRocmAudioSpec } elseif ($ROCmGfxArch -and $torchaudioFloorMap -and $torchaudioFloorMap.ContainsKey($ROCmGfxArch)) { $torchaudioFloorMap[$ROCmGfxArch] } else { "torchaudio" } substep "PyTorch flavor mismatch (installed $installedTorchTag, need ROCm) -- reinstalling correct build..." "Yellow" - $torchFixExit = Invoke-InstallCommand { uv pip install --python $VenvPython --force-reinstall --default-index $ROCmIndexUrl $rocmSpec $visionSpec $audioSpec } + $torchFixExit = Invoke-InstallCommand -Label "reinstall PyTorch (ROCm)" { uv pip install --python $VenvPython --force-reinstall --default-index $ROCmIndexUrl $rocmSpec $visionSpec $audioSpec } if ($torchFixExit -ne 0) { Write-Host "[ERROR] Failed to reinstall PyTorch with the correct ROCm build (exit code $torchFixExit)" -ForegroundColor Red return (Exit-InstallFailure "Failed to reinstall PyTorch (ROCm) (exit code $torchFixExit)" $torchFixExit) @@ -2544,7 +2560,7 @@ exit 0 } elseif ($expectedTorchTag -ne 'rocm') { # CUDA: stale +cpu (or wrong cuXXX) against a CUDA index -> reinstall triplet. substep "PyTorch flavor mismatch (installed $installedTorchTag, need $expectedTorchTag) -- reinstalling correct build..." "Yellow" - $torchFixExit = Invoke-InstallCommand { uv pip install --python $VenvPython "torch>=2.4,<2.11.0" "torchvision>=0.19,<0.26.0" "torchaudio>=2.4,<2.11.0" --default-index $TorchIndexUrl --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio } + $torchFixExit = Invoke-InstallCommand -Label "reinstall PyTorch ($expectedTorchTag)" { uv pip install --python $VenvPython "torch>=2.4,<2.11.0" "torchvision>=0.19,<0.26.0" "torchaudio>=2.4,<2.11.0" --default-index $TorchIndexUrl --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio } if ($torchFixExit -ne 0) { Write-Host "[ERROR] Failed to reinstall PyTorch with the correct CUDA build (exit code $torchFixExit)" -ForegroundColor Red return (Exit-InstallFailure "Failed to reinstall PyTorch ($expectedTorchTag) (exit code $torchFixExit)" $torchFixExit) @@ -2645,6 +2661,9 @@ exit 0 # an inherited value would put llama.cpp in the wrong place. $previousUnslothStudioHome = $env:UNSLOTH_STUDIO_HOME $hadPreviousUnslothStudioHome = ($null -ne $previousUnslothStudioHome) + $previousTauriMode = $env:UNSLOTH_TAURI_MODE + $hadPreviousTauriMode = ($null -ne $previousTauriMode) + $env:UNSLOTH_TAURI_MODE = if ($TauriMode) { "1" } else { "0" } if ($StudioRedirectMode -eq 'env') { $env:UNSLOTH_STUDIO_HOME = $StudioHome } else { @@ -2674,14 +2693,22 @@ exit 0 } else { Remove-Item Env:UNSLOTH_STUDIO_HOME -ErrorAction SilentlyContinue } + if ($hadPreviousTauriMode) { + $env:UNSLOTH_TAURI_MODE = $previousTauriMode + } else { + Remove-Item Env:UNSLOTH_TAURI_MODE -ErrorAction SilentlyContinue + } Remove-Item Env:UNSLOTH_LOCAL_LLAMA_CPP_DIR -ErrorAction SilentlyContinue Remove-Item Env:UNSLOTH_INSTALL_ROLLBACK_MANAGED -ErrorAction SilentlyContinue Remove-Item Env:UNSLOTH_SETUP_PYTHON -ErrorAction SilentlyContinue } if ($setupExit -ne 0) { - Write-Host "[ERROR] unsloth studio setup failed (exit code $setupExit)" -ForegroundColor Red + if (-not $TauriMode) { + Write-Host "[ERROR] unsloth studio setup failed (exit code $setupExit)" -ForegroundColor Red + } return (Exit-InstallFailure "unsloth studio setup failed (exit code $setupExit)" $setupExit) } + Clear-TauriInstallError "studio setup completed" # ── Expose `unsloth` via a shim dir containing only unsloth.exe ── # We do NOT add the venv Scripts dir to PATH (it also holds python.exe diff --git a/install.sh b/install.sh index fece7b173b..376daa8fab 100755 --- a/install.sh +++ b/install.sh @@ -207,18 +207,37 @@ run_install_cmd() { # command's exit code across the pipe without relying on pipefail # (this script runs under plain sh). _rcf=$(mktemp) - { "$@" 2>&1; printf '%s' "$?" > "$_rcf"; } | _redact_install_output + tauri_stream_log stdout "OUTPUT_CLEAR" "$_label" + { + if "$@" 2>&1; then + _cmd_rc=0 + else + _cmd_rc=$? + fi + printf '%s' "$_cmd_rc" > "$_rcf" + } | _redact_install_output _rc=$(cat "$_rcf" 2>/dev/null || echo 1) rm -f "$_rcf" - [ "${_rc:-1}" -eq 0 ] 2>/dev/null && return 0 + _rc=${_rc:-1} + if [ "$_rc" -eq 0 ] 2>/dev/null; then + tauri_clear_install_error "$_label recovered" + return 0 + fi + tauri_stream_log stdout "ERROR_OUTPUT" "$_label failed (exit code $_rc)" step "error" "$_label failed (exit code $_rc)" "$C_ERR" >&2 return "$_rc" fi _log=$(mktemp) - "$@" >"$_log" 2>&1 && { rm -f "$_log"; return 0; } + tauri_stream_log stderr "OUTPUT_CLEAR" "$_label" + "$@" >"$_log" 2>&1 && { + rm -f "$_log" + tauri_clear_install_error "$_label recovered" + return 0 + } _rc=$? step "error" "$_label failed (exit code $_rc)" "$C_ERR" >&2 _redact_install_output "$_log" >&2 + tauri_stream_log stderr "ERROR_OUTPUT" "$_label failed (exit code $_rc)" rm -f "$_log" return $_rc } @@ -383,6 +402,34 @@ tauri_log() { fi } +tauri_stream_log() { + _tsl_stream="$1" + _tsl_tag="$2" + shift 2 + if [ "$TAURI_MODE" = true ]; then + if [ "$_tsl_stream" = stderr ]; then + printf '[TAURI:%s] %s\n' "$_tsl_tag" "$*" >&2 + else + printf '[TAURI:%s] %s\n' "$_tsl_tag" "$*" + fi + fi +} + +rollback_substep() { + if [ "$TAURI_MODE" = true ]; then + tauri_log "PROGRESS" "$1" + else + substep "$@" + fi +} + +tauri_clear_install_error() { + if [ "$TAURI_MODE" = true ]; then + tauri_log "ERROR_CLEAR" "$1" + printf '[TAURI:ERROR_CLEAR] %s\n' "$1" >&2 + fi +} + tauri_diag_marker() { _diag_gpu_branch="${1:-unknown}" _diag_torch_index_family="${2:-none}" @@ -543,10 +590,10 @@ _restore_studio_venv_replacement() { _VENV_ROLLBACK_ACTIVE=false return 0 } - substep "restoring previous environment after failed install..." "$C_WARN" + rollback_substep "restoring previous environment after failed install..." "$C_WARN" rm -rf "$_VENV_ROLLBACK_TARGET" if mv "$_VENV_ROLLBACK_DIR" "$_VENV_ROLLBACK_TARGET"; then - substep "restored previous environment" + rollback_substep "restored previous environment" _VENV_ROLLBACK_ACTIVE=false _VENV_ROLLBACK_DIR="" else @@ -4055,6 +4102,7 @@ if [ -n "$VENV_ABS_BIN" ]; then fi if ! command -v bash >/dev/null 2>&1; then + tauri_log "ERROR" "bash is required to run studio setup" step "setup" "bash is required to run studio setup" "$C_ERR" substep "Please install bash and re-run install.sh" exit 1 @@ -4093,6 +4141,7 @@ if [ "$STUDIO_LOCAL_INSTALL" = true ]; then STUDIO_LOCAL_REPO="$_REPO_ROOT" \ UNSLOTH_NO_TORCH="$SKIP_TORCH" \ UNSLOTH_LOCAL_LLAMA_CPP_DIR="$_WITH_LLAMA_CPP_DIR" \ + UNSLOTH_TAURI_MODE="$TAURI_MODE" \ bash "$SETUP_SH" =24.1.0", + # unsloth_cli/__init__.py reaches click via commands/start.py, so every + # command needs it. typer supplied it until 0.27 dropped the dependency. + "click>=8.0", ] [project.scripts] @@ -69,6 +75,33 @@ include = ["unsloth*", "unsloth_cli*", "studio", "studio.backend*"] exclude = ["images*", "tests*", "*.node_modules", "*.node_modules.*"] [project.optional-dependencies] +# Studio's server stack, mirroring studio/backend/requirements/studio.txt. +# test_studio_extra_matches_requirements.py catches drift. +studio = [ + "typer", + "fastapi", + "uvicorn", + "pydantic", + "packaging", + "matplotlib==3.10.9", + "pandas", + "nest_asyncio", + "datasets==4.3.0", + "pyjwt", + "huggingface-hub==0.36.2", + "structlog>=24.1.0", + "diceware", + "ddgs", + "cryptography>=42.0.0", + "boto3>=1.34.0", + "httpx>=0.27.0", + "fastmcp>=3.0.2", + "sqlite-vec==0.1.9", + "pymupdf==1.27.2.3", + "pymupdf4llm==0.3.4", + "python-docx==1.2.0", +] + triton = [ "triton>=3.0.0 ; ('linux' in sys_platform)", "triton-windows ; (sys_platform == 'win32') and (platform_machine == 'AMD64' or platform_machine == 'x86_64')", diff --git a/studio/backend/core/inference/api_monitor.py b/studio/backend/core/inference/api_monitor.py index ce32a6d3ef..b637ba56d1 100644 --- a/studio/backend/core/inference/api_monitor.py +++ b/studio/backend/core/inference/api_monitor.py @@ -5,6 +5,7 @@ from __future__ import annotations +import os import threading import time import uuid @@ -18,6 +19,14 @@ _MAX_PROMPT_CHARS = 12000 _MAX_REPLY_CHARS = 12000 _PREVIEW_CHARS = 360 +# Opt-in startup kill switch for Studio's in-memory API monitor. +_DISABLE_ENV = "UNSLOTH_STUDIO_DISABLE_API_MONITOR" +_TRUE_VALUES = frozenset({"1", "true", "yes", "on"}) + + +def _api_monitor_disabled() -> bool: + return os.environ.get(_DISABLE_ENV, "").strip().lower() in _TRUE_VALUES + def _trim(text: Optional[str], limit: int) -> str: if not text: @@ -104,10 +113,16 @@ class ApiMonitorEntry: class ApiMonitor: - def __init__(self, max_entries: int = _MAX_ENTRIES): + def __init__( + self, + max_entries: int = _MAX_ENTRIES, + *, + enabled: bool = True, + ): self._entries: deque[ApiMonitorEntry] = deque() self._max_entries = max(0, max_entries) self._lock = threading.Lock() + self._enabled = enabled def start( self, @@ -119,6 +134,8 @@ class ApiMonitor: context_length: Optional[int] = None, subject: Optional[str] = None, ) -> str: + if not self._enabled: + return "" now = time.time() entry = ApiMonitorEntry( id = f"apireq_{uuid.uuid4().hex[:12]}", @@ -152,6 +169,8 @@ class ApiMonitor: :meth:`fail`; an unload is terminal on arrival. Rows are shared (visible to every subject) and share the request retention budget. """ + if not self._enabled: + return "" now = time.time() entry = ApiMonitorEntry( id = f"apievt_{uuid.uuid4().hex[:12]}", @@ -392,4 +411,4 @@ class ApiMonitor: self._entries = kept -api_monitor = ApiMonitor() +api_monitor = ApiMonitor(enabled = not _api_monitor_disabled()) diff --git a/studio/backend/core/inference/chat_template_helpers.py b/studio/backend/core/inference/chat_template_helpers.py index 528c059fbc..3a8463855b 100644 --- a/studio/backend/core/inference/chat_template_helpers.py +++ b/studio/backend/core/inference/chat_template_helpers.py @@ -326,6 +326,58 @@ def _normalize_tool_call_arguments(messages: list) -> list: return out if mutated else messages +def _take_tool_result(pending: list, call_id) -> Optional[dict]: + if call_id: + for i, result in enumerate(pending): + if result.get("tool_call_id") == call_id: + return pending.pop(i) + for i, result in enumerate(pending): + if not result.get("tool_call_id"): + return pending.pop(i) + return None + + +def _split_parallel_tool_calls(messages: list) -> list: + """Llama 3.x templates render one call per message, so split parallel calls + into consecutive single-call messages, each followed by its own result.""" + if not any(isinstance(m, dict) and len(m.get("tool_calls") or ()) > 1 for m in messages): + return messages + + out: list = [] + i = 0 + total = len(messages) + while i < total: + msg = messages[i] + calls = msg.get("tool_calls") if isinstance(msg, dict) else None + if not calls or len(calls) <= 1: + out.append(msg) + i += 1 + continue + + # Tool results right after this message answer its calls. + j = i + 1 + pending: list = [] + while ( + j < total + and isinstance(messages[j], dict) + and messages[j].get("role") in ("tool", "ipython") + ): + pending.append(messages[j]) + j += 1 + + for idx, call in enumerate(calls): + piece = {**msg, "tool_calls": [call]} + if idx: + piece["content"] = "" + out.append(piece) + result = _take_tool_result(pending, call.get("id") if isinstance(call, dict) else None) + if result is not None: + out.append(result) + out.extend(pending) + i = j + return out + + def apply_chat_template_for_generation( tokenizer, messages: list, @@ -378,13 +430,21 @@ def apply_chat_template_for_generation( try: return _render(messages) except Exception: - # Strict tool templates reject the JSON-string ``arguments`` form via - # TypeError or a broad Jinja raise_exception, so retry with dicts coerced. - # Original messages render first, so working templates stay byte-identical. + # Retry with repairs applied cumulatively. Originals render first, so + # working templates stay byte-identical. + candidates: list = [] normalized = _normalize_tool_call_arguments(messages) - if normalized is messages: - raise - return _render(normalized) + if normalized is not messages: + candidates.append(normalized) + split = _split_parallel_tool_calls(normalized) + if split is not normalized: + candidates.append(split) + for candidate in candidates: + try: + return _render(candidate) + except Exception: + continue + raise def render_native_template( diff --git a/studio/backend/core/inference/worker.py b/studio/backend/core/inference/worker.py index 254eda40a3..3f32b3bd57 100644 --- a/studio/backend/core/inference/worker.py +++ b/studio/backend/core/inference/worker.py @@ -25,7 +25,7 @@ from pathlib import Path from typing import Any logger = get_logger(__name__) -from utils.hardware import apply_gpu_ids +from utils.hardware import apply_gpu_ids, is_apple_silicon _SHARE_OBJECT_MAX_BYTES = 1 << 20 _SHARE_OBJECT_ERROR_SIZE = -1 @@ -801,10 +801,7 @@ def run_inference_process( # ── 0. MLX fast-path — skip torch/transformers ── _ensure_backend_on_path() - from utils.hardware import hardware as _hw - - _hw.detect_hardware() - if _hw.DEVICE == _hw.DeviceType.MLX: + if is_apple_silicon(): # Non-fatal: fall through with the installed version, but log the cause # instead of swallowing it (issue #6103). try: @@ -816,6 +813,11 @@ def run_inference_process( model_name, exc, ) + + from utils.hardware import hardware as _hw + + _hw.detect_hardware() + if _hw.DEVICE == _hw.DeviceType.MLX: try: from core.inference.mlx_inference import MLXInferenceBackend, _init_mlx_distributed diff --git a/studio/backend/loggers/handlers.py b/studio/backend/loggers/handlers.py index 716c4f40d2..5d99ca85c6 100644 --- a/studio/backend/loggers/handlers.py +++ b/studio/backend/loggers/handlers.py @@ -8,12 +8,19 @@ filter_sensitive_data (structlog processor for sanitization), and get_logger (factory for structured loggers). """ +from __future__ import annotations + import os import re import time +from typing import TYPE_CHECKING import structlog -from starlette.types import ASGIApp, Message, Receive, Scope, Send + +# Annotations only: a runtime import makes the ASGI stack a hard dependency of +# every CLI command. +if TYPE_CHECKING: + from starlette.types import ASGIApp, Message, Receive, Scope, Send from utils.native_path_leases import redact_native_paths diff --git a/studio/backend/main.py b/studio/backend/main.py index ff09ccb36a..02f5a20106 100644 --- a/studio/backend/main.py +++ b/studio/backend/main.py @@ -1075,7 +1075,9 @@ async def liveness_check(): "status": "alive", "service": "Unsloth UI Backend", "desktop_protocol_version": 1, - "desktop_manageability_version": 1, + # Lockstep with DESKTOP_MANAGEABILITY_VERSION in + # studio/src-tauri/src/preflight/version.rs and `desktop-capabilities`. + "desktop_manageability_version": 2, "supports_desktop_auth": True, "supports_desktop_backend_ownership": True, "studio_root_id": _studio_root_id(), @@ -1098,7 +1100,8 @@ async def health_check(request: Request): "service": "Unsloth UI Backend", "chat_only": _hw_module.CHAT_ONLY, "desktop_protocol_version": 1, - "desktop_manageability_version": 1, + # Lockstep: see the note in /api/liveness above. + "desktop_manageability_version": 2, "supports_desktop_auth": True, "supports_desktop_backend_ownership": True, # Opaque per-install id; launchers reject sibling Studios on the same port. diff --git a/studio/backend/tests/test_api_monitor.py b/studio/backend/tests/test_api_monitor.py index 7dd4baa2dd..4602b5cf62 100644 --- a/studio/backend/tests/test_api_monitor.py +++ b/studio/backend/tests/test_api_monitor.py @@ -260,6 +260,63 @@ def test_api_monitor_append_reply_exact_cap_then_more_marks_truncated(): assert len(reply) == m._MAX_REPLY_CHARS and reply.endswith("...") +def test_api_monitor_disabled_is_noop(): + monitor = ApiMonitor(max_entries = 3, enabled = False) + + request_id = monitor.start( + endpoint = "/v1/chat/completions", + method = "POST", + model = "local-model", + prompt = "user: hello", + context_length = 100, + ) + load_id = monitor.record_lifecycle( + event = "load", + model = "local-model", + running = True, + ) + unload_id = monitor.record_lifecycle( + event = "unload", + model = "local-model", + ) + assert request_id == load_id == unload_id == "" + + # Every mutator must be a safe no-op on the falsy id. + monitor.append_reply(request_id, "hi") + monitor.set_reply(request_id, "hi") + monitor.set_usage(request_id, prompt_tokens = 4, completion_tokens = 6) + monitor.relabel(load_id, "renamed-model") + monitor.set_progress(load_id, 50) + monitor.finish(load_id) + monitor.fail_open(load_id, "boom") + monitor.fail(request_id, "boom") + monitor.discard(unload_id) + + assert monitor.snapshot() == [] + assert monitor.active_count() == 0 + assert monitor.get(request_id) is None + + +def test_api_monitor_disable_env_var_truthy(monkeypatch): + import core.inference.api_monitor as m + for value in ("1", "true", "yes", "on", "TRUE", "On", " yes "): + monkeypatch.setenv(m._DISABLE_ENV, value) + assert m._api_monitor_disabled() is True, value + + +def test_api_monitor_disable_env_var_falsy(monkeypatch): + import core.inference.api_monitor as m + for value in ("", "0", "false", "no", "off", "disabled"): + monkeypatch.setenv(m._DISABLE_ENV, value) + assert m._api_monitor_disabled() is False, value + + +def test_api_monitor_disable_env_var_unset(monkeypatch): + import core.inference.api_monitor as m + monkeypatch.delenv(m._DISABLE_ENV, raising = False) + assert m._api_monitor_disabled() is False + + # ── model lifecycle rows (load / unload) ──────────────────────────── diff --git a/studio/backend/tests/test_chat_template_tool_arguments.py b/studio/backend/tests/test_chat_template_tool_arguments.py index 13d1ecabaa..8a927ea93c 100644 --- a/studio/backend/tests/test_chat_template_tool_arguments.py +++ b/studio/backend/tests/test_chat_template_tool_arguments.py @@ -6,10 +6,14 @@ from the OpenAI JSON-string form to a dict before rendering. Strict tool templates (e.g. mlx-community Qwen3.5 checkpoints) iterate arguments.items() and raise "Can only get item pairs from a mapping." on the string form when a prior tool call is re-rendered on the next turn (MLX + transformers paths). + +It must likewise split parallel tool calls for templates that render only one +call per message (Llama 3.x). """ from __future__ import annotations +import json import sys from pathlib import Path @@ -21,6 +25,7 @@ if str(_BACKEND) not in sys.path: from core.inference.chat_template_helpers import ( # noqa: E402 _normalize_tool_call_arguments, + _split_parallel_tool_calls, apply_chat_template_for_generation, ) @@ -155,3 +160,152 @@ def test_unrelated_template_error_still_propagates_with_dict_args(): with pytest.raises(ValueError, match = "broken"): apply_chat_template_for_generation(_AlwaysRaises(), _conv({"query": "x"})) + + +def _parallel_conv( + *, + ids = ("c1", "c2"), + results_have_ids = True, + content = "sure", +): + a, b = ids + return [ + {"role": "user", "content": "search then render"}, + { + "role": "assistant", + "content": content, + "tool_calls": [ + { + "type": "function", + "id": a, + "function": {"name": "web_search", "arguments": {"query": "x"}}, + }, + { + "type": "function", + "id": b, + "function": {"name": "render_html", "arguments": {"html": ""}}, + }, + ], + }, + { + "role": "tool", + "name": "web_search", + **({"tool_call_id": a} if results_have_ids else {}), + "content": "no text", + }, + { + "role": "tool", + "name": "render_html", + **({"tool_call_id": b} if results_have_ids else {}), + "content": "ok", + }, + ] + + +class _SingleToolCallTokenizer: + """Mimics the Llama 3.x template: rejects >1 call per message.""" + + def apply_chat_template( + self, + messages, + *, + tokenize = False, + add_generation_prompt = True, + **kw, + ): + for msg in messages: + if len(msg.get("tool_calls") or ()) > 1: + raise ValueError("This model only supports single tool-calls at once!") + return "RENDERED" + + +def test_parallel_calls_split_into_sequential_single_call_turns(): + out = _split_parallel_tool_calls(_parallel_conv()) + assert [(m["role"], m.get("name")) for m in out] == [ + ("user", None), + ("assistant", None), + ("tool", "web_search"), + ("assistant", None), + ("tool", "render_html"), + ] + assert [len(m["tool_calls"]) for m in out if m.get("tool_calls")] == [1, 1] + assert out[1]["tool_calls"][0]["function"]["name"] == "web_search" + assert out[3]["tool_calls"][0]["function"]["name"] == "render_html" + + +def test_split_pairs_results_by_tool_call_id_not_position(): + conv = _parallel_conv() + conv[2], conv[3] = conv[3], conv[2] # results arrive out of order + out = _split_parallel_tool_calls(conv) + assert out[1]["tool_calls"][0]["id"] == "c1" and out[2]["tool_call_id"] == "c1" + assert out[3]["tool_calls"][0]["id"] == "c2" and out[4]["tool_call_id"] == "c2" + + +def test_split_falls_back_to_order_when_results_have_no_ids(): + out = _split_parallel_tool_calls(_parallel_conv(results_have_ids = False)) + assert [m["role"] for m in out] == ["user", "assistant", "tool", "assistant", "tool"] + assert out[2]["name"] == "web_search" and out[4]["name"] == "render_html" + + +def test_split_keeps_content_on_first_piece_only(): + out = _split_parallel_tool_calls(_parallel_conv(content = "sure")) + assert out[1]["content"] == "sure" + assert out[3]["content"] == "" + + +def test_split_keeps_unmatched_results_after_the_split(): + conv = _parallel_conv() + del conv[3] # second call never returned a result + out = _split_parallel_tool_calls(conv) + assert [m["role"] for m in out] == ["user", "assistant", "tool", "assistant"] + + +def test_split_leaves_later_turns_intact(): + conv = _parallel_conv() + [ + {"role": "assistant", "content": "done"}, + {"role": "user", "content": "thanks"}, + ] + out = _split_parallel_tool_calls(conv) + assert [m["role"] for m in out[-2:]] == ["assistant", "user"] + assert out[-2]["content"] == "done" + + +def test_single_call_and_plain_conversations_pass_through_unchanged(): + conv = _conv({"query": "x"}) + assert _split_parallel_tool_calls(conv) is conv + plain = [{"role": "user", "content": "hi"}] + assert _split_parallel_tool_calls(plain) is plain + + +def test_render_succeeds_on_single_call_template_with_parallel_calls(): + # Regression: two calls in one turn used to break every later render. + result = apply_chat_template_for_generation(_SingleToolCallTokenizer(), _parallel_conv()) + assert result == "RENDERED" + + +def test_string_arguments_and_parallel_calls_are_repaired_together(): + conv = _parallel_conv() + for call in conv[1]["tool_calls"]: + call["function"]["arguments"] = json.dumps(call["function"]["arguments"]) + + class _StrictAndSingleCall(_SingleToolCallTokenizer): + def apply_chat_template(self, messages, **kw): + for msg in messages: + for call in msg.get("tool_calls", []) or []: + if isinstance(call.get("function", {}).get("arguments"), str): + raise TypeError("Can only get item pairs from a mapping.") + return super().apply_chat_template(messages, **kw) + + assert apply_chat_template_for_generation(_StrictAndSingleCall(), conv) == "RENDERED" + + +def test_lenient_template_never_sees_a_split_conversation(): + seen = {} + + class _Lenient: + def apply_chat_template(self, messages, **kw): + seen["n"] = len(messages) + return "RENDERED" + + apply_chat_template_for_generation(_Lenient(), _parallel_conv()) + assert seen["n"] == 4 # unsplit diff --git a/studio/backend/tests/test_mlx_inference_backend.py b/studio/backend/tests/test_mlx_inference_backend.py index d49a2281a0..3d20dd4bcc 100644 --- a/studio/backend/tests/test_mlx_inference_backend.py +++ b/studio/backend/tests/test_mlx_inference_backend.py @@ -1,8 +1,11 @@ # SPDX-License-Identifier: AGPL-3.0-only +import json +import subprocess import sys import types from contextlib import contextmanager +from pathlib import Path from types import SimpleNamespace import pytest @@ -376,6 +379,128 @@ def test_worker_share_object_receives_distributed_payload(monkeypatch): assert response["object"] == shared_obj +def test_worker_activates_mlx_sidecar_before_hardware_detection(tmp_path): + backend_dir = Path(__file__).resolve().parent.parent + fake_modules = tmp_path / "base" + sidecar = tmp_path / ".venv_t5_530" + packages = { + fake_modules / "transformers" / "__init__.py": '__version__ = "4.57.6"\n', + fake_modules / "mlx" / "__init__.py": "", + fake_modules / "mlx" / "core.py": "", + fake_modules / "mlx_lm" / "__init__.py": "import transformers\n", + fake_modules / "mlx_lm" / "sample_utils.py": "", + fake_modules / "mlx_vlm" / "__init__.py": "", + sidecar / "transformers" / "__init__.py": '__version__ = "5.3.0"\n', + } + for path, contents in packages.items(): + path.parent.mkdir(parents = True, exist_ok = True) + path.write_text(contents) + + script = r""" +import json +import os +import sys + +sys.path.insert(0, os.environ["FAKE_MODULES"]) +from core.inference import worker +from utils.hardware import hardware +import utils.mlx_repair as mlx_repair +import utils.transformers_version as transformers_version + +bootstrap_roots = sorted( + { + name.split(".", 1)[0] + for name in sys.modules + if name.split(".", 1)[0] + in { + "huggingface_hub", + "mlx", + "mlx_lm", + "mlx_vlm", + "torch", + "transformers", + "unsloth", + "unsloth_zoo", + } + } +) +assert not bootstrap_roots, f"worker bootstrap imported ML modules: {bootstrap_roots}" + +worker.is_apple_silicon = lambda: True +hardware.is_apple_silicon = lambda: True +hardware._has_torch = lambda: False +mlx_repair._mlx_versions_satisfy_minimums = lambda: True +transformers_version._VENV_T5_530_DIR = os.environ["SIDECAR"] +transformers_version._ensure_venv_t5_530_exists = lambda: True + +observed = {"bootstrap_roots": bootstrap_roots} + +def capture_active_version(_backend, _config, _responses): + module = sys.modules["transformers"] + observed["active"] = module.__version__ + observed["file"] = module.__file__ + observed["device"] = hardware.DEVICE.value + +class CommandQueue: + def get(self, timeout): + return {"type": "shutdown"} + +class ResponseQueue: + def put(self, _response): + pass + +worker._handle_load = capture_active_version +worker.run_inference_process( + cmd_queue = CommandQueue(), + resp_queue = ResponseQueue(), + cancel_event = None, + config = { + "model_name": "Ministral-3-regression", + "hf_token": "", + "resolved_gpu_ids": None, + "device_backend": "mlx", + }, +) +observed["tier"] = transformers_version.get_transformers_tier( + "Ministral-3-regression" +) +print("RESULT " + json.dumps(observed, sort_keys = True)) +""" + result = subprocess.run( + [sys.executable, "-c", script], + cwd = backend_dir, + env = { + **__import__("os").environ, + "FAKE_MODULES": str(fake_modules), + "SIDECAR": str(sidecar), + "UNSLOTH_STUDIO_HOME": str(tmp_path), + "HF_HOME": str(tmp_path / "hf"), + "HF_HUB_CACHE": str(tmp_path / "hf" / "hub"), + "HF_HUB_OFFLINE": "1", + "TRANSFORMERS_OFFLINE": "1", + }, + capture_output = True, + text = True, + ) + + assert result.returncode == 0, result.stdout + result.stderr + result_line = next( + ( + line.removeprefix("RESULT ") + for line in result.stdout.splitlines() + if line.startswith("RESULT ") + ), + None, + ) + assert result_line is not None, result.stdout + result.stderr + observed = json.loads(result_line) + assert observed["bootstrap_roots"] == [] + assert observed["tier"] == "530" + assert observed["device"] == "mlx" + assert observed["active"] == "5.3.0" + assert observed["file"] == str(sidecar / "transformers" / "__init__.py") + + def test_worker_share_object_oversize_notifies_peers(monkeypatch): from core.inference import worker diff --git a/studio/frontend/src/components/ui/sonner.tsx b/studio/frontend/src/components/ui/sonner.tsx index aec1235b81..4c65edc4b7 100644 --- a/studio/frontend/src/components/ui/sonner.tsx +++ b/studio/frontend/src/components/ui/sonner.tsx @@ -8,8 +8,8 @@ import { MultiplicationSignCircleIcon, } from "@hugeicons/core-free-icons"; import { HugeiconsIcon } from "@hugeicons/react"; -import { Spinner } from "@/components/ui/spinner"; import { useTheme } from "@/features/settings/stores/theme-store"; +import { createLoadingToastIcon } from "@/lib/toast"; import { Toaster as Sonner, type ToasterProps } from "sonner"; // Make toast text selectable. Sonner's onPointerDown calls setPointerCapture(), @@ -78,7 +78,7 @@ const Toaster = ({ ...props }: ToasterProps) => { /> ), // App-wide arc spinner so loading toasts match the "Downloading model" toast. - loading: , + loading: createLoadingToastIcon(), }} style={ { diff --git a/studio/frontend/src/features/chat/api/chat-adapter.ts b/studio/frontend/src/features/chat/api/chat-adapter.ts index fb9331ecc2..a0be3ea640 100644 --- a/studio/frontend/src/features/chat/api/chat-adapter.ts +++ b/studio/frontend/src/features/chat/api/chat-adapter.ts @@ -6,7 +6,7 @@ import { resolveInitialConfig } from "@/features/model-picker"; import { projectHasSources } from "@/features/rag/api/rag-api"; import { apiUrl } from "@/lib/api-base"; import { parseParamCountB } from "@/lib/model-size"; -import { toast } from "@/lib/toast"; +import { createLoadingToastIcon, toast } from "@/lib/toast"; import type { MessageTiming, ToolCallMessagePart } from "@assistant-ui/core"; import type { ChatModelAdapter } from "@assistant-ui/react"; import { parsePartialJsonObject } from "assistant-stream/utils"; @@ -1512,13 +1512,38 @@ async function autoLoadSmallestModel(): Promise<{ const trustRemoteCode = store.params.trustRemoteCode ?? false; const specSettings = resolveSpeculativeSettingsForLoad(); const lastLoaded = readLastLocalModelLoad(); - const toastId = toast("Loading a model…", { + let autoLoadToastDismissed = false; + const toastId = toast.message("Loading a model…", { description: lastLoaded ? "Loading last used model." : "Auto-selecting the smallest downloaded model.", - duration: 5000, + duration: Number.POSITIVE_INFINITY, closeButton: true, + icon: createLoadingToastIcon(), + onDismiss: () => { + autoLoadToastDismissed = true; + }, }); + const updateAutoLoadToast = (message: string, description: string): void => { + if (autoLoadToastDismissed) return; + toast.message(message, { + id: toastId, + description, + duration: Number.POSITIVE_INFINITY, + }); + }; + const showAutoLoadSuccess = (message: string): void => { + const options = { + description: undefined, + duration: 5000, + icon: undefined, + }; + if (autoLoadToastDismissed) { + toast.success(message, options); + return; + } + toast.success(message, { ...options, id: toastId }); + }; let blockedByTrustRemoteCode = false; let hadNonTrustFailure = false; let loadAttempts = 0; @@ -1774,7 +1799,7 @@ async function autoLoadSmallestModel(): Promise<{ ggufVariant: candidate.ggufVariant, }); } - toast.success(candidate.successLabel, { id: toastId }); + showAutoLoadSuccess(candidate.successLabel); return true; } try { @@ -1800,11 +1825,10 @@ async function autoLoadSmallestModel(): Promise<{ isAutoLoadableGgufVariant(entry), ); if (variant) { - toast("Loading last used model…", { - id: toastId, - description: `${repo.repo_id} (${variant.quant})`, - duration: 5000, - }); + updateAutoLoadToast( + "Loading last used model…", + `${repo.repo_id} (${variant.quant})`, + ); if ( await loadAutoLoadCandidate({ id: repo.repo_id, @@ -1829,11 +1853,7 @@ async function autoLoadSmallestModel(): Promise<{ const repo = findCachedRepo(modelRepos, lastLoaded.id); if (repo) { try { - toast("Loading last used model…", { - id: toastId, - description: repo.repo_id, - duration: 5000, - }); + updateAutoLoadToast("Loading last used model…", repo.repo_id); if ( await loadAutoLoadCandidate({ id: repo.repo_id, @@ -1854,11 +1874,10 @@ async function autoLoadSmallestModel(): Promise<{ } } } - toast("Loading a model…", { - id: toastId, - description: "Auto-selecting the smallest downloaded model.", - duration: 5000, - }); + updateAutoLoadToast( + "Loading a model…", + "Auto-selecting the smallest downloaded model.", + ); } // GGUF first: smallest-total-size repo, then its smallest variant. @@ -1949,12 +1968,10 @@ async function autoLoadSmallestModel(): Promise<{ } // No cached models — try downloading a small default GGUF. - toast("Downloading a small model…", { - id: toastId, - description: - "No downloaded models found. Fetching Qwen3.5-4B-MTP (UD-Q4_K_XL).", - duration: 30000, - }); + updateAutoLoadToast( + "Downloading a small model…", + "No downloaded models found. Fetching Qwen3.5-4B-MTP (UD-Q4_K_XL).", + ); try { const rt = useChatRuntimeStore.getState(); if ( @@ -2050,7 +2067,7 @@ async function autoLoadSmallestModel(): Promise<{ kind: "gguf", ggufVariant: "UD-Q4_K_XL", }); - toast.success("Loaded Qwen3.5-4B-MTP (UD-Q4_K_XL)", { id: toastId }); + showAutoLoadSuccess("Loaded Qwen3.5-4B-MTP (UD-Q4_K_XL)"); return { loaded: true, blockedByTrustRemoteCode: false }; } catch { toast.dismiss(toastId); diff --git a/studio/frontend/src/features/recipe-studio/hooks/use-recipe-executions.ts b/studio/frontend/src/features/recipe-studio/hooks/use-recipe-executions.ts index a067855fd5..bc4bc2a391 100644 --- a/studio/frontend/src/features/recipe-studio/hooks/use-recipe-executions.ts +++ b/studio/frontend/src/features/recipe-studio/hooks/use-recipe-executions.ts @@ -2,7 +2,7 @@ // Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 import { getInferenceStatus, loadModel } from "@/features/chat"; -import { toast } from "@/lib/toast"; +import { createLoadingToastIcon, toast } from "@/lib/toast"; import { toastError } from "@/shared/toast"; import { useCallback, useEffect, useState } from "react"; import { useShallow } from "zustand/react/shallow"; @@ -238,8 +238,15 @@ async function loadLocalModelSelection( ): Promise { const { target, ggufVariant } = selection; const modelLabel = ggufVariant ? `${target} (${ggufVariant})` : target; - const toastId = toast.loading(`Loading ${modelLabel}...`, { + let loadToastDismissed = false; + const toastId = toast.message(`Loading ${modelLabel}...`, { description: "Starting the local inference server for this recipe.", + duration: Number.POSITIVE_INFINITY, + closeButton: true, + icon: createLoadingToastIcon(), + onDismiss: () => { + loadToastDismissed = true; + }, }); try { const isGguf = GGUF_MODEL_PATTERN.test(target) || Boolean(ggufVariant); @@ -267,7 +274,16 @@ async function loadLocalModelSelection( // biome-ignore lint/style/useNamingConvention: api schema tensor_parallel: false, }); - toast.success(`Loaded ${modelLabel}`, { id: toastId, duration: 2000 }); + const successOptions = { + description: undefined, + duration: 2000, + icon: undefined, + }; + if (loadToastDismissed) { + toast.success(`Loaded ${modelLabel}`, successOptions); + } else { + toast.success(`Loaded ${modelLabel}`, { ...successOptions, id: toastId }); + } return null; } catch (error) { toast.dismiss(toastId); diff --git a/studio/frontend/src/lib/toast.ts b/studio/frontend/src/lib/toast.ts index 6b1635b42e..b500c11ec2 100644 --- a/studio/frontend/src/lib/toast.ts +++ b/studio/frontend/src/lib/toast.ts @@ -4,5 +4,15 @@ // Re-export of sonner. Swipe blocking lives on the Toaster via // `swipeDirections={[]}`, so no per-toast dismissible override. +import { Spinner } from "@/components/ui/spinner"; +import { createElement } from "react"; + +function createLoadingToastIcon() { + return createElement(Spinner, { + className: "size-4 text-muted-foreground", + }); +} + export { toast } from "sonner"; export type { ExternalToast } from "sonner"; +export { createLoadingToastIcon }; diff --git a/studio/install_manifest.py b/studio/install_manifest.py new file mode 100644 index 0000000000..8f48dcf35d --- /dev/null +++ b/studio/install_manifest.py @@ -0,0 +1,305 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Install-completeness manifest for Unsloth Studio. + +install_python_stack.py drops the manifest before the dependency pass and writes +it back only after the last step, so its presence means "the install finished". +Read by `unsloth studio verify-install`, `desktop-capabilities` (and through it +the Tauri preflight) and setup.sh/setup.ps1's fast path. + +Without it an installer killed part-way leaves a venv with `unsloth` but not +studio.txt's dependencies, which still answers `-h` and so looked ready right up +until the backend died on `import structlog`. + +Must import inside that half-installed venv: stdlib only, `packaging` optional. +""" + +from __future__ import annotations + +import hashlib +import json +import os +import platform +import re +import sys +import time +from pathlib import Path +from typing import Dict, List, Optional, Tuple + +MANIFEST_NAME = "unsloth_install_manifest.json" +MANIFEST_SCHEMA = 1 + +# Fingerprinted into the manifest, relative to studio/backend/requirements/. +# Editing one (a --local install) invalidates it and forces a dependency pass. +TRACKED_REQUIREMENT_FILES: Tuple[str, ...] = ( + "studio.txt", + "base.txt", + "extras.txt", + "extras-no-deps.txt", + "no-torch-runtime.txt", + "single-env/data-designer-deps.txt", + "single-env/data-designer.txt", +) + +# The import chain studio/backend/run.py walks on startup. +BOOT_REQUIREMENT_FILE = "studio.txt" + + +def venv_root() -> Path: + """Directory holding pyvenv.cfg for the interpreter running this code.""" + return Path(sys.prefix) + + +def manifest_path(root: Optional[Path] = None) -> Path: + return (root or venv_root()) / MANIFEST_NAME + + +def requirements_root(script_dir: Optional[Path] = None) -> Path: + """studio/backend/requirements/ next to this module (or a given studio/ dir).""" + return (script_dir or Path(__file__).resolve().parent) / "backend" / "requirements" + + +def _sha256(path: Path) -> Optional[str]: + try: + return hashlib.sha256(path.read_bytes()).hexdigest() + except OSError: + return None + + +def requirement_digests(req_root: Optional[Path] = None) -> Dict[str, str]: + """sha256 of every tracked requirement file that exists.""" + root = req_root or requirements_root() + digests: Dict[str, str] = {} + for name in TRACKED_REQUIREMENT_FILES: + digest = _sha256(root / name) + if digest is not None: + digests[name] = digest + return digests + + +def _canonical(name: str) -> str: + """PEP 503 normalisation, so PyJWT / pyjwt / py_jwt compare equal.""" + return re.sub(r"[-_.]+", "-", name).lower() + + +def _installed_version(dist_name: str, installed: Optional[Dict[str, str]] = None) -> Optional[str]: + if installed is not None: + return installed.get(_canonical(dist_name)) + from importlib.metadata import PackageNotFoundError, version + try: + return version(dist_name) + except PackageNotFoundError: + return None + except Exception: + return None + + +def remove_manifest(root: Optional[Path] = None) -> bool: + """Called before the dependency pass so an aborted run cannot leave a valid one. + + True when no manifest remains. A surviving marker (Windows raises on a + read-only or locked file) still names this version and these digests, so a + pass killed afterwards would verify as complete. + """ + try: + manifest_path(root).unlink() + except FileNotFoundError: + return True + except OSError: + return False + return True + + +def write_manifest( + root: Optional[Path] = None, + req_root: Optional[Path] = None, + steps_total: int = 0, + package_name: str = "unsloth", +) -> Optional[Path]: + """Record a completed install. Never raises: no manifest reads as incomplete, + which is the safe answer.""" + payload = { + "schema": MANIFEST_SCHEMA, + "completed_at_ms": int(time.time() * 1000), + "package": package_name, + "package_version": _installed_version(package_name), + "python": platform.python_version(), + "platform": f"{sys.platform}-{platform.machine()}", + "prefix": str(venv_root()), + "steps_total": steps_total, + "requirement_files": requirement_digests(req_root), + } + path = manifest_path(root) + try: + tmp = path.with_suffix(".json.tmp") + tmp.write_text(json.dumps(payload, indent = 2, sort_keys = True), encoding = "utf-8") + os.replace(tmp, path) + return path + except OSError: + return None + + +def read_manifest(root: Optional[Path] = None) -> Optional[dict]: + try: + raw = manifest_path(root).read_text(encoding = "utf-8") + except OSError: + return None + try: + data = json.loads(raw) + except ValueError: + return None + return data if isinstance(data, dict) else None + + +def _parse_requirement_line(line: str) -> Optional[Tuple[str, str, str]]: + """(distribution name, marker, specifier) for a requirement, or None. + + Covers what studio.txt uses: names, specifiers, inline comments, markers. + pip flags are skipped. + """ + text = line.split("#", 1)[0].strip() + if not text or text.startswith("-"): + return None + try: + from packaging.requirements import Requirement + requirement = Requirement(text) + return ( + requirement.name, + str(requirement.marker or ""), + str(requirement.specifier), + ) + except Exception: + pass + marker = "" + if ";" in text: + text, marker = text.split(";", 1) + marker = marker.strip() + name = text.strip() + for sep in ("===", "==", ">=", "<=", "~=", "!=", ">", "<", "[", " "): + idx = name.find(sep) + if idx > 0: + name = name[:idx] + name = name.strip() + return (name, marker, "") if name else None + + +def _marker_applies(marker: str) -> bool: + """True when the environment marker matches (or cannot be evaluated).""" + if not marker: + return True + try: + from packaging.markers import Marker + except Exception: + # No packaging: assume it applies. Over-reporting costs one extra pass. + return True + try: + return bool(Marker(marker).evaluate()) + except Exception: + return True + + +def _version_satisfies(version: str, specifier: str) -> bool: + if not specifier: + return True + try: + from packaging.specifiers import SpecifierSet + return SpecifierSet(specifier).contains(version) + except Exception: + return False + + +def missing_requirements( + req_file: Optional[Path] = None, installed: Optional[Dict[str, str]] = None +) -> List[str]: + """Distribution names that are missing or outside their required versions. + + Checked via importlib.metadata, not import names, because studio.txt lists + PyJWT / python-docx / pymupdf whose import names (jwt, docx, fitz) differ. + + `installed` (canonical distribution name -> version) checks a venv other + than the one running this code, which importlib.metadata cannot see. + """ + from importlib.metadata import PackageNotFoundError, distribution + + path = req_file or (requirements_root() / BOOT_REQUIREMENT_FILE) + try: + lines = path.read_text(encoding = "utf-8").splitlines() + except OSError: + return [] + + missing: List[str] = [] + for line in lines: + parsed = _parse_requirement_line(line) + if parsed is None: + continue + name, marker, specifier = parsed + if not _marker_applies(marker): + continue + if installed is not None: + version = installed.get(_canonical(name)) + if version is None or not _version_satisfies(version, specifier): + missing.append(name) + continue + try: + dist = distribution(name) + except PackageNotFoundError: + missing.append(name) + except Exception: + missing.append(name) + else: + if not _version_satisfies(dist.version, specifier): + missing.append(name) + return missing + + +def verify_install( + root: Optional[Path] = None, + req_root: Optional[Path] = None, + package_name: str = "unsloth", + installed: Optional[Dict[str, str]] = None, +) -> dict: + """Report whether the managed install finished and can still boot. + + Reason strings are surfaced verbatim by the desktop preflight as its + staleness reason, so keep them stable. + + Pass `installed` (and the matching `root` / `req_root`) to describe a venv + other than this interpreter's; without it the version and dependency checks + would answer for the venv the caller happens to be running in. + """ + reqs = req_root or requirements_root() + missing = missing_requirements(reqs / BOOT_REQUIREMENT_FILE, installed = installed) + deps_ok = not missing + + manifest = read_manifest(root) + manifest_ok = False + reason: Optional[str] = None + + if manifest is None: + reason = "studio_install_incomplete" + elif manifest.get("schema") != MANIFEST_SCHEMA: + reason = "studio_install_manifest_schema" + else: + # `update --package X` records X, so comparing against unsloth would + # report a permanent version change. + current = _installed_version(manifest.get("package") or package_name, installed) + recorded = manifest.get("package_version") + if current and recorded and current != recorded: + reason = "studio_install_version_changed" + elif manifest.get("requirement_files") != requirement_digests(reqs): + reason = "studio_install_requirements_changed" + else: + manifest_ok = True + + if manifest_ok and not deps_ok: + # Install finished but the boot deps are gone: venv edited afterwards. + reason = "studio_deps_missing" + + return { + "ok": manifest_ok and deps_ok, + "manifest_ok": manifest_ok, + "deps_ok": deps_ok, + "missing": missing, + "reason": None if (manifest_ok and deps_ok) else (reason or "studio_deps_missing"), + } diff --git a/studio/install_python_stack.py b/studio/install_python_stack.py index 2883f30b20..4004a3b048 100644 --- a/studio/install_python_stack.py +++ b/studio/install_python_stack.py @@ -28,6 +28,9 @@ _BACKEND_DIR = Path(__file__).resolve().parent / "backend" if str(_BACKEND_DIR) not in sys.path: sys.path.insert(1, str(_BACKEND_DIR)) +# setup.sh/setup.ps1 invoke this by path, so its directory is sys.path[0]. +import install_manifest # noqa: E402 + from backend.utils.wheel_utils import ( flash_attn_package_version, flash_attn_wheel_url, @@ -2856,6 +2859,18 @@ def install_python_stack() -> int: base_total += 2 # flash-attn + torch final repair (step 13), Linux _TOTAL = (base_total - 1) if skip_base else base_total + # Drop it up front: a missing manifest is what tells the CLI, setup.sh and + # the preflight that an interrupted run left the venv half-built. Stop if it + # survives rather than mutate the venv behind a marker that still verifies. + if not install_manifest.remove_manifest(): + print( + f"error: could not remove the stale {install_manifest.MANIFEST_NAME} in " + f"{install_manifest.venv_root()}; refusing to install behind a marker " + "that would still report this venv as complete", + file = sys.stderr, + ) + return 1 + # 1. Try uv for faster installs (before pip upgrade -- uv venvs don't # include pip by default). USE_UV = _bootstrap_uv() @@ -3234,6 +3249,23 @@ def install_python_stack() -> int: **_windows_hidden_subprocess_kwargs(), ) + # 15. Record success. Written last so an earlier kill leaves none. Exiting 0 + # without it reports a finished install every later check calls unfinished. + if ( + install_manifest.write_manifest( + req_root = REQ_ROOT, + steps_total = _TOTAL, + package_name = package_name, + ) + is None + ): + print( + f"error: could not write {install_manifest.MANIFEST_NAME} to " + f"{install_manifest.venv_root()}", + file = sys.stderr, + ) + return 1 + _step(_LABEL, "installed") return 0 diff --git a/studio/setup.ps1 b/studio/setup.ps1 index bb21ce2da6..ea84068809 100644 --- a/studio/setup.ps1 +++ b/studio/setup.ps1 @@ -74,6 +74,20 @@ $script:NvccPath = $null $script:CudaToolkitRoot = $null $script:CudaArch = $null +function Exit-SetupFailure { + param( + [Parameter(Mandatory = $true)][string]$Message, + [int]$Code = 1 + ) + if ($Code -eq 0) { $Code = 1 } + if (@("1", "true") -contains $env:UNSLOTH_TAURI_MODE) { + $singleLine = ($Message -replace '[\r\n]+', ' ').Trim() + [Console]::Out.WriteLine("[TAURI:ERROR] $singleLine") + [Console]::Out.Flush() + } + exit $Code +} + # Detect if running from pip install (no frontend/ dir in studio) $FrontendDir = Join-Path $ScriptDir "frontend" $OxcValidatorDir = Join-Path $ScriptDir "backend\core\data_recipe\oxc-validator" @@ -851,7 +865,7 @@ function Ensure-BuildToolsForLlamaSourceBuild { Write-Host " Manual install:" -ForegroundColor Red Write-Host ' 1. winget install Microsoft.VisualStudio.2022.BuildTools --source winget' -ForegroundColor Yellow Write-Host ' 2. Open Visual Studio Installer -> Modify -> check "Desktop development with C++"' -ForegroundColor Yellow - exit 1 + Exit-SetupFailure "Visual Studio Build Tools are required for the llama.cpp source build" } } @@ -1652,7 +1666,7 @@ if (-not $HasGit) { if (-not $HasGit) { Write-Host "[ERROR] Git is required but could not be installed automatically." -ForegroundColor Red Write-Host " Install Git from https://git-scm.com/download/win and re-run." -ForegroundColor Red - exit 1 + Exit-SetupFailure "Git is required but could not be installed automatically" } step "git" "$(git --version)" } else { @@ -1821,7 +1835,7 @@ if (-not $NvccPath -and $IncompatibleToolkit) { Write-Host "========================================================================" -ForegroundColor Red Write-Host "[ERROR] CUDA source build cannot use the installed toolkit with this driver." -ForegroundColor Red Write-Host "========================================================================" -ForegroundColor Red - exit 1 + Exit-SetupFailure "The installed CUDA toolkit is incompatible with the current driver" } # -- No toolkit at all: install via winget (only when a source build needs it) -- @@ -1893,7 +1907,7 @@ if (-not $NvccPath) { } else { Write-Host " Install CUDA Toolkit from https://developer.nvidia.com/cuda-downloads" -ForegroundColor Yellow } - exit 1 + Exit-SetupFailure "A compatible CUDA Toolkit could not be found or installed" } # -- Set CUDA env vars so cmake AND MSBuild can find the toolkit -- @@ -2043,7 +2057,7 @@ if (-not $IsPipInstall) { if (-not (Test-Path -LiteralPath $NodeOverride -PathType Container)) { Write-Host "ERROR: UNSLOTH_STUDIO_HOME/STUDIO_HOME=$NodeOverride does not exist." -ForegroundColor Red Write-Host " Run install.ps1 to create the install root before 'unsloth studio update'." -ForegroundColor Red - exit 1 + Exit-SetupFailure "UNSLOTH_STUDIO_HOME/STUDIO_HOME=$NodeOverride does not exist" } $NodeParent = (Resolve-Path -LiteralPath $NodeOverride).Path # An override pointing at the legacy default maps to the legacy sibling @@ -2227,7 +2241,7 @@ if ($PythonOk) { if (-not $HasPython) { Write-Host "[ERROR] Python could not be installed automatically." -ForegroundColor Red Write-Host " Install Python 3.12 from https://python.org/downloads/" -ForegroundColor Yellow - exit 1 + Exit-SetupFailure "Python could not be installed automatically" } step "python" "$(python --version 2>&1)" $PythonOk = $true @@ -2237,7 +2251,7 @@ if ($PythonOk) { Write-Host "[ERROR] No supported Python (3.11-3.13) found on this system." -ForegroundColor Red Write-Host " py.exe could not locate -3.11/-3.12/-3.13 and `python` on PATH is unsupported." -ForegroundColor Yellow Write-Host " Install Python 3.12 from https://python.org/downloads/" -ForegroundColor Yellow - exit 1 + Exit-SetupFailure "No supported Python 3.11-3.13 was found" } # Add user-scheme Python Scripts dir to PATH (nt_user only, no venv fallback). @@ -2319,7 +2333,7 @@ if ($NeedNodeForSetup) { if (-not (Test-Path -LiteralPath $nodeOwnedMarker) -and -not (Test-Path -LiteralPath $nodeMeta)) { Write-Host "[ERROR] $NodeDir already exists and is not an Unsloth-owned Node install." -ForegroundColor Red Write-Host " Move it aside or choose an empty UNSLOTH_STUDIO_HOME before re-running." -ForegroundColor Yellow - exit 1 + Exit-SetupFailure "$NodeDir is not an Unsloth-owned Node install" } } substep "installing isolated Node (system Node/npm left untouched)..." @@ -2331,12 +2345,12 @@ if ($NeedNodeForSetup) { if ($nodeExit -eq 3) { Write-Host $nodeOut -ForegroundColor DarkGray step "node" "install blocked by another active Unsloth install" "Red" - exit 3 + Exit-SetupFailure "Node install is blocked by another active Unsloth install" 3 } elseif ($nodeExit -ne 0) { Write-Host $nodeOut -ForegroundColor DarkGray Write-Host "[ERROR] Could not install an isolated Node automatically." -ForegroundColor Red Write-Host " Install Node >= 20.19 (with npm >= 11) from https://nodejs.org/ and re-run, or check your network." -ForegroundColor Yellow - exit 1 + Exit-SetupFailure "Could not install an isolated Node runtime" } if ($NodeOverride -and (Test-Path -LiteralPath $NodeDir -PathType Container)) { New-Item -ItemType File -Force -Path (Join-Path $NodeDir ".unsloth-studio-owned") -ErrorAction SilentlyContinue | Out-Null @@ -2456,7 +2470,7 @@ if ($NeedFrontendBuild -and -not $IsPipInstall) { Write-Host "[ERROR] npm install failed (exit code $npmExit)" -ForegroundColor Red Write-Host " Try running 'npm install' manually in frontend/ to see errors" -ForegroundColor Yellow Show-NpmRegistryHint - exit 1 + Exit-SetupFailure "Frontend dependency installation failed (exit code $npmExit)" } } @@ -2467,7 +2481,7 @@ if ($NeedFrontendBuild -and -not $IsPipInstall) { $ErrorActionPreference = $prevEAP_npm foreach ($gi in $HiddenGitignores) { Rename-Item -Path "$gi._twbuild" -NewName (Split-Path $gi -Leaf) -Force -ErrorAction SilentlyContinue } Write-Host "[ERROR] npm run build failed (exit code $buildExit)" -ForegroundColor Red - exit 1 + Exit-SetupFailure "Frontend build failed (exit code $buildExit)" } Pop-Location $ErrorActionPreference = $prevEAP_npm @@ -2498,7 +2512,7 @@ if ((Test-Path $OxcValidatorDir) -and $NodeSource -ne "skip" -and (Get-Command n $ErrorActionPreference = $prevEAP_oxc Write-Host "[ERROR] OXC validator npm install failed (exit code $oxcInstallExit)" -ForegroundColor Red Show-NpmRegistryHint - exit 1 + Exit-SetupFailure "OXC validator dependency installation failed (exit code $oxcInstallExit)" } Pop-Location $ErrorActionPreference = $prevEAP_oxc @@ -2597,7 +2611,7 @@ if (-not $PythonCmd) { Write-Host "[ERROR] No standalone Python 3.11-3.13 found (conda Python is not supported)." -ForegroundColor Red Write-Host " Install Python from https://python.org/downloads/ or via:" -ForegroundColor Yellow Write-Host " winget install -e --id Python.Python.3.12" -ForegroundColor Yellow - exit 1 + Exit-SetupFailure "No standalone Python 3.11-3.13 was found" } substep "Python found: $PythonCmd" @@ -2630,12 +2644,12 @@ if ($_studioOverride) { Remove-Item -LiteralPath $_setupWriteProbe -Force -ErrorAction SilentlyContinue } catch { Write-Host "ERROR: $_studioOverrideVar=$StudioHome is not writable." -ForegroundColor Red - exit 1 + Exit-SetupFailure "$_studioOverrideVar=$StudioHome is not writable" } } else { Write-Host "ERROR: $_studioOverrideVar=$_studioOverride does not exist." -ForegroundColor Red Write-Host " Run install.ps1 to create the install root before 'unsloth studio update'." -ForegroundColor Red - exit 1 + Exit-SetupFailure "$_studioOverrideVar=$_studioOverride does not exist" } } else { $StudioHome = Join-Path $env:USERPROFILE ".unsloth\studio" @@ -2679,7 +2693,7 @@ function Assert-StudioOwnedOrAbsent { } Write-Host "[ERROR] $Path already exists and is not marked as an Unsloth-owned $Label." -ForegroundColor Red Write-Host " Move it aside or choose an empty UNSLOTH_STUDIO_HOME before re-running." -ForegroundColor Yellow - exit 1 + Exit-SetupFailure "$Label path is not an Unsloth-owned install: $Path" } } function Mark-StudioOwned { @@ -2815,7 +2829,7 @@ if ((Test-Path -LiteralPath $VenvDir -PathType Container) -and -not $NoTorchMode substep "Stale venv detected ($reason)." "Yellow" Write-Host " [ERROR] The existing Unsloth environment needs repair." -ForegroundColor Red Write-Host " Re-run install.ps1 so it can replace the environment safely with rollback." -ForegroundColor Yellow - exit 1 + Exit-SetupFailure "The existing Unsloth environment needs repair" } substep "Stale venv detected ($reason) -- rebuilding..." "Yellow" # why: mirror install.ps1 env-mode guard so an update against a custom @@ -2829,14 +2843,14 @@ if ((Test-Path -LiteralPath $VenvDir -PathType Container) -and -not $NoTorchMode ) { Write-Host "[ERROR] $VenvDir already exists but does not look like an Unsloth Studio install." -ForegroundColor Red Write-Host " Move it aside or choose an empty UNSLOTH_STUDIO_HOME before re-running." -ForegroundColor Yellow - exit 1 + Exit-SetupFailure "$VenvDir is not an Unsloth Studio environment" } try { Remove-Item -LiteralPath $VenvDir -Recurse -Force -ErrorAction Stop } catch { Write-Host " [ERROR] Could not remove stale venv: $($_.Exception.Message)" -ForegroundColor Red Write-Host " Close any running Unsloth/Python processes and re-run setup." -ForegroundColor Red - exit 1 + Exit-SetupFailure "Could not remove the stale environment at $VenvDir" } } } @@ -2845,7 +2859,7 @@ if (-not (Test-Path -LiteralPath $VenvDir)) { Write-Host "[ERROR] Virtual environment not found at $VenvDir" -ForegroundColor Red Write-Host " Run install.ps1 first to create the environment:" -ForegroundColor Yellow Write-Host " irm https://unsloth.ai/install.ps1 | iex" -ForegroundColor Yellow - exit 1 + Exit-SetupFailure "Virtual environment not found at $VenvDir" } else { substep "reusing existing virtual environment at $VenvDir" $_venvPyExe = Join-Path $VenvDir "Scripts\python.exe" @@ -2963,6 +2977,26 @@ sys.exit(0 if (major, minor) >= (4, 14) else 1) substep "anyio >=4.14 found (#6483) -- forcing dependency pass to repair..." "Cyan" $SkipPythonDeps = $false } + # An interrupted install leaves $_PkgName current while studio.txt + # never finished, so the compare above says "up to date" and update -- + # plus the desktop Repair button -- no-ops on a venv that cannot boot. + $_studioInstallIncomplete = $false + try { + & python -c " +import sys +sys.path.insert(0, sys.argv[1]) +try: + import install_manifest +except Exception: + sys.exit(0) # older tree without the manifest helper: leave the fast path alone +sys.exit(0 if install_manifest.verify_install()['ok'] else 1) +" "$PSScriptRoot" 2>$null + if ($LASTEXITCODE -ne 0) { $_studioInstallIncomplete = $true } + } catch {} + if ($_studioInstallIncomplete) { + substep "studio install incomplete -- forcing dependency pass to repair..." "Cyan" + $SkipPythonDeps = $false + } # ...but not if an AMD GPU is present and installed PyTorch is CPU-only # (host predates ROCm-wheel support, or GPU added later): the fast "up to # date" path would leave the user on CPU torch with Train/Export disabled. @@ -3009,6 +3043,28 @@ if ($script:PinChangedForceReinstall) { $SkipPythonDeps = $false } if (-not $SkipPythonDeps) { +# install_python_stack.py drops the manifest before its own dependency pass, but +# pip, torch and triton are replaced first here. Drop it now so a run killed in +# those leaves the venv marked half-built, not behind a marker that verifies. +$_ManifestDropped = $true +try { + & python -c " +import sys +sys.path.insert(0, sys.argv[1]) +try: + import install_manifest +except Exception: + sys.exit(0) # older tree without the manifest helper +sys.exit(0 if install_manifest.remove_manifest() else 1) +" "$PSScriptRoot" 2>$null + if ($LASTEXITCODE -ne 0) { $_ManifestDropped = $false } +} catch { $_ManifestDropped = $false } +if (-not $_ManifestDropped) { + Write-Host "[ERROR] Could not remove the stale unsloth_install_manifest.json." -ForegroundColor Red + Write-Host " Refusing to install behind a marker that still reports this venv as complete." -ForegroundColor Red + exit 1 +} + if ($script:UnslothVerbose) { Fast-Install --upgrade pip } else { @@ -3216,7 +3272,7 @@ if (-not $ROCmIndexUrl -and ($CuTag -eq "cpu" -or $ROCmCpuFallback)) { if ($torchInstallExit -ne 0) { Write-Host "[FAILED] PyTorch install failed (exit code $torchInstallExit)" -ForegroundColor Red Write-Host (Redact-InstallOutput $output) -ForegroundColor Red - exit 1 + Exit-SetupFailure "PyTorch installation failed (exit code $torchInstallExit)" } } elseif (-not $ROCmIndexUrl) { substep "installing PyTorch with CUDA support ($CuTag)..." @@ -3248,7 +3304,7 @@ if (-not $ROCmIndexUrl -and ($CuTag -eq "cpu" -or $ROCmCpuFallback)) { if ($torchInstallExit -ne 0) { Write-Host "[FAILED] PyTorch CUDA install failed (exit code $torchInstallExit)" -ForegroundColor Red Write-Host (Redact-InstallOutput $output) -ForegroundColor Red - exit 1 + Exit-SetupFailure "PyTorch CUDA installation failed (exit code $torchInstallExit)" } # Install Triton for Windows (enables torch.compile -- without it training can hang) @@ -3284,7 +3340,7 @@ $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 - exit 1 + Exit-SetupFailure "Python dependency installation failed (exit code $stackExit)" } } else { @@ -3362,7 +3418,7 @@ foreach ($pkg in @("transformers==5.3.0", "huggingface_hub==1.8.0", "hf_xet==1.4 Write-Host "[FAIL] Could not install $pkg into .venv_t5_530/" -ForegroundColor Red Write-Host (Redact-InstallOutput $output) -ForegroundColor Red $ErrorActionPreference = $prevEAP_t5 - exit 1 + Exit-SetupFailure "Could not install $pkg into .venv_t5_530" } } if ($script:UnslothVerbose) { @@ -3397,7 +3453,7 @@ foreach ($pkg in @("transformers==5.5.0", "huggingface_hub==1.8.0", "hf_xet==1.4 Write-Host "[FAIL] Could not install $pkg into .venv_t5_550/" -ForegroundColor Red Write-Host (Redact-InstallOutput $output) -ForegroundColor Red $ErrorActionPreference = $prevEAP_t5 - exit 1 + Exit-SetupFailure "Could not install $pkg into .venv_t5_550" } } if ($script:UnslothVerbose) { @@ -3432,7 +3488,7 @@ foreach ($pkg in @("transformers==5.10.2", "huggingface_hub==1.8.0", "hf_xet==1. Write-Host "[FAIL] Could not install $pkg into .venv_t5_510/" -ForegroundColor Red Write-Host (Redact-InstallOutput $output) -ForegroundColor Red $ErrorActionPreference = $prevEAP_t5 - exit 1 + Exit-SetupFailure "Could not install $pkg into .venv_t5_510" } } if ($script:UnslothVerbose) { @@ -3547,7 +3603,7 @@ if (-not $LlamaPr -and $LlamaPrForce -and $LlamaPrForce -match '^\d+$' -and [int if ($LlamaPr) { if ($LlamaPr -notmatch '^\d+$' -or [int]$LlamaPr -le 0) { Write-Host "[ERROR] UNSLOTH_LLAMA_PR=$LlamaPr is not a valid PR number" -ForegroundColor Red - exit 1 + Exit-SetupFailure "UNSLOTH_LLAMA_PR=$LlamaPr is not a valid PR number" } step "llama.cpp" "UNSLOTH_LLAMA_PR=$LlamaPr -- will build from PR head" "Yellow" $ResolvedLlamaTag = "pr-$LlamaPr" @@ -3563,7 +3619,7 @@ $LocalLlamaCppSrc = $env:UNSLOTH_LOCAL_LLAMA_CPP_DIR if ($LocalLlamaCppSrc) { if (-not (Test-Path -LiteralPath $LocalLlamaCppSrc -PathType Container)) { step "llama.cpp" "UNSLOTH_LOCAL_LLAMA_CPP_DIR does not exist: $LocalLlamaCppSrc" "Red" - exit 1 + Exit-SetupFailure "UNSLOTH_LOCAL_LLAMA_CPP_DIR does not exist: $LocalLlamaCppSrc" } $ResolvedLocal = (Resolve-Path -LiteralPath $LocalLlamaCppSrc).Path # Reusing a local dir disables both the prebuilt download and the source @@ -3594,7 +3650,7 @@ if ($LocalLlamaCppSrc) { # and leave Unsloth with no usable binary. if (-not $LocalLlamaServerFound) { step "llama.cpp" "no llama-server.exe under $ResolvedLocal (looked for .\llama-server.exe, .\build\bin and .\build\bin\Release) -- build llama.cpp there first, or drop --with-llama-cpp-dir" "Red" - exit 1 + Exit-SetupFailure "No llama-server.exe was found under $ResolvedLocal" } # If the target is already a junction/symlink (e.g. a previous # --with-llama-cpp-dir run), delete only the link via DirectoryInfo.Delete(). @@ -3620,7 +3676,7 @@ if ($LocalLlamaCppSrc) { if (Test-Path -LiteralPath $LlamaCppDir) { step "llama.cpp" "install blocked by active llama.cpp process" "Yellow" substep "Close Unsloth or other llama.cpp users and retry" "Yellow" - exit 3 + Exit-SetupFailure "llama.cpp install is blocked by an active llama.cpp process" 3 } } cmd /c "mklink /J `"$LlamaCppDir`" `"$ResolvedLocal`"" 2>&1 | Out-Null @@ -3762,7 +3818,7 @@ if ($LocalLlamaCppLinked) { substep "Existing install was restored" "Yellow" } substep "Close Unsloth or other llama.cpp users and retry" "Yellow" - exit 3 + Exit-SetupFailure "llama.cpp install is blocked by an active llama.cpp process" 3 } elseif ($prebuiltExit -eq 4) { step "llama.cpp" "not enough disk space to install llama.cpp" "Yellow" Write-LlamaFailureLog -Output $prebuiltOutput @@ -4025,7 +4081,7 @@ if ($LocalLlamaCppLinked) { } else { Write-Host "[ERROR] CMake 4.2+ is required to build llama.cpp with the Visual Studio 2026 generator, and no older Visual Studio toolchain was found to fall back to." -ForegroundColor Red Write-Host " Upgrade CMake from https://cmake.org/download/ and re-run, or use a prebuilt llama.cpp bundle." -ForegroundColor Red - exit 1 + Exit-SetupFailure "CMake cannot drive the Visual Studio 2026 generator" } } } @@ -4532,5 +4588,5 @@ Write-Host "" # failure. Direct 'unsloth studio update' does not set SKIP_STUDIO_BASE, # so it keeps degraded installs successful. if ($script:LlamaCppDegraded -and $env:SKIP_STUDIO_BASE -eq "1") { - exit 1 + Exit-SetupFailure "llama.cpp setup did not produce a usable server" } diff --git a/studio/setup.sh b/studio/setup.sh index b4088c15b6..f623c9ab0b 100755 --- a/studio/setup.sh +++ b/studio/setup.sh @@ -67,6 +67,18 @@ fi step() { printf " ${C_DIM}%-15.15s${C_RST}${3:-$C_OK}%s${C_RST}\n" "$1" "$2"; } substep() { printf " %-15s${2:-$C_DIM}%s${C_RST}\n" "" "$1"; } +setup_fail() { + local exit_code=$1 + shift + [ "$exit_code" -ne 0 ] || exit_code=1 + local message + message=$(printf '%s' "$*" | tr '\r\n' ' ') + case "${UNSLOTH_TAURI_MODE:-0}" in + 1|true) printf '[TAURI:ERROR] %s\n' "$message" ;; + esac + exit "$exit_code" +} + # ── Helper: can the controlling terminal actually be opened for reading? ── # `test -r` only checks permission bits, which look fine in containers and # systemd units where open() then fails with ENXIO. Probe with a real open. @@ -173,7 +185,7 @@ _run_quiet() { exit_code=$? step "error" "$label failed (exit code $exit_code)" "$C_ERR" >&2 if [ "$on_fail" = "exit" ]; then - exit "$exit_code" + setup_fail "$exit_code" "$label failed (exit code $exit_code)" else return "$exit_code" fi @@ -182,7 +194,10 @@ _run_quiet() { local tmplog tmplog=$(mktemp) || { step "error" "Failed to create temporary file" "$C_ERR" >&2 - [ "$on_fail" = "exit" ] && exit 1 || return 1 + if [ "$on_fail" = "exit" ]; then + setup_fail 1 "Failed to create temporary file for $label" + fi + return 1 } if "$@" >"$tmplog" 2>&1; then @@ -196,7 +211,7 @@ _run_quiet() { rm -f "$tmplog" if [ "$on_fail" = "exit" ]; then - exit "$exit_code" + setup_fail "$exit_code" "$label failed (exit code $exit_code)" else return "$exit_code" fi @@ -549,10 +564,14 @@ if [ -n "$_studio_override" ]; then if [ ! -d "$_studio_override" ]; then echo "ERROR: $_studio_override_var=$_studio_override does not exist." >&2 echo " Run install.sh to create the install root before 'unsloth studio update'." >&2 - exit 1 + setup_fail 1 "$_studio_override_var=$_studio_override does not exist" fi - [ -w "$_studio_override" ] || { echo "ERROR: $_studio_override_var=$_studio_override is not writable." >&2; exit 1; } - STUDIO_HOME="$(CDPATH= cd -P -- "$_studio_override" && pwd -P)" || exit 1 + if [ ! -w "$_studio_override" ]; then + echo "ERROR: $_studio_override_var=$_studio_override is not writable." >&2 + setup_fail 1 "$_studio_override_var=$_studio_override is not writable" + fi + STUDIO_HOME="$(CDPATH= cd -P -- "$_studio_override" && pwd -P)" || + setup_fail 1 "Could not resolve $_studio_override_var=$_studio_override" else STUDIO_HOME="$HOME/.unsloth/studio" fi @@ -598,7 +617,7 @@ _assert_studio_owned_or_absent() { fi echo "ERROR: $_aso_dir already exists and is not marked as an Unsloth-owned $_aso_label." >&2 echo " Move it aside or choose an empty UNSLOTH_STUDIO_HOME before re-running." >&2 - exit 1 + setup_fail 1 "$_aso_label path is not an Unsloth-owned install: $_aso_dir" fi } @@ -716,12 +735,12 @@ elif [ "$NODE_SOURCE" = bundled ]; then step "node" "install blocked by another active Unsloth install" "$C_ERR" sed 's/^/ | /' "$_NODE_LOG" >&2; rm -f "$_NODE_LOG" substep "close other Unsloth installs and retry" - exit 3 + setup_fail 3 "Node install is blocked by another active Unsloth install" elif [ "$_NODE_STATUS" -ne 0 ]; then step "node" "isolated Node install failed" "$C_ERR" sed 's/^/ | /' "$_NODE_LOG" >&2; rm -f "$_NODE_LOG" substep "install Node >= 20.19 (with npm >= 11) yourself and re-run, or check your network" - exit 1 + setup_fail 1 "Could not install an isolated Node runtime" fi grep -Fq "already matches" "$_NODE_LOG" && verbose_substep "isolated Node already up to date" rm -f "$_NODE_LOG" @@ -854,7 +873,7 @@ if [ "$_bun_install_ok" = false ]; then if [ "$_npm_install_rc" -ne 0 ]; then _suggest_npm_registry "$_FRONTEND_INSTALL_LOG" rm -f "$_FRONTEND_INSTALL_LOG" - exit "$_npm_install_rc" + setup_fail "$_npm_install_rc" "Frontend dependency installation failed (exit code $_npm_install_rc)" fi fi _CAPTURE_LOG="" @@ -894,7 +913,7 @@ if [ -d "$_OXC_DIR" ] && [ "${NODE_SOURCE:-}" != skip ] && command -v npm &>/dev if [ "$_oxc_install_rc" -ne 0 ]; then _suggest_npm_registry "$_OXC_INSTALL_LOG" rm -f "$_OXC_INSTALL_LOG" - exit "$_oxc_install_rc" + setup_fail "$_oxc_install_rc" "OXC validator dependency installation failed (exit code $_oxc_install_rc)" fi rm -f "$_OXC_INSTALL_LOG" cd "$SCRIPT_DIR" @@ -932,7 +951,7 @@ if [ ! -x "$VENV_DIR/bin/python" ]; then if ! run_quiet_no_exit "install Colab backend deps" pip install -q -r "$_COLAB_REQS_TMP"; then rm -f "$_COLAB_REQS_TMP" step "python" "Colab backend dependency install failed" "$C_ERR" - exit 1 + setup_fail 1 "Colab backend dependency installation failed" fi else step "python" "no Colab backend dependencies resolved from requirements file" "$C_WARN" @@ -943,7 +962,7 @@ if [ ! -x "$VENV_DIR/bin/python" ]; then step "python" "venv not found at $VENV_DIR" "$C_ERR" substep "Run install.sh first to create the environment:" substep "curl -fsSL https://unsloth.ai/install.sh | sh" - exit 1 + setup_fail 1 "Virtual environment not found at $VENV_DIR" fi else source "$VENV_DIR/bin/activate" @@ -1025,6 +1044,21 @@ sys.exit(0 if (major, minor) >= (4, 14) else 1) substep "anyio >=4.14 found (#6483) -- forcing dependency pass to repair..." _SKIP_PYTHON_DEPS=false fi + # An interrupted install leaves $_PKG_NAME current while studio.txt + # never finished, so the compare above says "up to date" and update -- + # plus the desktop Repair button -- no-ops on a venv that cannot boot. + if ! "$VENV_DIR/bin/python" -c " +import sys +sys.path.insert(0, sys.argv[1]) +try: + import install_manifest +except Exception: + sys.exit(0) # older tree without the manifest helper: leave the fast path alone +sys.exit(0 if install_manifest.verify_install()['ok'] else 1) +" "$SCRIPT_DIR" 2>/dev/null; then + substep "studio install incomplete -- forcing dependency pass to repair..." + _SKIP_PYTHON_DEPS=false + fi elif [ -n "$INSTALLED_VER" ] && [ -n "$LATEST_VER" ]; then substep "$_PKG_NAME $INSTALLED_VER -> $LATEST_VER available, updating..." elif [ -z "$LATEST_VER" ]; then @@ -1277,7 +1311,7 @@ fi if [ -n "$_LLAMA_PR" ]; then if ! [[ "$_LLAMA_PR" =~ ^[0-9]+$ ]] || [ "$_LLAMA_PR" -le 0 ]; then step "llama.cpp" "UNSLOTH_LLAMA_PR=$_LLAMA_PR is not a valid PR number" "$C_ERR" - exit 1 + setup_fail 1 "UNSLOTH_LLAMA_PR=$_LLAMA_PR is not a valid PR number" fi step "llama.cpp" "UNSLOTH_LLAMA_PR=$_LLAMA_PR -- will build from PR head" "$C_WARN" _RESOLVED_LLAMA_TAG="pr-$_LLAMA_PR" @@ -1313,7 +1347,7 @@ _LOCAL_LLAMA_CPP_LINKED=false if [ -n "${UNSLOTH_LOCAL_LLAMA_CPP_DIR:-}" ]; then if [ ! -d "$UNSLOTH_LOCAL_LLAMA_CPP_DIR" ]; then step "llama.cpp" "UNSLOTH_LOCAL_LLAMA_CPP_DIR does not exist: $UNSLOTH_LOCAL_LLAMA_CPP_DIR" "$C_ERR" - exit 1 + setup_fail 1 "UNSLOTH_LOCAL_LLAMA_CPP_DIR does not exist: $UNSLOTH_LOCAL_LLAMA_CPP_DIR" fi _RESOLVED_LOCAL="$(CDPATH= cd -P -- "$UNSLOTH_LOCAL_LLAMA_CPP_DIR" && pwd -P)" # Canonicalize the install path the same way before comparing: _RESOLVED_LOCAL @@ -1351,7 +1385,7 @@ if [ -n "${UNSLOTH_LOCAL_LLAMA_CPP_DIR:-}" ]; then # with no usable binary. if ! _has_local_llama_server "$_RESOLVED_LOCAL"; then step "llama.cpp" "no llama-server under $_RESOLVED_LOCAL (looked for ./llama-server and ./build/bin/llama-server) -- build llama.cpp there first, or drop --with-llama-cpp-dir" "$C_ERR" - exit 1 + setup_fail 1 "No llama-server was found under $_RESOLVED_LOCAL" fi # A stale link from a previous --with-llama-cpp-dir run isn't Unsloth-owned # content; drop it before the ownership check so re-runs stay idempotent @@ -1460,7 +1494,7 @@ else substep "existing install was restored" fi substep "close Unsloth or other llama.cpp users and retry" - exit 3 + setup_fail 3 "llama.cpp install is blocked by an active llama.cpp process" elif [ "$_PREBUILT_STATUS" -eq 4 ]; then step "llama.cpp" "not enough disk space to install llama.cpp" "$C_WARN" print_llama_error_log "$_PREBUILT_LOG" @@ -2217,5 +2251,5 @@ echo "" # successful -- the footer above already reports the limitation and Unsloth # is still usable for non-GGUF workflows. if [ "$_LLAMA_CPP_DEGRADED" = true ] && [ "${SKIP_STUDIO_BASE:-0}" = "1" ]; then - exit 1 + setup_fail 1 "llama.cpp setup did not produce a usable server" fi diff --git a/studio/src-tauri/src/desktop_backend_owner.rs b/studio/src-tauri/src/desktop_backend_owner.rs index c7d0a7b309..8a174a0731 100644 --- a/studio/src-tauri/src/desktop_backend_owner.rs +++ b/studio/src-tauri/src/desktop_backend_owner.rs @@ -528,7 +528,7 @@ fn lifecycle_control_block_reason(liveness: &DesktopLiveness) -> Option return Some("desktop_auth_unsupported".to_string()); } if liveness.desktop_manageability_version.unwrap_or(0) - < crate::preflight::DESKTOP_MANAGEABILITY_VERSION + < crate::preflight::DESKTOP_BACKEND_MANAGEABILITY_VERSION { return Some("desktop_manageability_unsupported".to_string()); } @@ -1014,14 +1014,12 @@ mod tests { assert!(!metadata_is_well_formed(&metadata)); } - #[test] - fn liveness_verification_requires_root_kind_and_token_sha() { - let metadata = metadata(1, Some(8888)); - let liveness = DesktopLiveness { + fn owned_liveness(manageability: u16) -> DesktopLiveness { + DesktopLiveness { status: Some("alive".to_string()), service: Some("Unsloth UI Backend".to_string()), desktop_protocol_version: Some(1), - desktop_manageability_version: Some(1), + desktop_manageability_version: Some(manageability), supports_desktop_auth: Some(true), supports_desktop_backend_ownership: Some(true), studio_root_id: Some(ROOT_ID.to_string()), @@ -1029,7 +1027,57 @@ mod tests { kind: Some(OWNER_KIND_TAURI.to_string()), token_sha256: Some(token_sha256(TOKEN)), }), - }; + } + } + + #[test] + fn legacy_manageability_backend_stays_lifecycle_controllable() { + // A backend from the previous app version reports manageability 1. + // studio_install_ok is CLI-side, not part of this backend's HTTP + // contract: blocking makes preflight answer ExternalConflict and never + // adopt a process the root id and token already prove is ours. + assert_eq!(lifecycle_control_block_reason(&owned_liveness(1)), None); + assert_eq!( + lifecycle_control_block_reason(&owned_liveness( + crate::preflight::DESKTOP_MANAGEABILITY_VERSION + )), + None + ); + + // The bits a live backend really must carry are still enforced. + let mut no_ownership = owned_liveness(1); + no_ownership.supports_desktop_backend_ownership = Some(false); + assert_eq!( + lifecycle_control_block_reason(&no_ownership).as_deref(), + Some("desktop_backend_ownership_unsupported") + ); + + let mut no_auth = owned_liveness(1); + no_auth.supports_desktop_auth = Some(false); + assert_eq!( + lifecycle_control_block_reason(&no_auth).as_deref(), + Some("desktop_auth_unsupported") + ); + + let mut old_protocol = owned_liveness(1); + old_protocol.desktop_protocol_version = Some(0); + assert_eq!( + lifecycle_control_block_reason(&old_protocol).as_deref(), + Some("desktop_protocol_incompatible") + ); + + let mut no_manageability = owned_liveness(1); + no_manageability.desktop_manageability_version = None; + assert_eq!( + lifecycle_control_block_reason(&no_manageability).as_deref(), + Some("desktop_manageability_unsupported") + ); + } + + #[test] + fn liveness_verification_requires_root_kind_and_token_sha() { + let metadata = metadata(1, Some(8888)); + let liveness = owned_liveness(1); assert!(liveness_verifies_metadata(&liveness, &metadata)); let mut wrong_root = liveness; diff --git a/studio/src-tauri/src/diagnostics/mod.rs b/studio/src-tauri/src/diagnostics/mod.rs index 998efc893e..3fdbac06dd 100644 --- a/studio/src-tauri/src/diagnostics/mod.rs +++ b/studio/src-tauri/src/diagnostics/mod.rs @@ -31,6 +31,10 @@ pub const TAIL_MAX_LINES: usize = 1000; pub const TAIL_MAX_BYTES: usize = 200 * 1024; pub const REPORT_MAX_BYTES: usize = 1024 * 1024; +pub(crate) fn redact_for_display(text: &str) -> String { + redaction::redact_text(text, &mut redaction::RedactionReport::default()) +} + pub(crate) const MAX_STATE_ITEMS: usize = 200; pub(crate) const MAX_PHASE_LINE_BYTES: usize = 16 * 1024; pub(crate) const FOOTER_BUDGET_BYTES: usize = 8 * 1024; diff --git a/studio/src-tauri/src/diagnostics/redaction.rs b/studio/src-tauri/src/diagnostics/redaction.rs index 0c7a92944a..bd36f9254e 100644 --- a/studio/src-tauri/src/diagnostics/redaction.rs +++ b/studio/src-tauri/src/diagnostics/redaction.rs @@ -16,6 +16,8 @@ pub(crate) fn redact_text(text: &str, report: &mut RedactionReport) -> String { } out = replace_regex(private_key_re(), &out, "", report); out = replace_regex(url_credentials_re(), &out, "$1@", report); + out = replace_regex(url_query_value_re(), &out, "$1=", report); + out = replace_regex(url_fragment_re(), &out, "$1#", report); out = replace_regex(auth_header_re(), &out, "$1: ", report); out = replace_regex(cookie_re(), &out, "$1: ", report); out = replace_regex(token_re(), &out, "", report); @@ -110,6 +112,16 @@ fn url_credentials_re() -> &'static Regex { RE.get_or_init(|| Regex::new(r"(?i)\b([a-z][a-z0-9+.-]*://)[^/\s:@]+(:[^/\s@]*)?@").unwrap()) } +fn url_query_value_re() -> &'static Regex { + static RE: OnceLock = OnceLock::new(); + RE.get_or_init(|| Regex::new(r"([?&][^=\s&`]+)=[^&#\s`]+").unwrap()) +} + +fn url_fragment_re() -> &'static Regex { + static RE: OnceLock = OnceLock::new(); + RE.get_or_init(|| Regex::new(r"(?i)(https?://[^\s`#]+)#[^\s`]+").unwrap()) +} + fn auth_header_re() -> &'static Regex { static RE: OnceLock = OnceLock::new(); RE.get_or_init(|| Regex::new(r"(?i)\b(authorization|proxy-authorization)\s*[:=]\s*(bearer|basic)?\s*[A-Za-z0-9._~+/=-]+").unwrap()) @@ -183,6 +195,7 @@ mod tests { "API_KEY=secret123\n", "native_path_lease=abc.DEF_123\n", "url=https://user:pass@example.com/path\n", + "signed=https://example.com/object?X-Amz-Signature=presignedvalue987&version=1#fragmentsecret\n", "email=alex@example.com\n", "path=/Users/alex/.unsloth/studio/logs/install.log\n", "win=C:\\Users\\Alex\\.unsloth\\studio\\logs\\install.log\n", @@ -198,6 +211,9 @@ mod tests { assert!(!redacted.contains("abc.DEF_123")); assert!(redacted.contains("native_path_lease=")); assert!(redacted.contains("https://@example.com/path")); + assert!(!redacted.contains("presignedvalue987")); + assert!(!redacted.contains("fragmentsecret")); + assert!(redacted.contains("?X-Amz-Signature=&version=#")); assert!(!redacted.contains("alex@example.com")); assert!(redacted.contains("")); assert!(!redacted.contains("PRIVATE KEY-----\nabc")); diff --git a/studio/src-tauri/src/install.rs b/studio/src-tauri/src/install.rs index 024b730735..d7226bf901 100644 --- a/studio/src-tauri/src/install.rs +++ b/studio/src-tauri/src/install.rs @@ -1,6 +1,7 @@ use crate::diagnostics::{self, AttemptLog, DiagnosticsState}; use log::{error, info, warn}; use process_wrap::std::*; +use std::collections::VecDeque; use std::io::BufRead; use std::path::{Path, PathBuf}; use std::process::{Command, ExitStatus, Stdio}; @@ -38,6 +39,160 @@ pub fn new_install_state() -> InstallState { use crate::process::trim_line_endings; +const FAILURE_CONTEXT_LINES: usize = 8; +const FAILURE_CONTEXT_LINE_BYTES: usize = 1_000; + +fn generic_failure_message(code: i32) -> String { + format!( + "Installation failed with exit code {}. Open the installer logs for details.", + code + ) +} + +#[derive(Clone, Copy, PartialEq, Eq)] +enum InstallOutputStream { + Stdout, + Stderr, +} + +struct InstallOutputLine { + stream: InstallOutputStream, + text: String, +} + +#[derive(Default)] +struct InstallFailureContext { + explicit_error: Option, + explicit_error_stream: Option, + default_error: Option, + output_tail: VecDeque, +} + +impl InstallFailureContext { + fn observe_stdout(&mut self, text: &str) -> bool { + if text.starts_with("[TAURI:ERROR_CLEAR] ") { + self.clear_failure(InstallOutputStream::Stdout); + return true; + } + if text.starts_with("[TAURI:OUTPUT_CLEAR] ") { + self.clear_stream(InstallOutputStream::Stdout); + return true; + } + if let Some(message) = text.strip_prefix("[TAURI:ERROR] ") { + let message = message.trim(); + if !message.is_empty() { + self.explicit_error = Some(Self::bounded_line(message)); + self.explicit_error_stream = Some(InstallOutputStream::Stdout); + } + return false; + } + if let Some(message) = text.strip_prefix("[TAURI:ERROR_OUTPUT] ") { + self.capture_output_error(InstallOutputStream::Stdout, message); + return true; + } + if let Some(message) = text.strip_prefix("[TAURI:ERROR_DEFAULT] ") { + let message = message.trim(); + if !message.is_empty() { + self.default_error = Some(Self::bounded_line(message)); + } + return true; + } + if !text.starts_with("[TAURI:") { + self.push_output(InstallOutputStream::Stdout, text); + } + false + } + + fn observe_stderr(&mut self, text: &str) -> bool { + if text.starts_with("[TAURI:ERROR_CLEAR] ") { + self.clear_failure(InstallOutputStream::Stderr); + return true; + } + if text.starts_with("[TAURI:OUTPUT_CLEAR] ") { + self.clear_stream(InstallOutputStream::Stderr); + return true; + } + if let Some(message) = text.strip_prefix("[TAURI:ERROR_OUTPUT] ") { + self.capture_output_error(InstallOutputStream::Stderr, message); + return true; + } + self.push_output(InstallOutputStream::Stderr, text); + false + } + + fn capture_output_error(&mut self, stream: InstallOutputStream, fallback: &str) { + let fallback = fallback.trim(); + let detail = self + .output_tail + .iter() + .rev() + .find(|line| line.stream == stream) + .map(|line| line.text.as_str()); + if let Some(error) = match (fallback.is_empty(), detail) { + (_, Some(detail)) if fallback == detail => Some(detail.to_owned()), + (false, Some(detail)) => Some(Self::bounded_line(&format!("{fallback}: {detail}"))), + (false, None) => Some(Self::bounded_line(fallback)), + (true, Some(detail)) => Some(detail.to_owned()), + (true, None) => None, + } { + self.explicit_error = Some(error); + self.explicit_error_stream = Some(stream); + } + } + + fn clear_failure(&mut self, stream: InstallOutputStream) { + if self.explicit_error_stream == Some(stream) { + self.explicit_error = None; + self.explicit_error_stream = None; + } + if stream == InstallOutputStream::Stdout { + self.default_error = None; + } + self.clear_stream(stream); + } + + fn clear_stream(&mut self, stream: InstallOutputStream) { + self.output_tail.retain(|line| line.stream != stream); + } + + fn push_output(&mut self, stream: InstallOutputStream, text: &str) { + let text = text.trim(); + if text.is_empty() { + return; + } + let text = Self::bounded_line(text); + self.output_tail + .push_back(InstallOutputLine { stream, text }); + while self.output_tail.len() > FAILURE_CONTEXT_LINES { + self.output_tail.pop_front(); + } + } + + fn bounded_line(text: &str) -> String { + let mut text = diagnostics::redact_for_display(text); + let boundary = + diagnostics::valid_utf8_boundary(&text, text.len().min(FAILURE_CONTEXT_LINE_BYTES)); + text.truncate(boundary); + text + } + + fn message(&self, code: i32) -> String { + let detail = self + .explicit_error + .as_deref() + .or(self.default_error.as_deref()) + .or_else(|| self.output_tail.back().map(|line| line.text.as_str())); + match detail { + Some(detail) => format!("Installation failed: {}", detail), + None => generic_failure_message(code), + } + } +} + +fn is_elevation_request(code: i32, packages: &[String]) -> bool { + code == 2 && !packages.is_empty() +} + // ── Script Resolution ── /// Returns (script_path, args) depending on dev vs production mode. @@ -232,7 +387,7 @@ fn spawn_script( // ── Stream ── /// Spawns reader threads for stdout/stderr. -/// Parses [TAURI:*] lines from stdout for structured events. +/// Parses structured events from stdout and failure controls from both streams. fn stream_output( app: &AppHandle, state: &InstallState, @@ -241,14 +396,19 @@ fn stream_output( attempt: AttemptLog, stdout: Option, stderr: Option, -) -> Vec> { +) -> ( + Vec>, + Arc>, +) { let mut threads = Vec::new(); + let failure_context = Arc::new(Mutex::new(InstallFailureContext::default())); if let Some(out) = stdout { let app_clone = app.clone(); let state_clone = Arc::clone(state); let diagnostics_clone = diagnostics.clone(); let attempt_clone = attempt.clone(); + let failure_context_clone = Arc::clone(&failure_context); threads.push(std::thread::spawn(move || { let mut reader = std::io::BufReader::new(out); let mut buf = Vec::new(); @@ -259,6 +419,14 @@ fn stream_output( Ok(_) => { let text = String::from_utf8_lossy(trim_line_endings(&buf)).into_owned(); diagnostics::append_phase_line(&attempt_clone.handle, "stdout", &text); + let is_failure_control = failure_context_clone + .lock() + .map(|mut context| context.observe_stdout(&text)) + .unwrap_or(false); + if is_failure_control { + info!("[install][stdout] {}", text); + continue; + } // Parse structured Tauri protocol lines if let Some(packages) = text.strip_prefix("[TAURI:NEED_SUDO] ") { let pkgs: Vec = @@ -314,6 +482,7 @@ fn stream_output( if let Some(err) = stderr { let app_clone = app.clone(); let attempt_clone = attempt.clone(); + let failure_context_clone = Arc::clone(&failure_context); threads.push(std::thread::spawn(move || { let mut reader = std::io::BufReader::new(err); let mut buf = Vec::new(); @@ -324,6 +493,14 @@ fn stream_output( Ok(_) => { let text = String::from_utf8_lossy(trim_line_endings(&buf)).into_owned(); diagnostics::append_phase_line(&attempt_clone.handle, "stderr", &text); + let is_failure_control = failure_context_clone + .lock() + .map(|mut context| context.observe_stderr(&text)) + .unwrap_or(false); + if is_failure_control { + info!("[install][stderr] {}", text); + continue; + } warn!("[install][stderr] {}", text); let _ = app_clone.emit(event_mode.progress_event(), &text); } @@ -336,7 +513,7 @@ fn stream_output( })); } - threads + (threads, failure_context) } // ── Wait & Finalize ── @@ -457,7 +634,7 @@ fn run_install_with_event_mode( return Err(msg); } }; - let threads = stream_output( + let (threads, failure_context) = stream_output( &app, &state, event_mode, @@ -490,12 +667,12 @@ fn run_install_with_event_mode( } Ok((status, intentional)) => { let code = status.code().unwrap_or(-1); - if code == 2 { + let packages = state + .lock() + .map(|install| install.needed_packages.clone()) + .unwrap_or_default(); + if is_elevation_request(code, &packages) { // Script needs elevated package install — report to frontend - let packages = state - .lock() - .map(|i| i.needed_packages.clone()) - .unwrap_or_default(); diagnostics::record_elevation_packages(&diagnostics, &attempt, &packages); diagnostics::finish_attempt( &diagnostics, @@ -508,7 +685,10 @@ fn run_install_with_event_mode( let _ = app.emit(event_mode.needs_elevation_event(), &packages); Err("NEEDS_ELEVATION".to_string()) } else { - let msg = format!("Installer exited with code {}", code); + let msg = failure_context + .lock() + .map(|context| context.message(code)) + .unwrap_or_else(|_| generic_failure_message(code)); diagnostics::finish_attempt( &diagnostics, &attempt, @@ -896,5 +1076,211 @@ mod tests { "repair-needs-elevation" ); assert!(!InstallEventMode::Repair.emit_terminal_events()); + assert!(!is_elevation_request(2, &[])); + assert!(is_elevation_request(2, &["cmake".to_string()])); + assert!(!is_elevation_request(1, &["cmake".to_string()])); + } + + #[test] + fn explicit_installer_error_beats_stderr_noise() { + let mut context = InstallFailureContext::default(); + context.observe_stdout("[TAURI:ERROR] Failed to install PyTorch"); + context.observe_stderr("rollback cleanup failed"); + assert_eq!( + context.message(7), + "Installation failed: Failed to install PyTorch" + ); + } + + #[test] + fn command_error_includes_preceding_output_from_the_same_stream() { + let mut context = InstallFailureContext::default(); + context.observe_stdout("unrelated stdout"); + context.observe_stderr("resolver error: no space left on device"); + assert!(context.observe_stderr("[TAURI:ERROR_OUTPUT] install unsloth failed (exit code 1)")); + context.observe_stdout("[TAURI:ERROR_DEFAULT] Failed to install unsloth"); + assert_eq!( + context.message(1), + "Installation failed: install unsloth failed (exit code 1): resolver error: no space left on device" + ); + } + + #[test] + fn command_error_without_output_uses_its_fallback() { + let mut context = InstallFailureContext::default(); + context.observe_stdout("unrelated output from an earlier step"); + assert!(context.observe_stdout("[TAURI:OUTPUT_CLEAR] create venv")); + assert!(context.observe_stdout("[TAURI:ERROR_OUTPUT] create venv failed (exit code 2)")); + assert_eq!( + context.message(2), + "Installation failed: create venv failed (exit code 2)" + ); + } + + #[test] + fn recovered_retry_clears_stale_installer_error() { + let mut context = InstallFailureContext::default(); + context.observe_stderr("ERROR: transient PyTorch download failure"); + assert!(context.observe_stderr("[TAURI:ERROR_OUTPUT] install PyTorch failed (exit code 1)")); + assert!(context.observe_stdout("[TAURI:ERROR_CLEAR] install PyTorch recovered after retry")); + assert!(context.observe_stderr("[TAURI:ERROR_CLEAR] install PyTorch recovered after retry")); + context.observe_stderr("ERROR: studio setup failed"); + let message = context.message(7); + assert!(message.contains("studio setup failed")); + assert!(!message.contains("install PyTorch")); + assert!(!message.contains("transient PyTorch")); + } + + #[test] + fn recovery_clear_is_order_independent_across_streams() { + let mut context = InstallFailureContext::default(); + assert!(context.observe_stdout("[TAURI:ERROR_CLEAR] install PyTorch recovered")); + context.observe_stderr("ERROR: transient PyTorch download failure"); + assert!(context.observe_stderr("[TAURI:ERROR_OUTPUT] install PyTorch failed (exit code 1)")); + assert!(context.observe_stderr("[TAURI:ERROR_CLEAR] install PyTorch recovered")); + context.observe_stdout("[TAURI:ERROR] later setup failure"); + assert!(context.observe_stderr("[TAURI:ERROR_CLEAR] delayed recovery clear")); + assert_eq!( + context.message(1), + "Installation failed: later setup failure" + ); + } + + #[test] + fn successful_fallback_clears_unstructured_stderr() { + let mut context = InstallFailureContext::default(); + context.observe_stderr("bitsandbytes pre-release install failed"); + assert!(context.observe_stderr("[TAURI:ERROR_CLEAR] bitsandbytes pypi fallback recovered")); + context.observe_stderr("mkdir: cannot create directory: Permission denied"); + assert_eq!( + context.message(1), + "Installation failed: mkdir: cannot create directory: Permission denied" + ); + } + + #[test] + fn setup_failure_uses_explicit_producer_error_before_default() { + let mut context = InstallFailureContext::default(); + context.observe_stdout("CMake not found -- installing via winget"); + context.observe_stdout("[TAURI:ERROR] UNSLOTH_LLAMA_PR=invalid is not a valid PR number"); + assert!(context.observe_stdout("[TAURI:ERROR_DEFAULT] studio setup failed (exit code 4)")); + assert_eq!( + context.message(4), + "Installation failed: UNSLOTH_LLAMA_PR=invalid is not a valid PR number" + ); + } + + #[test] + fn setup_failure_uses_default_without_specific_output() { + let mut context = InstallFailureContext::default(); + context.observe_stdout("Finishing setup"); + assert!(context.observe_stdout("[TAURI:ERROR_DEFAULT] studio setup failed (exit code 4)")); + context.observe_stderr("restored previous environment"); + assert_eq!( + context.message(4), + "Installation failed: studio setup failed (exit code 4)" + ); + } + + #[test] + fn explicit_setup_error_survives_optional_output_and_footer() { + let mut context = InstallFailureContext::default(); + context.observe_stdout("[TAURI:ERROR] llama.cpp setup did not produce a usable server"); + context.observe_stderr("whisper.cpp source build failed (exit code 1)"); + context.observe_stdout( + "whisper.cpp prebuilt install failed; browser and Transformers dictation remain available", + ); + for index in 0..10 { + context.observe_stdout(&format!("setup footer line {index}")); + } + assert!(context.observe_stdout("[TAURI:ERROR_DEFAULT] studio setup failed (exit code 1)")); + assert_eq!( + context.message(1), + "Installation failed: llama.cpp setup did not produce a usable server" + ); + } + + #[test] + fn setup_default_outranks_nonfatal_failure_output() { + let mut context = InstallFailureContext::default(); + context.observe_stdout("long paths failed to enable"); + context.observe_stderr("Triton install failed; torch.compile may not work"); + assert!(context.observe_stdout("[TAURI:ERROR_DEFAULT] studio setup failed (exit code 3)")); + assert_eq!( + context.message(3), + "Installation failed: studio setup failed (exit code 3)" + ); + } + + #[test] + fn latest_output_is_used_without_structured_context() { + let mut context = InstallFailureContext::default(); + context.observe_stderr("first diagnostic"); + context.observe_stderr("mv: cannot move build: Permission denied"); + assert_eq!( + context.message(1), + "Installation failed: mv: cannot move build: Permission denied" + ); + } + + #[test] + fn structured_rollback_progress_does_not_replace_failure_output() { + let mut context = InstallFailureContext::default(); + context.observe_stderr("ln: cannot create symbolic link: Permission denied"); + context.observe_stdout( + "[TAURI:PROGRESS] restoring previous environment after failed install...", + ); + context.observe_stdout("[TAURI:PROGRESS] restored previous environment"); + assert_eq!( + context.message(1), + "Installation failed: ln: cannot create symbolic link: Permission denied" + ); + } + + #[test] + fn installer_exit_code_is_not_duplicated() { + let mut context = InstallFailureContext::default(); + context.observe_stdout("[TAURI:ERROR] Failed to install PyTorch (exit code 7)"); + let message = context.message(7); + assert_eq!( + message, + "Installation failed: Failed to install PyTorch (exit code 7)" + ); + assert_eq!(message.matches("exit code 7").count(), 1); + } + + #[test] + fn stderr_fallback_redacts_secrets() { + let mut context = InstallFailureContext::default(); + context.observe_stderr("ERROR: download failed for https://user:pass@example.com/package"); + let message = context.message(1); + assert!(message.contains("ERROR: download failed")); + assert!(message.contains("https://@example.com/package")); + assert!(!message.contains("user:pass")); + } + + #[test] + fn failure_context_is_bounded_and_utf8_safe() { + let mut context = InstallFailureContext::default(); + context.observe_stdout(&format!( + "[TAURI:ERROR] {}https://user:secret@example.com/package", + "é".repeat(500) + )); + for index in 0..20 { + context.observe_stderr(&format!("{index}: {}", "é".repeat(1_000))); + } + let explicit_error = context.explicit_error.as_ref().unwrap(); + assert!(explicit_error.len() <= FAILURE_CONTEXT_LINE_BYTES); + assert!(explicit_error.is_char_boundary(explicit_error.len())); + assert!(!explicit_error.contains("secret")); + assert_eq!(context.output_tail.len(), FAILURE_CONTEXT_LINES); + assert!(context + .output_tail + .iter() + .all(|line| line.text.len() <= FAILURE_CONTEXT_LINE_BYTES)); + assert!(context + .output_tail + .iter() + .all(|line| line.text.is_char_boundary(line.text.len()))); } } diff --git a/studio/src-tauri/src/preflight.rs b/studio/src-tauri/src/preflight.rs index 7ef5244754..16ea8adf99 100644 --- a/studio/src-tauri/src/preflight.rs +++ b/studio/src-tauri/src/preflight.rs @@ -14,7 +14,8 @@ use std::path::PathBuf; use types::{BackendProbe, ManagedProbe}; pub use types::{DesktopPreflightDisposition, DesktopPreflightResult, ExternalBackendConflict}; pub(crate) use version::{ - backend_version_stale_reason, DESKTOP_MANAGEABILITY_VERSION, DESKTOP_PROTOCOL_VERSION, + backend_version_stale_reason, DESKTOP_BACKEND_MANAGEABILITY_VERSION, + DESKTOP_MANAGEABILITY_VERSION, DESKTOP_PROTOCOL_VERSION, }; #[cfg(test)] @@ -577,7 +578,7 @@ exit 1 r#"#!/bin/sh if [ "$1" = "-h" ]; then exit 0; fi if [ "$1" = "studio" ] && [ "$2" = "desktop-capabilities" ] && [ "$3" = "--json" ]; then - printf '{"desktop_protocol_version":1,"desktop_manageability_version":1,"supports_api_only":true,"supports_provision_desktop_auth":true,"supports_desktop_backend_ownership":true,"version":"2026.5.3"}' + printf '{"desktop_protocol_version":1,"desktop_manageability_version":2,"supports_api_only":true,"supports_provision_desktop_auth":true,"supports_desktop_backend_ownership":true,"studio_install_ok":true,"version":"2026.5.3"}' exit 0 fi exit 1 @@ -589,7 +590,7 @@ exit 1 r#"#!/bin/sh if [ "$1" = "-h" ]; then exit 0; fi if [ "$1" = "studio" ] && [ "$2" = "desktop-capabilities" ] && [ "$3" = "--json" ]; then - printf '{"desktop_protocol_version":1,"desktop_manageability_version":1,"supports_api_only":true,"supports_provision_desktop_auth":false,"supports_desktop_backend_ownership":true,"desktop_auth_stale_reason":"cap_false","version":"2026.5.3"}' + printf '{"desktop_protocol_version":1,"desktop_manageability_version":2,"supports_api_only":true,"supports_provision_desktop_auth":false,"supports_desktop_backend_ownership":true,"desktop_auth_stale_reason":"cap_false","studio_install_ok":true,"version":"2026.5.3"}' exit 0 fi if [ "$1" = "studio" ] && [ "$2" = "provision-desktop-auth" ] && [ "$3" = "--help" ]; then exit 0; fi @@ -642,7 +643,7 @@ if [ "$1" = "-h" ]; then fi if [ "$1" = "studio" ] && [ "$2" = "desktop-capabilities" ] && [ "$3" = "--json" ]; then if [ -f "$modecap" ]; then exit 42; fi - printf '{"desktop_protocol_version":1,"desktop_manageability_version":1,"supports_api_only":true,"supports_provision_desktop_auth":true,"supports_desktop_backend_ownership":true,"version":"2026.5.3"}' + printf '{"desktop_protocol_version":1,"desktop_manageability_version":2,"supports_api_only":true,"supports_provision_desktop_auth":true,"supports_desktop_backend_ownership":true,"studio_install_ok":true,"version":"2026.5.3"}' exit 0 fi exit 1 @@ -712,7 +713,7 @@ exit 1 fn desktop_ready_health_with_owner(root_id: &str, include_owner: bool) -> String { let owner = desktop_owner_json(include_owner); format!( - r#"{{"status":"healthy","service":"Unsloth UI Backend","version":"2026.5.3","desktop_protocol_version":1,"desktop_manageability_version":1,"supports_desktop_auth":true,"supports_desktop_backend_ownership":true,"studio_root_id":"{root_id}"{owner}}}"# + r#"{{"status":"healthy","service":"Unsloth UI Backend","version":"2026.5.3","desktop_protocol_version":1,"desktop_manageability_version":2,"supports_desktop_auth":true,"supports_desktop_backend_ownership":true,"studio_root_id":"{root_id}"{owner}}}"# ) } @@ -772,7 +773,7 @@ exit 1 async fn backend_with_auth_support_but_missing_protocol_is_old() { let probe = probe_test_backend( format!( - r#"{{"status":"healthy","service":"Unsloth UI Backend","version":"2026.5.3","desktop_manageability_version":1,"supports_desktop_auth":true,"supports_desktop_backend_ownership":true,"studio_root_id":"{EXPECTED_ROOT_ID}"{}}}"#, + r#"{{"status":"healthy","service":"Unsloth UI Backend","version":"2026.5.3","desktop_manageability_version":2,"supports_desktop_auth":true,"supports_desktop_backend_ownership":true,"studio_root_id":"{EXPECTED_ROOT_ID}"{}}}"#, desktop_owner_json(true) ), "401 Unauthorized", @@ -790,6 +791,41 @@ exit 1 assert!(matches!(probe, BackendProbe::Ready { .. })); } + #[tokio::test] + async fn legacy_manageability_same_root_backend_is_still_ready() { + // Same migration window as the owned-backend case: a server from the + // release before the CLI gained studio_install_ok reports manageability + // 1. That capability is CLI-side, so it must not turn a live, + // protocol-compatible backend into a conflict the user has to kill. + let probe = probe_test_backend( + format!( + r#"{{"status":"healthy","service":"Unsloth UI Backend","version":"2026.5.3","desktop_protocol_version":1,"desktop_manageability_version":1,"supports_desktop_auth":true,"supports_desktop_backend_ownership":true,"studio_root_id":"{EXPECTED_ROOT_ID}"{}}}"#, + desktop_owner_json(true) + ), + "401 Unauthorized", + ) + .await; + + assert!(matches!(probe, BackendProbe::Ready { .. })); + } + + #[tokio::test] + async fn backend_without_any_manageability_field_is_old() { + let probe = probe_test_backend( + format!( + r#"{{"status":"healthy","service":"Unsloth UI Backend","version":"2026.5.3","desktop_protocol_version":1,"supports_desktop_auth":true,"supports_desktop_backend_ownership":true,"studio_root_id":"{EXPECTED_ROOT_ID}"{}}}"#, + desktop_owner_json(true) + ), + "401 Unauthorized", + ) + .await; + + assert!(matches!( + probe, + BackendProbe::Old { reason, .. } if reason == "desktop_manageability_unsupported" + )); + } + #[tokio::test] async fn compatible_same_root_without_desktop_owner_is_ready() { let probe = probe_test_backend( @@ -805,7 +841,7 @@ exit 1 async fn stale_same_root_without_desktop_owner_is_external_conflict() { let probe = probe_test_backend( format!( - r#"{{"status":"healthy","service":"Unsloth UI Backend","version":"2026.5.1","desktop_protocol_version":1,"desktop_manageability_version":1,"supports_desktop_auth":true,"supports_desktop_backend_ownership":true,"studio_root_id":"{EXPECTED_ROOT_ID}"}}"#, + r#"{{"status":"healthy","service":"Unsloth UI Backend","version":"2026.5.1","desktop_protocol_version":1,"desktop_manageability_version":2,"supports_desktop_auth":true,"supports_desktop_backend_ownership":true,"studio_root_id":"{EXPECTED_ROOT_ID}"}}"#, ), "401 Unauthorized", ) @@ -885,7 +921,7 @@ exit 1 async fn backend_capability_false_is_old_even_when_route_401() { let probe = probe_test_backend( format!( - r#"{{"status":"healthy","service":"Unsloth UI Backend","version":"2026.5.3","desktop_protocol_version":1,"desktop_manageability_version":1,"supports_desktop_auth":false,"supports_desktop_backend_ownership":true,"desktop_auth_stale_reason":"cap_false","studio_root_id":"{EXPECTED_ROOT_ID}"{}}}"#, + r#"{{"status":"healthy","service":"Unsloth UI Backend","version":"2026.5.3","desktop_protocol_version":1,"desktop_manageability_version":2,"supports_desktop_auth":false,"supports_desktop_backend_ownership":true,"desktop_auth_stale_reason":"cap_false","studio_root_id":"{EXPECTED_ROOT_ID}"{}}}"#, desktop_owner_json(true) ), "401 Unauthorized", diff --git a/studio/src-tauri/src/preflight/backend.rs b/studio/src-tauri/src/preflight/backend.rs index 5a58142667..9813b76c9a 100644 --- a/studio/src-tauri/src/preflight/backend.rs +++ b/studio/src-tauri/src/preflight/backend.rs @@ -1,6 +1,6 @@ use super::types::BackendProbe; use super::version::{ - backend_version_stale_reason, DESKTOP_MANAGEABILITY_VERSION, DESKTOP_PROTOCOL_VERSION, + backend_version_stale_reason, DESKTOP_BACKEND_MANAGEABILITY_VERSION, DESKTOP_PROTOCOL_VERSION, }; use serde::{Deserialize, Serialize}; @@ -149,7 +149,7 @@ fn backend_capability_stale_reason(health: &BackendHealth) -> Option { .clone() .or_else(|| Some("desktop_auth_unsupported".to_string())); } - if health.desktop_manageability_version.unwrap_or(0) < DESKTOP_MANAGEABILITY_VERSION { + if health.desktop_manageability_version.unwrap_or(0) < DESKTOP_BACKEND_MANAGEABILITY_VERSION { return Some("desktop_manageability_unsupported".to_string()); } if health.supports_desktop_backend_ownership != Some(true) { diff --git a/studio/src-tauri/src/preflight/managed.rs b/studio/src-tauri/src/preflight/managed.rs index 0d20f271c5..0d67a1c3e6 100644 --- a/studio/src-tauri/src/preflight/managed.rs +++ b/studio/src-tauri/src/preflight/managed.rs @@ -11,13 +11,18 @@ use std::time::{Duration, Instant, UNIX_EPOCH}; use tokio::io::AsyncReadExt; use tokio::process::Command; -const MANAGED_CAPABILITY_CACHE_SCHEMA: u16 = 2; +// 3: the cached capability gained studio_install_ok / studio_install_reason. +const MANAGED_CAPABILITY_CACHE_SCHEMA: u16 = 3; const FNV64_OFFSET_BASIS: u64 = 0xcbf29ce484222325; const FNV64_PRIME: u64 = 0x100000001b3; const HASHED_MARKER_MAX_BYTES: u64 = 64 * 1024; const FALLBACK_MARKER_NAMES: &[&str] = &[ + // In the fingerprint, not just the cached answer: a repair touching only + // studio.txt leaves every other marker alone, so a cache entry written + // while healthy would outlive the dropped manifest. Mirrors MANIFEST_NAME. + "unsloth_install_manifest.json", "pyvenv.cfg", "uv.lock", "requirements.txt", @@ -33,6 +38,10 @@ struct DesktopCapability { supports_provision_desktop_auth: Option, supports_desktop_backend_ownership: Option, desktop_auth_stale_reason: Option, + // A part-way install leaves a CLI that answers `-h` and a backend that dies + // on `import structlog`, so a running CLI does not mean ready. + studio_install_ok: Option, + studio_install_reason: Option, version: Option, } @@ -94,6 +103,41 @@ fn marker_content_hash(path: &Path, metadata: &fs::Metadata) -> Option { .map(|bytes| hash_bytes(FNV64_OFFSET_BASIS, &bytes)) } +fn site_packages_dirs(venv_dir: &Path) -> Vec { + let mut out = Vec::new(); + #[cfg(unix)] + { + if let Ok(lib_dir) = fs::read_dir(venv_dir.join("lib")) { + for entry in lib_dir.flatten() { + out.push(entry.path().join("site-packages")); + } + } + } + out.push(venv_dir.join("Lib").join("site-packages")); + // read_dir order is unspecified and the hashes below fold in order. + out.sort(); + out +} + +/// Hash of the .dist-info / .egg-info names present, version included. +/// +/// pip uninstall rewrites nothing else that is fingerprinted, so a venv that +/// lost a studio.txt dependency would keep serving the healthy verdict. +fn installed_distributions_hash(site_packages: &Path) -> Option { + let mut names: Vec = fs::read_dir(site_packages) + .ok()? + .flatten() + .filter_map(|entry| { + let name = entry.file_name().to_string_lossy().into_owned(); + (name.ends_with(".dist-info") || name.ends_with(".egg-info")).then_some(name) + }) + .collect(); + names.sort(); + Some(names.iter().fold(FNV64_OFFSET_BASIS, |hash, name| { + hash_bytes(hash, name.as_bytes()) + })) +} + fn marker_candidates_for_bin(bin: &Path) -> Vec { let Some(scripts_dir) = bin.parent() else { return Vec::new(); @@ -103,34 +147,18 @@ fn marker_candidates_for_bin(bin: &Path) -> Vec { }; let mut out = Vec::new(); - #[cfg(unix)] - { - if let Ok(lib_dir) = fs::read_dir(venv_dir.join("lib")) { - for entry in lib_dir.flatten() { - out.push( - entry - .path() - .join("site-packages") - .join("unsloth_cli") - .join("commands") - .join("studio.py"), - ); - } - } + for site_packages in site_packages_dirs(venv_dir) { + out.push( + site_packages + .join("unsloth_cli") + .join("commands") + .join("studio.py"), + ); } for marker_name in FALLBACK_MARKER_NAMES { out.push(venv_dir.join(marker_name)); out.push(scripts_dir.join(marker_name)); } - - out.push( - venv_dir - .join("Lib") - .join("site-packages") - .join("unsloth_cli") - .join("commands") - .join("studio.py"), - ); out } @@ -160,7 +188,7 @@ fn managed_bin_fingerprint(bin: &Path) -> Option { }) .collect(); marker_entries.sort_by(|left, right| left.path.cmp(&right.path)); - let marker_hash = marker_entries + let mut marker_hash = marker_entries .iter() .fold(FNV64_OFFSET_BASIS, |hash, marker| { let next = hash_bytes(hash, marker.path.as_bytes()); @@ -172,9 +200,20 @@ fn managed_bin_fingerprint(bin: &Path) -> Option { next } }); - let marker_path = (!marker_entries.is_empty()).then(|| "markers".to_string()); - let marker_size = (!marker_entries.is_empty()).then(|| marker_entries.len() as u64); - let marker_mtime_ms = (!marker_entries.is_empty()).then_some(marker_hash); + let mut tracked = marker_entries.len(); + if let Some(venv_dir) = bin.parent().and_then(Path::parent) { + for site_packages in site_packages_dirs(venv_dir) { + let Some(dist_hash) = installed_distributions_hash(&site_packages) else { + continue; + }; + marker_hash = hash_bytes(marker_hash, site_packages.to_string_lossy().as_bytes()); + marker_hash = hash_bytes(marker_hash, &dist_hash.to_le_bytes()); + tracked += 1; + } + } + let marker_path = (tracked > 0).then(|| "markers".to_string()); + let marker_size = (tracked > 0).then_some(tracked as u64); + let marker_mtime_ms = (tracked > 0).then_some(marker_hash); Some(ManagedBinFingerprint { bin_path, @@ -401,6 +440,16 @@ fn desktop_capability_stale_reason(capability: &DesktopCapability) -> Option ManagedProbe { pub async fn managed_install_ready() -> bool { matches!(probe_managed_install().await, ManagedProbe::Ready { .. }) } + +#[cfg(test)] +mod tests { + use super::*; + + fn healthy_capability() -> DesktopCapability { + DesktopCapability { + desktop_protocol_version: Some(DESKTOP_PROTOCOL_VERSION), + desktop_manageability_version: Some(DESKTOP_MANAGEABILITY_VERSION), + supports_api_only: Some(true), + supports_provision_desktop_auth: Some(true), + supports_desktop_backend_ownership: Some(true), + desktop_auth_stale_reason: None, + studio_install_ok: Some(true), + studio_install_reason: None, + version: Some("2026.7.5".to_string()), + } + } + + #[test] + fn complete_install_is_ready() { + assert_eq!(desktop_capability_stale_reason(&healthy_capability()), None); + assert!(desktop_capability_ready(&healthy_capability())); + } + + #[test] + fn incomplete_install_is_stale_with_the_cli_reason() { + // The venv has the CLI but not structlog, so preflight must repair + // rather than spawn a backend that cannot import. + let mut capability = healthy_capability(); + capability.studio_install_ok = Some(false); + capability.studio_install_reason = Some("studio_install_incomplete".to_string()); + assert_eq!( + desktop_capability_stale_reason(&capability).as_deref(), + Some("studio_install_incomplete") + ); + assert!(!desktop_capability_ready(&capability)); + } + + #[test] + fn deps_removed_after_install_is_stale() { + let mut capability = healthy_capability(); + capability.studio_install_ok = Some(false); + capability.studio_install_reason = Some("studio_deps_missing".to_string()); + assert_eq!( + desktop_capability_stale_reason(&capability).as_deref(), + Some("studio_deps_missing") + ); + } + + #[test] + fn missing_install_field_falls_back_to_a_generic_reason() { + let mut capability = healthy_capability(); + capability.studio_install_ok = None; + capability.studio_install_reason = None; + assert_eq!( + desktop_capability_stale_reason(&capability).as_deref(), + Some("studio_install_incomplete") + ); + } + + #[test] + fn older_cli_is_rejected_on_manageability_before_the_install_check() { + // A CLI predating this feature cannot answer studio_install_ok, so the + // more specific manageability reason must win in the diagnostics. + let mut capability = healthy_capability(); + capability.desktop_manageability_version = Some(1); + capability.studio_install_ok = None; + assert_eq!( + desktop_capability_stale_reason(&capability).as_deref(), + Some("desktop_manageability_unsupported") + ); + } + + #[test] + fn a_stale_capability_is_never_served_from_cache() { + // write_cached_capability runs before the ready check, so an incomplete + // install does get cached; reusing it would outlive the repair. + let mut capability = healthy_capability(); + capability.studio_install_ok = Some(false); + let cache = ManagedCapabilityCache { + schema: MANAGED_CAPABILITY_CACHE_SCHEMA, + bin_path: "/managed/unsloth".to_string(), + bin_size: 1, + bin_mtime_ms: 1, + studio_root_id: None, + marker_path: None, + marker_size: None, + marker_mtime_ms: None, + desktop_protocol_version: DESKTOP_PROTOCOL_VERSION, + desktop_manageability_version: DESKTOP_MANAGEABILITY_VERSION, + capability, + }; + let fingerprint = ManagedBinFingerprint { + bin_path: "/managed/unsloth".to_string(), + bin_size: 1, + bin_mtime_ms: 1, + studio_root_id: None, + marker_path: None, + marker_size: None, + marker_mtime_ms: None, + }; + assert!(!cache_matches(&cache, &fingerprint)); + } + + #[test] + fn dropping_the_manifest_changes_the_fingerprint() { + // Otherwise a cache entry written while healthy outlives the manifest, + // and the probe returns Ready on the very venv this is meant to catch. + let venv = std::env::temp_dir().join(format!( + "unsloth-fingerprint-{}-{:?}", + std::process::id(), + std::thread::current().id() + )); + let scripts = venv.join("bin"); + fs::create_dir_all(&scripts).unwrap(); + let bin = scripts.join("unsloth"); + fs::write(&bin, "#!/bin/sh\nexit 0\n").unwrap(); + let manifest = venv.join("unsloth_install_manifest.json"); + fs::write(&manifest, "{}").unwrap(); + + let with_manifest = managed_bin_fingerprint(&bin).unwrap(); + fs::remove_file(&manifest).unwrap(); + let without_manifest = managed_bin_fingerprint(&bin).unwrap(); + + assert_ne!(with_manifest, without_manifest); + let _ = fs::remove_dir_all(&venv); + } + + fn cache_for(fingerprint: &ManagedBinFingerprint) -> ManagedCapabilityCache { + ManagedCapabilityCache { + schema: MANAGED_CAPABILITY_CACHE_SCHEMA, + bin_path: fingerprint.bin_path.clone(), + bin_size: fingerprint.bin_size, + bin_mtime_ms: fingerprint.bin_mtime_ms, + studio_root_id: fingerprint.studio_root_id.clone(), + marker_path: fingerprint.marker_path.clone(), + marker_size: fingerprint.marker_size, + marker_mtime_ms: fingerprint.marker_mtime_ms, + desktop_protocol_version: DESKTOP_PROTOCOL_VERSION, + desktop_manageability_version: DESKTOP_MANAGEABILITY_VERSION, + capability: healthy_capability(), + } + } + + #[test] + fn losing_a_studio_package_changes_the_fingerprint() { + // pip uninstall rewrites no fingerprinted file: the manifest, pyvenv.cfg + // and the launcher survive and `unsloth -h` still exits 0. Without the + // installed distributions in the fingerprint the healthy answer sticks. + let venv = std::env::temp_dir().join(format!( + "unsloth-fingerprint-deps-{}-{:?}", + std::process::id(), + std::thread::current().id() + )); + let _ = fs::remove_dir_all(&venv); + let scripts = venv.join("bin"); + fs::create_dir_all(&scripts).unwrap(); + let bin = scripts.join("unsloth"); + fs::write(&bin, "#!/bin/sh\nexit 0\n").unwrap(); + fs::write(venv.join("pyvenv.cfg"), "home = /usr/bin\n").unwrap(); + fs::write(venv.join("unsloth_install_manifest.json"), "{}").unwrap(); + + let site_packages = venv.join("lib").join("python3.11").join("site-packages"); + fs::create_dir_all(site_packages.join("unsloth_cli").join("commands")).unwrap(); + fs::write( + site_packages + .join("unsloth_cli") + .join("commands") + .join("studio.py"), + "# cli\n", + ) + .unwrap(); + let dist_info = site_packages.join("fastmcp-3.0.2.dist-info"); + fs::create_dir_all(&dist_info).unwrap(); + fs::write(dist_info.join("METADATA"), "Name: fastmcp\n").unwrap(); + + let with_dep = managed_bin_fingerprint(&bin).unwrap(); + let healthy_cache = cache_for(&with_dep); + // read_dir order is unspecified, so an unsorted walk would miss its own + // cache every launch and the entry would never be worth writing. + assert_eq!(with_dep, managed_bin_fingerprint(&bin).unwrap()); + assert!(cache_matches(&healthy_cache, &with_dep)); + + fs::remove_dir_all(&dist_info).unwrap(); + let without_dep = managed_bin_fingerprint(&bin).unwrap(); + + assert_ne!(with_dep, without_dep); + assert!( + !cache_matches(&healthy_cache, &without_dep), + "a removed studio package must not keep serving the cached Ready answer" + ); + let _ = fs::remove_dir_all(&venv); + } +} diff --git a/studio/src-tauri/src/preflight/version.rs b/studio/src-tauri/src/preflight/version.rs index b920c8744b..a39681416a 100644 --- a/studio/src-tauri/src/preflight/version.rs +++ b/studio/src-tauri/src/preflight/version.rs @@ -1,7 +1,15 @@ use std::cmp::Ordering; pub(crate) const DESKTOP_PROTOCOL_VERSION: u16 = 1; -pub(crate) const DESKTOP_MANAGEABILITY_VERSION: u16 = 1; +// 2: the CLI must report studio_install_ok from `studio desktop-capabilities`, +// so an interrupted install is caught before the backend is spawned. A CLI +// reporting 1 is Stale and gets repaired, which reinstalls what it missed. +pub(crate) const DESKTOP_MANAGEABILITY_VERSION: u16 = 2; +// What a RUNNING backend must report to be adopted and stopped. Not the +// constant above: studio_install_ok is CLI-side, so gating on 2 would only +// reject (and so never adopt, or stop) a backend the previous app version +// spawned. Bump only for a real backend contract change, keep it <= main.py's. +pub(crate) const DESKTOP_BACKEND_MANAGEABILITY_VERSION: u16 = 1; // Explicit backend package minimum, not the desktop app Cargo version: backend // and app releases can diverge. When bumping, verify this package exists on PyPI. pub(super) const MIN_DESKTOP_BACKEND_VERSION: &str = "2026.5.3"; diff --git a/tests/sh/test_install_rollback_lifecycle.sh b/tests/sh/test_install_rollback_lifecycle.sh index d1ccae8e19..b0e6183fb4 100644 --- a/tests/sh/test_install_rollback_lifecycle.sh +++ b/tests/sh/test_install_rollback_lifecycle.sh @@ -32,6 +32,7 @@ run_signal_case() { { printf '%s\n' 'set -e' printf '%s\n' 'substep() { :; }' + printf '%s\n' 'rollback_substep() { substep "$@"; }' printf '%s\n' 'C_WARN=""' printf "STUDIO_HOME='%s'\n" "$_case_dir" printf "VENV_DIR='%s/unsloth_studio'\n" "$_case_dir" @@ -77,6 +78,7 @@ START_BOUNDARY_HARNESS="$START_BOUNDARY_DIR/harness.sh" { printf '%s\n' 'set -e' printf '%s\n' 'substep() { :; }' + printf '%s\n' 'rollback_substep() { substep "$@"; }' printf '%s\n' 'C_WARN=""' printf "STUDIO_HOME='%s'\n" "$START_BOUNDARY_DIR" printf "VENV_DIR='%s/unsloth_studio'\n" "$START_BOUNDARY_DIR" @@ -102,6 +104,7 @@ COMMIT_BOUNDARY_HARNESS="$COMMIT_BOUNDARY_DIR/harness.sh" { printf '%s\n' 'set -e' printf '%s\n' 'substep() { :; }' + printf '%s\n' 'rollback_substep() { substep "$@"; }' printf '%s\n' 'C_WARN=""' printf "STUDIO_HOME='%s'\n" "$COMMIT_BOUNDARY_DIR" printf "VENV_DIR='%s/unsloth_studio'\n" "$COMMIT_BOUNDARY_DIR" @@ -131,6 +134,7 @@ PRUNE_HARNESS="$PRUNE_DIR/harness.sh" { printf '%s\n' 'set -e' printf '%s\n' 'substep() { :; }' + printf '%s\n' 'rollback_substep() { substep "$@"; }' printf '%s\n' 'C_WARN=""' printf "STUDIO_HOME='%s'\n" "$PRUNE_DIR" printf "VENV_DIR='%s/unsloth_studio'\n" "$PRUNE_DIR" diff --git a/tests/sh/test_tauri_retry_failure_context.sh b/tests/sh/test_tauri_retry_failure_context.sh new file mode 100755 index 0000000000..e9034e8e54 --- /dev/null +++ b/tests/sh/test_tauri_retry_failure_context.sh @@ -0,0 +1,307 @@ +#!/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 +set -e + +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" +INSTALL_SH="$SCRIPT_DIR/../../install.sh" +INSTALL_PS1="$SCRIPT_DIR/../../install.ps1" +SETUP_SH="$SCRIPT_DIR/../../studio/setup.sh" +SETUP_PS1="$SCRIPT_DIR/../../studio/setup.ps1" + +_FUNC_FILE=$(mktemp) +{ + sed -n '/^run_install_cmd()/,/^}/p' "$INSTALL_SH" + sed -n '/^run_install_cmd_retry()/,/^}/p' "$INSTALL_SH" + sed -n '/^tauri_log()/,/^}/p' "$INSTALL_SH" + sed -n '/^tauri_stream_log()/,/^}/p' "$INSTALL_SH" + sed -n '/^tauri_clear_install_error()/,/^}/p' "$INSTALL_SH" +} > "$_FUNC_FILE" +# shellcheck disable=SC1090 +. "$_FUNC_FILE" +rm -f "$_FUNC_FILE" + +substep() { + : +} + +step() { + : +} + +sleep() { + : +} + +_is_verbose() { + return 1 +} + +_redact_install_output() { + cat "$@" +} + +echo "=== run_install_cmd_retry Tauri failure context ===" + +TAURI_MODE=true +_test_attempt=0 +_test_command() { + _test_attempt=$((_test_attempt + 1)) + [ "$_test_attempt" -eq 2 ] +} + +UNSLOTH_INSTALL_RETRIES=3 +UNSLOTH_INSTALL_RETRY_DELAY=0 +_stdout_file=$(mktemp) +_stderr_file=$(mktemp) +trap 'rm -f "$_stdout_file" "$_stderr_file"' EXIT +run_install_cmd_retry "install PyTorch" _test_command >"$_stdout_file" 2>"$_stderr_file" +_stdout_clear_count=$(grep -c '^\[TAURI:ERROR_CLEAR\] install PyTorch recovered$' "$_stdout_file") +_stderr_clear_count=$(grep -c '^\[TAURI:ERROR_CLEAR\] install PyTorch recovered$' "$_stderr_file") +if [ "$_stdout_clear_count" -ne 1 ] || [ "$_stderr_clear_count" -ne 1 ]; then + echo " FAIL: recovered retry emitted $_stdout_clear_count stdout and $_stderr_clear_count stderr clear markers" + exit 1 +fi +echo " PASS: recovered retry clears stale context on both streams" + +_test_command() { + printf '%s\n' "resolver error: no space left on device" + return 9 +} + +if run_install_cmd_retry "install PyTorch" _test_command >"$_stdout_file" 2>"$_stderr_file"; then + echo " FAIL: permanent failure returned success" + exit 1 +else + _exit_code=$? +fi +if [ "$_exit_code" -ne 9 ]; then + echo " FAIL: permanent failure returned exit code $_exit_code" + exit 1 +fi +if grep -q '^\[TAURI:ERROR_CLEAR\]' "$_stdout_file" || + grep -q '^\[TAURI:ERROR_CLEAR\]' "$_stderr_file"; then + echo " FAIL: permanent failure cleared its failure context" + exit 1 +fi +if ! grep -qxF '[TAURI:OUTPUT_CLEAR] install PyTorch' "$_stderr_file" || + ! grep -qxF 'resolver error: no space left on device' "$_stderr_file" || + ! tail -n 1 "$_stderr_file" | + grep -qxF '[TAURI:ERROR_OUTPUT] install PyTorch failed (exit code 9)'; then + echo " FAIL: quiet failure did not bind its command output to the structured error" + exit 1 +fi +echo " PASS: permanent failure retains its command output and exit code" + +UNSLOTH_INSTALL_RETRIES=1 +if run_install_cmd_retry "preferred PyTorch build" _test_command >"$_stdout_file" 2>"$_stderr_file"; then + echo " FAIL: failed preferred build returned success" + exit 1 +fi +_test_command() { + return 0 +} +run_install_cmd_retry "fallback PyTorch build" _test_command >"$_stdout_file" 2>"$_stderr_file" +if ! grep -q '^\[TAURI:ERROR_CLEAR\] fallback PyTorch build recovered$' "$_stdout_file" || + ! grep -q '^\[TAURI:ERROR_CLEAR\] fallback PyTorch build recovered$' "$_stderr_file"; then + echo " FAIL: successful fallback retained the preferred build failure" + exit 1 +fi +echo " PASS: successful fallback clears an exhausted preferred failure" + +_test_command() { + return 0 +} +run_install_cmd "successful unstructured fallback" _test_command >"$_stdout_file" 2>"$_stderr_file" +if ! grep -q '^\[TAURI:ERROR_CLEAR\] successful unstructured fallback recovered$' "$_stdout_file" || + ! grep -q '^\[TAURI:ERROR_CLEAR\] successful unstructured fallback recovered$' "$_stderr_file"; then + echo " FAIL: initial success did not clear unstructured failure context" + exit 1 +fi +echo " PASS: every successful wrapped command clears unstructured failure context" + +_is_verbose() { + return 0 +} +_test_command() { + printf '%s\n' "resolver error: proxy authentication required" + return 7 +} +set +e +( + set -e + run_install_cmd "verbose failure" _test_command +) >"$_stdout_file" 2>"$_stderr_file" +_exit_code=$? +set -e +if [ "$_exit_code" -ne 7 ]; then + echo " FAIL: verbose failure returned exit code $_exit_code instead of 7" + exit 1 +fi +if ! grep -qxF '[TAURI:OUTPUT_CLEAR] verbose failure' "$_stdout_file" || + ! grep -qxF 'resolver error: proxy authentication required' "$_stdout_file" || + ! tail -n 1 "$_stdout_file" | + grep -qxF '[TAURI:ERROR_OUTPUT] verbose failure failed (exit code 7)'; then + echo " FAIL: verbose failure did not bind its command output and exit code" + exit 1 +fi +echo " PASS: verbose failure retains its command output and exit code under set -e" + +_test_command() { + _cmd_rc=9 + return 0 +} +run_install_cmd "verbose clobbering success" _test_command >"$_stdout_file" 2>"$_stderr_file" +if ! grep -q '^\[TAURI:ERROR_CLEAR\] verbose clobbering success recovered$' "$_stdout_file"; then + echo " FAIL: verbose success inherited a status variable written by the wrapped function" + exit 1 +fi +echo " PASS: verbose success records its status after the wrapped function returns" + +_test_command() { + return 7 +} +_missing_status_parent=$(mktemp -d) +rmdir "$_missing_status_parent" +_missing_status_path="$_missing_status_parent/status" +set +e +( + set -e + mktemp() { + printf '%s\n' "$_missing_status_path" + } + run_install_cmd "missing status file" _test_command +) >"$_stdout_file" 2>"$_stderr_file" +_exit_code=$? +set -e +if [ "$_exit_code" -ne 1 ] || + ! grep -q '^\[TAURI:ERROR_OUTPUT\] missing status file failed (exit code 1)$' "$_stdout_file"; then + echo " FAIL: missing verbose status defaulted to an empty or invalid exit code" + exit 1 +fi +echo " PASS: missing verbose status defaults before reporting the failure" + +_SETUP_FUNC_FILE=$(mktemp) +sed -n '/^setup_fail()/,/^}/p' "$SETUP_SH" > "$_SETUP_FUNC_FILE" +# shellcheck disable=SC1090 +. "$_SETUP_FUNC_FILE" +rm -f "$_SETUP_FUNC_FILE" + +set +e +( + UNSLOTH_TAURI_MODE=1 + setup_fail 7 "specific setup failure" +) >"$_stdout_file" 2>"$_stderr_file" +_exit_code=$? +set -e +if [ "$_exit_code" -ne 7 ] || + ! grep -qxF '[TAURI:ERROR] specific setup failure' "$_stdout_file"; then + echo " FAIL: Tauri setup failure did not emit its explicit error and exit code" + exit 1 +fi + +set +e +( + UNSLOTH_TAURI_MODE=0 + setup_fail 7 "specific setup failure" +) >"$_stdout_file" 2>"$_stderr_file" +_exit_code=$? +set -e +if [ "$_exit_code" -ne 7 ] || [ -s "$_stdout_file" ] || [ -s "$_stderr_file" ]; then + echo " FAIL: non-Tauri setup failure emitted desktop protocol output" + exit 1 +fi +echo " PASS: setup failures emit explicit context only in Tauri mode" + +_setup_mode_count=$(grep -c 'UNSLOTH_TAURI_MODE="$TAURI_MODE"' "$INSTALL_SH") +if [ "$_setup_mode_count" -ne 2 ]; then + echo " FAIL: Unix installer does not pass Tauri mode to both setup invocations" + exit 1 +fi + +_setup_exit_count=$(grep -Ec '^[[:space:]]*exit[[:space:]]+' "$SETUP_SH") +if [ "$_setup_exit_count" -ne 1 ] || + ! grep -q '^[[:space:]]*exit "\$exit_code"$' "$SETUP_SH"; then + echo " FAIL: Unix setup has explicit exits outside setup_fail" + exit 1 +fi +echo " PASS: Unix setup routes explicit exits through setup_fail" + +_rollback_block=$(sed -n \ + '/^_restore_studio_venv_replacement()/,/^}/p' \ + "$INSTALL_SH") +_rollback_progress_count=$(printf '%s\n' "$_rollback_block" | + grep -c 'rollback_substep') +if [ "$_rollback_progress_count" -ne 2 ]; then + echo " FAIL: successful Unix rollback output can replace failure context" + exit 1 +fi +echo " PASS: successful Unix rollback remains structured progress" + +_setup_success_block=$(sed -n \ + '/^if \[ "$_SETUP_EXIT" -eq 0 \]; then$/,/^mkdir -p "\$_LOCAL_BIN"$/p' \ + "$INSTALL_SH") +if ! printf '%s\n' "$_setup_success_block" | + grep -q 'tauri_clear_install_error "studio setup completed"'; then + echo " FAIL: successful studio setup does not clear recovered setup errors before post-setup work" + exit 1 +fi + +_setup_failure_block=$(sed -n \ + '/^# If setup.sh failed, report and exit now\.$/,/^fi$/p' \ + "$INSTALL_SH") +if ! printf '%s\n' "$_setup_failure_block" | + grep -q 'tauri_log "ERROR_DEFAULT" "studio setup failed'; then + echo " FAIL: failed studio setup does not preserve output before its generic fallback" + exit 1 +fi +echo " PASS: studio setup success clears recovered errors and failure preserves specific output" + +_ps_setup_block=$(sed -n \ + '/if (\$setupExit -ne 0) {/,/# ── Expose `unsloth` via a shim dir/p' \ + "$INSTALL_PS1") +if ! printf '%s\n' "$_ps_setup_block" | grep -q 'Exit-InstallFailure' || + ! printf '%s\n' "$_ps_setup_block" | + grep -q 'Clear-TauriInstallError "studio setup completed"'; then + echo " FAIL: Windows setup does not preserve failed output and clear successful output" + exit 1 +fi +echo " PASS: Windows setup uses the same failure-context boundaries" + +if ! grep -q '\$env:UNSLOTH_TAURI_MODE = if (\$TauriMode)' "$INSTALL_PS1"; then + echo " FAIL: Windows installer does not pass Tauri mode to setup" + exit 1 +fi + +_ps_setup_exit_count=$(grep -Ec '^[[:space:]]*exit[[:space:]]+' "$SETUP_PS1") +if [ "$_ps_setup_exit_count" -ne 1 ] || + ! grep -q '^[[:space:]]*exit \$Code$' "$SETUP_PS1"; then + echo " FAIL: Windows setup has explicit exits outside Exit-SetupFailure" + exit 1 +fi +echo " PASS: Windows setup routes explicit exits through Exit-SetupFailure" + +_ps_command_block=$(sed -n \ + '/function Invoke-InstallCommand {/,/function New-StudioShortcuts {/p' \ + "$INSTALL_PS1") +if ! printf '%s\n' "$_ps_command_block" | + grep -q 'Write-TauriLog "ERROR_OUTPUT" "$Label failed' || + ! printf '%s\n' "$_ps_command_block" | + grep -q 'Write-TauriLog "OUTPUT_CLEAR" \$Label' || + ! printf '%s\n' "$_ps_command_block" | + grep -q 'Clear-TauriInstallError "$Label recovered"' || + ! printf '%s\n' "$_ps_command_block" | + grep -q 'Invoke-InstallCommand -Command \$Command -Label \$Label'; then + echo " FAIL: Windows command output is not attributed and cleared at the command boundary" + exit 1 +fi + +_ps_exit_block=$(sed -n \ + '/function Exit-InstallFailure {/,/# ── Parse flags/p' \ + "$INSTALL_PS1") +if ! printf '%s\n' "$_ps_exit_block" | + grep -q 'Write-TauriLog "ERROR_DEFAULT" \$Message'; then + echo " FAIL: Windows finalization can overwrite producer-owned failure context" + exit 1 +fi +echo " PASS: Windows command failures preserve output through retries and finalization" diff --git a/tests/sh/test_with_llama_cpp_dir_link_behavior.sh b/tests/sh/test_with_llama_cpp_dir_link_behavior.sh index fb09e56c51..2d87b7f541 100644 --- a/tests/sh/test_with_llama_cpp_dir_link_behavior.sh +++ b/tests/sh/test_with_llama_cpp_dir_link_behavior.sh @@ -37,6 +37,7 @@ case "$block" in *'_has_local_llama_server'*) : ;; PREAMBLE=' set -u step() { :; }; substep() { :; }; verbose_substep() { :; } +setup_fail() { exit "$1"; } _assert_studio_owned_or_absent() { :; } C_ERR="" _STUDIO_HOME_IS_CUSTOM=false diff --git a/tests/studio/install/test_install_manifest.py b/tests/studio/install/test_install_manifest.py new file mode 100644 index 0000000000..79b2c1db50 --- /dev/null +++ b/tests/studio/install/test_install_manifest.py @@ -0,0 +1,203 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Coverage for studio/install_manifest.py. + +The manifest separates "the install finished" from "the installer was killed +part-way and the venv only looks fine". The CLI, setup.sh's fast path and the +Tauri preflight all read it, so a wrong answer either crashes the backend on +launch or forces needless reinstalls for everyone. +""" + +from __future__ import annotations + +import importlib.util +import json +import pathlib + +import pytest + +REPO_ROOT = pathlib.Path(__file__).resolve().parents[3] +MODULE_PATH = REPO_ROOT / "studio" / "install_manifest.py" + + +def _load_module(): + spec = importlib.util.spec_from_file_location("studio_install_manifest_under_test", MODULE_PATH) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +im = _load_module() + + +@pytest.fixture +def req_root(tmp_path: pathlib.Path) -> pathlib.Path: + """A requirements tree whose studio.txt names one installed and one absent dist.""" + root = tmp_path / "requirements" + root.mkdir() + (root / "studio.txt").write_text( + "# comment line\n\npytest\nunsloth-definitely-not-a-real-package\n", + encoding = "utf-8", + ) + return root + + +@pytest.fixture +def install_root(tmp_path: pathlib.Path) -> pathlib.Path: + root = tmp_path / "venv" + root.mkdir() + return root + + +def test_parse_requirement_line_handles_the_shapes_studio_txt_uses(): + assert im._parse_requirement_line("structlog>=24.1.0") == ("structlog", "", ">=24.1.0") + assert im._parse_requirement_line("matplotlib==3.10.9") == ("matplotlib", "", "==3.10.9") + assert im._parse_requirement_line("boto3>=1.34.0 # optional: S3") == ( + "boto3", + "", + ">=1.34.0", + ) + assert im._parse_requirement_line("uvicorn[standard]") == ("uvicorn", "", "") + assert im._parse_requirement_line("# just a comment") is None + assert im._parse_requirement_line("") is None + assert im._parse_requirement_line("--index-url https://example.invalid") is None + name, marker, specifier = im._parse_requirement_line("pywin32 ; sys_platform == 'win32'") + assert name == "pywin32" + assert "sys_platform" in marker + assert specifier == "" + + +def test_missing_requirements_rejects_an_incompatible_installed_version(tmp_path): + req = tmp_path / "studio.txt" + req.write_text( + "matplotlib==3.10.9\nstructlog>=24.1.0\n", + encoding = "utf-8", + ) + installed = { + "matplotlib": "3.9.0", + "structlog": "24.1.0", + } + assert im.missing_requirements(req, installed = installed) == ["matplotlib"] + + +def test_platform_gated_lines_are_skipped_when_the_marker_does_not_apply(tmp_path): + req = tmp_path / "studio.txt" + req.write_text( + "unsloth-not-real-a ; sys_platform == 'definitely-not-this-platform'\n" + "unsloth-not-real-b\n", + encoding = "utf-8", + ) + missing = im.missing_requirements(req) + assert missing == ["unsloth-not-real-b"], ( + "a requirement gated to another OS must not be reported missing, or every " + "install would look broken on the platforms that legitimately skip it" + ) + + +def test_missing_requirements_matches_on_distribution_not_import_name(tmp_path): + # studio.txt lists PyJWT / python-docx / pymupdf, whose import names are + # jwt / docx / fitz, so matching on imports would look missing. + req = tmp_path / "studio.txt" + req.write_text("pytest\n", encoding = "utf-8") + assert im.missing_requirements(req) == [] + + +def test_complete_install_verifies_ok(install_root, req_root): + im.write_manifest(root = install_root, req_root = req_root, package_name = "pytest") + state = im.verify_install(root = install_root, req_root = req_root, package_name = "pytest") + assert state["manifest_ok"] is True + assert state["deps_ok"] is False # the fake dist is intentionally absent + assert state["reason"] == "studio_deps_missing" + assert "unsloth-definitely-not-a-real-package" in state["missing"] + + +def test_missing_manifest_reports_incomplete(install_root, req_root): + state = im.verify_install(root = install_root, req_root = req_root, package_name = "pytest") + assert state["ok"] is False + assert state["manifest_ok"] is False + assert state["reason"] == "studio_install_incomplete" + + +def test_interrupted_install_leaves_no_manifest(install_root, req_root): + # remove_manifest() runs before the dependency pass, so a later kill cannot + # leave a stale-but-valid manifest behind. + im.write_manifest(root = install_root, req_root = req_root, package_name = "pytest") + assert im.manifest_path(install_root).is_file() + assert im.remove_manifest(install_root) is True + assert not im.manifest_path(install_root).is_file() + state = im.verify_install(root = install_root, req_root = req_root, package_name = "pytest") + assert state["reason"] == "studio_install_incomplete" + + +def test_remove_manifest_reports_whether_the_marker_is_really_gone( + install_root, req_root, monkeypatch +): + # Nothing to remove is success: a first install has no manifest yet. + assert im.remove_manifest(install_root) is True + + # A surviving marker must be reported, not swallowed: the dependency pass + # would then run behind a manifest that still verifies, so a part-way kill + # looks complete. + im.write_manifest(root = install_root, req_root = req_root, package_name = "pytest") + path = im.manifest_path(install_root) + + def _refuse(*_args, **_kwargs): + raise PermissionError(13, "Access is denied") + + monkeypatch.setattr(pathlib.Path, "unlink", _refuse) + assert im.remove_manifest(install_root) is False + monkeypatch.undo() + + # The stale marker still verifies, which is why the installer has to stop. + assert path.is_file() + state = im.verify_install(root = install_root, req_root = req_root, package_name = "pytest") + assert state["manifest_ok"] is True + + +def test_schema_bump_invalidates_an_old_manifest(install_root, req_root): + im.write_manifest(root = install_root, req_root = req_root, package_name = "pytest") + path = im.manifest_path(install_root) + data = json.loads(path.read_text(encoding = "utf-8")) + data["schema"] = im.MANIFEST_SCHEMA + 1 + path.write_text(json.dumps(data), encoding = "utf-8") + state = im.verify_install(root = install_root, req_root = req_root, package_name = "pytest") + assert state["reason"] == "studio_install_manifest_schema" + + +def test_package_upgrade_invalidates_the_manifest(install_root, req_root): + im.write_manifest(root = install_root, req_root = req_root, package_name = "pytest") + path = im.manifest_path(install_root) + data = json.loads(path.read_text(encoding = "utf-8")) + data["package_version"] = "0.0.0-not-the-installed-version" + path.write_text(json.dumps(data), encoding = "utf-8") + state = im.verify_install(root = install_root, req_root = req_root, package_name = "pytest") + assert state["reason"] == "studio_install_version_changed" + + +def test_verify_follows_the_package_the_manifest_names(install_root, req_root): + # `studio update --package X` records X. Checking unsloth's version instead + # would report a change on every probe and repair for ever. + im.write_manifest(root = install_root, req_root = req_root, package_name = "pytest") + state = im.verify_install( + root = install_root, + req_root = req_root, + package_name = "unsloth-definitely-not-a-real-package", + ) + assert state["manifest_ok"] is True + + +def test_edited_requirements_invalidate_the_manifest(install_root, req_root): + # The --local dev path: an edited studio.txt must re-run the dependency + # pass, not sit behind setup.sh's "up to date" fast path. + im.write_manifest(root = install_root, req_root = req_root, package_name = "pytest") + (req_root / "studio.txt").write_text("pytest\nrich\n", encoding = "utf-8") + state = im.verify_install(root = install_root, req_root = req_root, package_name = "pytest") + assert state["reason"] == "studio_install_requirements_changed" + + +def test_unwritable_root_degrades_to_incomplete(tmp_path, req_root): + missing_root = tmp_path / "does" / "not" / "exist" + assert im.write_manifest(root = missing_root, req_root = req_root) is None + state = im.verify_install(root = missing_root, req_root = req_root, package_name = "pytest") + assert state["ok"] is False diff --git a/tests/studio/install/test_setup_fast_path_guard.py b/tests/studio/install/test_setup_fast_path_guard.py new file mode 100644 index 0000000000..da6f9aa7d8 --- /dev/null +++ b/tests/studio/install/test_setup_fast_path_guard.py @@ -0,0 +1,99 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""setup.sh / setup.ps1 must not skip the dependency pass on a half-built venv. + +Both short-circuit all dependency work when the installed unsloth version equals +PyPI's latest, which is true on an interrupted install: unsloth goes in early and +studio.txt never finishes. So update, and the desktop Repair button behind it, +said "up to date" while the server kept dying on `import structlog`. + +That branch only runs for a non-local update, which reinstalls from PyPI and +clobbers the tree under test, so assert the guard structurally instead. +""" + +from __future__ import annotations + +import pathlib +import re + +import pytest + +REPO_ROOT = pathlib.Path(__file__).resolve().parents[3] +SETUP_SH = REPO_ROOT / "studio" / "setup.sh" +SETUP_PS1 = REPO_ROOT / "studio" / "setup.ps1" + + +@pytest.mark.parametrize("script", [SETUP_SH, SETUP_PS1], ids = ["setup.sh", "setup.ps1"]) +def test_fast_path_consults_the_install_manifest(script: pathlib.Path): + text = script.read_text(encoding = "utf-8") + assert "install_manifest" in text, ( + f"{script.name} no longer consults studio/install_manifest.py. Without it " + "the 'up to date' fast path skips the dependency pass on an interrupted " + "install, and `unsloth studio update` becomes a silent no-op." + ) + assert "verify_install" in text, ( + f"{script.name} must call install_manifest.verify_install() so the check " + "matches what `unsloth studio verify-install` and the desktop preflight use." + ) + + +@pytest.mark.parametrize("script", [SETUP_SH, SETUP_PS1], ids = ["setup.sh", "setup.ps1"]) +def test_guard_can_still_force_the_dependency_pass(script: pathlib.Path): + """The guard has to clear the skip flag, not merely log a warning.""" + text = script.read_text(encoding = "utf-8") + if script.name.endswith(".ps1"): + pattern = r"studio install incomplete[\s\S]{0,200}?\$SkipPythonDeps\s*=\s*\$false" + else: + pattern = r"studio install incomplete[\s\S]{0,200}?_SKIP_PYTHON_DEPS=false" + assert re.search(pattern, text), ( + f"{script.name} detects an incomplete install but does not clear the " + "skip flag, so the dependency pass would still be skipped." + ) + + +def test_ps1_drops_the_manifest_before_its_first_install(): + """Nothing may mutate the venv while the marker still says "install finished". + + install_python_stack.py drops it before its own dependency pass, which is + enough for setup.sh: the stack is the first thing that pass runs. setup.ps1 + replaces pip, torch and triton first, so a run killed there would leave a + manifest that still verifies and a venv with half a PyTorch. + """ + text = SETUP_PS1.read_text(encoding = "utf-8") + pass_start = text.index("if (-not $SkipPythonDeps) {") + removal = text.find("remove_manifest", pass_start) + first_install = text.index("Fast-Install", pass_start) + stack = text.index(r'python "$PSScriptRoot\install_python_stack.py"', pass_start) + + assert removal != -1, ( + "setup.ps1 never drops the install manifest; install_python_stack.py " + "only does so after setup.ps1 has already replaced pip and torch" + ) + assert removal < first_install < stack, ( + "setup.ps1 must invalidate the install manifest before its first " + "Fast-Install, not leave it to install_python_stack.py" + ) + + +def test_sh_dependency_pass_mutates_nothing_before_the_stack(): + """setup.sh relies on install_python_stack.py dropping the marker, which only + holds while the stack is the first thing its dependency pass runs.""" + text = SETUP_SH.read_text(encoding = "utf-8") + pass_start = text.index('if [ "$_SKIP_PYTHON_DEPS" = false ]') + body = text[pass_start : text.index("install_python_stack", pass_start)] + assert "fast_install" not in body and "pip install" not in body, ( + "setup.sh installs something before install_python_stack.py drops the " + "manifest, so an interrupted run would keep a marker that verifies" + ) + + +def test_sh_guard_runs_before_the_skip_decision(): + text = SETUP_SH.read_text(encoding = "utf-8") + guard = text.find("studio install incomplete") + decision = text.find('if [ "$_SKIP_PYTHON_DEPS" = false ]') + assert guard != -1 and decision != -1 + assert guard < decision, ( + "the incomplete-install guard must run before setup.sh acts on " + "_SKIP_PYTHON_DEPS, otherwise it can never change the outcome" + ) diff --git a/tests/studio/install/test_studio_deps_cli.py b/tests/studio/install/test_studio_deps_cli.py new file mode 100644 index 0000000000..85d1212b61 --- /dev/null +++ b/tests/studio/install/test_studio_deps_cli.py @@ -0,0 +1,273 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Coverage for unsloth_cli/_studio_deps.py. + +Two things have to be right for the CLI half of the install check. + +It must describe the venv it was *asked* about. The wheel ships studio/, so a +CLI installed outside the managed venv always finds its own copy of the manifest +helper, and would otherwise report on its own prefix: a healthy managed install +comes back "incomplete", a broken one comes back with the wrong missing list. + +And it must name the *distribution* to install rather than the import that +failed. `pip install jwt` / `docx` / `fitz` all succeed and install unrelated +PyPI projects, leaving the backend just as broken as before. +""" + +from __future__ import annotations + +import importlib.util +import io +import json +import contextlib +import pathlib +import shutil +import sys + +import pytest +import typer + +REPO_ROOT = pathlib.Path(__file__).resolve().parents[3] +DEPS_PATH = REPO_ROOT / "unsloth_cli" / "_studio_deps.py" +MANIFEST_PATH = REPO_ROOT / "studio" / "install_manifest.py" +REQUIREMENTS = REPO_ROOT / "studio" / "backend" / "requirements" + + +def _load(path: pathlib.Path, name: str): + spec = importlib.util.spec_from_file_location(name, path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +_MANIFEST = _load(MANIFEST_PATH, "install_manifest_for_deps_test") + + +def _studio_distributions() -> list: + lines = (REQUIREMENTS / "studio.txt").read_text(encoding = "utf-8").splitlines() + parsed = [_MANIFEST._parse_requirement_line(line) for line in lines] + return [name for name, _, _ in (p for p in parsed if p is not None)] + + +def _studio_distribution_versions() -> dict: + versions = {} + lines = (REQUIREMENTS / "studio.txt").read_text(encoding = "utf-8").splitlines() + for parsed in (_MANIFEST._parse_requirement_line(line) for line in lines): + if parsed is None: + continue + name, _marker, specifier = parsed + version = "1.0.0" + for part in specifier.split(","): + if part.startswith("=="): + version = part[2:] + break + if part.startswith(">="): + version = part[2:] + versions[name] = version + return versions + + +def _make_venv( + root: pathlib.Path, + *, + unsloth_version: str, + distributions, + extra_requirement = "", +): + """A venv tree: pyvenv.cfg, the shipped studio/ package and .dist-info dirs.""" + site_packages = root / "lib" / "python3.11" / "site-packages" + (site_packages / "studio" / "backend").mkdir(parents = True) + shutil.copy(MANIFEST_PATH, site_packages / "studio" / "install_manifest.py") + shutil.copytree(REQUIREMENTS, site_packages / "studio" / "backend" / "requirements") + if extra_requirement: + studio_txt = site_packages / "studio" / "backend" / "requirements" / "studio.txt" + studio_txt.write_text( + studio_txt.read_text(encoding = "utf-8") + extra_requirement, encoding = "utf-8" + ) + (root / "pyvenv.cfg").write_text("home = /usr/bin\n", encoding = "utf-8") + studio_versions = _studio_distribution_versions() + for name in [*distributions, "unsloth"]: + version = unsloth_version if name == "unsloth" else studio_versions.get(name, "1.0.0") + dist_info = site_packages / f"{name.replace('-', '_')}-{version}.dist-info" + dist_info.mkdir() + (dist_info / "METADATA").write_text( + f"Metadata-Version: 2.1\nName: {name}\nVersion: {version}\n", + encoding = "utf-8", + ) + return site_packages + + +def _write_manifest(root: pathlib.Path, site_packages: pathlib.Path, version: str): + (root / _MANIFEST.MANIFEST_NAME).write_text( + json.dumps( + { + "schema": _MANIFEST.MANIFEST_SCHEMA, + "package": "unsloth", + "package_version": version, + "requirement_files": _MANIFEST.requirement_digests( + site_packages / "studio" / "backend" / "requirements", + ), + } + ), + encoding = "utf-8", + ) + + +@pytest.fixture +def cross_venv(tmp_path, monkeypatch): + """`unsloth studio verify-install` run from a CLI outside the managed venv. + + Returns a callable: build the managed venv, then ask about it. + """ + + def build( + *, + managed_version = "2026.6.1", + caller_version = "2026.7.9", + managed_distributions = None, + extra_requirement = "", + with_manifest = True, + ): + caller = tmp_path / "caller_venv" + caller_site = _make_venv(caller, unsloth_version = caller_version, distributions = []) + managed = tmp_path / "studio_home" / "unsloth_studio" + managed_site = _make_venv( + managed, + unsloth_version = managed_version, + distributions = _studio_distributions() + if managed_distributions is None + else managed_distributions, + extra_requirement = extra_requirement, + ) + if with_manifest: + _write_manifest(managed, managed_site, managed_version) + + (caller_site / "unsloth_cli").mkdir(parents = True) + shutil.copy(DEPS_PATH, caller_site / "unsloth_cli" / "_studio_deps.py") + monkeypatch.setattr(sys, "prefix", str(caller)) + deps = _load(caller_site / "unsloth_cli" / "_studio_deps.py", "studio_deps_cross_venv") + return deps.install_state(extra_roots = (managed,)) + + return build + + +def test_a_healthy_managed_venv_is_not_reported_incomplete(cross_venv): + """The caller's own prefix has no manifest and none of studio.txt, so + describing it instead sends a working install through a needless repair.""" + state = cross_venv() + assert state["ok"] is True, state + assert state["reason"] is None + assert state["missing"] == [] + + +def test_a_newer_caller_does_not_look_like_a_changed_managed_version(cross_venv): + """The version and requirement digests must come from the managed venv too: + reading them here compares two unrelated installs.""" + state = cross_venv(managed_version = "2026.1.1", caller_version = "2026.12.31") + assert state["ok"] is True, state + + +def test_a_managed_venv_missing_a_boot_dep_names_that_dep(cross_venv): + """The other direction: report what is actually absent over there.""" + state = cross_venv( + managed_distributions = [d for d in _studio_distributions() if d != "fastmcp"], + ) + assert state["ok"] is False + assert state["reason"] == "studio_deps_missing" + assert state["missing"] == ["fastmcp"], state + + +def test_an_unfinished_managed_install_is_still_reported_incomplete(cross_venv): + state = cross_venv(with_manifest = False) + assert state["ok"] is False + assert state["reason"] == "studio_install_incomplete" + + +# ── import name vs distribution name ───────────────────────────────── + + +@pytest.fixture +def deps(): + return _load(DEPS_PATH, "studio_deps_under_test") + + +def _remediation(deps, trigger: str, studio_missing) -> str: + deps._missing_studio_packages = lambda: list(studio_missing) + stderr = io.StringIO() + with contextlib.redirect_stderr(stderr), pytest.raises(typer.Exit): + with deps.studio_backend_imports("unsloth studio"): + raise ModuleNotFoundError(f"No module named '{trigger}'", name = trigger) + return stderr.getvalue() + + +@pytest.mark.parametrize( + "trigger, distribution", + [("jwt", "pyjwt"), ("docx", "python-docx"), ("fitz", "pymupdf")], +) +def test_a_missing_studio_package_is_named_by_its_distribution(deps, trigger, distribution): + """`pip install jwt` installs a different JWT library and repairs nothing.""" + output = _remediation(deps, trigger, [distribution]) + assert f"pip install {trigger}" not in output, output + assert distribution in output + assert "unsloth studio update" in output + + +def test_a_normalised_name_still_counts_as_a_studio_dependency(deps): + """studio.txt writes huggingface-hub; the import is huggingface_hub.""" + output = _remediation(deps, "huggingface_hub", ["huggingface-hub"]) + assert "Install it:" not in output, output + assert "also missing:" not in output, output + + +def test_a_missing_submodule_is_traced_to_its_installable_package(deps): + """exc.name is dotted when the top level survived a partial install, and + `pip install fastmcp.server` is not a package name at all.""" + output = _remediation(deps, "fastmcp.server", ["fastmcp"]) + assert "fastmcp.server" not in output.split("Install it:")[-1], output + assert "Install it:" not in output, output + assert "unsloth studio update" in output + + +def test_a_non_studio_dependency_keeps_its_own_install_line(deps): + """train reaches torch through the same wrapped import and the studio extra + does not carry it.""" + output = _remediation(deps, "torch", ["pyjwt"]) + assert "pip install torch" in output + assert "also missing: pyjwt" in output + + +def test_studio_only_guard_preserves_non_studio_failures(deps): + deps._missing_studio_packages = lambda: ["pyjwt"] + with pytest.raises(ModuleNotFoundError): + with deps.studio_backend_imports("unsloth inference", studio_only = True): + raise ModuleNotFoundError("No module named 'mlx'", name = "mlx") + + +def test_the_import_map_only_names_studio_distributions(): + """Drift guard: an entry pointing at a dropped requirement is dead advice.""" + known = {deps_name.lower() for deps_name in _studio_distributions()} + module = _load(DEPS_PATH, "studio_deps_map_check") + for import_name, distribution in module._IMPORT_TO_DISTRIBUTION.items(): + assert distribution.lower() in known, ( + f"_IMPORT_TO_DISTRIBUTION maps {import_name} to {distribution}, " + "which studio.txt no longer requires" + ) + + +def test_a_torn_tree_without_the_manifest_helper_is_incomplete(tmp_path, monkeypatch): + """studio/install_manifest.py ships in the same wheel as _studio_deps.py, so + only a torn install has one without the other. Answering yes here launches a + backend whose own files may be just as absent.""" + caller = tmp_path / "caller_venv" + site_packages = caller / "lib" / "python3.11" / "site-packages" + (site_packages / "unsloth_cli").mkdir(parents = True) + shutil.copy(DEPS_PATH, site_packages / "unsloth_cli" / "_studio_deps.py") + (caller / "pyvenv.cfg").write_text("home = /usr/bin\n", encoding = "utf-8") + monkeypatch.setattr(sys, "prefix", str(caller)) + + deps = _load(site_packages / "unsloth_cli" / "_studio_deps.py", "studio_deps_torn_tree") + state = deps.install_state() + + assert state["ok"] is False, state + assert state["reason"] == "studio_install_manifest_missing" diff --git a/tests/studio/install/test_studio_extra_matches_requirements.py b/tests/studio/install/test_studio_extra_matches_requirements.py new file mode 100644 index 0000000000..83988ce64e --- /dev/null +++ b/tests/studio/install/test_studio_extra_matches_requirements.py @@ -0,0 +1,85 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""The studio extra must mirror studio/backend/requirements/studio.txt. + +Nothing else keeps them in sync, and drift reintroduces #4701 / #5260 / #7147. +""" + +from __future__ import annotations + +import pathlib +import sys + +import pytest + +REPO_ROOT = pathlib.Path(__file__).resolve().parents[3] +PYPROJECT = REPO_ROOT / "pyproject.toml" +STUDIO_TXT = REPO_ROOT / "studio" / "backend" / "requirements" / "studio.txt" + +# Imported at module scope by the chain every CLI command walks: structlog via +# studio.backend, click via unsloth_cli/commands/start.py. +CORE_RUNTIME_PACKAGES = ("structlog", "click") + + +def _load_pyproject() -> dict: + if sys.version_info >= (3, 11): + import tomllib + else: + tomllib = pytest.importorskip("tomli") + return tomllib.loads(PYPROJECT.read_text(encoding = "utf-8")) + + +def _requirement_lines(path: pathlib.Path) -> list[str]: + out = [] + for line in path.read_text(encoding = "utf-8").splitlines(): + text = line.split("#", 1)[0].strip() + if text and not text.startswith("-"): + out.append(text) + return out + + +def _normalise(name: str) -> str: + """PEP 503 normalisation, so PyJWT/pyjwt and nest_asyncio/nest-asyncio match.""" + head = name + for sep in ("===", "==", ">=", "<=", "~=", "!=", ">", "<", "[", ";", " "): + idx = head.find(sep) + if idx > 0: + head = head[:idx] + return head.strip().lower().replace("_", "-").replace(".", "-") + + +def test_studio_extra_exists(): + extras = _load_pyproject()["project"]["optional-dependencies"] + assert "studio" in extras, ( + "pyproject.toml has no `studio` extra. The wheel ships studio/ and " + "studio.backend*, so their dependencies need a pip-installable home." + ) + + +def test_studio_extra_matches_requirements_file(): + extras = _load_pyproject()["project"]["optional-dependencies"] + extra = sorted(_normalise(entry) for entry in extras["studio"]) + required = sorted(_normalise(entry) for entry in _requirement_lines(STUDIO_TXT)) + + missing = sorted(set(required) - set(extra)) + surplus = sorted(set(extra) - set(required)) + assert not missing, ( + f"studio.txt lists {missing} but the `studio` extra does not. " + '`pip install "unsloth[studio]"` would build a venv the Studio server ' + "cannot boot in. Add them to [project.optional-dependencies] studio." + ) + assert not surplus, ( + f"The `studio` extra lists {surplus} but studio.txt does not. " + "Remove them, or add them to studio.txt if install.sh needs them too." + ) + + +@pytest.mark.parametrize("package", CORE_RUNTIME_PACKAGES) +def test_cli_runtime_packages_are_core_dependencies(package): + core = [_normalise(entry) for entry in _load_pyproject()["project"]["dependencies"]] + assert _normalise(package) in core, ( + f"{package} is imported at module scope by the studio.backend chain " + f"`unsloth train` / `unsloth export` walk, so a plain `pip install " + f"unsloth` must provide it or they die with ModuleNotFoundError." + ) diff --git a/tests/studio/test_model_picker_contracts.py b/tests/studio/test_model_picker_contracts.py index e1aba66b1b..00ee83efc7 100644 --- a/tests/studio/test_model_picker_contracts.py +++ b/tests/studio/test_model_picker_contracts.py @@ -73,6 +73,57 @@ def test_autoload_records_backend_loaded_model_identity(): assert "m.id === loadedModelId" in autoload +def test_chat_autoload_toast_is_persistent_and_dismissible(): + """Send-triggered autoload stays visible until it settles but remains + dismissible, matching the explicit model-loading toast's lifetime.""" + src = _read("features/chat/api/chat-adapter.ts") + auto_load = src.split("async function autoLoadSmallestModel", 1)[1] + auto_load = auto_load.split("export function createOpenAIStreamAdapter", 1)[0] + assert "toast.loading(" not in auto_load + assert "const updateAutoLoadToast =" in auto_load + assert "if (autoLoadToastDismissed) return;" in auto_load + assert auto_load.count("toast.message(") == 2 + assert auto_load.count("updateAutoLoadToast(") >= 4 + assert "duration: Number.POSITIVE_INFINITY" in auto_load + assert "closeButton: true" in auto_load + assert "icon: createLoadingToastIcon()" in auto_load + assert "onDismiss:" in auto_load + # Terminal success uses a fresh finite toast after manual progress dismissal. + assert "showAutoLoadSuccess" in auto_load + assert "description: undefined" in auto_load + assert "icon: undefined" in auto_load + assert "duration: 5000" in auto_load + assert "duration: 30000" not in auto_load + assert auto_load.count("toast.dismiss(toastId)") >= 4 + + explicit_load = _read("features/chat/hooks/use-chat-model-runtime.ts") + assert "duration: Infinity" in explicit_load + + +def test_recipe_model_load_toast_is_persistent_and_dismissible(): + """Recipe model loading uses the same dismissible persistent lifecycle as + chat loading because both call the non-abortable loadModel API.""" + src = _read("features/recipe-studio/hooks/use-recipe-executions.ts") + model_load = src.split("async function loadLocalModelSelection", 1)[1] + model_load = model_load.split("function getLocalModelLoadPlanForPayload", 1)[0] + assert "toast.loading(" not in model_load + assert "toast.message(" in model_load + assert "duration: Number.POSITIVE_INFINITY" in model_load + assert "closeButton: true" in model_load + assert "icon: createLoadingToastIcon()" in model_load + assert "onDismiss:" in model_load + assert "description: undefined" in model_load + assert "icon: undefined" in model_load + assert "duration: 2000" in model_load + + toast_lib = _read("lib/toast.ts") + assert "createElement(Spinner" in toast_lib + assert 'className: "size-4 text-muted-foreground"' in toast_lib + + sonner = _read("components/ui/sonner.tsx") + assert "loading: createLoadingToastIcon()" in sonner + + def test_rollback_restores_native_lease_expiry_with_token(): """A failed model switch that rolls back to a previously loaded picked GGUF must restore the lease expiry paired with the token, never the token alone diff --git a/unsloth_cli/_inference.py b/unsloth_cli/_inference.py index 3ad901542c..43268423c6 100644 --- a/unsloth_cli/_inference.py +++ b/unsloth_cli/_inference.py @@ -61,15 +61,20 @@ def ensure_studio_backend_path() -> None: def configure_quiet_logging() -> None: import logging - import structlog - # The CLI never configures structlog, so without this every backend INFO # line prints. LOG_LEVEL is exported so the worker subprocess inherits it. level_name = os.environ.setdefault("LOG_LEVEL", "WARNING").upper() level = getattr(logging, level_name, logging.WARNING) - structlog.configure(wrapper_class = structlog.make_filtering_bound_logger(level)) os.environ.setdefault("HF_HUB_DISABLE_PROGRESS_BARS", "1") + # Quieting logs must not fail a command before the import that really needs + # structlog gets to report itself. + try: + import structlog + except ModuleNotFoundError: + return + structlog.configure(wrapper_class = structlog.make_filtering_bound_logger(level)) + def _parse_nonnegative_int(value: Optional[str]) -> Optional[int]: if value is None: @@ -433,7 +438,9 @@ def load_chat_backend( fresh_backend uses a private orchestrator so a second model (compare's base column) can run alongside the main one. """ - with quiet_if_nonzero_mlx_rank(): + from unsloth_cli._studio_deps import studio_backend_imports + + with studio_backend_imports("unsloth inference", studio_only = True), quiet_if_nonzero_mlx_rank(): is_mlx_distributed, rank, _world_size = mlx_distributed_info() if model_config is None: model_config = resolve_model_config(model, hf_token = hf_token) diff --git a/unsloth_cli/_studio_deps.py b/unsloth_cli/_studio_deps.py new file mode 100644 index 0000000000..fa02aabefc --- /dev/null +++ b/unsloth_cli/_studio_deps.py @@ -0,0 +1,255 @@ +# SPDX-License-Identifier: AGPL-3.0-only +# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0 + +"""Studio dependency checks shared by the CLI commands. + +The wheel ships studio/ and studio.backend*, so train / export / chat / +inference / studio all work after a plain `pip install unsloth` right up to the +point they import the backend. studio_backend_imports() turns the resulting +traceback into one sentence and the two commands that fix it. + +Also loads studio/install_manifest.py for `unsloth studio verify-install`. +""" + +from __future__ import annotations + +import contextlib +import importlib.util +import inspect +import re +import sys +from pathlib import Path +from typing import Dict, Iterable, List, Optional, Sequence + +import typer + +# One parent up is the package root: site-packages, or the repo root if editable. +_PACKAGE_ROOT = Path(__file__).resolve().parent.parent + +_MANIFEST_MODULE = None +_MANIFEST_LOADED = False + + +def _manifest_candidates(extra_roots: Sequence[Path] = ()) -> Iterable[Path]: + yield _PACKAGE_ROOT / "studio" / "install_manifest.py" + roots: List[Path] = [Path(sys.prefix), *extra_roots] + for root in roots: + for pattern in ( + "lib/python*/site-packages/studio/install_manifest.py", + "Lib/site-packages/studio/install_manifest.py", + ): + yield from root.glob(pattern) + + +def load_install_manifest_module(extra_roots: Sequence[Path] = ()): + """Load studio/install_manifest.py by file path, or None if unavailable. + + By path for the same reason as studio.backend.run: a partial + site-packages/studio/ tree can shadow an editable install, which is exactly + what this check exists to detect. + """ + global _MANIFEST_MODULE, _MANIFEST_LOADED + if _MANIFEST_LOADED: + return _MANIFEST_MODULE + + _MANIFEST_LOADED = True + for path in _manifest_candidates(extra_roots): + if not path.is_file(): + continue + spec = importlib.util.spec_from_file_location("studio.install_manifest", path) + if spec is None or spec.loader is None: + continue + module = importlib.util.module_from_spec(spec) + try: + spec.loader.exec_module(module) + except Exception: + continue + _MANIFEST_MODULE = module + return _MANIFEST_MODULE + return None + + +def _venv_root_for_module(module) -> Optional[Path]: + """Prefix owning a manifest module, which may be a venv other than ours.""" + path = Path(getattr(module, "__file__", "") or "") + for parent in path.parents: + if (parent / "pyvenv.cfg").is_file(): + return parent + return None + + +def _canonical(name: str) -> str: + """PEP 503 normalisation, so PyJWT / pyjwt / py_jwt compare equal.""" + return re.sub(r"[-_.]+", "-", name).lower() + + +def _resolved(path: Path) -> Path: + try: + return path.resolve() + except OSError: + return path + + +def _venv_site_packages(root: Path) -> List[Path]: + out: List[Path] = [] + for pattern in ("lib/python*/site-packages", "Lib/site-packages"): + out.extend(sorted(root.glob(pattern))) + return out + + +def _managed_root(extra_roots: Sequence[Path]) -> Optional[Path]: + """A requested venv that is not the one this CLI runs in. + + The wheel ships studio/, so a CLI installed outside the managed venv always + finds its own copy of the helper first; without this it would then verify + its own prefix instead of the venv it was asked about. + """ + running = _resolved(Path(sys.prefix)) + for root in extra_roots: + if (root / "pyvenv.cfg").is_file() and _resolved(root) != running: + return root + return None + + +def _distributions_in(root: Path) -> Optional[Dict[str, str]]: + """Canonical distribution name -> version inside another venv. + + importlib.metadata reports the running interpreter only, so a foreign + site-packages has to be handed to the finder explicitly. + """ + paths = [str(path) for path in _venv_site_packages(root)] + if not paths: + return None + from importlib.metadata import Distribution, DistributionFinder + + found: Dict[str, str] = {} + try: + for dist in Distribution.discover(context = DistributionFinder.Context(path = paths)): + name = getattr(dist, "name", None) or dist.metadata["Name"] + if name: + found.setdefault(_canonical(name), dist.version or "") + except Exception: + return None + return found + + +def _requirements_root_in(root: Path) -> Optional[Path]: + for path in _venv_site_packages(root): + reqs = path / "studio" / "backend" / "requirements" + if reqs.is_dir(): + return reqs + return None + + +def _supports_foreign_root(module) -> bool: + """A manifest helper predating the installed= parameter cannot describe another venv.""" + try: + return "installed" in inspect.signature(module.verify_install).parameters + except (TypeError, ValueError): + return False + + +def install_state(extra_roots: Sequence[Path] = ()) -> dict: + """verify_install() result, or incomplete when the helper cannot be loaded. + + studio/install_manifest.py ships in the same wheel as this file, so a tree + that has one without the other is a torn install, not an old one: a CLI + predating both never reaches this code, and the desktop already calls it + stale on desktop_manageability_version. Answering yes here would launch a + backend whose own files may be just as absent. + """ + module = load_install_manifest_module(extra_roots) + if module is None: + return { + "ok": False, + "manifest_ok": False, + "deps_ok": False, + "missing": [], + "reason": "studio_install_manifest_missing", + } + # The requested managed venv is the subject, even though the helper above + # came from this CLI's own tree. + root = _managed_root(extra_roots) or _venv_root_for_module(module) + foreign = root is not None and _resolved(root) != _resolved(Path(sys.prefix)) + installed = _distributions_in(root) if foreign else None + req_root = _requirements_root_in(root) if foreign else None + try: + if installed is not None and req_root is not None and _supports_foreign_root(module): + # That venv's own metadata: unreadable through this interpreter. + return module.verify_install(root = root, req_root = req_root, installed = installed) + state = module.verify_install(root = root) + if foreign and not state["deps_ok"]: + # The manifest came from another venv but the dependency walk ran + # here, so it says nothing about that venv. + state = dict(state, deps_ok = True, missing = []) + state["ok"] = state["manifest_ok"] + state["reason"] = None if state["ok"] else state["reason"] + return state + except Exception as exc: + return { + "ok": False, + "manifest_ok": False, + "deps_ok": False, + "missing": [], + "reason": f"studio_install_check_failed:{type(exc).__name__}", + } + + +def _missing_studio_packages() -> List[str]: + """Studio packages studio.txt asks for and the venv does not have.""" + module = load_install_manifest_module() + if module is None: + return [] + try: + return list(module.missing_requirements()) + except Exception: + return [] + + +# studio.txt names distributions, ModuleNotFoundError names the import. Only +# pairs differing by more than PEP 503 normalisation need an entry, and each +# import name below is itself a real but unrelated PyPI project. +_IMPORT_TO_DISTRIBUTION = { + "jwt": "pyjwt", + "docx": "python-docx", + "fitz": "pymupdf", +} + + +@contextlib.contextmanager +def studio_backend_imports(feature: str = "This command", *, studio_only: bool = False): + """Report a missing dependency as a message instead of a traceback. + + Only ModuleNotFoundError is intercepted; any other ImportError from the + backend is a real bug and keeps its traceback. + """ + try: + yield + except ModuleNotFoundError as exc: + studio_missing = _missing_studio_packages() + # The failed import may not be a studio dependency at all: `train` + # reaches torch through the same wrapper, so only offer the extra when + # it helps. + trigger = exc.name or "" + # Match on the owning distribution, never the import: `pip install jwt` + # (or fastmcp.server) installs the wrong thing or nothing at all. + top = trigger.split(".", 1)[0] + needed = _IMPORT_TO_DISTRIBUTION.get(top, top) + wanted = _canonical(needed) + from_studio = not trigger or any(_canonical(name) == wanted for name in studio_missing) + if studio_only and not from_studio: + raise + typer.echo( + f"Error: {feature} needs {needed or 'a dependency'}, which is not installed.", + err = True, + ) + others = [name for name in studio_missing if _canonical(name) != wanted] + if others: + typer.echo(f" also missing: {', '.join(others)}", err = True) + typer.echo("", err = True) + if not from_studio: + typer.echo(f" Install it: pip install {needed}", err = True) + if from_studio or others: + typer.echo(" Studio install: unsloth studio update", err = True) + typer.echo(' Plain pip: pip install "unsloth[studio]"', err = True) + raise typer.Exit(code = 1) from None diff --git a/unsloth_cli/commands/export.py b/unsloth_cli/commands/export.py index a72ceb5f27..651657abc6 100644 --- a/unsloth_cli/commands/export.py +++ b/unsloth_cli/commands/export.py @@ -6,6 +6,8 @@ from typing import Optional import typer +from unsloth_cli._studio_deps import studio_backend_imports + EXPORT_FORMATS = ["merged-16bit", "merged-4bit", "gguf", "lora"] GGUF_QUANTS = ["q4_k_m", "q5_k_m", "q8_0", "f16"] @@ -17,7 +19,8 @@ def list_checkpoints( ), ): """List checkpoints detected in the outputs directory.""" - from studio.backend.core.export import ExportBackend + with studio_backend_imports("unsloth list-checkpoints"): + from studio.backend.core.export import ExportBackend backend = ExportBackend() checkpoints = backend.scan_checkpoints(outputs_dir = str(outputs_dir)) @@ -72,7 +75,8 @@ def export( typer.echo("Error: --repo-id required when using --push-to-hub", err = True) raise typer.Exit(code = 2) - from studio.backend.core.export import ExportBackend + with studio_backend_imports("unsloth export"): + from studio.backend.core.export import ExportBackend backend = ExportBackend() diff --git a/unsloth_cli/commands/studio.py b/unsloth_cli/commands/studio.py index 1bb0f42016..c2bdbbc915 100644 --- a/unsloth_cli/commands/studio.py +++ b/unsloth_cli/commands/studio.py @@ -22,6 +22,7 @@ from pathlib import Path from typing import List, Literal, Optional import typer +from unsloth_cli import _studio_deps from unsloth_cli.commands import _password_prompt studio_app = typer.Typer(help = "Unsloth Studio commands.") @@ -229,6 +230,15 @@ def _find_run_py() -> Optional[Path]: return None +def _install_state() -> dict: + """verify_install() result for this install root. + + STUDIO_HOME is an extra search root so a CLI installed outside the managed + venv still inspects the venv the desktop app launches. + """ + return _studio_deps.install_state(extra_roots = (STUDIO_HOME / "unsloth_studio",)) + + _RUN_MODULE = None @@ -1555,7 +1565,8 @@ def studio_default( typer.echo("Unsloth Studio not set up. Run install.sh first.") raise typer.Exit(1) - run_mod = _load_run_module() + with _studio_deps.studio_backend_imports("unsloth studio"): + run_mod = _load_run_module() run_server = run_mod.run_server if not silent: @@ -2201,7 +2212,8 @@ def run( os.environ.pop(_START_API_KEY_MARKER_ENV, None) # ── 2. Start server (always suppress built-in banner) ───────────── - run_mod = _load_run_module() + with _studio_deps.studio_backend_imports("unsloth studio"): + run_mod = _load_run_module() run_server = run_mod.run_server # Match the route handlers' import path: run.py adds studio/backend/ to @@ -2804,12 +2816,18 @@ def desktop_capabilities( help = "Emit machine-readable JSON.", ), ): + state = _install_state() payload = { "desktop_protocol_version": 1, - "desktop_manageability_version": 1, + # 2 adds studio_install_ok; the desktop treats < 2 as stale rather than + # guess at an absent field. + "desktop_manageability_version": 2, "supports_provision_desktop_auth": True, "supports_api_only": True, "supports_desktop_backend_ownership": True, + # Did the install finish and are the backend's boot deps still there. + "studio_install_ok": bool(state["ok"]), + "studio_install_reason": state["reason"], "version": "unknown", } try: @@ -2826,6 +2844,36 @@ def desktop_capabilities( typer.echo(f"{key}: {value}") +@studio_app.command("verify-install") +def verify_install( + json_output: bool = typer.Option( + False, + "--json", + help = "Emit machine-readable JSON.", + ), +): + """Check that the Unsloth Studio dependency install completed. + + Exits 0 when complete, 1 otherwise. setup.sh / setup.ps1 use the exit code + to decide whether the "already up to date" fast path may be taken. + """ + state = _install_state() + + if json_output: + typer.echo(json.dumps(state, sort_keys = True)) + raise typer.Exit(0 if state["ok"] else 1) + + if state["ok"]: + typer.echo("Unsloth Studio install is complete.") + raise typer.Exit(0) + + typer.echo(f"Unsloth Studio install is incomplete ({state['reason']}).") + if state["missing"]: + typer.echo(f" missing packages: {', '.join(state['missing'])}") + typer.echo(" repair with: unsloth studio update") + raise typer.Exit(1) + + @studio_app.command("provision-desktop-auth", hidden = True) def provision_desktop_auth(): """Create/repair desktop auth state for the local machine.""" diff --git a/unsloth_cli/commands/train.py b/unsloth_cli/commands/train.py index c52c2344b7..4bff2f2876 100644 --- a/unsloth_cli/commands/train.py +++ b/unsloth_cli/commands/train.py @@ -8,13 +8,15 @@ from typing import Optional import typer from unsloth_cli._inference import ensure_studio_backend_path +from unsloth_cli._studio_deps import studio_backend_imports from unsloth_cli.config import Config, load_config from unsloth_cli.options import add_options_from_config def _should_use_mlx_backend_for_cli() -> bool: ensure_studio_backend_path() - from studio.backend.core.training.training import should_use_mlx_training_backend + with studio_backend_imports("unsloth train"): + from studio.backend.core.training.training import should_use_mlx_training_backend return should_use_mlx_training_backend() @@ -33,12 +35,14 @@ def _create_cli_trainer(model_name: str, hf_token: Optional[str]): _activate_mlx_transformers(model_name, hf_token) # MLX is torch-free: use the lightweight adapter, not trainer.py (imports torch/unsloth/trl at load). ensure_studio_backend_path() - from studio.backend.core.training.training import create_mlx_trainer_adapter + with studio_backend_imports("unsloth train"): + from studio.backend.core.training.training import create_mlx_trainer_adapter return create_mlx_trainer_adapter() ensure_studio_backend_path() - from studio.backend.core.training.trainer import UnslothTrainer + with studio_backend_imports("unsloth train"): + from studio.backend.core.training.trainer import UnslothTrainer return UnslothTrainer()