mirror of
https://github.com/unslothai/unsloth.git
synced 2026-04-28 03:19:57 +00:00
fix failed to start on docker
This commit is contained in:
Roland Tannous
parent
bcf5c59d5c
commit
9092544099
1 changed files with 51 additions and 370 deletions
|
|
@ -1,19 +1,11 @@
|
|||
# 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 secrets
|
||||
import sqlite3
|
||||
import subprocess
|
||||
import sys
|
||||
import tempfile
|
||||
import time
|
||||
import types
|
||||
from datetime import datetime, timezone
|
||||
from pathlib import Path
|
||||
from typing import Optional
|
||||
import typer
|
||||
|
|
@ -21,52 +13,12 @@ import typer
|
|||
studio_app = typer.Typer(help = "Unsloth Studio commands.")
|
||||
|
||||
STUDIO_HOME = Path.home() / ".unsloth" / "studio"
|
||||
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 _studio_venv_python() -> Optional[Path]:
|
||||
"""Return the studio venv Python binary, or None if not set up."""
|
||||
if platform.system() == "Windows":
|
||||
|
|
@ -145,228 +97,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,
|
||||
|
|
@ -417,11 +153,6 @@ def studio_default(
|
|||
host: str = typer.Option("0.0.0.0", "--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."""
|
||||
if ctx.invoked_subcommand is not None:
|
||||
|
|
@ -433,49 +164,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)
|
||||
|
|
@ -488,7 +219,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)
|
||||
|
|
@ -760,16 +491,9 @@ def _run_setup_script(*, verbose: bool = False) -> None:
|
|||
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"]
|
||||
)
|
||||
powershell_args.extend(["-ExecutionPolicy", "Bypass", "-File", str(script)])
|
||||
result = subprocess.run(
|
||||
powershell_args,
|
||||
["powershell", "-ExecutionPolicy", "Bypass", "-File", str(script)],
|
||||
env = env,
|
||||
**_windows_hidden_subprocess_kwargs(),
|
||||
)
|
||||
else:
|
||||
result = subprocess.run(["bash", str(script)], env = env)
|
||||
|
|
@ -825,44 +549,6 @@ def update(
|
|||
# ── 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,
|
||||
"supports_provision_desktop_auth": True,
|
||||
"supports_api_only": 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.
|
||||
|
|
@ -873,18 +559,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.")
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue