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/pyproject.toml b/pyproject.toml index 0f57ecf4df..62623499d6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -30,6 +30,12 @@ dependencies = [ "pydantic", "pyyaml", "nest-asyncio", + # Every CLI command imports studio.backend.*, which reaches structlog at + # module level. The rest of the server stack lives in the studio extra. + "structlog>=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] @@ -68,6 +74,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/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/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 4a87fd79e0..ea84068809 100644 --- a/studio/setup.ps1 +++ b/studio/setup.ps1 @@ -2977,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. @@ -3023,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 { diff --git a/studio/setup.sh b/studio/setup.sh index 1cad0e2dbe..f623c9ab0b 100755 --- a/studio/setup.sh +++ b/studio/setup.sh @@ -1044,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 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/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/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/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()