From 3357d746f3053707d09754835d5d2daa78446bd5 Mon Sep 17 00:00:00 2001 From: Roland Tannous Date: Thu, 23 Apr 2026 20:48:58 +0000 Subject: [PATCH] fix failed to start on docker --- unsloth_cli/commands/studio.py | 333 +++++---------------------------- 1 file changed, 52 insertions(+), 281 deletions(-) diff --git a/unsloth_cli/commands/studio.py b/unsloth_cli/commands/studio.py index 52c84ebdf9..2635b4933e 100644 --- a/unsloth_cli/commands/studio.py +++ b/unsloth_cli/commands/studio.py @@ -1,9 +1,6 @@ # 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 @@ -11,7 +8,6 @@ import secrets import sqlite3 import subprocess import sys -import tempfile import time import types import urllib.error @@ -98,6 +94,7 @@ 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 +STUDIO_HOME = Path.home() / ".unsloth" / "studio" # __file__ is unsloth_cli/commands/studio.py -- two parents up is the package root # (either site-packages or the repo root for editable installs). @@ -234,228 +231,12 @@ def _create_api_key_inprocess(name: str) -> str: ``POST /api/auth/api-keys`` on fresh installs. Safe because the CLI already has filesystem access to ``~/.unsloth/studio``. """ - storage = _load_backend_auth_storage() + from auth.storage import create_api_key, DEFAULT_ADMIN_USERNAME - raw_key, _row = storage.create_api_key( - username = storage.DEFAULT_ADMIN_USERNAME, - name = name, - ) + raw_key, _row = create_api_key(username = 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, @@ -509,11 +290,6 @@ def studio_default( 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. @@ -527,49 +303,49 @@ def 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() - 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 + 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") + # 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) + proc = _sp.Popen(args) + 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) @@ -582,7 +358,7 @@ def studio_default( 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) + run_kwargs = dict(host = host, port = port, silent = silent) if frontend is not None: run_kwargs["frontend_path"] = frontend run_server(**run_kwargs) @@ -1053,7 +829,7 @@ def _run_setup_script(*, verbose: bool = False) -> None: # CREATE_NO_WINDOW. Empty update.log on the windows-latest # CI was the smoking gun (run 25533694490 and 25534292239). result = subprocess.run( - powershell_args, + ["powershell", "-ExecutionPolicy", "Bypass", "-File", str(script)], env = env, stdin = _stream_for_subprocess(sys.stdin), stdout = _stream_for_subprocess(sys.stdout), @@ -1412,18 +1188,13 @@ def reset_password(): """ 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() + pw_file = auth_dir / ".bootstrap_password" - db_file.unlink(missing_ok = True) - for path in stale_files: - path.unlink(missing_ok = True) - - if not had_db: + if not db_file.exists(): typer.echo("No auth database found -- nothing to reset.") raise typer.Exit(0) + db_file.unlink(missing_ok = True) + pw_file.unlink(missing_ok = True) + typer.echo("Auth database deleted. Restart Unsloth Studio to get a new password.")