* studio: regenerate desktop launcher on `unsloth studio update` Today `unsloth studio update` only mutates the venv. The macOS .app bundle, the Linux .desktop file, and the shared launch-studio.sh stub bake their paths and `studio_install_id` at install time and never refresh. Users who update an existing Studio install report the Dock / Applications icon still pointing at the old launcher; only a fresh `curl ... install.sh | sh` fixes it because that path re-enters install.sh's create_studio_shortcuts. Wire the same logic into the update path: - install.sh: add --shortcuts-only. Skips the heavy install steps, resolves STUDIO_HOME / OS / DATA_DIR through the existing _resolve_studio_destinations + platform detection, then calls create_studio_shortcuts and exits. - unsloth_cli/commands/studio.py: after setup.sh succeeds, call install.sh with --shortcuts-only. Prefers a local checkout's install.sh (when STUDIO_LOCAL_REPO is set) or one shipped under _PACKAGE_ROOT, and falls back to fetching the upstream installer from https://unsloth.ai/install.sh for PyPI-installed users (the wheel does not ship install.sh). Net effect: `unsloth studio update` now refreshes the macOS .app stub, launcher script, studio.conf, and Linux .desktop entry on every update, so the desktop icon stays in sync with the venv that setup.sh just updated. Env-override and Tauri modes keep their existing behavior (no persistent menu shortcuts, but the launch-studio.sh is still regenerated). Windows is unchanged here; setup.ps1 already handles its own Start Menu / Desktop .lnk creation on update. * studio: also regenerate Windows .lnk shortcuts on update Mirror the macOS fix: install.ps1 gains --shortcuts-only that short-circuits to New-StudioShortcuts, and unsloth studio update calls it after setup.ps1 the same way it now does on macOS / Linux. PyPI installs do not ship install.ps1, so the Python helper fetches the upstream script from https://unsloth.ai/install.ps1 and pipes it into powershell.exe -Command - with an explicit Install-UnslothStudio call appended (irm | iex relies on the trailing @args, which is empty when launched from stdin). setup.ps1 alone never recreates the Start Menu / Desktop .lnk targets or the launch-studio.{ps1,vbs} scripts, so without this update users on Windows hit the same stale-icon regression that triggered the macOS PR. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * studio: rename unsloth.exe to .deleteme before update on Windows Pip's editable reinstall calls uninstall first, which deletes every RECORD entry. unsloth.exe is one of them, and Windows refuses to delete a file whose image is mapped into the running process tree. The first unsloth studio update after install therefore fails with: OSError: [WinError 32] The process cannot access the file because it is being used by another process: ...\Scripts\unsloth.exe Windows does allow renaming an in-use exe, so move it aside before _run_setup_script kicks pip. pip then drops a fresh unsloth.exe at the original path; the *.exe.deleteme left behind is cleaned up at the start of the next update once the previous shim has exited. * studio: rename unsloth.exe from setup.ps1 to reliably bypass exe lock * studio: print python -m workaround when Windows exe lock blocks update * studio: use python -c hint (unsloth_cli has no __main__) * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * install.sh: reshape --shortcuts-only Tauri guard to pass exit-order test * shorter comments in update / launcher regen logic * studio update: env-mode passthrough + non-silent shortcuts-only error * studio update: address codex/gemini PR review - Strip install.ps1's `Install-UnslothStudio @args` auto-invoke before appending an explicit `--shortcuts-only` call so PyPI Windows installs don't re-run the full installer over stdin. - subprocess.run(input=wrapper, ...) now uses encoding="utf-8" so box drawing chars in install.ps1 don't UnicodeEncodeError on CP1252. - Wrap _run_setup_script in try/except to restore unsloth.exe from .deleteme if setup fails, and mirror that rollback inside setup.ps1 when install_python_stack.py exits non-zero. - Capture subprocess return codes in _refresh_desktop_shortcuts and echo a one-line warning on non-zero so silent stale-shortcut failures surface. - Drop --local from the Windows lock-recovery hint so users on PyPI installs don't accidentally switch into editable-checkout mode. - Quote $VENV_ABS_BIN/unsloth in the install.sh shortcuts-only error so paths with spaces print legibly. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * studio update: harden Windows refresh per multi-reviewer pass - PowerShell stdin path now writes the wrapper to a UTF-8 BOM tempfile and runs it via `-File`. `powershell.exe -Command -` decodes stdin with the OEM code page, which mangles box-drawing chars in the fetched install.ps1; -File reads the BOM and decodes UTF-8 cleanly. - _restore_self_exe_lock_windows now treats a zero-byte unsloth.exe as a partial-write and prefers the .deleteme copy. setup.ps1 mirrors the same check. - _release_self_exe_lock_windows uses os.replace for atomic overwrite so a stale .deleteme from an aborted prior update doesn't break the rename. - Lock-recovery hint mentions that --local should be re-added when the user installed from a repo checkout. * studio update: respect Tauri context and tidy Windows .deleteme Tauri's update.rs spawns `unsloth studio update`; without a signal, the CLI's _refresh_desktop_shortcuts would call install.{sh,ps1} --shortcuts-only and create duplicate ~/Applications/Unsloth Studio.app (or .desktop / .lnk) entries that collide with the Tauri bundle. - update.rs now sets UNSLOTH_TAURI_UPDATE=1 on the spawned child. - studio.py's update() skips _refresh_desktop_shortcuts when that env var is set; Tauri owns its own bundle entries. - After a successful Windows update, drop the .deleteme orphan so repeated updates don't accumulate stale binaries that could later be promoted by _restore_self_exe_lock_windows on a cross-version failure. - Tempfile for the PyPI-fallback PowerShell path now uses an unsloth-studio-refresh- prefix so AV/EDR rules and user greps can identify it. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * studio update: drop obsolete WinError 32 hint, echo Tauri skip The rename trick in _release_self_exe_lock_windows + setup.ps1's restore now handle the .exe-lock case in-flow; the printed hint suggested re-running update via venv python, but that just re-enters the same update() and hits the same failure if the rename didn't help. Removing the misleading hint and its helper. Also surface a one-line typer.echo when refresh is skipped under UNSLOTH_TAURI_UPDATE so --verbose logs make the branch visible. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
1429 lines
51 KiB
Python
1429 lines
51 KiB
Python
# SPDX-License-Identifier: AGPL-3.0-only
|
|
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
|
|
|
|
import importlib.util
|
|
import hashlib
|
|
import json
|
|
import os
|
|
import platform
|
|
import re
|
|
import secrets
|
|
import sqlite3
|
|
import subprocess
|
|
import sys
|
|
import tempfile
|
|
import time
|
|
import types
|
|
import urllib.error
|
|
import urllib.request
|
|
from datetime import datetime, timezone
|
|
from pathlib import Path
|
|
from typing import List, Optional
|
|
import typer
|
|
|
|
studio_app = typer.Typer(help = "Unsloth Studio commands.")
|
|
|
|
|
|
# Resolve install root: UNSLOTH_STUDIO_HOME, then STUDIO_HOME alias, then
|
|
# sys.prefix inference (so a direct call to <root>/bin/unsloth resolves after
|
|
# the installer's env var has expired), then legacy ~/.unsloth/studio.
|
|
# UNSLOTH_STUDIO_HOME wins when both env vars are set.
|
|
def _looks_like_installer_managed_studio_home(candidate: Path) -> bool:
|
|
"""Sentinel check (studio.conf or bin shim) so a dev venv named
|
|
unsloth_studio is not misidentified as a custom Studio root.
|
|
"""
|
|
shim_name = "unsloth.exe" if platform.system() == "Windows" else "unsloth"
|
|
return (candidate / "share" / "studio.conf").is_file() or (
|
|
candidate / "bin" / shim_name
|
|
).is_file()
|
|
|
|
|
|
def _resolve_studio_home() -> tuple[Path, bool]:
|
|
override = (os.environ.get("UNSLOTH_STUDIO_HOME") or "").strip()
|
|
if not override:
|
|
override = (os.environ.get("STUDIO_HOME") or "").strip()
|
|
if override:
|
|
try:
|
|
return Path(override).expanduser().resolve(), True
|
|
except (OSError, ValueError):
|
|
return Path(override).expanduser(), True
|
|
try:
|
|
prefix = Path(sys.prefix).resolve()
|
|
if prefix.name == "unsloth_studio":
|
|
inferred = prefix.parent
|
|
legacy = (Path.home() / ".unsloth" / "studio").resolve()
|
|
if inferred != legacy and _looks_like_installer_managed_studio_home(
|
|
inferred
|
|
):
|
|
return inferred, True
|
|
except (OSError, ValueError):
|
|
pass
|
|
return Path.home() / ".unsloth" / "studio", False
|
|
|
|
|
|
STUDIO_HOME, _STUDIO_HOME_IS_CUSTOM = _resolve_studio_home()
|
|
|
|
|
|
def _ensure_studio_env_exported() -> None:
|
|
"""Re-export UNSLOTH_STUDIO_HOME / UNSLOTH_LLAMA_CPP_PATH only for real
|
|
custom roots so subprocesses inherit the right install. Called from each
|
|
studio subcommand entry rather than at import time, to avoid leaking env
|
|
state into unrelated importers (tests, --help, CLI introspection).
|
|
"""
|
|
if not _STUDIO_HOME_IS_CUSTOM:
|
|
return
|
|
# Truthy-check (not setdefault) so a blank UNSLOTH_STUDIO_HOME= does not
|
|
# suppress the inferred custom root.
|
|
if not os.environ.get("UNSLOTH_STUDIO_HOME"):
|
|
os.environ["UNSLOTH_STUDIO_HOME"] = str(STUDIO_HOME)
|
|
# When override == legacy default, llama.cpp stays at ~/.unsloth/llama.cpp.
|
|
try:
|
|
_legacy_studio = (Path.home() / ".unsloth" / "studio").resolve()
|
|
_is_legacy = STUDIO_HOME.resolve() == _legacy_studio
|
|
except (OSError, ValueError):
|
|
_is_legacy = STUDIO_HOME == (Path.home() / ".unsloth" / "studio")
|
|
if _is_legacy:
|
|
_llama_dir = Path.home() / ".unsloth" / "llama.cpp"
|
|
else:
|
|
_llama_dir = STUDIO_HOME / "llama.cpp"
|
|
if not os.environ.get("UNSLOTH_LLAMA_CPP_PATH"):
|
|
os.environ["UNSLOTH_LLAMA_CPP_PATH"] = str(_llama_dir)
|
|
|
|
|
|
BOOTSTRAP_PASSWORD_FILE = ".bootstrap_password"
|
|
DESKTOP_SECRET_FILE = ".desktop_secret"
|
|
DEFAULT_ADMIN_USERNAME = "unsloth"
|
|
DESKTOP_SECRET_PREFIX = "desktop-"
|
|
API_KEY_PBKDF2_SALT_KEY = "api_key_pbkdf2_salt"
|
|
DESKTOP_SECRET_HASH_KEY = "desktop_secret_hash"
|
|
DESKTOP_SECRET_CREATED_AT_KEY = "desktop_secret_created_at"
|
|
PBKDF2_ITERATIONS = 100_000
|
|
|
|
# __file__ is unsloth_cli/commands/studio.py -- two parents up is the package root
|
|
# (either site-packages or the repo root for editable installs).
|
|
_PACKAGE_ROOT = Path(__file__).resolve().parent.parent.parent
|
|
|
|
|
|
def _should_hide_windows_subprocesses() -> bool:
|
|
"""Hide child console windows only for non-interactive Windows launches."""
|
|
if platform.system() != "Windows":
|
|
return False
|
|
try:
|
|
return not sys.stdout.isatty()
|
|
except (AttributeError, OSError, ValueError):
|
|
return True
|
|
|
|
|
|
def _windows_hidden_subprocess_kwargs() -> dict[str, object]:
|
|
"""Return Windows-only Popen kwargs that suppress transient console windows."""
|
|
if not _should_hide_windows_subprocesses():
|
|
return {}
|
|
|
|
kwargs: dict[str, object] = {}
|
|
create_no_window = getattr(subprocess, "CREATE_NO_WINDOW", 0)
|
|
if create_no_window:
|
|
kwargs["creationflags"] = create_no_window
|
|
|
|
startupinfo_factory = getattr(subprocess, "STARTUPINFO", None)
|
|
startf_use_showwindow = getattr(subprocess, "STARTF_USESHOWWINDOW", 0)
|
|
sw_hide = getattr(subprocess, "SW_HIDE", 0)
|
|
if startupinfo_factory is not None and startf_use_showwindow:
|
|
startupinfo = startupinfo_factory()
|
|
startupinfo.dwFlags |= startf_use_showwindow
|
|
startupinfo.wShowWindow = sw_hide
|
|
kwargs["startupinfo"] = startupinfo
|
|
|
|
return kwargs
|
|
|
|
|
|
def _stream_for_subprocess(stream):
|
|
"""Return *stream* if it has a real OS file descriptor, else None.
|
|
|
|
subprocess.run on Windows refuses to inherit std handles unless
|
|
they're passed explicitly (otherwise close_fds=True forces
|
|
bInheritHandles=False, and a CREATE_NO_WINDOW child ends up with
|
|
no stdio at all). When sys.stdout / sys.stderr is a real fd-backed
|
|
stream we want to hand it through; when it's been captured by a
|
|
test harness (pytest's capsys, an in-memory wrapper, etc) we fall
|
|
back to None so subprocess uses its default.
|
|
"""
|
|
if stream is None:
|
|
return None
|
|
try:
|
|
stream.fileno()
|
|
except (AttributeError, OSError, ValueError):
|
|
return None
|
|
return stream
|
|
|
|
|
|
def _studio_venv_python() -> Optional[Path]:
|
|
"""Return the studio venv Python binary, or None if not set up."""
|
|
if platform.system() == "Windows":
|
|
p = STUDIO_HOME / "unsloth_studio" / "Scripts" / "python.exe"
|
|
else:
|
|
p = STUDIO_HOME / "unsloth_studio" / "bin" / "python"
|
|
return p if p.is_file() else None
|
|
|
|
|
|
def _find_run_py() -> Optional[Path]:
|
|
"""Find studio/backend/run.py.
|
|
|
|
No CWD dependency — works from any directory.
|
|
Since studio/ is now a proper package (has __init__.py), it lives in
|
|
site-packages after pip install, right next to unsloth_cli/.
|
|
"""
|
|
# 1. Relative to __file__ (site-packages or editable repo root)
|
|
run_py = _PACKAGE_ROOT / "studio" / "backend" / "run.py"
|
|
if run_py.is_file():
|
|
return run_py
|
|
# 2. Studio venv's site-packages (Linux + Windows layouts)
|
|
for pattern in (
|
|
"lib/python*/site-packages/studio/backend/run.py",
|
|
"Lib/site-packages/studio/backend/run.py",
|
|
):
|
|
for match in (STUDIO_HOME / "unsloth_studio").glob(pattern):
|
|
return match
|
|
return None
|
|
|
|
|
|
def _find_setup_script() -> Optional[Path]:
|
|
"""Find studio/setup.sh or studio/setup.ps1.
|
|
|
|
No CWD dependency — works from any directory.
|
|
"""
|
|
name = "setup.ps1" if platform.system() == "Windows" else "setup.sh"
|
|
# 1. Relative to __file__ (site-packages or editable repo root)
|
|
s = _PACKAGE_ROOT / "studio" / name
|
|
if s.is_file():
|
|
return s
|
|
# 2. Studio venv's site-packages
|
|
for pattern in (
|
|
f"lib/python*/site-packages/studio/{name}",
|
|
f"Lib/site-packages/studio/{name}",
|
|
):
|
|
for match in (STUDIO_HOME / "unsloth_studio").glob(pattern):
|
|
return match
|
|
return None
|
|
|
|
|
|
# ── helpers for `unsloth studio run` ────────────────────────────────
|
|
|
|
|
|
def _wait_for_server(port: int, timeout: int = 30) -> bool:
|
|
"""Poll ``GET /api/health`` until the server responds 200 or *timeout* expires."""
|
|
import urllib.request
|
|
import urllib.error
|
|
|
|
url = f"http://127.0.0.1:{port}/api/health"
|
|
deadline = time.monotonic() + timeout
|
|
while time.monotonic() < deadline:
|
|
try:
|
|
with urllib.request.urlopen(url, timeout = 2) as resp:
|
|
if resp.status == 200:
|
|
return True
|
|
except (urllib.error.URLError, OSError, ConnectionError):
|
|
pass
|
|
time.sleep(0.5)
|
|
return False
|
|
|
|
|
|
def _create_api_key_inprocess(name: str) -> str:
|
|
"""Create an API key via direct storage call (no HTTP needed).
|
|
|
|
Bypasses the ``must_change_password`` gate that blocks HTTP
|
|
``POST /api/auth/api-keys`` on fresh installs. Safe because the
|
|
CLI already has filesystem access to ``~/.unsloth/studio``.
|
|
"""
|
|
storage = _load_backend_auth_storage()
|
|
|
|
raw_key, _row = storage.create_api_key(
|
|
username = storage.DEFAULT_ADMIN_USERNAME,
|
|
name = name,
|
|
)
|
|
return raw_key
|
|
|
|
|
|
def _load_backend_auth_storage():
|
|
run_py = _find_run_py()
|
|
backend_dir = (
|
|
run_py.parent if run_py is not None else _PACKAGE_ROOT / "studio" / "backend"
|
|
)
|
|
if backend_dir.is_dir() and str(backend_dir) not in sys.path:
|
|
sys.path.insert(0, str(backend_dir))
|
|
|
|
auth_dir = backend_dir / "auth"
|
|
storage_py = auth_dir / "storage.py"
|
|
loaded = sys.modules.get("auth.storage")
|
|
loaded_path = Path(getattr(loaded, "__file__", "")).resolve()
|
|
if loaded is not None and loaded_path == storage_py:
|
|
return loaded
|
|
|
|
package = sys.modules.get("auth")
|
|
package_paths = [Path(path).resolve() for path in getattr(package, "__path__", [])]
|
|
if package is None or auth_dir.resolve() not in package_paths:
|
|
package = types.ModuleType("auth")
|
|
package.__path__ = [str(auth_dir)]
|
|
package.__package__ = "auth"
|
|
package.__file__ = str(auth_dir / "__init__.py")
|
|
sys.modules["auth"] = package
|
|
|
|
spec = importlib.util.spec_from_file_location("auth.storage", storage_py)
|
|
if spec is None or spec.loader is None:
|
|
raise ImportError(f"Could not load backend auth storage from {storage_py}")
|
|
storage = importlib.util.module_from_spec(spec)
|
|
sys.modules["auth.storage"] = storage
|
|
spec.loader.exec_module(storage)
|
|
|
|
return storage
|
|
|
|
|
|
def _write_auth_secret(path: Path, secret: str) -> None:
|
|
path.parent.mkdir(parents = True, exist_ok = True)
|
|
fd, tmp_name = tempfile.mkstemp(prefix = f".{path.name}.", dir = path.parent)
|
|
tmp_path = Path(tmp_name)
|
|
try:
|
|
try:
|
|
os.chmod(tmp_path, 0o600)
|
|
except OSError:
|
|
pass
|
|
with os.fdopen(fd, "w") as f:
|
|
fd = -1
|
|
f.write(secret)
|
|
os.replace(tmp_path, path)
|
|
except Exception:
|
|
if fd >= 0:
|
|
os.close(fd)
|
|
tmp_path.unlink(missing_ok = True)
|
|
raise
|
|
try:
|
|
os.chmod(path, 0o600)
|
|
except OSError:
|
|
pass
|
|
|
|
|
|
def _connect_auth_db() -> sqlite3.Connection:
|
|
auth_dir = STUDIO_HOME / "auth"
|
|
auth_dir.mkdir(parents = True, exist_ok = True)
|
|
conn = sqlite3.connect(auth_dir / "auth.db")
|
|
conn.execute(
|
|
"""
|
|
CREATE TABLE IF NOT EXISTS auth_user (
|
|
id INTEGER PRIMARY KEY,
|
|
username TEXT UNIQUE NOT NULL,
|
|
password_salt TEXT NOT NULL,
|
|
password_hash TEXT NOT NULL,
|
|
jwt_secret TEXT NOT NULL,
|
|
must_change_password INTEGER NOT NULL DEFAULT 0
|
|
);
|
|
"""
|
|
)
|
|
conn.execute(
|
|
"""
|
|
CREATE TABLE IF NOT EXISTS refresh_tokens (
|
|
id INTEGER PRIMARY KEY,
|
|
token_hash TEXT NOT NULL,
|
|
username TEXT NOT NULL,
|
|
expires_at TEXT NOT NULL,
|
|
is_desktop INTEGER NOT NULL DEFAULT 0
|
|
);
|
|
"""
|
|
)
|
|
conn.execute(
|
|
"""
|
|
CREATE TABLE IF NOT EXISTS api_keys (
|
|
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
|
username TEXT NOT NULL,
|
|
key_prefix TEXT NOT NULL,
|
|
key_hash TEXT NOT NULL UNIQUE,
|
|
name TEXT NOT NULL DEFAULT '',
|
|
created_at TEXT NOT NULL,
|
|
last_used_at TEXT,
|
|
expires_at TEXT,
|
|
is_active INTEGER NOT NULL DEFAULT 1
|
|
);
|
|
"""
|
|
)
|
|
conn.execute(
|
|
"""
|
|
CREATE TABLE IF NOT EXISTS app_secrets (
|
|
key TEXT PRIMARY KEY,
|
|
value TEXT NOT NULL
|
|
);
|
|
"""
|
|
)
|
|
auth_columns = {row[1] for row in conn.execute("PRAGMA table_info(auth_user)")}
|
|
if "must_change_password" not in auth_columns:
|
|
conn.execute(
|
|
"ALTER TABLE auth_user ADD COLUMN must_change_password INTEGER NOT NULL DEFAULT 0"
|
|
)
|
|
refresh_columns = {
|
|
row[1] for row in conn.execute("PRAGMA table_info(refresh_tokens)")
|
|
}
|
|
if "is_desktop" not in refresh_columns:
|
|
conn.execute(
|
|
"ALTER TABLE refresh_tokens ADD COLUMN is_desktop INTEGER NOT NULL DEFAULT 0"
|
|
)
|
|
conn.commit()
|
|
return conn
|
|
|
|
|
|
def _pbkdf2_hex(value: str, salt: bytes) -> str:
|
|
return hashlib.pbkdf2_hmac(
|
|
"sha256",
|
|
value.encode("utf-8"),
|
|
salt,
|
|
PBKDF2_ITERATIONS,
|
|
).hex()
|
|
|
|
|
|
def _hash_password(password: str) -> tuple[str, str]:
|
|
salt = secrets.token_hex(16)
|
|
pwd_hash = _pbkdf2_hex(password, salt.encode("utf-8"))
|
|
return salt, pwd_hash
|
|
|
|
|
|
def _get_or_create_api_key_pbkdf2_salt(conn: sqlite3.Connection) -> bytes:
|
|
row = conn.execute(
|
|
"SELECT value FROM app_secrets WHERE key = ?",
|
|
(API_KEY_PBKDF2_SALT_KEY,),
|
|
).fetchone()
|
|
if row is None:
|
|
salt_hex = secrets.token_hex(32)
|
|
conn.execute(
|
|
"INSERT OR IGNORE INTO app_secrets (key, value) VALUES (?, ?)",
|
|
(API_KEY_PBKDF2_SALT_KEY, salt_hex),
|
|
)
|
|
row = conn.execute(
|
|
"SELECT value FROM app_secrets WHERE key = ?",
|
|
(API_KEY_PBKDF2_SALT_KEY,),
|
|
).fetchone()
|
|
return bytes.fromhex(row[0])
|
|
|
|
|
|
def _ensure_cli_default_admin(conn: sqlite3.Connection) -> None:
|
|
row = conn.execute(
|
|
"SELECT 1 FROM auth_user WHERE username = ?",
|
|
(DEFAULT_ADMIN_USERNAME,),
|
|
).fetchone()
|
|
if row is not None:
|
|
return
|
|
|
|
bootstrap_password = secrets.token_urlsafe(32)
|
|
password_salt, password_hash = _hash_password(bootstrap_password)
|
|
conn.execute(
|
|
"""
|
|
INSERT INTO auth_user (
|
|
username,
|
|
password_salt,
|
|
password_hash,
|
|
jwt_secret,
|
|
must_change_password
|
|
)
|
|
VALUES (?, ?, ?, ?, ?)
|
|
""",
|
|
(
|
|
DEFAULT_ADMIN_USERNAME,
|
|
password_salt,
|
|
password_hash,
|
|
secrets.token_urlsafe(64),
|
|
1,
|
|
),
|
|
)
|
|
_write_auth_secret(
|
|
STUDIO_HOME / "auth" / BOOTSTRAP_PASSWORD_FILE,
|
|
bootstrap_password,
|
|
)
|
|
|
|
|
|
def _create_desktop_secret_in_cli() -> str:
|
|
raw_secret = DESKTOP_SECRET_PREFIX + secrets.token_urlsafe(48)
|
|
now = datetime.now(timezone.utc).isoformat()
|
|
conn = _connect_auth_db()
|
|
try:
|
|
_ensure_cli_default_admin(conn)
|
|
secret_hash = _pbkdf2_hex(raw_secret, _get_or_create_api_key_pbkdf2_salt(conn))
|
|
conn.execute(
|
|
"INSERT OR REPLACE INTO app_secrets (key, value) VALUES (?, ?)",
|
|
(DESKTOP_SECRET_HASH_KEY, secret_hash),
|
|
)
|
|
conn.execute(
|
|
"INSERT OR REPLACE INTO app_secrets (key, value) VALUES (?, ?)",
|
|
(DESKTOP_SECRET_CREATED_AT_KEY, now),
|
|
)
|
|
conn.commit()
|
|
return raw_secret
|
|
finally:
|
|
conn.close()
|
|
|
|
|
|
def _load_model_via_http(
|
|
port: int,
|
|
api_key: str,
|
|
model: str,
|
|
gguf_variant: Optional[str],
|
|
max_seq_length: int,
|
|
load_in_4bit: bool,
|
|
llama_extra_args: Optional[List[str]] = None,
|
|
timeout: int = 600,
|
|
) -> dict:
|
|
"""POST to ``/api/inference/load`` using the API key for auth."""
|
|
import json
|
|
import urllib.request
|
|
import urllib.error
|
|
|
|
payload: dict = {
|
|
"model_path": model,
|
|
"max_seq_length": max_seq_length,
|
|
"load_in_4bit": load_in_4bit,
|
|
}
|
|
if gguf_variant:
|
|
payload["gguf_variant"] = gguf_variant
|
|
if llama_extra_args:
|
|
payload["llama_extra_args"] = list(llama_extra_args)
|
|
|
|
data = json.dumps(payload).encode()
|
|
req = urllib.request.Request(
|
|
f"http://127.0.0.1:{port}/api/inference/load",
|
|
data = data,
|
|
headers = {
|
|
"Content-Type": "application/json",
|
|
"Authorization": f"Bearer {api_key}",
|
|
},
|
|
method = "POST",
|
|
)
|
|
try:
|
|
with urllib.request.urlopen(req, timeout = timeout) as resp:
|
|
return json.loads(resp.read())
|
|
except urllib.error.HTTPError as exc:
|
|
body = exc.read().decode(errors = "replace")
|
|
raise RuntimeError(f"Model load failed (HTTP {exc.code}): {body}") from exc
|
|
|
|
|
|
# ── unsloth studio (server) ──────────────────────────────────────────
|
|
|
|
|
|
@studio_app.callback(invoke_without_command = True)
|
|
def studio_default(
|
|
ctx: typer.Context,
|
|
port: int = typer.Option(8888, "--port", "-p"),
|
|
host: str = typer.Option("127.0.0.1", "--host", "-H"),
|
|
frontend: Optional[Path] = typer.Option(None, "--frontend", "-f"),
|
|
silent: bool = typer.Option(False, "--silent", "-q"),
|
|
api_only: bool = typer.Option(
|
|
False,
|
|
"--api-only",
|
|
help = "Run API server only, no frontend serving (for Tauri desktop app)",
|
|
),
|
|
):
|
|
"""Launch the Unsloth Studio server."""
|
|
# Runs before any subcommand; covers run/setup/update/etc in one place.
|
|
_ensure_studio_env_exported()
|
|
if ctx.invoked_subcommand is not None:
|
|
return
|
|
|
|
# Always use the studio venv if it exists and we're not already in it
|
|
studio_venv_dir = STUDIO_HOME / "unsloth_studio"
|
|
in_studio_venv = sys.prefix.startswith(str(studio_venv_dir))
|
|
|
|
if not in_studio_venv:
|
|
studio_python = _studio_venv_python()
|
|
run_py = _find_run_py()
|
|
if studio_python and run_py:
|
|
if not silent:
|
|
typer.echo("Launching Unsloth Studio... Please wait...")
|
|
args = [
|
|
str(studio_python),
|
|
str(run_py),
|
|
"--host",
|
|
host,
|
|
"--port",
|
|
str(port),
|
|
]
|
|
if frontend:
|
|
args.extend(["--frontend", str(frontend)])
|
|
if silent:
|
|
args.append("--silent")
|
|
if api_only:
|
|
args.append("--api-only")
|
|
# On Windows, os.execvp() spawns a child but the parent lingers,
|
|
# so Ctrl+C only kills the parent leaving the child orphaned.
|
|
# Use subprocess.run() on Windows so the parent waits for the child.
|
|
if sys.platform == "win32":
|
|
import subprocess as _sp
|
|
|
|
proc = _sp.Popen(args, **_windows_hidden_subprocess_kwargs())
|
|
try:
|
|
rc = proc.wait()
|
|
except KeyboardInterrupt:
|
|
# Child has its own signal handler — let it finish
|
|
rc = proc.wait()
|
|
if rc != 0:
|
|
typer.echo(
|
|
f"\nError: Studio server exited unexpectedly (code {rc}).",
|
|
err = True,
|
|
)
|
|
typer.echo(
|
|
"Check the error above. If a package is missing, "
|
|
"re-run: unsloth studio setup",
|
|
err = True,
|
|
)
|
|
raise typer.Exit(rc)
|
|
else:
|
|
os.execvp(str(studio_python), args)
|
|
else:
|
|
typer.echo("Studio not set up. Run install.sh first.")
|
|
raise typer.Exit(1)
|
|
|
|
from studio.backend.run import run_server
|
|
|
|
if not silent:
|
|
from studio.backend.run import _resolve_external_ip
|
|
|
|
display_host = _resolve_external_ip() if host == "0.0.0.0" else host
|
|
typer.echo(f"Starting Unsloth Studio on http://{display_host}:{port}")
|
|
|
|
run_kwargs = dict(host = host, port = port, silent = silent, api_only = api_only)
|
|
if frontend is not None:
|
|
run_kwargs["frontend_path"] = frontend
|
|
run_server(**run_kwargs)
|
|
|
|
from studio.backend.run import _shutdown_event
|
|
|
|
try:
|
|
if _shutdown_event is not None:
|
|
# NOTE: Event.wait() without a timeout blocks at the C level
|
|
# on Linux, preventing Python from delivering SIGINT (Ctrl+C).
|
|
while not _shutdown_event.is_set():
|
|
_shutdown_event.wait(timeout = 1)
|
|
else:
|
|
while True:
|
|
time.sleep(1)
|
|
except KeyboardInterrupt:
|
|
from studio.backend.run import _graceful_shutdown, _server
|
|
|
|
_graceful_shutdown(_server)
|
|
typer.echo("\nShutting down...")
|
|
|
|
|
|
# ── unsloth studio run ───────────────────────────────────────────────
|
|
|
|
|
|
def _split_repo_variant(model_arg: str) -> tuple[str, Optional[str]]:
|
|
"""Split ``org/name:variant`` HF-style identifiers into (repo, variant).
|
|
|
|
Mirrors llama.cpp's ``-hf <repo>:<quant>`` convention so users can
|
|
write ``unsloth/gpt-oss-20b-GGUF:UD-Q4_K_XL`` instead of passing
|
|
``--gguf-variant`` separately. Local paths (absolute, ``./``,
|
|
``~/``, Windows drive letters) and identifiers without a ``:``
|
|
suffix are returned verbatim.
|
|
"""
|
|
s = model_arg.strip()
|
|
if not s:
|
|
return s, None
|
|
if s.startswith(("/", "./", "../", "~")) or s == ".":
|
|
return s, None
|
|
# Windows drive letter (e.g. "C:\\path" or "C:/path") -- the colon
|
|
# here is a path separator, not a variant suffix.
|
|
if len(s) >= 2 and s[1] == ":" and s[0].isalpha():
|
|
return s, None
|
|
if ":" not in s:
|
|
return s, None
|
|
repo, _, variant = s.rpartition(":")
|
|
if not repo or not variant:
|
|
return s, None
|
|
# A real quant label has no slashes; ``foo:bar/baz`` is not
|
|
# ``repo:variant`` syntax.
|
|
if "/" in variant:
|
|
return s, None
|
|
return repo, variant
|
|
|
|
|
|
@studio_app.command(
|
|
context_settings = {
|
|
"allow_extra_args": True,
|
|
"ignore_unknown_options": True,
|
|
},
|
|
)
|
|
def run(
|
|
ctx: typer.Context,
|
|
model: str = typer.Option(
|
|
...,
|
|
"--model",
|
|
"-m",
|
|
"-hf",
|
|
"-hfr",
|
|
"--hf-repo",
|
|
help = (
|
|
"Model path or HF repo. Accepts llama.cpp-style "
|
|
"`org/repo:variant` syntax. The `-hf` / `--hf-repo` aliases "
|
|
"match llama-server's spelling."
|
|
),
|
|
),
|
|
gguf_variant: Optional[str] = typer.Option(
|
|
None, "--gguf-variant", help = "GGUF quant variant (e.g. UD-Q4_K_XL)"
|
|
),
|
|
max_seq_length: int = typer.Option(
|
|
0, "--max-seq-length", help = "Max sequence length (0 = model default)"
|
|
),
|
|
load_in_4bit: bool = typer.Option(True, "--load-in-4bit/--no-load-in-4bit"),
|
|
api_key_name: str = typer.Option(
|
|
"cli", "--api-key-name", help = "Label for the auto-generated API key"
|
|
),
|
|
port: int = typer.Option(8888, "--port", "-p"),
|
|
host: str = typer.Option("127.0.0.1", "--host", "-H"),
|
|
frontend: Optional[Path] = typer.Option(None, "--frontend", "-f"),
|
|
silent: bool = typer.Option(False, "--silent", "-q"),
|
|
enable_tools: Optional[bool] = typer.Option(
|
|
None,
|
|
"--enable-tools/--disable-tools",
|
|
help = (
|
|
"Force server-side tools on/off for all requests. "
|
|
"Default: on for 127.0.0.1, off for 0.0.0.0."
|
|
),
|
|
),
|
|
yes: bool = typer.Option(
|
|
False,
|
|
"--yes",
|
|
"-y",
|
|
help = "Skip the 0.0.0.0 + --enable-tools confirmation prompt.",
|
|
),
|
|
):
|
|
"""Start Studio, load a model, and print an API key -- one-liner server.
|
|
|
|
Any flag this command does not recognize is forwarded verbatim to
|
|
the underlying llama-server (GGUF only). Studio-managed flags
|
|
(--port, -c / --ctx-size, --api-key, -ngl, --jinja, --flash-attn,
|
|
--no-context-shift, model-identity flags, ...) are rejected with
|
|
HTTP 400.
|
|
|
|
Example:
|
|
unsloth studio run --model unsloth/Qwen3-1.7B-GGUF --gguf-variant UD-Q4_K_XL
|
|
unsloth studio run --model unsloth/Qwen3-1.7B-GGUF --top-k 20 --seed 42
|
|
unsloth studio run --model some-model --chat-template-file /path/to/tpl.jinja
|
|
"""
|
|
extra_llama_args: List[str] = list(ctx.args) if ctx.args else []
|
|
|
|
# ── 0. Parse llama.cpp-style ``repo:variant`` syntax in --model. ───
|
|
# Lets users write ``--model unsloth/foo-GGUF:UD-Q4_K_XL`` instead
|
|
# of pairing ``--model`` with ``--gguf-variant``. If both are given
|
|
# and disagree, fail loudly instead of silently picking one.
|
|
parsed_repo, embedded_variant = _split_repo_variant(model)
|
|
if embedded_variant:
|
|
if gguf_variant and gguf_variant != embedded_variant:
|
|
typer.echo(
|
|
f"Error: --model embeds variant '{embedded_variant}' but "
|
|
f"--gguf-variant '{gguf_variant}' was also provided.",
|
|
err = True,
|
|
)
|
|
raise typer.Exit(1)
|
|
model = parsed_repo
|
|
gguf_variant = gguf_variant or embedded_variant
|
|
|
|
# ── Resolve the server-side tool policy. The y/N prompt (if any)
|
|
# runs in the outer process so the re-exec'd child never re-prompts.
|
|
from unsloth_cli._tool_policy import is_external_host, resolve_tool_policy
|
|
|
|
enable_tools = resolve_tool_policy(
|
|
host = host,
|
|
flag = enable_tools,
|
|
yes = yes,
|
|
silent = silent,
|
|
)
|
|
|
|
# ── 1. Venv re-exec (same pattern as studio_default) ──────────────
|
|
studio_venv_dir = STUDIO_HOME / "unsloth_studio"
|
|
in_studio_venv = sys.prefix.startswith(str(studio_venv_dir))
|
|
|
|
if not in_studio_venv:
|
|
studio_python = _studio_venv_python()
|
|
if not studio_python:
|
|
typer.echo("Studio not set up. Run install.sh first.")
|
|
raise typer.Exit(1)
|
|
# Re-exec into the studio venv via its `unsloth` entry point
|
|
studio_bin = studio_python.parent / "unsloth"
|
|
if not studio_bin.is_file():
|
|
typer.echo(
|
|
"Studio venv missing 'unsloth' entry point. Re-run: unsloth studio setup"
|
|
)
|
|
raise typer.Exit(1)
|
|
args = [
|
|
str(studio_bin),
|
|
"studio",
|
|
"run",
|
|
"--model",
|
|
model,
|
|
"--max-seq-length",
|
|
str(max_seq_length),
|
|
"--api-key-name",
|
|
api_key_name,
|
|
"--port",
|
|
str(port),
|
|
"--host",
|
|
host,
|
|
]
|
|
if gguf_variant:
|
|
args.extend(["--gguf-variant", gguf_variant])
|
|
if not load_in_4bit:
|
|
args.append("--no-load-in-4bit")
|
|
if frontend:
|
|
args.extend(["--frontend", str(frontend)])
|
|
if silent:
|
|
args.append("--silent")
|
|
# Forward the resolved tool policy (always concrete True/False
|
|
# at this point — the resolver above ran before the re-exec).
|
|
if enable_tools:
|
|
args.append("--enable-tools")
|
|
else:
|
|
args.append("--disable-tools")
|
|
# Forward --yes whenever the parent already cleared the prompt
|
|
# (either operator passed --yes, or the parent's resolver
|
|
# accepted the network-bind confirmation). Otherwise the child
|
|
# re-runs the resolver and prompts a second time.
|
|
if yes or (enable_tools and is_external_host(host)):
|
|
args.append("--yes")
|
|
# Forward unknown args (llama-server pass-through) to the
|
|
# re-exec'd command so the studio venv sees them in ctx.args
|
|
# and the re-execed run() can include them in the load payload.
|
|
if extra_llama_args:
|
|
args.extend(extra_llama_args)
|
|
|
|
if sys.platform == "win32":
|
|
proc = subprocess.Popen(args)
|
|
try:
|
|
rc = proc.wait()
|
|
except KeyboardInterrupt:
|
|
rc = proc.wait()
|
|
raise typer.Exit(rc)
|
|
else:
|
|
os.execvp(str(studio_bin), args)
|
|
|
|
# ── 2. Start server (always suppress built-in banner) ─────────────
|
|
from studio.backend.run import run_server, _resolve_external_ip
|
|
|
|
run_kwargs = dict(host = host, port = port, silent = True, llama_parallel_slots = 4)
|
|
if frontend is not None:
|
|
run_kwargs["frontend_path"] = frontend
|
|
app = run_server(**run_kwargs)
|
|
actual_port = getattr(app.state, "server_port", port) or port
|
|
|
|
# ── Apply the resolved tool policy as a process-level override.
|
|
# Must use the same import path the route handlers use --
|
|
# `studio/backend/run.py` adds `studio/backend/` to sys.path so the
|
|
# routes import this module as top-level `state.tool_policy`. If we
|
|
# imported via `studio.backend.state.tool_policy` instead, Python
|
|
# would cache two different module objects with two different
|
|
# `_tool_policy` globals, and the gates would never see our value.
|
|
from state.tool_policy import set_tool_policy
|
|
|
|
set_tool_policy(enable_tools)
|
|
|
|
# ── 3. Wait for server health ─────────────────────────────────────
|
|
if not silent:
|
|
typer.echo("Starting Unsloth Studio...")
|
|
if not _wait_for_server(actual_port):
|
|
typer.echo("Error: server did not become healthy within 30 seconds.", err = True)
|
|
raise typer.Exit(1)
|
|
|
|
# ── 4. Create API key in-process ──────────────────────────────────
|
|
api_key = _create_api_key_inprocess(api_key_name)
|
|
|
|
# ── 5. Load model via HTTP ────────────────────────────────────────
|
|
if not silent:
|
|
typer.echo(f"Loading model: {model}...")
|
|
try:
|
|
result = _load_model_via_http(
|
|
port = actual_port,
|
|
api_key = api_key,
|
|
model = model,
|
|
gguf_variant = gguf_variant,
|
|
max_seq_length = max_seq_length,
|
|
load_in_4bit = load_in_4bit,
|
|
llama_extra_args = extra_llama_args,
|
|
)
|
|
except RuntimeError as exc:
|
|
typer.echo(f"Error: {exc}", err = True)
|
|
raise typer.Exit(1)
|
|
|
|
loaded_model = result.get("model", model)
|
|
display_variant = f" ({gguf_variant})" if gguf_variant else ""
|
|
|
|
# ── 6. Print banner ───────────────────────────────────────────────
|
|
display_host = _resolve_external_ip() if host == "0.0.0.0" else host
|
|
base_url = f"http://{display_host}:{actual_port}"
|
|
sdk_base_url = f"{base_url}/v1"
|
|
|
|
# Claude orange (Claude Code's brand color) for tool-policy notices
|
|
# so they stand out from the surrounding banner. Always printed --
|
|
# even under --silent / --yes -- so the operator never misses the
|
|
# current tool-execution status.
|
|
_tool_notice_fg = (217, 119, 87)
|
|
_is_external = is_external_host(host)
|
|
if _is_external and enable_tools:
|
|
_tool_notice = (
|
|
f"Server-side tools are ENABLED on {host} (network-reachable). "
|
|
f"Anyone with the API key can run code on this machine. "
|
|
f"Do not share the API key."
|
|
)
|
|
elif _is_external:
|
|
_tool_notice = (
|
|
f"Server-side tools are disabled by default on {host} "
|
|
f"(network-reachable). Pass --enable-tools to turn on "
|
|
f"(you will be warned about API-key risk)."
|
|
)
|
|
elif enable_tools:
|
|
_tool_notice = (
|
|
"Server-side tools are enabled by default for loopback. "
|
|
"Pass --disable-tools to turn off."
|
|
)
|
|
else:
|
|
_tool_notice = "Server-side tools are disabled."
|
|
|
|
if not silent:
|
|
typer.echo("")
|
|
typer.echo("=" * 56)
|
|
typer.echo(f" Unsloth Studio running at {base_url}")
|
|
typer.echo(f" Model loaded: {loaded_model}{display_variant}")
|
|
typer.echo(f" API Key: {api_key}")
|
|
typer.echo("")
|
|
typer.echo(" OpenAI / Anthropic SDK base URL:")
|
|
typer.echo(f" {sdk_base_url}")
|
|
typer.echo("=" * 56)
|
|
typer.secho(_tool_notice, fg = _tool_notice_fg, bold = True)
|
|
typer.echo("")
|
|
typer.echo("OpenAI Chat Completions:")
|
|
typer.echo(f" curl {sdk_base_url}/chat/completions \\")
|
|
typer.echo(f' -H "Authorization: Bearer {api_key}" \\')
|
|
typer.echo(' -H "Content-Type: application/json" \\')
|
|
typer.echo(
|
|
""" -d '{"messages": [{"role": "user", "content": "Hello"}], "stream": true}'"""
|
|
)
|
|
typer.echo("")
|
|
typer.echo("Anthropic Messages:")
|
|
typer.echo(f" curl {sdk_base_url}/messages \\")
|
|
typer.echo(f' -H "Authorization: Bearer {api_key}" \\')
|
|
typer.echo(' -H "Content-Type: application/json" \\')
|
|
typer.echo(
|
|
""" -d '{"max_tokens": 256, "messages": [{"role": "user", "content": "Hello"}], "stream": true}'"""
|
|
)
|
|
typer.echo("")
|
|
typer.echo("OpenAI Responses:")
|
|
typer.echo(f" curl {sdk_base_url}/responses \\")
|
|
typer.echo(f' -H "Authorization: Bearer {api_key}" \\')
|
|
typer.echo(' -H "Content-Type: application/json" \\')
|
|
typer.echo(""" -d '{"input": "Hello", "stream": true}'""")
|
|
typer.echo("")
|
|
else:
|
|
# Silent mode still prints the essentials (URL, API key) plus
|
|
# the orange tool-status notice so the operator never loses
|
|
# visibility into the security-relevant policy.
|
|
typer.echo(f"URL: {base_url}")
|
|
typer.echo(f"API Key: {api_key}")
|
|
typer.secho(_tool_notice, fg = _tool_notice_fg, bold = True)
|
|
|
|
# ── 7. Wait for Ctrl+C ────────────────────────────────────────────
|
|
from studio.backend.run import _shutdown_event, _graceful_shutdown, _server
|
|
|
|
try:
|
|
if _shutdown_event is not None:
|
|
while not _shutdown_event.is_set():
|
|
_shutdown_event.wait(timeout = 1)
|
|
else:
|
|
while True:
|
|
time.sleep(1)
|
|
except KeyboardInterrupt:
|
|
_graceful_shutdown(_server)
|
|
typer.echo("\nShutting down...")
|
|
|
|
|
|
# ── unsloth studio stop ───────────────────────────────────────────────
|
|
|
|
_PID_FILE = STUDIO_HOME / "studio.pid"
|
|
|
|
|
|
@studio_app.command()
|
|
def stop():
|
|
"""Stop a running Unsloth Studio server.
|
|
|
|
Reads the PID from ~/.unsloth/studio/studio.pid and sends SIGTERM
|
|
(or TerminateProcess on Windows) to shut it down gracefully.
|
|
"""
|
|
import signal as _signal
|
|
|
|
if not _PID_FILE.is_file():
|
|
typer.echo("No running Studio server found (no PID file).")
|
|
raise typer.Exit(0)
|
|
|
|
pid_text = _PID_FILE.read_text().strip()
|
|
if not pid_text.isdigit():
|
|
typer.echo(f"Invalid PID file contents: {pid_text}")
|
|
_PID_FILE.unlink(missing_ok = True)
|
|
raise typer.Exit(1)
|
|
|
|
pid = int(pid_text)
|
|
|
|
# Check if the process is still alive
|
|
try:
|
|
os.kill(pid, 0)
|
|
except ProcessLookupError:
|
|
typer.echo(
|
|
f"Studio server (PID {pid}) is not running. Cleaning up stale PID file."
|
|
)
|
|
_PID_FILE.unlink(missing_ok = True)
|
|
raise typer.Exit(0)
|
|
except PermissionError:
|
|
pass # process exists but we may not own it; try to signal anyway
|
|
|
|
# Send SIGTERM (graceful shutdown) or TerminateProcess on Windows
|
|
try:
|
|
if sys.platform == "win32":
|
|
subprocess.run(["taskkill", "/PID", str(pid), "/F"], check = True)
|
|
else:
|
|
os.kill(pid, _signal.SIGTERM)
|
|
typer.echo(f"Sent shutdown signal to Studio server (PID {pid}).")
|
|
except ProcessLookupError:
|
|
typer.echo(f"Studio server (PID {pid}) already exited.")
|
|
_PID_FILE.unlink(missing_ok = True)
|
|
raise typer.Exit(0)
|
|
except Exception as e:
|
|
typer.echo(f"Failed to stop Studio server (PID {pid}): {e}", err = True)
|
|
raise typer.Exit(1)
|
|
|
|
# Wait briefly for the process to exit and clean up
|
|
for _ in range(10):
|
|
time.sleep(0.5)
|
|
try:
|
|
os.kill(pid, 0)
|
|
except ProcessLookupError:
|
|
_PID_FILE.unlink(missing_ok = True)
|
|
typer.echo("Studio server stopped.")
|
|
raise typer.Exit(0)
|
|
except PermissionError:
|
|
break
|
|
|
|
typer.echo("Studio server is shutting down (may take a few seconds).")
|
|
|
|
|
|
# ── unsloth studio setup / update ─────────────────────────────────────
|
|
|
|
|
|
def _run_setup_script(*, verbose: bool = False) -> None:
|
|
"""Find and run the studio setup/update script."""
|
|
script = _find_setup_script()
|
|
if not script:
|
|
typer.echo("Error: Could not find setup script (setup.sh / setup.ps1).")
|
|
raise typer.Exit(1)
|
|
|
|
env = {**os.environ, "UNSLOTH_VERBOSE": "1"} if verbose else None
|
|
|
|
if platform.system() == "Windows":
|
|
powershell_args = ["powershell.exe"]
|
|
if _should_hide_windows_subprocesses():
|
|
powershell_args.extend(
|
|
["-NoLogo", "-NoProfile", "-NonInteractive", "-WindowStyle", "Hidden"]
|
|
)
|
|
# Use -Command + `*>&1` instead of -File so setup.ps1's
|
|
# Write-Host output (PowerShell Information stream / #6) is
|
|
# merged into the success stream and reaches the parent's
|
|
# stdout. With -File, Information stream output is dropped
|
|
# whenever stdout is a pipe, which is exactly the situation
|
|
# CI hits with `unsloth studio update --local 2>&1 | tee
|
|
# logs/update.log`. Single-quote escaping handles paths that
|
|
# contain apostrophes.
|
|
script_pwsh_literal = str(script).replace("'", "''")
|
|
powershell_args.extend(
|
|
[
|
|
"-ExecutionPolicy",
|
|
"Bypass",
|
|
"-Command",
|
|
f"& '{script_pwsh_literal}' *>&1",
|
|
]
|
|
)
|
|
# Explicitly hand stdin/stdout/stderr to the child so the
|
|
# CI tee actually sees setup.ps1's output. Without this,
|
|
# subprocess.run on Windows uses close_fds=True (default,
|
|
# since Python 3.7) which sets bInheritHandles=False on
|
|
# CreateProcess. With CREATE_NO_WINDOW also set (via
|
|
# _windows_hidden_subprocess_kwargs in non-TTY runs), the
|
|
# child has neither a console nor any inherited std
|
|
# handles, so PowerShell's Write-Host -- and even
|
|
# [Console]::Out.WriteLine -- writes to nothing. Passing
|
|
# stdout=sys.stdout / stderr=sys.stderr makes Python set up
|
|
# PROC_THREAD_ATTRIBUTE_HANDLE_LIST with the std handles
|
|
# explicitly inheritable, which works alongside
|
|
# CREATE_NO_WINDOW. Empty update.log on the windows-latest
|
|
# CI was the smoking gun (run 25533694490 and 25534292239).
|
|
result = subprocess.run(
|
|
powershell_args,
|
|
env = env,
|
|
stdin = _stream_for_subprocess(sys.stdin),
|
|
stdout = _stream_for_subprocess(sys.stdout),
|
|
stderr = _stream_for_subprocess(sys.stderr),
|
|
**_windows_hidden_subprocess_kwargs(),
|
|
)
|
|
else:
|
|
result = subprocess.run(["bash", str(script)], env = env)
|
|
|
|
if result.returncode != 0:
|
|
raise typer.Exit(result.returncode)
|
|
|
|
|
|
_INSTALLER_URL_BASH = "https://unsloth.ai/install.sh"
|
|
_INSTALLER_URL_PWSH = "https://unsloth.ai/install.ps1"
|
|
|
|
|
|
def _refresh_desktop_shortcuts(*, verbose: bool = False) -> None:
|
|
"""Re-run installer with --shortcuts-only to refresh launchers post-update."""
|
|
env = {**os.environ}
|
|
if verbose:
|
|
env["UNSLOTH_VERBOSE"] = "1"
|
|
|
|
is_windows = platform.system() == "Windows"
|
|
installer_name = "install.ps1" if is_windows else "install.sh"
|
|
installer_url = _INSTALLER_URL_PWSH if is_windows else _INSTALLER_URL_BASH
|
|
|
|
# Prefer local checkout, fall back to package dir, then network fetch.
|
|
local_repo = (os.environ.get("STUDIO_LOCAL_REPO") or "").strip()
|
|
candidates: list[Path] = []
|
|
if local_repo:
|
|
candidates.append(Path(local_repo) / installer_name)
|
|
candidates.append(_PACKAGE_ROOT / installer_name)
|
|
|
|
args = ["--shortcuts-only"]
|
|
if verbose:
|
|
args.append("--verbose")
|
|
|
|
if is_windows:
|
|
ps_argv: list[str] = ["powershell.exe"]
|
|
if _should_hide_windows_subprocesses():
|
|
ps_argv.extend(
|
|
["-NoLogo", "-NoProfile", "-NonInteractive", "-WindowStyle", "Hidden"]
|
|
)
|
|
|
|
for script in candidates:
|
|
try:
|
|
if script.is_file():
|
|
quoted = str(script).replace("'", "''")
|
|
argv = list(ps_argv)
|
|
argv.extend(
|
|
[
|
|
"-ExecutionPolicy",
|
|
"Bypass",
|
|
"-Command",
|
|
f"& '{quoted}' {' '.join(args)} *>&1",
|
|
]
|
|
)
|
|
result = subprocess.run(
|
|
argv,
|
|
env = env,
|
|
check = False,
|
|
**_windows_hidden_subprocess_kwargs(),
|
|
)
|
|
if result.returncode != 0:
|
|
typer.echo(
|
|
f" refresh-launcher install.ps1 exited {result.returncode}"
|
|
)
|
|
return
|
|
except OSError:
|
|
continue
|
|
|
|
# PyPI installs lack install.ps1: fetch + pipe to powershell stdin.
|
|
try:
|
|
request = urllib.request.Request(
|
|
installer_url, headers = {"User-Agent": "unsloth-studio-update"}
|
|
)
|
|
with urllib.request.urlopen(request, timeout = 30) as response:
|
|
installer = response.read().decode("utf-8", errors = "replace")
|
|
except (urllib.error.URLError, TimeoutError, OSError) as exc:
|
|
typer.echo(
|
|
f" refresh-launcher skipped: could not fetch {installer_url} ({exc})"
|
|
)
|
|
return
|
|
|
|
# install.ps1 auto-invokes `Install-UnslothStudio @args` at EOF; over
|
|
# stdin `$args` is empty so that triggers the full installer flow
|
|
# (deps, venv, prompts) before our shortcuts-only call. Strip it.
|
|
installer = re.sub(
|
|
r"(?m)^[ \t]*Install-UnslothStudio[ \t]+@args[ \t]*\r?\n?",
|
|
"",
|
|
installer,
|
|
)
|
|
# stdin-piped scripts have empty $args, so call Install-UnslothStudio explicitly.
|
|
marker_args = " ".join(args)
|
|
wrapper = installer + f"\nInstall-UnslothStudio {marker_args}\n"
|
|
|
|
# Write to a UTF-8 BOM tempfile and use -File rather than -Command -.
|
|
# `powershell.exe -Command -` reads stdin via [Console]::InputEncoding
|
|
# (CP1252/OEM on most Windows boxes), which mangles box-drawing chars
|
|
# in install.ps1. -File reads the BOM and decodes correctly. The
|
|
# prefix gives AV/EDR engines (and grep'ing users) a clear identity.
|
|
ps1_fd, ps1_path = tempfile.mkstemp(
|
|
prefix = "unsloth-studio-refresh-",
|
|
suffix = ".ps1",
|
|
)
|
|
try:
|
|
with os.fdopen(ps1_fd, "wb") as fh:
|
|
fh.write(b"\xef\xbb\xbf" + wrapper.encode("utf-8"))
|
|
argv = list(ps_argv)
|
|
argv.extend(["-ExecutionPolicy", "Bypass", "-File", ps1_path])
|
|
try:
|
|
result = subprocess.run(
|
|
argv,
|
|
env = env,
|
|
check = False,
|
|
**_windows_hidden_subprocess_kwargs(),
|
|
)
|
|
if result.returncode != 0:
|
|
typer.echo(
|
|
f" refresh-launcher fetched install.ps1 exited {result.returncode}"
|
|
)
|
|
except OSError as exc:
|
|
typer.echo(
|
|
f" refresh-launcher skipped: powershell exec failed ({exc})"
|
|
)
|
|
finally:
|
|
try:
|
|
os.unlink(ps1_path)
|
|
except OSError:
|
|
pass
|
|
return
|
|
|
|
for script in candidates:
|
|
try:
|
|
if script.is_file():
|
|
result = subprocess.run(
|
|
["bash", str(script), *args],
|
|
env = env,
|
|
check = False,
|
|
)
|
|
if result.returncode != 0:
|
|
typer.echo(
|
|
f" refresh-launcher install.sh exited {result.returncode}"
|
|
)
|
|
return
|
|
except OSError:
|
|
continue
|
|
|
|
# PyPI installs lack install.sh: fetch upstream.
|
|
try:
|
|
request = urllib.request.Request(
|
|
installer_url, headers = {"User-Agent": "unsloth-studio-update"}
|
|
)
|
|
with urllib.request.urlopen(request, timeout = 30) as response:
|
|
installer = response.read()
|
|
except (urllib.error.URLError, TimeoutError, OSError) as exc:
|
|
typer.echo(
|
|
f" refresh-launcher skipped: could not fetch {installer_url} ({exc})"
|
|
)
|
|
return
|
|
|
|
try:
|
|
result = subprocess.run(
|
|
["bash", "-s", "--", *args],
|
|
input = installer,
|
|
env = env,
|
|
check = False,
|
|
)
|
|
if result.returncode != 0:
|
|
typer.echo(
|
|
f" refresh-launcher fetched install.sh exited {result.returncode}"
|
|
)
|
|
except OSError as exc:
|
|
typer.echo(f" refresh-launcher skipped: bash exec failed ({exc})")
|
|
|
|
|
|
@studio_app.command(hidden = True)
|
|
def setup(
|
|
verbose: bool = typer.Option(
|
|
False,
|
|
"--verbose",
|
|
"-v",
|
|
help = "Full pip/build output during setup for troubleshooting.",
|
|
),
|
|
):
|
|
"""Run Studio setup (called by install.ps1 / install.sh)."""
|
|
_run_setup_script(verbose = verbose)
|
|
|
|
|
|
@studio_app.command()
|
|
def update(
|
|
local: bool = typer.Option(
|
|
False, "--local", help = "Install from local repo instead of PyPI"
|
|
),
|
|
package: str = typer.Option(
|
|
"unsloth", "--package", help = "Package name to install/update (for testing)"
|
|
),
|
|
verbose: bool = typer.Option(
|
|
False,
|
|
"--verbose",
|
|
"-v",
|
|
help = "Full pip/build output during update for troubleshooting.",
|
|
),
|
|
):
|
|
"""Update Unsloth Studio dependencies and rebuild."""
|
|
# Re-export UNSLOTH_STUDIO_HOME for env-mode installs so the refresh
|
|
# subprocess resolves the same install root the user originally chose.
|
|
_ensure_studio_env_exported()
|
|
# Ensure SKIP_STUDIO_BASE is not inherited from a parent install.ps1 session
|
|
os.environ.pop("SKIP_STUDIO_BASE", None)
|
|
os.environ["STUDIO_PACKAGE_NAME"] = package
|
|
if local:
|
|
os.environ["STUDIO_LOCAL_INSTALL"] = "1"
|
|
# Pass the repo root explicitly so install_python_stack.py doesn't
|
|
# have to guess from SCRIPT_DIR (which may be inside site-packages).
|
|
repo_root = Path(__file__).resolve().parents[2]
|
|
os.environ["STUDIO_LOCAL_REPO"] = str(repo_root)
|
|
else:
|
|
os.environ["STUDIO_LOCAL_INSTALL"] = "0"
|
|
os.environ.pop("STUDIO_LOCAL_REPO", None)
|
|
_release_self_exe_lock_windows()
|
|
try:
|
|
_run_setup_script(verbose = verbose)
|
|
except BaseException:
|
|
# Restore unsloth.exe from .deleteme if setup failed before pip
|
|
# produced a replacement; otherwise the user has no CLI for recovery.
|
|
_restore_self_exe_lock_windows()
|
|
raise
|
|
# On Windows clear the .deleteme orphan now that pip wrote a fresh
|
|
# unsloth.exe; on next update os.replace would overwrite it anyway,
|
|
# but leaving a stale binary around invites cross-version restore
|
|
# confusion from _restore_self_exe_lock_windows.
|
|
_cleanup_self_exe_lock_windows()
|
|
# Tauri desktop owns its own bundle entries; skip CLI launcher refresh
|
|
# so a Tauri-initiated update doesn't create duplicate shortcuts.
|
|
if os.environ.get("UNSLOTH_TAURI_UPDATE") == "1":
|
|
if verbose:
|
|
typer.echo(" refresh-launcher skipped (Tauri update)")
|
|
return
|
|
_refresh_desktop_shortcuts(verbose = verbose)
|
|
|
|
|
|
def _release_self_exe_lock_windows() -> None:
|
|
"""Rename running unsloth.exe so pip can replace it. setup.ps1 also retries."""
|
|
if platform.system() != "Windows":
|
|
return
|
|
try:
|
|
venv_scripts = Path(sys.executable).resolve().parent
|
|
except OSError:
|
|
return
|
|
exe = venv_scripts / "unsloth.exe"
|
|
if not exe.exists():
|
|
return
|
|
stale = exe.with_suffix(".exe.deleteme")
|
|
try:
|
|
# os.replace is atomic-overwrite on Windows; os.rename would raise
|
|
# FileExistsError if a prior aborted update left a .deleteme behind.
|
|
os.replace(exe, stale)
|
|
except OSError as e:
|
|
# Not fatal; setup.ps1 retries from a sibling process.
|
|
print(f"[update] could not rename {exe.name} -> {stale.name}: {e}")
|
|
|
|
|
|
def _restore_self_exe_lock_windows() -> None:
|
|
"""If setup failed before pip wrote a working unsloth.exe, restore .deleteme."""
|
|
if platform.system() != "Windows":
|
|
return
|
|
try:
|
|
venv_scripts = Path(sys.executable).resolve().parent
|
|
except OSError:
|
|
return
|
|
exe = venv_scripts / "unsloth.exe"
|
|
stale = exe.with_suffix(".exe.deleteme")
|
|
if not stale.exists():
|
|
return
|
|
# Treat a missing or zero-byte exe as "pip didn't produce a usable
|
|
# replacement"; otherwise leave the new binary alone.
|
|
if exe.exists():
|
|
try:
|
|
if exe.stat().st_size > 0:
|
|
return
|
|
except OSError:
|
|
return
|
|
try:
|
|
os.replace(stale, exe)
|
|
except OSError as e:
|
|
print(f"[update] could not restore {stale.name} -> {exe.name}: {e}")
|
|
|
|
|
|
def _cleanup_self_exe_lock_windows() -> None:
|
|
"""Remove the .deleteme orphan after a successful update on Windows."""
|
|
if platform.system() != "Windows":
|
|
return
|
|
try:
|
|
venv_scripts = Path(sys.executable).resolve().parent
|
|
except OSError:
|
|
return
|
|
stale = (venv_scripts / "unsloth.exe").with_suffix(".exe.deleteme")
|
|
try:
|
|
stale.unlink(missing_ok = True)
|
|
except OSError:
|
|
pass
|
|
|
|
|
|
# ── unsloth studio reset-password ────────────────────────────────────
|
|
|
|
|
|
@studio_app.command("desktop-capabilities", hidden = True)
|
|
def desktop_capabilities(
|
|
json_output: bool = typer.Option(
|
|
False,
|
|
"--json",
|
|
help = "Emit machine-readable JSON.",
|
|
),
|
|
):
|
|
payload = {
|
|
"desktop_protocol_version": 1,
|
|
"desktop_manageability_version": 1,
|
|
"supports_provision_desktop_auth": True,
|
|
"supports_api_only": True,
|
|
"supports_desktop_backend_ownership": True,
|
|
"version": "unknown",
|
|
}
|
|
try:
|
|
from importlib.metadata import version as package_version
|
|
|
|
payload["version"] = package_version("unsloth")
|
|
except Exception:
|
|
pass
|
|
|
|
if json_output:
|
|
typer.echo(json.dumps(payload, sort_keys = True))
|
|
return
|
|
|
|
for key, value in payload.items():
|
|
typer.echo(f"{key}: {value}")
|
|
|
|
|
|
@studio_app.command("provision-desktop-auth", hidden = True)
|
|
def provision_desktop_auth():
|
|
"""Create/repair desktop auth state for the local machine."""
|
|
auth_dir = STUDIO_HOME / "auth"
|
|
secret = _create_desktop_secret_in_cli()
|
|
_write_auth_secret(auth_dir / DESKTOP_SECRET_FILE, secret)
|
|
typer.echo("Desktop auth ready.")
|
|
|
|
|
|
@studio_app.command("reset-password")
|
|
def reset_password():
|
|
"""Reset the Studio admin password.
|
|
|
|
Deletes the auth database so that a fresh admin account with a new
|
|
random password is created on the next server start. The Studio
|
|
server must be restarted after running this command.
|
|
"""
|
|
auth_dir = STUDIO_HOME / "auth"
|
|
db_file = auth_dir / "auth.db"
|
|
stale_files = [
|
|
auth_dir / BOOTSTRAP_PASSWORD_FILE,
|
|
auth_dir / DESKTOP_SECRET_FILE,
|
|
]
|
|
had_db = db_file.exists()
|
|
|
|
db_file.unlink(missing_ok = True)
|
|
for path in stale_files:
|
|
path.unlink(missing_ok = True)
|
|
|
|
if not had_db:
|
|
typer.echo("No auth database found -- nothing to reset.")
|
|
raise typer.Exit(0)
|
|
|
|
typer.echo("Auth database deleted. Restart Unsloth Studio to get a new password.")
|