fix failed to start on docker
This commit is contained in:
Roland Tannous
parent
5d3fea9cdf
commit
3357d746f3
1 changed files with 53 additions and 282 deletions
|
|
@ -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.")
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue