Compare commits

...
Sign in to create a new pull request.

6 commits

Author SHA1 Message Date
Roland Tannous
799f239ce8 remove model load timeouts, add elapsed-time logging 2026-03-19 18:08:13 +00:00
Roland Tannous
4dbbc0945c replace setup.sh and install python stack files 2026-03-19 16:43:02 +00:00
Roland Tannous
a2c3f39e12 setup: support CUDA_ARCHS_EXTRA env for Docker builds 2026-03-19 16:33:51 +00:00
Roland Tannous
d447e80620 compile for cuda compatibility 86 for A10 2026-03-19 16:33:51 +00:00
Lee Jackson
0e66986d55 studio: enable training progress stream without JWT when auth disabled (#4422)
* fix(studio): enable training progress stream without JWT when auth disabled

* fix(studio): support auth-disabled training runtime sync

* fix(studio): harden auth-disabled training runtime hydration and retry logic
2026-03-19 16:33:51 +00:00
Lee Jackson
7b73882c68 feat(studio): disable auth flow for HF Spaces deployment (#4375) 2026-03-19 16:33:51 +00:00
11 changed files with 403 additions and 349 deletions

View file

@ -1,6 +1,7 @@
# 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 os
import secrets
from datetime import datetime, timedelta, timezone
from typing import Optional, Tuple
@ -21,7 +22,28 @@ ALGORITHM = "HS256"
ACCESS_TOKEN_EXPIRE_MINUTES = 60
REFRESH_TOKEN_EXPIRE_DAYS = 7
security = HTTPBearer() # Reads Authorization: Bearer <token>
security = HTTPBearer(auto_error = False) # Reads Authorization: Bearer <token>
def is_auth_disabled() -> bool:
"""
Return True when auth should be bypassed.
For the HF Spaces branch, auth is disabled by default so users can access
Studio without login/signup. You can explicitly re-enable auth with:
UNSLOTH_STUDIO_ENABLE_AUTH=1
"""
raw_enable = os.getenv("UNSLOTH_STUDIO_ENABLE_AUTH", "").strip().lower()
if raw_enable in {"1", "true", "yes", "on"}:
return False
raw = os.getenv("UNSLOTH_STUDIO_DISABLE_AUTH", "").strip().lower()
if raw in {"1", "true", "yes", "on"}:
return True
if raw in {"0", "false", "no", "off"}:
return False
return True
def _get_secret_for_subject(subject: str) -> str:
@ -103,7 +125,7 @@ def reload_secret() -> None:
async def get_current_subject(
credentials: HTTPAuthorizationCredentials = Depends(security),
credentials: Optional[HTTPAuthorizationCredentials] = Depends(security),
) -> str:
"""Validate JWT and require the password-change flow to be completed."""
return await _get_current_subject(
@ -113,7 +135,7 @@ async def get_current_subject(
async def get_current_subject_allow_password_change(
credentials: HTTPAuthorizationCredentials = Depends(security),
credentials: Optional[HTTPAuthorizationCredentials] = Depends(security),
) -> str:
"""Validate JWT but allow access to the password-change endpoint."""
return await _get_current_subject(
@ -123,7 +145,7 @@ async def get_current_subject_allow_password_change(
async def _get_current_subject(
credentials: HTTPAuthorizationCredentials,
credentials: Optional[HTTPAuthorizationCredentials],
*,
allow_password_change: bool,
) -> str:
@ -136,6 +158,15 @@ async def _get_current_subject(
async def secure_endpoint(current_subject: str = Depends(get_current_subject)):
...
"""
if is_auth_disabled():
return "hf-space-user"
if credentials is None:
raise HTTPException(
status_code = status.HTTP_401_UNAUTHORIZED,
detail = "Missing bearer token",
)
token = credentials.credentials
subject = _decode_subject_without_verification(token)
if subject is None:

View file

@ -1010,8 +1010,8 @@ class LlamaCppBackend:
self._is_vision = is_vision
self._model_identifier = model_identifier
# Wait for llama-server to become healthy
if not self._wait_for_health(timeout = 120.0):
# Wait for llama-server to become healthy (no timeout — let it take as long as needed)
if not self._wait_for_health():
self._kill_process()
raise RuntimeError(
"llama-server failed to start. "
@ -1129,39 +1129,49 @@ class LlamaCppBackend:
"""atexit handler to ensure llama-server is terminated."""
self._kill_process()
def _wait_for_health(self, timeout: float = 120.0, interval: float = 0.5) -> bool:
def _wait_for_health(self, interval: float = 0.5) -> bool:
"""
Poll llama-server's /health endpoint until it responds 200.
No timeout waits indefinitely so large models on slow hardware
(e.g. HF Spaces) have time to load. Logs elapsed time every 30s.
Also monitors subprocess for early exit/crash.
"""
deadline = time.monotonic() + timeout
start = time.monotonic()
last_log = start
url = f"http://127.0.0.1:{self._port}/health"
while time.monotonic() < deadline:
while True:
# Check if process crashed
if self._process.poll() is not None:
# Give the drain thread a moment to collect final output
if self._stdout_thread is not None:
self._stdout_thread.join(timeout = 2)
output = "\n".join(self._stdout_lines[-50:])
elapsed = time.monotonic() - start
logger.error(
f"llama-server exited with code {self._process.returncode}. "
f"Output: {output[:2000]}"
f"llama-server exited with code {self._process.returncode} "
f"after {elapsed:.1f}s. Output: {output[:2000]}"
)
return False
try:
resp = httpx.get(url, timeout = 2.0)
if resp.status_code == 200:
elapsed = time.monotonic() - start
logger.info(f"llama-server healthy after {elapsed:.1f}s")
return True
except (httpx.ConnectError, httpx.TimeoutException):
pass
time.sleep(interval)
# Periodic progress logging
now = time.monotonic()
if now - last_log >= 30.0:
elapsed = now - start
logger.info(f"Waiting for llama-server health check... {elapsed:.0f}s elapsed")
last_log = now
logger.error(f"llama-server health check timed out after {timeout}s")
return False
time.sleep(interval)
# ── Message building (OpenAI format) ──────────────────────────

View file

@ -262,28 +262,56 @@ class InferenceOrchestrator:
except (EOFError, OSError, ValueError):
return None
def _wait_response(self, expected_type: str, timeout: float = 120.0) -> dict:
def _wait_response(self, expected_type: str, timeout: Optional[float] = None) -> dict:
"""Block until a response of the expected type arrives.
Also handles 'status' and 'error' events during the wait.
Returns the matching response dict.
Raises RuntimeError on timeout or subprocess crash.
Raises RuntimeError on subprocess crash.
If timeout is None, waits indefinitely (logs progress every 30s).
"""
deadline = time.monotonic() + timeout
start = time.monotonic()
last_log = start
while time.monotonic() < deadline:
remaining = max(0.1, deadline - time.monotonic())
resp = self._read_resp(timeout = min(remaining, 1.0))
while True:
if timeout is not None:
remaining = timeout - (time.monotonic() - start)
if remaining <= 0:
raise RuntimeError(
f"Timeout waiting for '{expected_type}' response after {timeout}s"
)
poll_timeout = min(remaining, 1.0)
else:
poll_timeout = 1.0
resp = self._read_resp(timeout = poll_timeout)
if resp is None:
# Check subprocess health
if not self._ensure_subprocess_alive():
raise RuntimeError("Inference subprocess crashed during wait")
elapsed = time.monotonic() - start
raise RuntimeError(
f"Inference subprocess crashed during wait after {elapsed:.1f}s"
)
# Periodic progress logging
now = time.monotonic()
if now - last_log >= 30.0:
elapsed = now - start
logger.info(
"Waiting for '%s' response... %.0fs elapsed",
expected_type,
elapsed,
)
last_log = now
continue
rtype = resp.get("type", "")
if rtype == expected_type:
elapsed = time.monotonic() - start
logger.info(
"Received '%s' response after %.1fs", expected_type, elapsed
)
return resp
if rtype == "error":
@ -301,10 +329,6 @@ class InferenceOrchestrator:
expected_type,
)
raise RuntimeError(
f"Timeout waiting for '{expected_type}' response after {timeout}s"
)
def _drain_queue(self) -> list:
"""Drain all pending responses."""
events = []
@ -614,7 +638,7 @@ class InferenceOrchestrator:
needed_major,
)
self._spawn_subprocess(sub_config)
resp = self._wait_response("loaded", timeout = 180)
resp = self._wait_response("loaded")
# Update local state from response
if resp.get("success"):

View file

@ -156,6 +156,7 @@ def _handle_load(backend, config: dict, resp_queue: Any) -> None:
except Exception as e:
logger.warning("Could not read adapter_config.json: %s", e)
load_start = time.monotonic()
success = backend.load_model(
config = mc,
max_seq_length = config.get("max_seq_length", 2048),
@ -163,6 +164,8 @@ def _handle_load(backend, config: dict, resp_queue: Any) -> None:
hf_token = hf_token,
trust_remote_code = config.get("trust_remote_code", False),
)
load_elapsed = time.monotonic() - load_start
logger.info("Inference model load took %.1fs (success=%s)", load_elapsed, success)
if success:
# Build model_info for the parent to mirror

View file

@ -415,6 +415,7 @@ def run_training_process(
# ── 4c. Load training model (uses VRAM — dataset already formatted) ──
_send_status(event_queue, "Loading model...")
load_start = time.monotonic()
success = trainer.load_model(
model_name = model_name,
max_seq_length = config["max_seq_length"],
@ -425,6 +426,8 @@ def run_training_process(
is_dataset_audio = config.get("is_dataset_audio", False),
trust_remote_code = config.get("trust_remote_code", False),
)
load_elapsed = time.monotonic() - load_start
logger.info("Training model load took %.1fs (success=%s)", load_elapsed, success)
if not success or trainer.should_stop:
if trainer.should_stop:
event_queue.put(

View file

@ -34,6 +34,10 @@ class AuthStatusResponse(BaseModel):
...,
description = "True if the seeded admin must still change the default password",
)
auth_disabled: bool = Field(
default = False,
description = "True when auth is bypassed (HF Spaces mode).",
)
class ChangePasswordRequest(BaseModel):

View file

@ -20,6 +20,7 @@ from auth.authentication import (
create_refresh_token,
get_current_subject,
get_current_subject_allow_password_change,
is_auth_disabled,
refresh_access_token,
)
@ -34,6 +35,14 @@ async def auth_status() -> AuthStatusResponse:
- initialized = False -> frontend should wait for the seeded admin bootstrap.
- initialized = True -> frontend should show login or force the first password change.
"""
if is_auth_disabled():
return AuthStatusResponse(
initialized = True,
default_username = storage.DEFAULT_ADMIN_USERNAME,
requires_password_change = False,
auth_disabled = True,
)
return AuthStatusResponse(
initialized = storage.is_initialized(),
default_username = storage.DEFAULT_ADMIN_USERNAME,
@ -42,6 +51,7 @@ async def auth_status() -> AuthStatusResponse:
)
if storage.is_initialized()
else True,
auth_disabled = False,
)

View file

@ -10,6 +10,22 @@ import {
refreshSession,
} from "@/features/auth";
type AuthStatus = {
initialized: boolean;
requires_password_change: boolean;
auth_disabled?: boolean;
};
async function getAuthStatus(): Promise<AuthStatus | null> {
try {
const res = await fetch("/api/auth/status");
if (!res.ok) return null;
return (await res.json()) as AuthStatus;
} catch {
return null;
}
}
async function hasActiveSession(): Promise<boolean> {
if (hasAuthToken()) return true;
if (!hasRefreshToken()) return false;
@ -17,28 +33,22 @@ async function hasActiveSession(): Promise<boolean> {
}
async function checkAuthInitialized(): Promise<boolean> {
try {
const res = await fetch("/api/auth/status");
if (!res.ok) return true; // fallback to login on error
const data = (await res.json()) as { initialized: boolean };
return data.initialized;
} catch {
return true; // fallback to login on error
}
const status = await getAuthStatus();
if (status?.auth_disabled) return true;
return status?.initialized ?? true; // fallback to login on error
}
async function checkPasswordChangeRequired(): Promise<boolean> {
try {
const res = await fetch("/api/auth/status");
if (!res.ok) return mustChangePassword();
const data = (await res.json()) as { requires_password_change: boolean };
return data.requires_password_change || mustChangePassword();
} catch {
return mustChangePassword();
}
const status = await getAuthStatus();
if (status?.auth_disabled) return false;
if (!status) return mustChangePassword();
return status.requires_password_change || mustChangePassword();
}
export async function requireAuth(): Promise<void> {
const status = await getAuthStatus();
if (status?.auth_disabled) return;
if (await hasActiveSession()) {
if (await checkPasswordChangeRequired()) {
throw redirect({ to: "/change-password" });
@ -52,11 +62,21 @@ export async function requireAuth(): Promise<void> {
}
export async function requireGuest(): Promise<void> {
const status = await getAuthStatus();
if (status?.auth_disabled) {
throw redirect({ to: getPostAuthRoute() });
}
if (!(await hasActiveSession())) return;
throw redirect({ to: getPostAuthRoute() });
}
export async function requirePasswordChangeFlow(): Promise<void> {
const status = await getAuthStatus();
if (status?.auth_disabled) {
throw redirect({ to: getPostAuthRoute() });
}
const requiresPasswordChange = await checkPasswordChangeRequired();
if (requiresPasswordChange) return;

View file

@ -15,6 +15,15 @@ import type { TrainingRuntimeStore } from "../types/runtime";
const STATUS_POLL_INTERVAL_MS = 3000;
const METRICS_POLL_INTERVAL_MS = 5000;
const STREAM_RECONNECT_DELAY_MS = 1500;
const AUTH_STATUS_RETRY_INTERVAL_MS = 3000;
const AUTH_STATUS_TIMEOUT_MS = 3000;
const INITIAL_HYDRATE_TIMEOUT_MS = 4000;
function wait(ms: number): Promise<void> {
return new Promise((resolve) => {
setTimeout(resolve, ms);
});
}
function shouldUseLiveSync(state: TrainingRuntimeStore): boolean {
return state.isTrainingRunning || state.phase === "training";
@ -26,9 +35,46 @@ export function useTrainingRuntimeLifecycle(): void {
let openingStream = false;
let streamController: AbortController | null = null;
let reconnectTimer: ReturnType<typeof setTimeout> | null = null;
/** HF Spaces / auth-disabled: no JWT, but train APIs allow anonymous access. */
let authDisabled = false;
const runtimeStore = useTrainingRuntimeStore;
const canUseTrainApi = () => hasAuthToken() || authDisabled;
let authProbeInFlight = false;
let lastAuthProbeStartedAt = 0;
const maybeRefreshAuthMode = async (force = false) => {
if (disposed || hasAuthToken() || authDisabled) return;
if (authProbeInFlight) return;
const now = Date.now();
if (
!force &&
now - lastAuthProbeStartedAt < AUTH_STATUS_RETRY_INTERVAL_MS
) {
return;
}
authProbeInFlight = true;
lastAuthProbeStartedAt = now;
const controller = new AbortController();
const timeout = setTimeout(() => {
controller.abort();
}, AUTH_STATUS_TIMEOUT_MS);
try {
const res = await fetch("/api/auth/status", { signal: controller.signal });
if (!res.ok) return;
const data = (await res.json()) as { auth_disabled?: boolean };
authDisabled = Boolean(data.auth_disabled);
} catch {
// Keep previous mode and retry later.
} finally {
clearTimeout(timeout);
authProbeInFlight = false;
}
};
const clearReconnect = () => {
if (reconnectTimer) {
clearTimeout(reconnectTimer);
@ -46,7 +92,10 @@ export function useTrainingRuntimeLifecycle(): void {
};
const pollMetrics = async () => {
if (!hasAuthToken()) return;
if (!canUseTrainApi()) {
void maybeRefreshAuthMode();
return;
}
const gen = runtimeStore.getState().resetGeneration;
try {
const metrics = await getTrainingMetrics();
@ -55,14 +104,17 @@ export function useTrainingRuntimeLifecycle(): void {
}
runtimeStore.getState().applyMetrics(metrics);
} catch (error) {
if (!isAbortError(error) && !disposed && hasAuthToken()) {
if (!isAbortError(error) && !disposed && canUseTrainApi()) {
runtimeStore.getState().setSseConnected(false);
}
}
};
const pollStatus = async () => {
if (!hasAuthToken()) return;
if (!canUseTrainApi()) {
void maybeRefreshAuthMode();
return;
}
const gen = runtimeStore.getState().resetGeneration;
try {
const status = await getTrainingStatus();
@ -79,7 +131,7 @@ export function useTrainingRuntimeLifecycle(): void {
stopStream();
}
} catch (error) {
if (!isAbortError(error) && !disposed && hasAuthToken()) {
if (!isAbortError(error) && !disposed && canUseTrainApi()) {
runtimeStore.getState().setSseConnected(false);
}
}
@ -153,7 +205,11 @@ export function useTrainingRuntimeLifecycle(): void {
const hydrate = async () => {
runtimeStore.getState().setHydrating(true);
try {
await Promise.all([pollStatus(), pollMetrics()]);
await maybeRefreshAuthMode(true);
await Promise.race([
Promise.allSettled([pollStatus(), pollMetrics()]).then(() => undefined),
wait(INITIAL_HYDRATE_TIMEOUT_MS),
]);
} finally {
if (!disposed) {
runtimeStore.getState().setHydrating(false);
@ -162,13 +218,14 @@ export function useTrainingRuntimeLifecycle(): void {
}
};
void hydrate();
let statusTimer: ReturnType<typeof setInterval> | null = null;
let metricsTimer: ReturnType<typeof setInterval> | null = null;
const statusTimer = setInterval(() => {
void hydrate();
statusTimer = setInterval(() => {
void pollStatus();
}, STATUS_POLL_INTERVAL_MS);
const metricsTimer = setInterval(() => {
metricsTimer = setInterval(() => {
const state = runtimeStore.getState();
if (shouldUseLiveSync(state) || state.currentStep > 0) {
void pollMetrics();
@ -177,8 +234,8 @@ export function useTrainingRuntimeLifecycle(): void {
return () => {
disposed = true;
clearInterval(statusTimer);
clearInterval(metricsTimer);
if (statusTimer) clearInterval(statusTimer);
if (metricsTimer) clearInterval(metricsTimer);
stopStream();
};
}, []);

View file

@ -140,15 +140,15 @@ def _bootstrap_uv() -> bool:
global UV_NEEDS_SYSTEM
if not shutil.which("uv"):
return False
# Probe: try a dry-run install targeting the current Python explicitly.
# Without --python, uv can ignore the activated venv on some platforms.
# Probe: try a dry-run install without --system.
# If uv can't find a venv it exits with code 2.
probe = subprocess.run(
["uv", "pip", "install", "--dry-run", "--python", sys.executable, "pip"],
["uv", "pip", "install", "--dry-run", "pip"],
stdout = subprocess.PIPE,
stderr = subprocess.STDOUT,
)
if probe.returncode != 0:
# Retry with --system (some envs need it when uv can't find a venv)
# Retry with --system to confirm it works
probe_sys = subprocess.run(
["uv", "pip", "install", "--dry-run", "--system", "pip"],
stdout = subprocess.PIPE,
@ -204,10 +204,6 @@ def _build_uv_cmd(args: tuple[str, ...]) -> list[str]:
cmd = ["uv", "pip", "install"]
if UV_NEEDS_SYSTEM:
cmd.append("--system")
# Always pass --python so uv targets the correct environment.
# Without this, uv can ignore an activated venv and install into
# the system Python (observed on Colab and similar environments).
cmd.extend(["--python", sys.executable])
cmd.extend(_translate_pip_args_for_uv(args))
cmd.append("--torch-backend=auto")
return cmd

View file

@ -41,26 +41,13 @@ if [[ "$keynames" == *$'\nCOLAB_'* ]]; then
fi
# ── Detect whether frontend needs building ──
# Skip if dist/ exists AND no tracked input is newer than dist/.
# Checks top-level config/entry files and src/, public/ recursively.
# This handles: PyPI installs (dist/ bundled), repeat runs (no changes),
# and upgrades/pulls (source newer than dist/ triggers rebuild).
_NEED_FRONTEND_BUILD=true
if [ -d "$SCRIPT_DIR/frontend/dist" ]; then
# Check all top-level files (package.json, bun.lock, vite.config.ts, index.html, etc.)
_changed=$(find "$SCRIPT_DIR/frontend" -maxdepth 1 -type f \
-newer "$SCRIPT_DIR/frontend/dist" -print -quit 2>/dev/null)
# Check src/ and public/ recursively (|| true guards against set -e when dirs are missing)
if [ -z "$_changed" ]; then
_changed=$(find "$SCRIPT_DIR/frontend/src" "$SCRIPT_DIR/frontend/public" \
-type f -newer "$SCRIPT_DIR/frontend/dist" -print -quit 2>/dev/null) || true
fi
if [ -z "$_changed" ]; then
_NEED_FRONTEND_BUILD=false
fi
fi
if [ "$_NEED_FRONTEND_BUILD" = false ]; then
echo "✅ Frontend already built and up to date -- skipping Node/npm check."
# Only skip when BOTH conditions are true:
# 1. We're inside site-packages (PyPI / pip install, not editable)
# 2. dist/ already exists (pre-built in the wheel)
# Otherwise always (re)build — handles upgrades, editable installs, and
# pip-from-source where dist/ was never built.
if [[ "$SCRIPT_DIR" == */site-packages/* ]] && [ -d "$SCRIPT_DIR/frontend/dist" ]; then
echo "✅ Frontend pre-built (PyPI) — skipping Node/npm check."
else
NEED_NODE=true
if command -v node &>/dev/null && command -v npm &>/dev/null; then
@ -159,27 +146,12 @@ run_quiet "npm run build" npm run build
_restore_gitignores
trap - EXIT
# Validate CSS output -- catch truncated Tailwind builds
_MAX_CSS=$(find "$SCRIPT_DIR/frontend/dist/assets" -name '*.css' -exec wc -c {} + 2>/dev/null | sort -n | tail -1 | awk '{print $1}')
if [ -z "$_MAX_CSS" ]; then
echo "⚠️ WARNING: No CSS files were emitted. The frontend build may have failed."
elif [ "$_MAX_CSS" -lt 100000 ]; then
echo "⚠️ WARNING: Largest CSS file is only $((_MAX_CSS / 1024))KB (expected >100KB)."
echo " Tailwind may not have scanned all source files. Check for .gitignore interference."
fi
cd "$SCRIPT_DIR/backend/core/data_recipe/oxc-validator"
run_quiet "npm install (oxc validator runtime)" npm install
cd "$SCRIPT_DIR"
echo "✅ Frontend built to frontend/dist"
fi # end frontend build check
# ── oxc-validator runtime (needs npm -- skip if not available) ──
if [ -d "$SCRIPT_DIR/backend/core/data_recipe/oxc-validator" ] && command -v npm &>/dev/null; then
cd "$SCRIPT_DIR/backend/core/data_recipe/oxc-validator"
run_quiet "npm install (oxc validator runtime)" npm install
cd "$SCRIPT_DIR"
fi
fi # end frontend dist check
# ── 6. Python venv + deps ──
@ -251,261 +223,188 @@ install_python_stack() {
python "$SCRIPT_DIR/install_python_stack.py"
}
# Create venv under ~/.unsloth/studio/ (shared location, not in repo).
# All platforms (including Colab) use the same isolated venv so that
# studio dependencies are never installed into the system Python.
STUDIO_HOME="$HOME/.unsloth/studio"
VENV_DIR="$STUDIO_HOME/.venv"
VENV_T5_DIR="$STUDIO_HOME/.venv_t5"
mkdir -p "$STUDIO_HOME"
# Clean up legacy in-repo venvs if they exist
[ -d "$REPO_ROOT/.venv" ] && rm -rf "$REPO_ROOT/.venv"
[ -d "$REPO_ROOT/.venv_overlay" ] && rm -rf "$REPO_ROOT/.venv_overlay"
[ -d "$REPO_ROOT/.venv_t5" ] && rm -rf "$REPO_ROOT/.venv_t5"
rm -rf "$VENV_DIR"
rm -rf "$VENV_T5_DIR"
# Try creating venv with pip; fall back to --without-pip + bootstrap
# (some environments like Colab have broken ensurepip)
if ! "$BEST_PY" -m venv "$VENV_DIR" 2>/dev/null; then
"$BEST_PY" -m venv --without-pip "$VENV_DIR"
source "$VENV_DIR/bin/activate"
curl -sS https://bootstrap.pypa.io/get-pip.py | python > /dev/null
if [ "$IS_COLAB" = true ]; then
# Colab: install packages directly without venv
install_python_stack
else
# Local: create venv under studio home (shared location, not in repo)
# Configurable via UNSLOTH_STUDIO_HOME; defaults to ~/.unsloth/studio
STUDIO_HOME="${UNSLOTH_STUDIO_HOME:-$HOME/.unsloth/studio}"
echo " Studio home: $STUDIO_HOME"
# Persist for future `unsloth studio` runs (survives shell restarts)
mkdir -p "$HOME/.unsloth"
echo "$STUDIO_HOME" > "$HOME/.unsloth/studio_home"
VENV_DIR="$STUDIO_HOME/.venv"
VENV_T5_DIR="$STUDIO_HOME/.venv_t5"
mkdir -p "$STUDIO_HOME"
# Clean up legacy in-repo venvs if they exist
[ -d "$REPO_ROOT/.venv" ] && rm -rf "$REPO_ROOT/.venv"
[ -d "$REPO_ROOT/.venv_overlay" ] && rm -rf "$REPO_ROOT/.venv_overlay"
[ -d "$REPO_ROOT/.venv_t5" ] && rm -rf "$REPO_ROOT/.venv_t5"
rm -rf "$VENV_DIR"
rm -rf "$VENV_T5_DIR"
"$BEST_PY" -m venv "$VENV_DIR"
source "$VENV_DIR/bin/activate"
fi
cd "$SCRIPT_DIR"
install_python_stack
# ── Ensure uv is available (much faster than pip) ──
USE_UV=false
if command -v uv &>/dev/null; then
USE_UV=true
elif curl -LsSf https://astral.sh/uv/install.sh | sh > /dev/null 2>&1; then
export PATH="$HOME/.local/bin:$PATH"
command -v uv &>/dev/null && USE_UV=true
fi
# Helper: install a package, preferring uv with pip fallback
fast_install() {
if [ "$USE_UV" = true ]; then
uv pip install --python "$(command -v python)" "$@" && return 0
fi
python -m pip install "$@"
}
cd "$SCRIPT_DIR"
install_python_stack
# ── 6b. Pre-install transformers 5.x into .venv_t5/ ──
# Models like GLM-4.7-Flash need transformers>=5.3.0. Instead of pip-installing
# at runtime (slow, ~10-15s), we pre-install into a separate directory.
# The training subprocess just prepends .venv_t5/ to sys.path -- instant switch.
echo ""
echo " Pre-installing transformers 5.x for newer model support..."
mkdir -p "$VENV_T5_DIR"
run_quiet "install transformers 5.x" fast_install --target "$VENV_T5_DIR" --no-deps "transformers==5.3.0"
run_quiet "install huggingface_hub for t5" fast_install --target "$VENV_T5_DIR" --no-deps "huggingface_hub==1.7.1"
run_quiet "install hf_xet for t5" fast_install --target "$VENV_T5_DIR" --no-deps "hf_xet==1.4.2"
# tiktoken is needed by Qwen-family tokenizers. Install with deps since
# regex/requests may be missing on Windows.
run_quiet "install tiktoken for t5" fast_install --target "$VENV_T5_DIR" "tiktoken"
echo "✅ Transformers 5.x pre-installed to $VENV_T5_DIR/"
# ── 7. WSL: pre-install GGUF build dependencies ──
# On WSL, sudo requires a password and can't be entered during GGUF export
# (runs in a non-interactive subprocess). Install build deps here instead.
if grep -qi microsoft /proc/version 2>/dev/null; then
# ── 6b. Pre-install transformers 5.x into .venv_t5/ ──
# Models like GLM-4.7-Flash need transformers>=5.3.0. Instead of pip-installing
# at runtime (slow, ~10-15s), we pre-install into a separate directory.
# The training subprocess just prepends .venv_t5/ to sys.path — instant switch.
echo ""
echo "⚠️ WSL detected -- installing build dependencies for GGUF export..."
_GGUF_DEPS="pciutils build-essential cmake curl git libcurl4-openssl-dev"
echo " Pre-installing transformers 5.x for newer model support..."
mkdir -p "$VENV_T5_DIR"
run_quiet "pip install transformers 5.x" pip install --target "$VENV_T5_DIR" --no-deps "transformers==5.3.0"
run_quiet "pip install huggingface_hub for t5" pip install --target "$VENV_T5_DIR" --no-deps "huggingface_hub==1.3.0"
echo "✅ Transformers 5.x pre-installed to $VENV_T5_DIR/"
# Try without sudo first (works when already root)
apt-get update -y >/dev/null 2>&1 || true
apt-get install -y $_GGUF_DEPS >/dev/null 2>&1 || true
# Check which packages are still missing
_STILL_MISSING=""
for _pkg in $_GGUF_DEPS; do
case "$_pkg" in
build-essential) command -v gcc >/dev/null 2>&1 || _STILL_MISSING="$_STILL_MISSING $_pkg" ;;
pciutils) command -v lspci >/dev/null 2>&1 || _STILL_MISSING="$_STILL_MISSING $_pkg" ;;
libcurl4-openssl-dev) dpkg -s "$_pkg" >/dev/null 2>&1 || _STILL_MISSING="$_STILL_MISSING $_pkg" ;;
*) command -v "$_pkg" >/dev/null 2>&1 || _STILL_MISSING="$_STILL_MISSING $_pkg" ;;
esac
done
_STILL_MISSING=$(echo "$_STILL_MISSING" | sed 's/^ *//')
if [ -z "$_STILL_MISSING" ]; then
# ── 7. WSL: pre-install GGUF build dependencies ──
# On WSL, sudo requires a password and can't be entered during GGUF export
# (runs in a non-interactive subprocess). Install build deps here instead.
if grep -qi microsoft /proc/version 2>/dev/null; then
echo ""
echo "⚠️ WSL detected — installing build dependencies for GGUF export..."
echo " You may be prompted for your password."
sudo apt-get update -y
sudo apt-get install -y build-essential cmake curl git libcurl4-openssl-dev
echo "✅ GGUF build dependencies installed"
elif command -v sudo >/dev/null 2>&1; then
echo ""
echo " !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!"
echo " WARNING: We require sudo elevated permissions to install:"
echo " $_STILL_MISSING"
echo " If you accept, we'll run sudo now, and it'll prompt your password."
echo " !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!"
echo ""
printf " Accept? [Y/n] "
if [ -r /dev/tty ]; then
read -r REPLY </dev/tty || REPLY="y"
else
REPLY="y"
fi
case "$REPLY" in
[nN]*)
echo ""
echo " Please install these packages first, then re-run Unsloth Studio setup:"
echo " sudo apt-get update -y && sudo apt-get install -y $_STILL_MISSING"
_SKIP_GGUF_BUILD=true
;;
*)
sudo apt-get update -y
sudo apt-get install -y $_STILL_MISSING
echo "✅ GGUF build dependencies installed"
;;
esac
else
echo " sudo is not available on this system."
echo " Please install as root, then re-run setup:"
echo " apt-get install -y $_STILL_MISSING"
_SKIP_GGUF_BUILD=true
fi
fi
# ── 8. Build llama.cpp binaries for GGUF inference + export ──
# Builds at ~/.unsloth/llama.cpp — a single shared location under the user's
# home directory. This is used by both the inference server and the GGUF
# export pipeline (unsloth-zoo).
# - llama-server: for GGUF model inference
# - llama-quantize: for GGUF export quantization (symlinked to root for check_llama_cpp())
UNSLOTH_HOME="$HOME/.unsloth"
mkdir -p "$UNSLOTH_HOME"
LLAMA_CPP_DIR="$UNSLOTH_HOME/llama.cpp"
LLAMA_SERVER_BIN="$LLAMA_CPP_DIR/build/bin/llama-server"
if [ "${_SKIP_GGUF_BUILD:-}" = true ]; then
echo ""
echo "Skipping llama-server build (missing dependencies)"
echo " Install the missing packages and re-run setup to enable GGUF inference."
else
rm -rf "$LLAMA_CPP_DIR"
{
# Check prerequisites
if ! command -v cmake &>/dev/null; then
echo ""
echo "⚠️ cmake not found — skipping llama-server build (GGUF inference won't be available)"
echo " Install cmake and re-run setup.sh to enable GGUF inference."
elif ! command -v git &>/dev/null; then
echo ""
echo "⚠️ git not found — skipping llama-server build (GGUF inference won't be available)"
else
echo ""
echo "Building llama-server for GGUF inference..."
BUILD_OK=true
run_quiet "clone llama.cpp" git clone --depth 1 https://github.com/ggml-org/llama.cpp.git "$LLAMA_CPP_DIR" || BUILD_OK=false
if [ "$BUILD_OK" = true ]; then
# Skip tests/examples we don't need (faster build)
CMAKE_ARGS="-DLLAMA_BUILD_TESTS=OFF -DLLAMA_BUILD_EXAMPLES=OFF -DLLAMA_BUILD_SERVER=ON -DGGML_NATIVE=ON"
# Use ccache if available (dramatically faster rebuilds)
if command -v ccache &>/dev/null; then
CMAKE_ARGS="$CMAKE_ARGS -DCMAKE_C_COMPILER_LAUNCHER=ccache -DCMAKE_CXX_COMPILER_LAUNCHER=ccache -DCMAKE_CUDA_COMPILER_LAUNCHER=ccache"
echo " Using ccache for faster compilation"
fi
# Detect CUDA: check nvcc on PATH, then common install locations
NVCC_PATH=""
if command -v nvcc &>/dev/null; then
NVCC_PATH="$(command -v nvcc)"
elif [ -x /usr/local/cuda/bin/nvcc ]; then
NVCC_PATH="/usr/local/cuda/bin/nvcc"
export PATH="/usr/local/cuda/bin:$PATH"
elif ls /usr/local/cuda-*/bin/nvcc &>/dev/null 2>&1; then
# Pick the newest cuda-XX.X directory
NVCC_PATH="$(ls -d /usr/local/cuda-*/bin/nvcc 2>/dev/null | sort -V | tail -1)"
export PATH="$(dirname "$NVCC_PATH"):$PATH"
fi
if [ -n "$NVCC_PATH" ]; then
echo " Building with CUDA support (nvcc: $NVCC_PATH)..."
CMAKE_ARGS="$CMAKE_ARGS -DGGML_CUDA=ON"
# Detect GPU compute capability and limit CUDA architectures
# Without this, cmake builds for ALL default archs (very slow)
CUDA_ARCHS=""
if command -v nvidia-smi &>/dev/null; then
# Read all GPUs, deduplicate (handles mixed-GPU hosts)
_raw_caps=$(nvidia-smi --query-gpu=compute_cap --format=csv,noheader 2>/dev/null || true)
while IFS= read -r _cap; do
_cap=$(echo "$_cap" | tr -d '[:space:]')
if [[ "$_cap" =~ ^([0-9]+)\.([0-9]+)$ ]]; then
_arch="${BASH_REMATCH[1]}${BASH_REMATCH[2]}"
# Append if not already present
case ";$CUDA_ARCHS;" in
*";$_arch;"*) ;;
*) CUDA_ARCHS="${CUDA_ARCHS:+$CUDA_ARCHS;}$_arch" ;;
esac
fi
done <<< "$_raw_caps"
fi
if [ -n "$CUDA_ARCHS" ]; then
echo " GPU compute capabilities: ${CUDA_ARCHS//;/, } -- limiting build to detected archs"
CMAKE_ARGS="$CMAKE_ARGS -DCMAKE_CUDA_ARCHITECTURES=${CUDA_ARCHS}"
else
echo " Could not detect GPU arch -- building for all default CUDA architectures (slower)"
fi
# Multi-threaded nvcc compilation (uses all CPU cores per .cu file)
CMAKE_ARGS="$CMAKE_ARGS -DCMAKE_CUDA_FLAGS=--threads=0"
elif [ -d /usr/local/cuda ] || nvidia-smi &>/dev/null; then
echo " CUDA driver detected but nvcc not found — building CPU-only"
echo " To enable GPU: install cuda-toolkit or add nvcc to PATH"
else
echo " Building CPU-only (no CUDA detected)..."
fi
NCPU=$(nproc 2>/dev/null || sysctl -n hw.ncpu 2>/dev/null || echo 4)
# Use Ninja if available (faster parallel builds than Make)
CMAKE_GENERATOR_ARGS=""
if command -v ninja &>/dev/null; then
CMAKE_GENERATOR_ARGS="-G Ninja"
fi
run_quiet "cmake llama.cpp" cmake $CMAKE_GENERATOR_ARGS -S "$LLAMA_CPP_DIR" -B "$LLAMA_CPP_DIR/build" $CMAKE_ARGS || BUILD_OK=false
fi
if [ "$BUILD_OK" = true ]; then
run_quiet "build llama-server" cmake --build "$LLAMA_CPP_DIR/build" --config Release --target llama-server -j"$NCPU" || BUILD_OK=false
fi
# Also build llama-quantize (needed by unsloth-zoo's GGUF export pipeline)
if [ "$BUILD_OK" = true ]; then
run_quiet "build llama-quantize" cmake --build "$LLAMA_CPP_DIR/build" --config Release --target llama-quantize -j"$NCPU" || true
# Symlink to llama.cpp root — check_llama_cpp() looks for the binary there
QUANTIZE_BIN="$LLAMA_CPP_DIR/build/bin/llama-quantize"
if [ -f "$QUANTIZE_BIN" ]; then
ln -sf build/bin/llama-quantize "$LLAMA_CPP_DIR/llama-quantize"
fi
fi
if [ "$BUILD_OK" = true ]; then
if [ -f "$LLAMA_SERVER_BIN" ]; then
echo "✅ llama-server built at $LLAMA_SERVER_BIN"
else
echo "⚠️ llama-server binary not found after build — GGUF inference won't be available"
fi
if [ -f "$LLAMA_CPP_DIR/llama-quantize" ]; then
echo "✅ llama-quantize available for GGUF export"
fi
else
echo "⚠️ llama-server build failed — GGUF inference won't be available, but everything else works"
fi
fi
}
fi # end _SKIP_GGUF_BUILD check
# Disabled: llama.cpp build is commented out for now.
# UNCOMMENT the block below to re-enable.
#
# # Builds at ~/.unsloth/llama.cpp — a single shared location under the user's
# # home directory. This is used by both the inference server and the GGUF
# # export pipeline (unsloth-zoo).
# # - llama-server: for GGUF model inference
# # - llama-quantize: for GGUF export quantization (symlinked to root for check_llama_cpp())
# UNSLOTH_HOME="$HOME/.unsloth"
# mkdir -p "$UNSLOTH_HOME"
# LLAMA_CPP_DIR="$UNSLOTH_HOME/llama.cpp"
# LLAMA_SERVER_BIN="$LLAMA_CPP_DIR/build/bin/llama-server"
# rm -rf "$LLAMA_CPP_DIR"
# {
# # Check prerequisites
# if ! command -v cmake &>/dev/null; then
# echo ""
# echo "⚠️ cmake not found — skipping llama-server build (GGUF inference won't be available)"
# echo " Install cmake and re-run setup.sh to enable GGUF inference."
# elif ! command -v git &>/dev/null; then
# echo ""
# echo "⚠️ git not found — skipping llama-server build (GGUF inference won't be available)"
# else
# echo ""
# echo "Building llama-server for GGUF inference..."
#
# BUILD_OK=true
# run_quiet "clone llama.cpp" git clone --depth 1 https://github.com/ggml-org/llama.cpp.git "$LLAMA_CPP_DIR" || BUILD_OK=false
#
# if [ "$BUILD_OK" = true ]; then
# # Skip tests/examples we don't need (faster build)
# CMAKE_ARGS="-DLLAMA_BUILD_TESTS=OFF -DLLAMA_BUILD_EXAMPLES=OFF -DLLAMA_BUILD_SERVER=ON -DGGML_NATIVE=ON"
#
# # Use ccache if available (dramatically faster rebuilds)
# if command -v ccache &>/dev/null; then
# CMAKE_ARGS="$CMAKE_ARGS -DCMAKE_C_COMPILER_LAUNCHER=ccache -DCMAKE_CXX_COMPILER_LAUNCHER=ccache -DCMAKE_CUDA_COMPILER_LAUNCHER=ccache"
# echo " Using ccache for faster compilation"
# fi
#
# # Detect CUDA: check nvcc on PATH, then common install locations
# NVCC_PATH=""
# if command -v nvcc &>/dev/null; then
# NVCC_PATH="$(command -v nvcc)"
# elif [ -x /usr/local/cuda/bin/nvcc ]; then
# NVCC_PATH="/usr/local/cuda/bin/nvcc"
# export PATH="/usr/local/cuda/bin:$PATH"
# elif ls /usr/local/cuda-*/bin/nvcc &>/dev/null 2>&1; then
# # Pick the newest cuda-XX.X directory
# NVCC_PATH="$(ls -d /usr/local/cuda-*/bin/nvcc 2>/dev/null | sort -V | tail -1)"
# export PATH="$(dirname "$NVCC_PATH"):$PATH"
# fi
#
# if [ -n "$NVCC_PATH" ]; then
# echo " Building with CUDA support (nvcc: $NVCC_PATH)..."
# CMAKE_ARGS="$CMAKE_ARGS -DGGML_CUDA=ON"
#
# # Detect GPU compute capability and limit CUDA architectures
# # Without this, cmake builds for ALL default archs (very slow)
# CUDA_ARCHS=""
# if command -v nvidia-smi &>/dev/null; then
# # Read all GPUs, deduplicate (handles mixed-GPU hosts)
# _raw_caps=$(nvidia-smi --query-gpu=compute_cap --format=csv,noheader 2>/dev/null || true)
# while IFS= read -r _cap; do
# _cap=$(echo "$_cap" | tr -d '[:space:]')
# if [[ "$_cap" =~ ^([0-9]+)\.([0-9]+)$ ]]; then
# _arch="${BASH_REMATCH[1]}${BASH_REMATCH[2]}"
# # Append if not already present
# case ";$CUDA_ARCHS;" in
# *";$_arch;"*) ;;
# *) CUDA_ARCHS="${CUDA_ARCHS:+$CUDA_ARCHS;}$_arch" ;;
# esac
# fi
# done <<< "$_raw_caps"
# fi
#
# if [ -n "$CUDA_ARCHS" ]; then
# echo " GPU compute capabilities: ${CUDA_ARCHS//;/, } -- limiting build to detected archs"
# CMAKE_ARGS="$CMAKE_ARGS -DCMAKE_CUDA_ARCHITECTURES=${CUDA_ARCHS}"
# else
# echo " Could not detect GPU arch -- building for all default CUDA architectures (slower)"
# fi
#
# # Multi-threaded nvcc compilation (uses all CPU cores per .cu file)
# CMAKE_ARGS="$CMAKE_ARGS -DCMAKE_CUDA_FLAGS=--threads=0"
# elif [ -d /usr/local/cuda ] || nvidia-smi &>/dev/null; then
# echo " CUDA driver detected but nvcc not found — building CPU-only"
# echo " To enable GPU: install cuda-toolkit or add nvcc to PATH"
# else
# echo " Building CPU-only (no CUDA detected)..."
# fi
#
# NCPU=$(nproc 2>/dev/null || sysctl -n hw.ncpu 2>/dev/null || echo 4)
#
# # Use Ninja if available (faster parallel builds than Make)
# CMAKE_GENERATOR_ARGS=""
# if command -v ninja &>/dev/null; then
# CMAKE_GENERATOR_ARGS="-G Ninja"
# fi
#
# run_quiet "cmake llama.cpp" cmake $CMAKE_GENERATOR_ARGS -S "$LLAMA_CPP_DIR" -B "$LLAMA_CPP_DIR/build" $CMAKE_ARGS || BUILD_OK=false
# fi
#
# if [ "$BUILD_OK" = true ]; then
# run_quiet "build llama-server" cmake --build "$LLAMA_CPP_DIR/build" --config Release --target llama-server -j"$NCPU" || BUILD_OK=false
# fi
#
# # Also build llama-quantize (needed by unsloth-zoo's GGUF export pipeline)
# if [ "$BUILD_OK" = true ]; then
# run_quiet "build llama-quantize" cmake --build "$LLAMA_CPP_DIR/build" --config Release --target llama-quantize -j"$NCPU" || true
# # Symlink to llama.cpp root — check_llama_cpp() looks for the binary there
# QUANTIZE_BIN="$LLAMA_CPP_DIR/build/bin/llama-quantize"
# if [ -f "$QUANTIZE_BIN" ]; then
# ln -sf build/bin/llama-quantize "$LLAMA_CPP_DIR/llama-quantize"
# fi
# fi
#
# if [ "$BUILD_OK" = true ]; then
# if [ -f "$LLAMA_SERVER_BIN" ]; then
# echo "✅ llama-server built at $LLAMA_SERVER_BIN"
# else
# echo "⚠️ llama-server binary not found after build — GGUF inference won't be available"
# fi
# if [ -f "$LLAMA_CPP_DIR/llama-quantize" ]; then
# echo "✅ llama-quantize available for GGUF export"
# fi
# else
# echo "⚠️ llama-server build failed — GGUF inference won't be available, but everything else works"
# fi
# fi
# }
echo ""
if [ "$IS_COLAB" = true ]; then
@ -514,9 +413,6 @@ if [ "$IS_COLAB" = true ]; then
echo "╠══════════════════════════════════════╣"
echo "║ Unsloth Studio is ready to start ║"
echo "║ in your Colab notebook! ║"
echo "║ ║"
echo "║ from colab import start ║"
echo "║ start() ║"
echo "╚══════════════════════════════════════╝"
else
echo "╔══════════════════════════════════════╗"
@ -524,6 +420,6 @@ else
echo "╠══════════════════════════════════════╣"
echo "║ Launch with: ║"
echo "║ ║"
echo "║ unsloth studio -H 0.0.0.0 -p 8888 ║"
echo "║ unsloth studio -H 0.0.0.0 -p 8000 ║"
echo "╚══════════════════════════════════════╝"
fi