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

Merge branch 'diffusion-train-precision' into diffusion-train-tab-2

Browse source
This commit is contained in:
Daniel Han 2026-07-04 02:21:11 +00:00
commit f1dbb74308
131 changed files with 19488 additions and 3092 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]
}
}
}
@ -2430,6 +2439,13 @@ exit 0
}
$studioArgs = @('studio', 'setup')
if ($script:UnslothVerbose) { $studioArgs += '--verbose' }
if ($WithLlamaCppDir) {
if (-not (Test-Path -LiteralPath $WithLlamaCppDir -PathType Container)) {
Write-Host "[ERROR] --with-llama-cpp-dir path does not exist: $WithLlamaCppDir" -ForegroundColor Red
return (Exit-InstallFailure "--with-llama-cpp-dir path does not exist.")
}
$env:UNSLOTH_LOCAL_LLAMA_CPP_DIR = (Resolve-Path -LiteralPath $WithLlamaCppDir).Path
}
$env:UNSLOTH_INSTALL_ROLLBACK_MANAGED = "1"
# Hand the venv interpreter to setup.ps1 so it reuses the Python we already
# resolved and built the venv with, instead of re-probing the system (which
@ -2445,6 +2461,7 @@ exit 0
} else {
Remove-Item Env:UNSLOTH_STUDIO_HOME -ErrorAction SilentlyContinue
}
Remove-Item Env:UNSLOTH_LOCAL_LLAMA_CPP_DIR -ErrorAction SilentlyContinue
Remove-Item Env:UNSLOTH_INSTALL_ROLLBACK_MANAGED -ErrorAction SilentlyContinue
Remove-Item Env:UNSLOTH_SETUP_PYTHON -ErrorAction SilentlyContinue
}

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).
@ -3023,6 +3038,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" \
@ -3031,6 +3053,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
@ -3045,6 +3068,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

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

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."
),
},
]

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

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

@ -55,9 +55,15 @@ 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
@ -114,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.
@ -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

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.

139
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"
@ -36,6 +38,10 @@ if sys.platform == "win32":
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.
@ -226,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
@ -1078,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
@ -1088,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(),
}
@ -1135,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:

830
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
@ -62,25 +63,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]:
@ -1183,6 +1210,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] = []
@ -1198,6 +1226,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())
@ -1317,6 +1347,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()
@ -1367,8 +1399,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,
@ -1416,7 +1458,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()
@ -1469,6 +1512,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,
@ -1522,6 +1567,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())

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:

191
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
@ -47,6 +47,15 @@ from utils.preview_sharing_settings import (
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()
@ -229,6 +238,186 @@ def update_openai_auto_switch_override(
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

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

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

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

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

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

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

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:

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

@ -11,7 +11,6 @@ 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 imagesRoute } from "./routes/images";
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";
@ -26,7 +25,6 @@ const routeTree = rootRoute.addChildren([
onboardingRoute,
loginRoute,
changePasswordRoute,
gridTestRoute,
hubRoute,
settingsRoute,
studioRoute,

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

@ -78,6 +78,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

@ -291,13 +291,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"
@ -1216,7 +1215,7 @@ export function AppSidebar() {
pathname === "/studio" || pathname.startsWith("/studio/")
}
disabled={chatOnly}
tooltip={trainExportDisabledHint}
tooltip={trainDisabledHint}
spinner={trainingInProgress}
onClick={() => {
if (chatOnly) return;
@ -1245,7 +1244,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;
@ -1266,11 +1265,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 {
@ -103,6 +104,7 @@ import {
type FormatFilter,
estimateQuantBytes,
fitsDevice,
hfModelFitsDevice,
isMlxId,
isMobileVariant,
isRecommendableFormat,
@ -1414,6 +1416,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>>(
@ -1805,34 +1810,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,
@ -2111,23 +2101,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(() => {
@ -2138,6 +2111,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],
@ -2148,6 +2157,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/"))
@ -2174,6 +2189,9 @@ export function HubModelPicker({
isKnownGgufRepo,
isChatSupported,
formatFilter,
fitOnDeviceOnly,
downloadedSet,
gpu,
isMac,
task,
]);
@ -2463,6 +2481,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
@ -2473,6 +2520,7 @@ export function HubModelPicker({
align="end"
className={sortTriggerClassName}
contentClassName={sortMenuContentClassName}
footer={fitOnDeviceFooter}
/>
) : section === "downloaded" ? (
<HubOptionMenu
@ -2483,6 +2531,7 @@ export function HubModelPicker({
align="end"
className={sortTriggerClassName}
contentClassName={sortMenuContentClassName}
footer={fitOnDeviceFooter}
/>
) : (
<HubOptionMenu
@ -2493,6 +2542,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());
}

73
studio/frontend/src/features/settings/components/agent-command.ts Normal file
View file

@ -0,0 +1,73 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
// Build the `unsloth start <agent>` command for the API-keys panel.
// `unsloth start` reads UNSLOTH_STUDIO_URL (default 127.0.0.1:8888) and only
// auto-mints a key for a loopback server, so the bare command is correct only for
// the default local server. For a non-default port or tunnel/remote base, emit the
// URL (plus a key for non-loopback) so the copy targets what the UI shows.
const DEFAULT_STUDIO_PORT = "8888";
const DEFAULT_AGENT = "claude";
// URL.hostname brackets IPv6 literals (`new URL("http://[::1]:8888").hostname` is
// "[::1]"), so strip the brackets before matching the bare "::1" loopback rules below.
function normalizeHost(host: string): string {
const lower = host.toLowerCase();
return lower.startsWith("[") && lower.endsWith("]") ? lower.slice(1, -1) : lower;
}
// The bare `unsloth start` probes exactly http://127.0.0.1:8888, so only that literal
// host earns the bare command. `localhost` can resolve to ::1 (and `::1` is never
// probed), so both keep an explicit UNSLOTH_STUDIO_URL -- harmless when they alias
// 127.0.0.1, correct when they don't.
function isDefaultLocalHost(host: string): boolean {
return host === "127.0.0.1";
}
// Match the CLI auto-mint rule (is_loopback_url): localhost, ::1, and all of 127.0.0.0/8.
function isLoopbackHost(host: string): boolean {
if (host === "localhost" || host === "::1") return true;
const octets = host.split(".");
return (
octets.length === 4 &&
octets[0] === "127" &&
octets.every((o) => /^\d{1,3}$/.test(o) && Number(o) <= 255)
);
}
export function buildAgentCommand(
base: string | null | undefined,
key: string | null | undefined,
os: "unix" | "windows",
agent: string = DEFAULT_AGENT,
): string {
const bare = `unsloth start ${agent}`;
let url: URL | null = null;
try {
if (base) url = new URL(base);
} catch {
url = null;
}
// Unknown base: fall back to the bare default-local command.
if (!url) return bare;
const host = normalizeHost(url.hostname);
const loopback = isLoopbackHost(host);
// Default local server (http://127.0.0.1/localhost:8888): bare command
// auto-discovers it. The CLI's bare default probes plain HTTP, so an HTTPS
// loopback on the same port must keep its explicit UNSLOTH_STUDIO_URL.
if (url.protocol === "http:" && isDefaultLocalHost(host) && url.port === DEFAULT_STUDIO_PORT) {
return bare;
}
// Non-default server: set the URL; non-loopback also needs an explicit key.
let cmd = bare;
if (!loopback && key) cmd += ` --api-key ${key}`;
const studioUrl = url.origin;
return os === "windows"
? `$env:UNSLOTH_STUDIO_URL="${studioUrl}"; ${cmd}`
: `UNSLOTH_STUDIO_URL=${studioUrl} ${cmd}`;
}

134
studio/frontend/src/features/settings/components/embedding-model-combobox.tsx Normal file
View file

@ -0,0 +1,134 @@
// 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 {
Combobox,
ComboboxContent,
ComboboxEmpty,
ComboboxInput,
ComboboxItem,
ComboboxList,
} from "@/components/ui/combobox";
import { Spinner } from "@/components/ui/spinner";
import type { PipelineType } from "@huggingface/hub";
import { useHubModelSearch } from "@/features/hub/hooks/use-hub-model-search";
import { useDebouncedValue } from "@/hooks";
import { type ReactElement, useMemo, useRef } from "react";
// HF pipeline filter for embedding models; matches the backend's
// is_embedding_model signals (sentence-similarity / feature-extraction).
const EMBEDDING_TASKS: readonly PipelineType[] = [
"sentence-similarity",
"feature-extraction",
];
type EmbeddingModelComboboxProps = {
value: string;
/** Fires on typing, selection, and Enter with the current text. */
onChange: (value: string) => void;
accessToken?: string;
disabled?: boolean;
placeholder?: string;
ariaLabel?: string;
className?: string;
};
export function EmbeddingModelCombobox({
value,
onChange,
accessToken,
disabled,
placeholder,
ariaLabel,
className,
}: EmbeddingModelComboboxProps): ReactElement {
const selectingRef = useRef(false);
const anchorRef = useRef<HTMLDivElement>(null);
// Fully controlled: the parent updates value on every keystroke, so the
// prop itself is the search query.
const debouncedQuery = useDebouncedValue(value);
const { results, isLoading } = useHubModelSearch(debouncedQuery, {
task: EMBEDDING_TASKS,
accessToken,
excludeGguf: true,
enabled: !disabled,
// Curated unsloth listing when empty (the global top-downloads page holds
// no unsloth mirrors to float); a typed query searches the whole Hub.
ownerScope: debouncedQuery.trim() ? "all" : "unsloth",
});
const items = useMemo(() => {
const ids = results.map((item) => item.id);
const selected = value.trim();
if (selected && !ids.includes(selected)) {
ids.push(selected);
}
return ids;
}, [results, value]);
return (
<div
ref={anchorRef}
className={className}
onKeyDown={(event) => {
if (event.key !== "Enter") return;
if (!(event.target instanceof HTMLInputElement)) return;
event.preventDefault();
const typed = event.target.value.trim();
if (typed) {
onChange(typed);
} else if (items.length > 0) {
onChange(items[0]);
}
}}
>
<Combobox
items={items}
filteredItems={items}
filter={null}
value={value.trim() ? value : null}
onValueChange={(next) => onChange(next ?? "")}
onInputValueChange={(next) => {
if (selectingRef.current) {
selectingRef.current = false;
return;
}
onChange(next);
}}
itemToStringValue={(item) => item}
autoHighlight={true}
>
<ComboboxInput
className="h-8 w-full font-mono [&_input]:text-[11px]"
placeholder={placeholder}
aria-label={ariaLabel}
disabled={disabled}
/>
<ComboboxContent anchor={anchorRef}>
{isLoading ? (
<div className="flex items-center gap-2 px-2 py-3 text-xs text-muted-foreground">
<Spinner className="size-3.5" />
Searching...
</div>
) : (
<ComboboxEmpty>No embedding models found</ComboboxEmpty>
)}
<ComboboxList>
{(id: string) => (
<ComboboxItem
key={id}
value={id}
onPointerDown={() => {
selectingRef.current = true;
}}
>
<span className="truncate font-mono text-[11px]">{id}</span>
</ComboboxItem>
)}
</ComboboxList>
</ComboboxContent>
</Combobox>
</div>
);
}

190
studio/frontend/src/features/settings/components/usage-examples.tsx
View file

@ -32,45 +32,60 @@ import {
loadOpenAIAutoSwitchSettings,
updateOpenAIAutoSwitchSettings,
} from "../api/openai-auto-switch";
import { buildAgentCommand } from "./agent-command";
// API call type; OS axis applies to curl only (Python is OS-identical).
type ExampleType =
| "curl"
| "python"
| "javascript"
| "curlTools"
| "pythonTools"
| "javascriptTools"
| "curlAdvanced"
| "pythonAdvanced";
| "pythonAdvanced"
| "javascriptAdvanced";
type Os = "unix" | "windows";
// plain = bare call; tools = server-side tools; advanced = sampling + thinking + tools.
type Variant = "plain" | "tools" | "advanced";
const TYPE_TABS: { id: ExampleType; label: string }[] = [
{ id: "curl", label: "curl" },
{ id: "python", label: "Python" },
{ id: "javascript", label: "JavaScript" },
{ id: "curlTools", label: "curl + tools" },
{ id: "pythonTools", label: "Python + tools" },
{ id: "javascriptTools", label: "JavaScript + tools" },
{ id: "curlAdvanced", label: "curl + advanced" },
{ id: "pythonAdvanced", label: "Python + advanced" },
{ id: "javascriptAdvanced", label: "JavaScript + advanced" },
];
const TYPE_LABEL_KEY: Partial<Record<ExampleType, TranslationKey>> = {
curlTools: "settings.apiKeys.exampleCurlTools",
pythonTools: "settings.apiKeys.examplePythonTools",
javascriptTools: "settings.apiKeys.exampleJavaScriptTools",
curlAdvanced: "settings.apiKeys.exampleCurlAdvanced",
pythonAdvanced: "settings.apiKeys.examplePythonAdvanced",
javascriptAdvanced: "settings.apiKeys.exampleJavaScriptAdvanced",
};
const OS_AWARE: Record<ExampleType, boolean> = {
curl: true,
python: false,
javascript: false,
curlTools: true,
pythonTools: false,
javascriptTools: false,
curlAdvanced: true,
pythonAdvanced: false,
javascriptAdvanced: false,
};
const CURL_TYPES = new Set<ExampleType>(["curl", "curlTools", "curlAdvanced"]);
const JAVASCRIPT_TYPES = new Set<ExampleType>([
"javascript",
"javascriptTools",
"javascriptAdvanced",
]);
const PROMPT = "Can Unsloth Studio do API calling?";
// Auto-switch demo: a second call naming a different downloaded GGUF so the
@ -82,7 +97,6 @@ const SWITCH_MODEL = "your-other-downloaded-GGUF";
const SWITCH_PROMPT = "Now answer as a different model.";
// web_search + python + terminal are the reliable built-in tools.
const TOOLS = ["web_search", "python", "terminal"];
// Sampling/thinking knobs for the "+ advanced" examples.
const ADV = {
temperature: 0.7,
top_p: 0.8,
@ -93,37 +107,18 @@ const ADV = {
} as const;
const DOC_LINKS = [
{
label: "Claude Code",
href: "https://unsloth.ai/docs/basics/claude-code",
},
{
label: "Codex",
href: "https://unsloth.ai/docs/basics/codex",
},
{
label: "OpenClaw",
href: "https://unsloth.ai/docs/integrations/openclaw",
},
{
label: "OpenCode",
href: "https://unsloth.ai/docs/integrations/opencode",
},
{
label: "Hermes Agent",
href: "https://unsloth.ai/docs/integrations/hermes-agent",
},
{ label: "Claude Code", href: "https://unsloth.ai/docs/basics/claude-code" },
{ label: "Codex", href: "https://unsloth.ai/docs/basics/codex" },
{ label: "OpenClaw", href: "https://unsloth.ai/docs/integrations/openclaw" },
{ label: "OpenCode", href: "https://unsloth.ai/docs/integrations/opencode" },
{ label: "Hermes Agent", href: "https://unsloth.ai/docs/integrations/hermes-agent" },
];
// JSON-encode; also a valid Python literal, so odd model names never break output.
const j = (s: string): string => JSON.stringify(s);
// Embed in a POSIX single-quoted string: close, escaped quote, reopen.
const shSingle = (s: string): string => s.replace(/'/g, "'\\''");
// Embed in a PowerShell single-quoted string: '' is a literal quote.
const psSingle = (s: string): string => s.replace(/'/g, "''");
const toolsJson = TOOLS.map(j).join(", ");
// Shared body fields (after model/messages, before stream) per variant.
function bodyExtraLines(variant: Variant, indent: string): string[] {
const lines: string[] = [];
if (variant === "advanced") {
@ -152,7 +147,6 @@ function curlBodyPretty(model: string, variant: Variant): string {
return `{\n${lines.join("\n")}\n }`;
}
// One-line JSON for the Windows body file (PowerShell mangles inline quotes to curl.exe).
function winBody(model: string, variant: Variant): string {
const body: Record<string, unknown> = {
model,
@ -172,7 +166,7 @@ function winBody(model: string, variant: Variant): string {
body.enabled_tools = TOOLS;
}
body.stream = true;
return JSON.stringify(body);
return JSON.stringify(body, null, 2);
}
// A leading comment (valid in both bash and PowerShell) noting the model field
@ -193,7 +187,6 @@ function curlUnix(
-d '${shSingle(curlBodyPretty(model, variant))}'`;
}
// Windows PowerShell: curl aliases to Invoke-WebRequest, so use curl.exe + body file.
function curlWindows(
base: string,
key: string,
@ -233,7 +226,6 @@ function pythonSnippet(
variant: Variant,
autoSwitch: boolean,
): string {
// Standard OpenAI args are named; Unsloth extensions go through extra_body.
const named =
variant === "advanced"
? `
@ -258,7 +250,6 @@ function pythonSnippet(
${extra.join("\n")}
},`
: "";
// With tools, some chunks are tool-lifecycle events with no choices; guard it.
const loop =
variant !== "plain"
? `for chunk in response:
@ -281,6 +272,70 @@ response = client.chat.completions.create(
${loop}${autoSwitch ? pythonSwitchDemo() : ""}`;
}
function javascriptSnippet(
base: string,
key: string,
model: string,
variant: Variant,
autoSwitch: boolean,
): string {
const options: string[] = [];
if (variant === "advanced") {
options.push(` temperature: ${ADV.temperature},`);
options.push(` top_p: ${ADV.top_p},`);
options.push(` max_tokens: ${ADV.max_tokens},`);
}
// The JS SDK forwards unknown options into the request body, so these go at the
// top level (the Python SDK needs them under extra_body instead).
if (variant === "advanced") {
options.push(` top_k: ${ADV.top_k},`);
options.push(` min_p: ${ADV.min_p},`);
options.push(` repetition_penalty: ${ADV.repetition_penalty},`);
options.push(` enable_thinking: true,`);
}
if (variant !== "plain") {
options.push(` enable_tools: true,`);
options.push(` enabled_tools: [${toolsJson}],`);
}
const trailingOptions = options.length ? `\n${options.join("\n")}` : "";
return `import OpenAI from "openai";
const client = new OpenAI({
baseURL: ${j(`${base}/v1`)},
apiKey: ${j(key)},
});
const response = await client.chat.completions.create({
model: ${j(model)},
messages: [{ role: "user", content: ${j(PROMPT)} }],${trailingOptions}
stream: true,
});
for await (const chunk of response) {
process.stdout.write(chunk.choices?.[0]?.delta?.content || "");
}${autoSwitch ? javascriptSwitchDemo() : ""}`;
}
function javascriptSwitchDemo(): string {
return `
// "Switch model by request" is on: replace the model below with another GGUF you
// have downloaded and Studio loads it before serving. Unknown names keep serving
// the current model.
const switchResponse = await client.chat.completions.create({
model: ${j(SWITCH_MODEL)},
messages: [{ role: "user", content: ${j(SWITCH_PROMPT)} }],
stream: true,
});
for await (const chunk of switchResponse) {
process.stdout.write(chunk.choices?.[0]?.delta?.content || "");
}`;
}
function buildSnippets(
base: string,
key: string,
@ -292,17 +347,24 @@ function buildSnippets(
return {
curl: curl(base, key, model, "plain", autoSwitch),
python: pythonSnippet(base, key, model, "plain", autoSwitch),
javascript: javascriptSnippet(base, key, model, "plain", autoSwitch),
curlTools: curl(base, key, model, "tools", autoSwitch),
pythonTools: pythonSnippet(base, key, model, "tools", autoSwitch),
javascriptTools: javascriptSnippet(base, key, model, "tools", autoSwitch),
curlAdvanced: curl(base, key, model, "advanced", autoSwitch),
pythonAdvanced: pythonSnippet(base, key, model, "advanced", autoSwitch),
javascriptAdvanced: javascriptSnippet(
base,
key,
model,
"advanced",
autoSwitch,
),
};
}
const KEY_PLACEHOLDER = "sk-unsloth-YOUR_KEY";
const MODEL_FALLBACK = "unsloth/gemma-4-E4B-it-GGUF:UD-Q5_K_XL";
// Default ON: when a tunnel exists, examples should show the public base_url.
const USE_TUNNEL_KEY = "unsloth_api_use_tunnel";
function readUseTunnelPref(): boolean {
@ -319,11 +381,10 @@ function writeUseTunnelPref(value: boolean): void {
try {
window.localStorage.setItem(USE_TUNNEL_KEY, value ? "true" : "false");
} catch {
// Non-fatal: the toggle still applies for this session.
// Non-fatal
}
}
// Active local checkpoint as repo[:variant]; external/none falls back to a default.
function useLoadedModelName(): string {
const checkpoint = useChatRuntimeStore((s) => s.params.checkpoint);
const ggufVariant = useChatRuntimeStore((s) => s.activeGgufVariant);
@ -338,7 +399,6 @@ function useLoadedModelName(): string {
}, [checkpoint, ggufVariant]);
}
// shiki highlighting via the app's shared code plugin + themes (same as chat).
const SHIKI_THEMES = [unslothLightTheme, unslothDarkTheme] as [
typeof unslothLightTheme,
typeof unslothDarkTheme,
@ -352,7 +412,6 @@ function HighlightedCode({
code: string;
language: string;
}) {
// Fence so Streamdown's shiki plugin highlights it (no markdown inside a fence).
const markdown = useMemo(
() => `\`\`\`${language}\n${code}\n\`\`\``,
[code, language],
@ -383,6 +442,7 @@ export function UsageExamples({ apiKey }: { apiKey?: string | null }) {
);
const [copied, setCopied] = useState(false);
const [copiedUrl, setCopiedUrl] = useState(false);
const [copiedAgent, setCopiedAgent] = useState(false);
const [useTunnel, setUseTunnel] = useState<boolean>(readUseTunnelPref);
// null while loading; the same setting the General tab exposes (shared cache).
const [autoSwitch, setAutoSwitch] = useState<OpenAIAutoSwitchSettings | null>(
@ -390,7 +450,6 @@ export function UsageExamples({ apiKey }: { apiKey?: string | null }) {
);
const [savingAutoSwitch, setSavingAutoSwitch] = useState(false);
// Tunnel may start after the first /api/health read; refresh so it surfaces here.
useEffect(() => {
void fetchDeviceType({ force: true });
}, []);
@ -410,10 +469,7 @@ export function UsageExamples({ apiKey }: { apiKey?: string | null }) {
}, []);
const model = useLoadedModelName();
// Real key while revealed (before "Done"); otherwise a placeholder.
const key = apiKey || KEY_PLACEHOLDER;
// Toggle on + tunnel up: public tunnel URL. Off: backend direct host:port
// (origin is only a last-resort fallback).
const origin = typeof window !== "undefined" ? window.location.origin : "";
const base =
useTunnel && cloudflareUrl ? cloudflareUrl : (serverUrl ?? origin);
@ -423,13 +479,20 @@ export function UsageExamples({ apiKey }: { apiKey?: string | null }) {
() => buildSnippets(base, key, model, os, autoSwitchOn),
[base, key, model, os, autoSwitchOn],
);
// Agent command must target the server the panel shows, not the :8888 default.
const agentCommand = useMemo(
() => buildAgentCommand(base, key, os),
[base, key, os],
);
const osAware = OS_AWARE[lang];
const shikiLang = CURL_TYPES.has(lang)
? os === "windows"
? "powershell"
: "bash"
: "python";
: JAVASCRIPT_TYPES.has(lang)
? "javascript"
: "python";
const handleCopy = async () => {
if (await copyToClipboard(snippets[lang])) {
@ -464,6 +527,13 @@ export function UsageExamples({ apiKey }: { apiKey?: string | null }) {
}
};
const handleCopyAgent = async () => {
if (await copyToClipboard(agentCommand)) {
setCopiedAgent(true);
setTimeout(() => setCopiedAgent(false), 1800);
}
};
return (
<section className="flex min-w-0 max-w-full flex-col">
<h2 className="mb-2 text-sm font-semibold text-foreground">
@ -539,8 +609,6 @@ export function UsageExamples({ apiKey }: { apiKey?: string | null }) {
</Tooltip>
)}
</div>
{/* Always rendered (dimmed when off) so toggling never changes the
row height and shifts the code block below. */}
<button
type="button"
onClick={handleCopyUrl}
@ -629,15 +697,39 @@ export function UsageExamples({ apiKey }: { apiKey?: string | null }) {
/>
{copied ? t("settings.apiKeys.copied") : t("settings.apiKeys.copy")}
</button>
{/* key on the snippet so Streamdown remounts and re-highlights when
only a substring (e.g. the base URL) changes; its block memo
otherwise keeps the stale render. */}
<HighlightedCode
key={snippets[lang]}
code={snippets[lang]}
language={shikiLang}
/>
</div>
<div className="flex min-w-0 flex-col gap-1.5 border-t border-border px-3 py-2.5">
<span className="text-[11px] font-semibold text-foreground">
{t("settings.apiKeys.codingAgents")}
</span>
<span className="text-[11px] leading-snug text-muted-foreground">
{t("settings.apiKeys.codingAgentsHint")}
</span>
<div className="relative mt-0.5 min-w-0">
<code className="block min-w-0 overflow-x-auto rounded border border-border bg-muted/30 px-2 py-1.5 pr-14 font-mono text-[11px] text-foreground">
{agentCommand}
</code>
<button
type="button"
onClick={handleCopyAgent}
className="absolute right-1.5 top-1/2 flex -translate-y-1/2 items-center gap-1 rounded border border-border bg-background/80 px-1.5 py-0.5 text-[11px] text-muted-foreground backdrop-blur transition-colors hover:text-foreground focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring"
aria-label={t("settings.apiKeys.copySnippet")}
>
<HugeiconsIcon
icon={copiedAgent ? Tick02Icon : Copy01Icon}
className={cn("size-3.5", copiedAgent && "text-emerald-600")}
/>
</button>
</div>
<span className="text-[11px] leading-snug text-muted-foreground">
{t("settings.apiKeys.codingAgentsSwap")}
</span>
</div>
<div className="flex flex-wrap items-center gap-x-2 gap-y-1 border-t border-border px-3 py-2 text-[11px] text-muted-foreground">
<span>{t("settings.apiKeys.setupDocs")}</span>
{DOC_LINKS.map((link) => (

208
studio/frontend/src/features/settings/settings-dialog.tsx
View file

@ -12,6 +12,7 @@ import { cn } from "@/lib/utils";
import {
Cancel01Icon,
CloudIcon,
CpuIcon,
Globe02Icon,
HelpCircleIcon,
Message01Icon,
@ -33,6 +34,8 @@ import { ChatTab } from "./tabs/chat-tab";
import { ConnectionsTab } from "./tabs/connections-tab";
import { GeneralTab } from "./tabs/general-tab";
import { ProfileTab } from "./tabs/profile-tab";
import { ResourcesTab } from "./tabs/resources-tab";
import { FloatingMonitor } from "@/components/floating-monitor";
interface TabDef {
id: SettingsTab;
@ -49,6 +52,11 @@ const TABS: TabDef[] = [
labelKey: "settings.tabs.appearance",
icon: PaintBrush02Icon,
},
{
id: "resources",
labelKey: "settings.tabs.resources",
icon: CpuIcon,
},
{
id: "chat",
labelKey: "settings.tabs.chat",
@ -77,6 +85,8 @@ function renderTab(tab: SettingsTab) {
return <ProfileTab />;
case "appearance":
return <AppearanceTab />;
case "resources":
return <ResourcesTab />;
case "chat":
return <ChatTab />;
case "connections":
@ -100,6 +110,7 @@ export function SettingsDialog() {
general: null,
profile: null,
appearance: null,
resources: null,
chat: null,
connections: null,
"api-keys": null,
@ -115,110 +126,113 @@ export function SettingsDialog() {
}, [open, activeTab]);
return (
<Dialog open={open} onOpenChange={(o) => !o && closeDialog()}>
<DialogContent
showCloseButton={false}
overlayClassName="bg-black/30 supports-backdrop-filter:backdrop-blur-[2px]"
onCloseAutoFocus={(e) => {
// Restore focus to the element that triggered openDialog(). Radix's
// FocusScope races our rAF-scheduled tab focus and loses the
// previous-focus reference, so restore it by hand.
if (opener && opener.isConnected) {
e.preventDefault();
opener.focus({ preventScroll: true });
}
}}
className={cn(
// Cap at 820px but shrink to the viewport so it doesn't clip on
// iPad-portrait widths (640-820px) where fixed `w-[820px]` overflows.
"settings-surface !max-w-[min(820px,calc(100vw-2rem))] h-[560px] w-[min(820px,calc(100vw-2rem))] p-0 overflow-hidden",
// Soft shadow, no outline ring. Pin --radius to the light value so
// corner rounding matches in dark mode.
"shadow-border rounded-xl ring-0 [--radius:1.1rem]",
"max-sm:h-dvh max-sm:w-dvw max-sm:!max-w-none max-sm:rounded-none",
)}
>
<DialogTitle className="sr-only">
{t("settings.dialog.title")}
</DialogTitle>
<DialogDescription className="sr-only">
{t("settings.dialog.description")}
</DialogDescription>
<div className="flex h-full min-h-0 max-sm:flex-col">
<aside className="font-heading flex w-[216px] shrink-0 flex-col border-r border-sidebar-border bg-muted/20 p-2 dark:border-r-0 max-sm:w-full max-sm:border-r-0 max-sm:border-b max-sm:border-sidebar-border">
<h2 className="pl-3 pr-2.5 pt-3.5 pb-3.5 text-[19px] font-semibold text-foreground max-sm:hidden">
{t("settings.dialog.title")}
</h2>
<nav className="flex flex-col gap-0.5 max-sm:flex-row max-sm:overflow-x-auto">
{TABS.map((tab) => {
const active = activeTab === tab.id;
return (
<button
key={tab.id}
ref={(node) => {
tabButtonRefs.current[tab.id] = node;
}}
type="button"
onClick={() => setActiveTab(tab.id)}
className={cn(
"relative flex h-[32px] items-center gap-2.5 rounded-full pl-3 pr-2.5 text-[14.5px] leading-[19px] tracking-nav font-medium transition-colors",
"max-sm:shrink-0",
"focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring focus-visible:ring-offset-1 focus-visible:ring-offset-background",
active
? "text-black dark:text-white"
: "text-[#383835] dark:text-[#c7c7c4] hover:bg-[#ececec] dark:hover:bg-[#3a3d43] hover:text-black dark:hover:text-white",
)}
>
{active && (
<motion.span
layoutId="settings-active-pill"
className="absolute inset-0 rounded-full bg-[#ececec] dark:bg-[#3a3d43]"
transition={
reduced
? { duration: 0 }
: {
<>
<Dialog open={open} onOpenChange={(o) => !o && closeDialog()}>
<DialogContent
showCloseButton={false}
overlayClassName="bg-black/30 supports-backdrop-filter:backdrop-blur-[2px]"
onCloseAutoFocus={(e) => {
// Restore focus to the element that triggered openDialog(). Radix's
// FocusScope races our rAF-scheduled tab focus and loses the
// previous-focus reference, so restore it by hand.
if (opener && opener.isConnected) {
e.preventDefault();
opener.focus({ preventScroll: true });
}
}}
className={cn(
// Cap at 820px but shrink to the viewport so it doesn't clip on
// iPad-portrait widths (640-820px) where fixed `w-[820px]` overflows.
"settings-surface !max-w-[min(820px,calc(100vw-2rem))] h-[560px] w-[min(820px,calc(100vw-2rem))] p-0 overflow-hidden",
// Soft shadow, no outline ring. Pin --radius to the light value so
// corner rounding matches in dark mode.
"shadow-border rounded-xl ring-0 [--radius:1.1rem]",
"max-sm:h-dvh max-sm:w-dvw max-sm:!max-w-none max-sm:rounded-none",
)}
>
<DialogTitle className="sr-only">
{t("settings.dialog.title")}
</DialogTitle>
<DialogDescription className="sr-only">
{t("settings.dialog.description")}
</DialogDescription>
<div className="flex h-full min-h-0 max-sm:flex-col">
<aside className="font-heading flex w-[216px] shrink-0 flex-col border-r border-sidebar-border bg-muted/20 p-2 dark:border-r-0 max-sm:w-full max-sm:border-r-0 max-sm:border-b max-sm:border-sidebar-border">
<h2 className="pl-3 pr-2.5 pt-3.5 pb-3.5 text-[19px] font-semibold text-foreground max-sm:hidden">
{t("settings.dialog.title")}
</h2>
<nav className="flex flex-col gap-0.5 max-sm:flex-row max-sm:overflow-x-auto">
{TABS.map((tab) => {
const active = activeTab === tab.id;
return (
<button
key={tab.id}
ref={(node) => {
tabButtonRefs.current[tab.id] = node;
}}
type="button"
onClick={() => setActiveTab(tab.id)}
className={cn(
"relative flex h-[32px] items-center gap-2.5 rounded-full pl-3 pr-2.5 text-[14.5px] leading-[19px] tracking-nav font-medium transition-colors",
"max-sm:shrink-0",
"focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring focus-visible:ring-offset-1 focus-visible:ring-offset-background",
active
? "text-black dark:text-white"
: "text-[#383835] dark:text-[#c7c7c4] hover:bg-[#ececec] dark:hover:bg-[#3a3d43] hover:text-black dark:hover:text-white",
)}
>
{active && (
<motion.span
layoutId="settings-active-pill"
className="absolute inset-0 rounded-full bg-[#ececec] dark:bg-[#3a3d43]"
transition={
reduced
? { duration: 0 }
: {
type: "spring",
stiffness: 500,
damping: 35,
mass: 0.5,
}
}
}
/>
)}
<HugeiconsIcon
icon={tab.icon}
strokeWidth={1.75}
className="relative z-10 size-icon"
/>
)}
<HugeiconsIcon
icon={tab.icon}
strokeWidth={1.75}
className="relative z-10 size-icon"
/>
<span className="relative z-10 min-w-0 truncate">
{t(tab.labelKey)}
</span>
{tab.badgeKey ? (
<span className="relative z-10 ml-auto rounded-full bg-emerald-500/10 px-2 py-1 text-[10px] leading-none font-semibold text-emerald-700 dark:text-emerald-300">
{t(tab.badgeKey)}
<span className="relative z-10 min-w-0 truncate">
{t(tab.labelKey)}
</span>
) : null}
</button>
);
})}
</nav>
</aside>
{tab.badgeKey ? (
<span className="relative z-10 ml-auto rounded-full bg-emerald-500/10 px-2 py-1 text-[10px] leading-none font-semibold text-emerald-700 dark:text-emerald-300">
{t(tab.badgeKey)}
</span>
) : null}
</button>
);
})}
</nav>
</aside>
<main className="relative flex min-h-0 min-w-0 flex-1 flex-col">
<button
type="button"
onClick={closeDialog}
className="absolute top-3 right-3 z-10 flex size-7 items-center justify-center rounded-full text-[#383835] dark:text-[#c7c7c4] transition-colors hover:bg-[#ececec] dark:hover:bg-[#3a3d43] hover:text-black dark:hover:text-white focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring"
aria-label={t("settings.dialog.closeAriaLabel")}
>
<HugeiconsIcon icon={Cancel01Icon} className="size-4" />
</button>
<div className="hover-scrollbar flex min-h-0 min-w-0 flex-1 flex-col overflow-y-auto p-6 [scrollbar-gutter:stable]">
{renderTab(activeTab)}
</div>
</main>
</div>
</DialogContent>
</Dialog>
<main className="relative flex min-h-0 min-w-0 flex-1 flex-col">
<button
type="button"
onClick={closeDialog}
className="absolute top-3 right-3 z-10 flex size-7 items-center justify-center rounded-full text-[#383835] dark:text-[#c7c7c4] transition-colors hover:bg-[#ececec] dark:hover:bg-[#3a3d43] hover:text-black dark:hover:text-white focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring"
aria-label={t("settings.dialog.closeAriaLabel")}
>
<HugeiconsIcon icon={Cancel01Icon} className="size-4" />
</button>
<div className="hover-scrollbar flex min-h-0 min-w-0 flex-1 flex-col overflow-y-auto p-6 [scrollbar-gutter:stable]">
{renderTab(activeTab)}
</div>
</main>
</div>
</DialogContent>
</Dialog>
<FloatingMonitor />
</>
);
}

24
studio/frontend/src/features/settings/stores/monitor-overlay-store.ts Normal file
View file

@ -0,0 +1,24 @@
// 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 { create } from "zustand";
import { persist } from "zustand/middleware";
interface MonitorOverlayState {
isOpen: boolean;
isMinimized: boolean;
setIsOpen: (open: boolean) => void;
toggleMinimized: () => void;
}
export const useMonitorOverlayStore = create<MonitorOverlayState>()(
persist(
(set) => ({
isOpen: false,
isMinimized: false,
setIsOpen: (isOpen) => set({ isOpen }),
toggleMinimized: () => set((state) => ({ isMinimized: !state.isMinimized })),
}),
{ name: "unsloth_monitor_overlay" }
)
);

2
studio/frontend/src/features/settings/stores/settings-dialog-store.ts
View file

@ -7,6 +7,7 @@ export type SettingsTab =
| "general"
| "profile"
| "appearance"
| "resources"
| "chat"
| "connections"
| "api-keys"
@ -60,6 +61,7 @@ function loadInitialTab(): SettingsTab {
"general",
"profile",
"appearance",
"resources",
"chat",
"connections",
"api-keys",

3
studio/frontend/src/features/settings/tabs/api-keys-tab.tsx
View file

@ -17,6 +17,7 @@ import { fetchApiKeys, revokeApiKey, type ApiKey } from "../api/api-keys";
import { ApiMonitorConsole } from "../components/api-monitor-console";
import { ApiKeyRow } from "../components/api-key-row";
import { CreateKeyForm } from "../components/create-key-form";
import { ModelAutoSwitchSection } from "../components/model-auto-switch-section";
import { KeyRevealCard } from "../components/key-reveal-card";
import { UsageExamples } from "../components/usage-examples";
@ -171,6 +172,8 @@ export function ApiKeysTab() {
<UsageExamples apiKey={revealed} />
<ModelAutoSwitchSection />
<Dialog open={revokeTarget !== null} onOpenChange={(o) => !o && setRevokeTarget(null)}>
<DialogContent className="max-w-md">
<DialogHeader>

227
studio/frontend/src/features/settings/tabs/general-tab.tsx
View file

@ -41,6 +41,13 @@ import {
rotatePreviewLinks,
updatePreviewSharing,
} from "../api/preview-sharing";
import {
type EmbeddingModelSettings,
EmbeddingModelVerificationError,
loadEmbeddingModelSettings,
resetEmbeddingModelSettings,
updateEmbeddingModelSettings,
} from "../api/embedding-model";
import {
DEFAULT_UPLOAD_LIMIT_MB,
type UploadLimitSettings,
@ -48,7 +55,7 @@ import {
updateUploadLimitSettings,
} from "../api/upload-limit";
import { ChangePasswordDialog } from "../components/change-password-dialog";
import { ModelAutoSwitchSection } from "../components/model-auto-switch-section";
import { EmbeddingModelCombobox } from "../components/embedding-model-combobox";
import { SettingsRow } from "../components/settings-row";
import { SettingsSection } from "../components/settings-section";
import { StudioVersionSection } from "../components/studio-version-section";
@ -81,6 +88,7 @@ const PREFS_KEYS: string[] = [
"unsloth_chat_load_on_selection",
"unsloth_chat_expand_quantizations",
"unsloth_chat_show_all_quantizations",
"unsloth_models_fit_on_device_only",
// Chat presets
"unsloth_chat_custom_presets",
"unsloth_chat_active_preset",
@ -96,6 +104,7 @@ const PREFS_KEYS: string[] = [
"tour:studio:v1",
// Update notifications
"unsloth_show_llama_update_banner",
"unsloth_monitor_overlay",
];
// Set by resetAllPrefs so the unmount-commit effect skips writing back the
@ -163,6 +172,16 @@ export function GeneralTab() {
const [revokePreviewOpen, setRevokePreviewOpen] = useState(false);
const [isRevokingPreview, setIsRevokingPreview] = useState(false);
const [modelsFolder, setModelsFolder] = useState<ModelsFolder | null>(null);
const [embeddingModel, setEmbeddingModel] =
useState<EmbeddingModelSettings | null>(null);
const [draftEmbeddingModel, setDraftEmbeddingModel] = useState("");
const [embeddingModelError, setEmbeddingModelError] = useState<string | null>(
null,
);
// Set after a 409 (unverifiable model); offers "Save anyway".
const [embeddingModelNeedsForce, setEmbeddingModelNeedsForce] =
useState(false);
const [isSavingEmbeddingModel, setIsSavingEmbeddingModel] = useState(false);
const draftRef = useRef(draftToken);
useEffect(() => {
@ -257,6 +276,27 @@ export function GeneralTab() {
};
}, [t]);
useEffect(() => {
let cancelled = false;
void loadEmbeddingModelSettings()
.then((settings) => {
if (cancelled) return;
setEmbeddingModel(settings);
setDraftEmbeddingModel(settings.embeddingModel);
})
.catch((error) => {
if (cancelled) return;
setEmbeddingModelError(
error instanceof Error
? error.message
: t("settings.general.rag.loadError"),
);
});
return () => {
cancelled = true;
};
}, [t]);
useEffect(() => {
let cancelled = false;
void loadModelsFolder()
@ -349,6 +389,58 @@ export function GeneralTab() {
}
};
const saveEmbeddingModel = async (force: boolean) => {
const trimmed = draftEmbeddingModel.trim();
if (!trimmed) {
setEmbeddingModelError(t("settings.general.rag.emptyError"));
return;
}
setIsSavingEmbeddingModel(true);
setEmbeddingModelError(null);
try {
const settings = await updateEmbeddingModelSettings(trimmed, {
hfToken: hfToken || undefined,
force,
});
setEmbeddingModel(settings);
setDraftEmbeddingModel(settings.embeddingModel);
setEmbeddingModelNeedsForce(false);
toast.success(t("settings.general.rag.saved"), {
description: t("settings.general.rag.reindexWarning"),
});
} catch (error) {
if (error instanceof EmbeddingModelVerificationError) {
setEmbeddingModelNeedsForce(true);
}
setEmbeddingModelError(
error instanceof Error
? error.message
: t("settings.general.rag.saveError"),
);
} finally {
setIsSavingEmbeddingModel(false);
}
};
const resetEmbeddingModel = async () => {
setIsSavingEmbeddingModel(true);
setEmbeddingModelError(null);
setEmbeddingModelNeedsForce(false);
try {
const settings = await resetEmbeddingModelSettings();
setEmbeddingModel(settings);
setDraftEmbeddingModel(settings.embeddingModel);
} catch (error) {
setEmbeddingModelError(
error instanceof Error
? error.message
: t("settings.general.rag.saveError"),
);
} finally {
setIsSavingEmbeddingModel(false);
}
};
const saveUploadLimit = async () => {
const parsed = Number(draftUploadLimit);
if (!Number.isInteger(parsed)) {
@ -499,38 +591,6 @@ export function GeneralTab() {
</SettingsRow>
</SettingsSection>
<SettingsSection title={t("settings.general.helperLlm.sectionTitle")}>
<SettingsRow
label={t("settings.general.helperLlm.preloadOnStartup")}
description={t(
"settings.general.helperLlm.preloadOnStartupDescription",
)}
>
<div className="flex flex-col items-end gap-1">
<Switch
checked={helperPrecache?.enabled ?? false}
disabled={
!helperPrecache ||
isSavingHelperPrecache ||
helperPrecache.disabledByEnv
}
onCheckedChange={(enabled) => void saveHelperPrecache(enabled)}
/>
{helperPrecache?.disabledByEnv ? (
<span className="max-w-[260px] text-right text-xs text-muted-foreground">
{t("settings.general.helperLlm.disabledByEnv")}
</span>
) : helperPrecacheError ? (
<span className="max-w-[260px] text-right text-xs text-destructive">
{helperPrecacheError}
</span>
) : null}
</div>
</SettingsRow>
</SettingsSection>
<ModelAutoSwitchSection />
<SettingsSection
title={t("settings.general.previewSharing.sectionTitle")}
>
@ -567,6 +627,77 @@ export function GeneralTab() {
</SettingsRow>
</SettingsSection>
<SettingsSection title={t("settings.general.rag.sectionTitle")}>
<SettingsRow
label={t("settings.general.rag.embeddingModel")}
description={t("settings.general.rag.embeddingModelDescription", {
defaultModel: embeddingModel?.defaultEmbeddingModel ?? "",
})}
>
<div className="flex flex-col items-end gap-1">
<div className="flex items-center gap-2">
<EmbeddingModelCombobox
value={draftEmbeddingModel}
onChange={(next) => {
setDraftEmbeddingModel(next);
setEmbeddingModelNeedsForce(false);
setEmbeddingModelError(null);
}}
accessToken={hfToken || undefined}
disabled={!embeddingModel}
placeholder={embeddingModel?.defaultEmbeddingModel ?? ""}
ariaLabel={t("settings.general.rag.embeddingModel")}
className="w-[220px]"
/>
<Button
variant="outline"
size="sm"
disabled={
!embeddingModel ||
isSavingEmbeddingModel ||
draftEmbeddingModel.trim() === embeddingModel.embeddingModel
}
onClick={() => void saveEmbeddingModel(false)}
>
{isSavingEmbeddingModel
? t("common.saving")
: t("common.save")}
</Button>
</div>
{embeddingModelError ? (
<span className="max-w-[300px] text-right text-xs text-destructive">
{embeddingModelError}
</span>
) : null}
<div className="flex items-center gap-2">
{embeddingModelNeedsForce ? (
<Button
variant="outline"
size="sm"
disabled={isSavingEmbeddingModel}
onClick={() => void saveEmbeddingModel(true)}
>
{t("settings.general.rag.saveAnyway")}
</Button>
) : null}
{embeddingModel?.isCustom ? (
<Button
variant="ghost"
size="sm"
disabled={isSavingEmbeddingModel}
onClick={() => void resetEmbeddingModel()}
>
{t("settings.general.rag.resetAction")}
</Button>
) : null}
</div>
<span className="max-w-[300px] text-right text-xs text-muted-foreground">
{t("settings.general.rag.reindexWarning")}
</span>
</div>
</SettingsRow>
</SettingsSection>
<SettingsSection title={t("settings.general.uploads.sectionTitle")}>
<SettingsRow
label={t("settings.general.uploads.maxUploadSize")}
@ -632,6 +763,36 @@ export function GeneralTab() {
</SettingsSection>
)}
<SettingsSection title={t("settings.general.helperLlm.sectionTitle")}>
<SettingsRow
label={t("settings.general.helperLlm.preloadOnStartup")}
description={t(
"settings.general.helperLlm.preloadOnStartupDescription",
)}
>
<div className="flex flex-col items-end gap-1">
<Switch
checked={helperPrecache?.enabled ?? false}
disabled={
!helperPrecache ||
isSavingHelperPrecache ||
helperPrecache.disabledByEnv
}
onCheckedChange={(enabled) => void saveHelperPrecache(enabled)}
/>
{helperPrecache?.disabledByEnv ? (
<span className="max-w-[260px] text-right text-xs text-muted-foreground">
{t("settings.general.helperLlm.disabledByEnv")}
</span>
) : helperPrecacheError ? (
<span className="max-w-[260px] text-right text-xs text-destructive">
{helperPrecacheError}
</span>
) : null}
</div>
</SettingsRow>
</SettingsSection>
<SettingsSection
title={t("settings.general.resetPreferences.sectionTitle")}
>

477
studio/frontend/src/features/settings/tabs/resources-tab.tsx Normal file
View file

@ -0,0 +1,477 @@
// 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 { Switch } from "@/components/ui/switch";
import { openModelsDir } from "@/features/native-intents";
import { useSystemInfo, type GpuDevice } from "@/hooks/use-system";
import { isTauri } from "@/lib/api-base";
import { copyToClipboard } from "@/lib/copy-to-clipboard";
import { toast } from "@/lib/toast";
import { cn } from "@/lib/utils";
import { useT } from "@/i18n";
import { useEffect, useMemo, useState } from "react";
import { loadModelsFolder, type ModelsFolder } from "../api/models-folder";
import { SettingsRow } from "../components/settings-row";
import { SettingsSection } from "../components/settings-section";
import { useMonitorOverlayStore } from "../stores/monitor-overlay-store";
import { LayersIcon } from "lucide-react";
const POLL_MS = 3000;
function isFiniteNumber(value: number | null | undefined): value is number {
return typeof value === "number" && Number.isFinite(value);
}
function clampPercent(value: number | null | undefined): number {
if (!isFiniteNumber(value)) return 0;
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 | null | undefined): string {
const safe = isFiniteNumber(value) ? Math.max(0, value) : 0;
const digits = safe >= 10 ? 1 : 2;
return `${safe.toFixed(digits)} GB`;
}
function formatMb(value: number | null | undefined): string {
const safe = isFiniteNumber(value) ? Math.max(0, value) : 0;
return `${Math.round(safe).toLocaleString()} MB`;
}
function formatPercent(value: number | null | undefined): string {
return `${Math.round(clampPercent(value))}%`;
}
function formatFrequency(mhz: number | null | undefined): string | null {
if (!isFiniteNumber(mhz) || mhz <= 0) return null;
if (mhz >= 1000) return `${(mhz / 1000).toFixed(2)} GHz`;
return `${Math.round(mhz)} MHz`;
}
function formatUptime(seconds: number | null | undefined): string {
if (!isFiniteNumber(seconds) || seconds <= 0) return "0m";
const minutes = Math.floor(seconds / 60);
const hours = Math.floor(minutes / 60);
const days = Math.floor(hours / 24);
if (days > 0) return `${days}d ${hours % 24}h`;
if (hours > 0) return `${hours}h ${minutes % 60}m`;
return `${Math.max(1, minutes)}m`;
}
function MetricTile({
label,
value,
detail,
percent,
}: {
label: string;
value: string;
detail: string;
percent: number;
}) {
const safePercent = clampPercent(percent);
return (
<div className="flex min-w-0 flex-col gap-2 rounded-md border border-border/60 bg-muted/20 p-3">
<div className="flex items-center justify-between gap-3">
<span className="truncate text-[11px] font-semibold uppercase tracking-[0.08em] text-muted-foreground">
{label}
</span>
<span
className={cn(
"shrink-0 font-mono text-xs tabular-nums",
usageTextClass(safePercent),
)}
>
{formatPercent(safePercent)}
</span>
</div>
<div className="min-w-0">
<div className="truncate font-mono text-sm tabular-nums text-foreground">
{value}
</div>
<div className="mt-0.5 truncate text-xs text-muted-foreground">
{detail}
</div>
</div>
<Progress
value={safePercent}
aria-label={label}
className="h-1.5 rounded-full bg-muted"
indicatorClassName={usageIndicatorClass(safePercent)}
/>
</div>
);
}
function InfoRow({
label,
value,
detail,
}: {
label: string;
value: string;
detail?: string;
}) {
return (
<div className="flex min-w-0 items-center justify-between gap-4 py-2.5">
<span className="min-w-0 truncate text-sm font-medium text-foreground">
{label}
</span>
<span
title={detail ?? value}
className="min-w-0 max-w-[60%] truncate text-right font-mono text-xs tabular-nums text-muted-foreground"
>
{detail ? `${value} (${detail})` : value}
</span>
</div>
);
}
function deviceOrdinal(device: GpuDevice): number | undefined {
return device.visible_ordinal ?? device.index;
}
export function ResourcesTab() {
const t = useT();
const [liveUpdates, setLiveUpdates] = useState(true);
const { isOpen, setIsOpen } = useMonitorOverlayStore();
const systemInfo = useSystemInfo({
enabled: liveUpdates,
pollMs: liveUpdates ? POLL_MS : undefined,
});
const [modelsFolder, setModelsFolder] = useState<ModelsFolder | null>(null);
const [modelsFolderLoaded, setModelsFolderLoaded] = useState(false);
useEffect(() => {
let cancelled = false;
void loadModelsFolder()
.then((folder) => {
if (cancelled) return;
setModelsFolder(folder);
setModelsFolderLoaded(true);
})
.catch(() => {
if (cancelled) return;
setModelsFolderLoaded(true);
});
return () => {
cancelled = true;
};
}, []);
const metrics = useMemo(() => {
const devices = systemInfo.gpu?.devices ?? [];
const ramTotal = systemInfo.memory?.total_gb ?? 0;
const ramAvailable = systemInfo.memory?.available_gb ?? 0;
const ramUsed = Math.max(0, ramTotal - ramAvailable);
const diskTotal = systemInfo.disk?.total_gb ?? 0;
const diskFree = systemInfo.disk?.free_gb ?? 0;
const diskUsed = Math.max(0, diskTotal - diskFree);
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 vramFree = devices.reduce(
(sum, device) =>
sum +
(device.vram_free_gb ??
Math.max(0, (device.memory_total_gb ?? 0) - (device.vram_used_gb ?? 0))),
0,
);
const vramPercent = vramTotal > 0 ? (vramUsed / vramTotal) * 100 : 0;
return {
devices,
ramTotal,
ramUsed,
diskTotal,
diskFree,
diskUsed,
vramTotal,
vramUsed,
vramFree,
vramPercent,
};
}, [systemInfo]);
const handleModelsFolder = async () => {
const folder = modelsFolder;
if (!folder) return;
if (isTauri) {
try {
await openModelsDir(folder.path);
} catch (error) {
toast.error(t("settings.resources.storage.openError"), {
description: error instanceof Error ? error.message : undefined,
});
}
return;
}
if (await copyToClipboard(folder.path)) {
toast.success(t("settings.resources.storage.copied"));
} else {
toast.error(t("settings.resources.storage.copyError"));
}
};
const cpuCoresLabel =
systemInfo.cpu?.logical_count && systemInfo.cpu?.physical_count
? t("settings.resources.liveMonitor.cpuCores", {
logical: systemInfo.cpu.logical_count,
physical: systemInfo.cpu.physical_count,
})
: t("settings.resources.environment.unknown");
const cpuFrequencyLabel = formatFrequency(systemInfo.cpu?.frequency_mhz);
const hasGpu =
(systemInfo.gpu?.available ?? false) && metrics.devices.length > 0;
const backendLabel = (
systemInfo.gpu?.backend ?? systemInfo.device_backend ?? "cpu"
).toUpperCase();
const modelsFolderPath = modelsFolder
? modelsFolder.path
: modelsFolderLoaded
? t("settings.resources.environment.unknown")
: t("common.loading");
return (
<div className="flex flex-col gap-6">
<header className="flex flex-wrap items-start justify-between gap-3">
<div className="flex min-w-0 flex-col gap-1">
<h1 className="text-xl font-semibold font-heading">
{t("settings.resources.title")}
</h1>
<p className="text-xs text-muted-foreground">
{t("settings.resources.description")}
</p>
</div>
<div className="flex items-center gap-3">
<Button
variant={isOpen ? "secondary" : "outline"}
size="sm"
className="gap-1.5 h-8 text-xs rounded-full px-3"
onClick={() => setIsOpen(!isOpen)}
>
<LayersIcon className="size-3.5" />
{isOpen
? t("settings.resources.disableOverlay")
: t("settings.resources.floatingWindow")}
</Button>
<div className="flex shrink-0 items-center gap-2 rounded-full border border-border/60 px-2.5 py-1.5 text-xs font-medium text-foreground">
<span>{t("settings.resources.liveUpdates")}</span>
<Switch
aria-label={t("settings.resources.liveUpdates")}
checked={liveUpdates}
onCheckedChange={setLiveUpdates}
/>
</div>
</div>
</header>
<SettingsSection title={t("settings.resources.liveMonitor.title")}>
<div className="grid gap-2 py-3 sm:grid-cols-2">
<MetricTile
label={t("settings.resources.liveMonitor.cpu")}
value={cpuFrequencyLabel ?? cpuCoresLabel}
detail={
cpuFrequencyLabel
? cpuCoresLabel
: t("settings.resources.liveMonitor.currentLoad")
}
percent={systemInfo.cpu?.usage_percent ?? 0}
/>
<MetricTile
label={t("settings.resources.liveMonitor.ram")}
value={`${formatGb(metrics.ramUsed)} / ${formatGb(metrics.ramTotal)}`}
detail={t("settings.resources.liveMonitor.free", {
value: formatGb(systemInfo.memory?.available_gb),
})}
percent={systemInfo.memory?.percent_used ?? 0}
/>
<MetricTile
label={t("settings.resources.liveMonitor.disk")}
value={`${formatGb(metrics.diskUsed)} / ${formatGb(metrics.diskTotal)}`}
detail={t("settings.resources.liveMonitor.free", {
value: formatGb(metrics.diskFree),
})}
percent={systemInfo.disk?.percent_used ?? 0}
/>
<MetricTile
label={t("settings.resources.liveMonitor.vram")}
value={
hasGpu
? `${formatGb(metrics.vramUsed)} / ${formatGb(metrics.vramTotal)}`
: t("settings.resources.liveMonitor.noGpu")
}
detail={
hasGpu
? t("settings.resources.liveMonitor.free", {
value: formatGb(metrics.vramFree),
})
: backendLabel
}
percent={metrics.vramPercent}
/>
</div>
</SettingsSection>
<SettingsSection title={t("settings.resources.gpu.title")}>
{hasGpu ? (
metrics.devices.map((device, index) => {
const ordinal = deviceOrdinal(device);
const total = device.memory_total_gb ?? 0;
const used = device.vram_used_gb ?? 0;
const free = device.vram_free_gb ?? Math.max(0, total - used);
const percent =
device.vram_utilization_pct ??
(total > 0 ? (used / total) * 100 : null);
const safePercent = clampPercent(percent);
return (
<div
key={`${device.index ?? index}-${device.name ?? "gpu"}`}
className="flex min-w-0 flex-col gap-2 py-3"
>
<div className="flex min-w-0 items-start justify-between gap-4">
<div className="min-w-0 flex-1">
<div className="truncate text-sm font-medium text-foreground">
{device.name ??
t("settings.resources.gpu.unknownDevice")}
</div>
<div className="mt-0.5 truncate text-xs text-muted-foreground">
{ordinal === undefined
? backendLabel
: `${t("settings.resources.gpu.deviceWithIndex", {
index: ordinal,
})}, ${backendLabel}`}
</div>
</div>
<div className="shrink-0 font-mono text-xs tabular-nums text-muted-foreground">
<span>
{formatPercent(safePercent)}{" "}
{t("settings.resources.gpu.vramUtilization")}
</span>
</div>
</div>
<div className="grid gap-1 text-xs text-muted-foreground sm:grid-cols-3 sm:gap-2">
<span className="min-w-0 truncate font-mono tabular-nums">
{t("settings.resources.gpu.used", {
value: formatGb(used),
})}
</span>
<span className="min-w-0 truncate font-mono tabular-nums sm:text-center">
{t("settings.resources.gpu.free", {
value: formatGb(free),
})}
</span>
<span className="min-w-0 truncate font-mono tabular-nums sm:text-right">
{t("settings.resources.gpu.total", {
value: formatGb(total),
})}
</span>
</div>
<Progress
value={safePercent}
aria-label={device.name ?? "GPU"}
className="h-1.5 rounded-full bg-muted"
indicatorClassName={usageIndicatorClass(safePercent)}
/>
</div>
);
})
) : (
<div className="py-3 text-sm text-muted-foreground">
{t("settings.resources.gpu.noGpu")}
</div>
)}
</SettingsSection>
<SettingsSection title={t("settings.resources.storage.title")}>
<InfoRow
label={t("settings.resources.storage.systemDisk")}
value={t("settings.resources.storage.diskUsage", {
used: formatGb(metrics.diskUsed),
total: formatGb(metrics.diskTotal),
})}
detail={t("settings.resources.storage.diskFree", {
free: formatGb(metrics.diskFree),
})}
/>
<SettingsRow
label={t("settings.resources.storage.modelsFolder")}
description={t("settings.resources.storage.modelsFolderDescription")}
className="max-sm:flex-col max-sm:items-start max-sm:gap-2"
>
<div className="flex min-w-0 items-center gap-2 max-sm:max-w-[calc(100vw-5rem)]">
<span
title={modelsFolder?.path}
className="min-w-0 max-w-[280px] truncate font-mono text-xs text-muted-foreground max-sm:max-w-[180px]"
>
{modelsFolderPath}
</span>
<Button
variant="outline"
size="sm"
disabled={!modelsFolder}
onClick={() => void handleModelsFolder()}
>
{isTauri
? t("settings.resources.storage.openAction")
: t("settings.resources.storage.copyAction")}
</Button>
</div>
</SettingsRow>
</SettingsSection>
<SettingsSection title={t("settings.resources.environment.title")}>
<InfoRow
label={t("settings.resources.environment.backend")}
value={backendLabel}
/>
<InfoRow
label={t("settings.resources.environment.python")}
value={systemInfo.python_version}
/>
<InfoRow
label={t("settings.resources.environment.torch")}
value={
systemInfo.ml_packages.torch ??
t("settings.resources.environment.notInstalled")
}
/>
<InfoRow
label={t("settings.resources.environment.transformers")}
value={
systemInfo.ml_packages.transformers ??
t("settings.resources.environment.notInstalled")
}
/>
<InfoRow
label={t("settings.resources.environment.uptime")}
value={formatUptime(systemInfo.uptime_seconds)}
/>
<InfoRow
label={t("settings.resources.environment.processMemory")}
value={formatMb(systemInfo.memory?.process_used_mb)}
/>
</SettingsSection>
</div>
);
}

128
studio/frontend/src/features/studio/sections/progress-section.tsx
View file

@ -29,6 +29,7 @@ import {
import { getTrainingMethodLabel } from "@/features/training/lib/training-methods";
import type { TrainingViewData } from "@/features/training";
import { useGpuUtilization } from "@/hooks";
import type { GpuUtilization } from "@/hooks/use-gpu-utilization";
import { cn } from "@/lib/utils";
import {
ChartAverageIcon,
@ -42,7 +43,7 @@ import {
} from "@hugeicons/core-free-icons";
import { HugeiconsIcon } from "@hugeicons/react";
import { Link, useNavigate } from "@tanstack/react-router";
import { type ReactElement, type ReactNode, useState } from "react";
import { type ReactElement, type ReactNode, useEffect, useState } from "react";
import { useShallow } from "zustand/react/shallow";
import { ChartSettingsSheet } from "./charts/chart-settings-sheet";
import {
@ -123,18 +124,17 @@ export function ProgressSection({
const [stopDialogOpen, setStopDialogOpen] = useState(false);
const [stopRequestedLocal, setStopRequestedLocal] = useState(false);
// Auto-resets when training stops; no useEffect needed
const stopRequested = data.isTrainingRunning && stopRequestedLocal;
const pct =
data.totalSteps > 0
? Math.min(
100,
Math.max(
0,
Math.round((data.currentStep / data.totalSteps) * 100),
),
)
100,
Math.max(
0,
Math.round((data.currentStep / data.totalSteps) * 100),
),
)
: Math.round(data.progressPercent);
const elapsed = data.elapsedSeconds;
@ -214,16 +214,16 @@ export function ProgressSection({
},
...(data.trainingMethod !== "full"
? [
{
section: "LoRA",
rows: [
configRow(t("studio.progress.rank"), cfgLoraRank),
configRow(t("studio.progress.alpha"), cfgLoraAlpha),
configRow(t("studio.progress.dropout"), cfgLoraDropout),
configRow(t("studio.progress.variant"), cfgLoraVariant),
],
},
]
{
section: "LoRA",
rows: [
configRow(t("studio.progress.rank"), cfgLoraRank),
configRow(t("studio.progress.alpha"), cfgLoraAlpha),
configRow(t("studio.progress.dropout"), cfgLoraDropout),
configRow(t("studio.progress.variant"), cfgLoraVariant),
],
},
]
: []),
];
@ -350,8 +350,8 @@ export function ProgressSection({
{stepsPerSecond == null
? t("studio.progress.noStepsPerSecond")
: t("studio.progress.stepsPerSecond", {
value: stepsPerSecond.toFixed(2),
})}
value: stepsPerSecond.toFixed(2),
})}
</span>
{data.currentNumTokens != null && (
<span>{t("studio.progress.tokens", { value: data.currentNumTokens })}</span>
@ -373,14 +373,50 @@ function LiveGpuPanel({
isTrainingRunning: boolean;
}): ReactElement {
const t = useT();
const gpu = useGpuUtilization(isTrainingRunning);
const [selectedGpu, setSelectedGpu] = useState(0);
const gpuData = useGpuUtilization(isTrainingRunning);
const gpus: GpuUtilization[] =
Array.isArray(gpuData?.devices) && gpuData.devices.length > 0
? gpuData.devices
: gpuData && Object.keys(gpuData).length > 0
? [gpuData]
: [];
useEffect(() => {
if (selectedGpu > 0 && selectedGpu >= gpus.length) {
setSelectedGpu(0);
}
}, [gpus.length, selectedGpu]);
const gpuCount = gpus.length;
const currentGpu: Partial<GpuUtilization> = gpus[selectedGpu] || gpus[0] || {};
return (
<div className="flex flex-col gap-3">
<div className="flex items-center justify-between">
<p className="text-xs font-medium text-muted-foreground">
{t("studio.progress.gpuMonitor")}
</p>
<div className="flex items-center justify-between gap-4">
<div className="flex items-center gap-2">
<p className="text-xs font-medium text-muted-foreground">
{t("studio.progress.gpuMonitor")}
</p>
{gpuCount > 1 && (
<select
value={selectedGpu}
onChange={(e) => setSelectedGpu(Number(e.target.value))}
className="h-6 cursor-pointer rounded-md border border-border bg-popover px-1.5 py-0.5 text-[11px] text-popover-foreground outline-none hover:bg-muted focus:border-primary transition-colors font-medium appearance-none"
title="Select GPU"
>
{gpus.map((device, index) => (
<option
key={device.index ?? index}
value={index}
className="bg-popover text-popover-foreground dark:bg-zinc-900 dark:text-zinc-100"
>
GPU {device.visible_ordinal ?? index} - {device.backend} ({device.vram_total_gb ? `${Math.round(device.vram_total_gb)}GB` : "N/A"})
</option>
))}
</select>
)}
</div>
<span className="text-[11px] text-muted-foreground">
{t("studio.progress.live")}
</span>
@ -388,51 +424,44 @@ function LiveGpuPanel({
<div className="grid grid-cols-2 gap-2.5">
<GpuStat
label={t("studio.progress.utilization")}
icon={
<HugeiconsIcon
icon={DashboardSpeed01Icon}
className="size-3.5"
/>
}
icon={<HugeiconsIcon icon={DashboardSpeed01Icon} className="size-3.5" />}
value={
gpu.gpu_utilization_pct != null
? `${gpu.gpu_utilization_pct}%`
currentGpu.gpu_utilization_pct != null
? `${currentGpu.gpu_utilization_pct}%`
: "--"
}
pct={gpu.gpu_utilization_pct ?? 0}
pct={currentGpu.gpu_utilization_pct ?? 0}
/>
<GpuStat
label={t("studio.progress.temperature")}
icon={
<HugeiconsIcon icon={TemperatureIcon} className="size-3.5" />
}
icon={<HugeiconsIcon icon={TemperatureIcon} className="size-3.5" />}
value={
gpu.temperature_c != null ? `${gpu.temperature_c}°C` : "--"
currentGpu.temperature_c != null ? `${currentGpu.temperature_c}°C` : "--"
}
pct={gpu.temperature_c ?? 0}
pct={currentGpu.temperature_c ?? 0}
max={100}
/>
<GpuStat
label={t("studio.progress.vram")}
icon={<HugeiconsIcon icon={RamMemoryIcon} className="size-3.5" />}
value={
gpu.vram_used_gb != null && gpu.vram_total_gb != null
? `${gpu.vram_used_gb} / ${gpu.vram_total_gb} GB`
currentGpu.vram_used_gb != null && currentGpu.vram_total_gb != null
? `${currentGpu.vram_used_gb} / ${currentGpu.vram_total_gb} GB`
: "--"
}
pct={gpu.vram_utilization_pct ?? 0}
pct={currentGpu.vram_utilization_pct ?? 0}
/>
<GpuStat
label={t("studio.progress.power")}
icon={<HugeiconsIcon icon={ZapIcon} className="size-3.5" />}
value={
gpu.power_draw_w != null
? gpu.power_limit_w != null
? `${gpu.power_draw_w} / ${gpu.power_limit_w} W`
: `${gpu.power_draw_w} W`
currentGpu.power_draw_w != null
? currentGpu.power_limit_w != null
? `${currentGpu.power_draw_w} / ${currentGpu.power_limit_w} W`
: `${currentGpu.power_draw_w} W`
: "--"
}
pct={gpu.power_utilization_pct ?? 0}
pct={currentGpu.power_utilization_pct ?? 0}
/>
</div>
</div>
@ -560,7 +589,10 @@ function TrainingHeaderActions({
<HugeiconsIcon icon={StopIcon} className="size-3" />
{stopRequested ? t("studio.training.stopping") : t("studio.training.stopAction")}
</Button>
<AlertDialogContent overlayClassName="bg-background/40 supports-backdrop-filter:backdrop-blur-[1px]">
<AlertDialogContent
className="w-max max-w-[95vw]"
overlayClassName="bg-background/40 supports-backdrop-filter:backdrop-blur-[1px]"
>
<AlertDialogHeader>
<AlertDialogTitle>{t("studio.training.stopTitle")}</AlertDialogTitle>
<AlertDialogDescription>

2
studio/frontend/src/hooks/index.ts
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
export { useDebouncedValue } from "./use-debounced-value";
export { useGpuInfo } from "./use-gpu-info";
export { useGpuUtilization } from "./use-gpu-utilization";
@ -9,3 +10,4 @@ export { useHfDatasetSplits } from "./use-hf-dataset-splits";
export { useHfTokenValidation } from "./use-hf-token-validation";
export { useTauriBackend } from "./use-tauri-backend";
export { useCollapseScrollLock } from "./use-collapse-scroll-lock";
export { useSystemInfo } from "./use-system";

47
studio/frontend/src/hooks/use-gpu-info.ts
View file

@ -3,19 +3,26 @@
import { authFetch } from "@/features/auth";
import { useEffect, useState } from "react";
import type { SystemInfoResponse } from "./use-system";
export interface GpuInfo {
available: boolean;
name: string;
memoryTotalGb: number;
cpuCore: number;
cpuThread: number;
systemRamAvailableGb: number;
systemRamTotalGb: number
}
const DEFAULT_GPU: GpuInfo = {
available: false,
name: "Unknown",
memoryTotalGb: 0,
cpuCore: 0,
cpuThread: 0,
systemRamAvailableGb: 0,
systemRamTotalGb: 0
};
// Module-level cache so multiple components share one fetch.
@ -30,24 +37,30 @@ async function fetchGpuOnce(): Promise<GpuInfo> {
try {
const res = await authFetch("/api/system");
if (!res.ok) throw new Error(`HTTP ${res.status}`);
const data = await res.json();
const ramAvailableGb = data?.memory?.available_gb ?? 0;
const data = await res.json() as SystemInfoResponse;
const gpuData = data?.gpu;
if (!gpuData?.available || !gpuData.devices?.length) {
// No discrete GPU (e.g. Mac): still surface system RAM so memory math
// (unified memory) has a budget to work with.
const info: GpuInfo = { ...DEFAULT_GPU, systemRamAvailableGb: ramAvailableGb };
cachedGpu = info;
return info;
}
const devices = gpuData.devices as Array<{ name?: string; memory_total_gb?: number }>;
const totalGb = devices.reduce((sum, d) => sum + (d.memory_total_gb ?? 0), 0);
const info: GpuInfo = {
available: true,
name: devices[0]?.name ?? "Unknown",
memoryTotalGb: totalGb,
systemRamAvailableGb: ramAvailableGb,
// CPU/RAM exist even on hosts without a GPU, so populate them on every path.
// No discrete GPU (e.g. Mac): still surface system RAM so memory math
// (unified memory) has a budget to work with.
const base = {
cpuCore: data?.cpu?.physical_count ?? 0,
cpuThread: data?.cpu?.logical_count ?? 0,
systemRamAvailableGb: data?.memory?.available_gb ?? 0,
systemRamTotalGb: data?.memory?.total_gb ?? 0,
};
const devices = gpuData?.devices ?? [];
const info: GpuInfo =
gpuData?.available && devices.length
? {
...base,
available: true,
name: devices[0]?.name ?? "Unknown",
memoryTotalGb: devices.reduce((sum, d) => sum + (d.memory_total_gb ?? 0), 0),
}
: { ...DEFAULT_GPU, ...base };
cachedGpu = info;
return info;
} catch {
@ -78,4 +91,4 @@ export function useGpuInfo(): GpuInfo {
}, []);
return gpu;
}
}

6
studio/frontend/src/hooks/use-gpu-utilization.ts
View file

@ -7,6 +7,9 @@ import { useEffect, useRef, useState } from "react";
export interface GpuUtilization {
available: boolean;
backend: string | null;
devices?: GpuUtilization[];
index?: number;
visible_ordinal?: number;
gpu_utilization_pct: number | null;
temperature_c: number | null;
vram_used_gb: number | null;
@ -57,11 +60,10 @@ export function useGpuUtilization(
const json = (await res.json()) as GpuUtilization;
if (!cancelled) setData(json);
} catch {
// Silently ignore — next poll will retry
// Retry on the next poll.
}
}
// Fetch immediately, then set up interval
void poll();
timerRef.current = setInterval(() => void poll(), intervalMs);

15
studio/frontend/src/hooks/use-hardware-info.ts
View file

@ -25,6 +25,13 @@ export interface HardwareInfo {
transformers: string | null;
unsloth: string | null;
llamaCpp: string | null;
// Whether export can run here (true only on a supported accelerator), with a torch-aware
// reason. `null` until the authoritative response lands, so callers don't briefly enable
// export; `loaded` flips true once a real (non-error) response arrives.
exportSupported: boolean | null;
exportUnsupportedReason: string | null;
exportUnsupportedMessage: string | null;
loaded: boolean;
}
const DEFAULT: HardwareInfo = {
@ -38,6 +45,10 @@ const DEFAULT: HardwareInfo = {
transformers: null,
unsloth: null,
llamaCpp: null,
exportSupported: null,
exportUnsupportedReason: null,
exportUnsupportedMessage: null,
loaded: false,
};
// Module-level cache so multiple components share one fetch.
@ -87,6 +98,10 @@ async function fetchOnce(): Promise<HardwareInfo> {
transformers: data?.versions?.transformers ?? null,
unsloth: data?.versions?.unsloth ?? null,
llamaCpp: data?.llama_cpp ?? null,
exportSupported: data?.export_supported ?? null,
exportUnsupportedReason: data?.export_unsupported_reason ?? null,
exportUnsupportedMessage: data?.export_unsupported_message ?? null,
loaded: true,
};
if (generation === cacheGeneration) {
cached = info;

130
studio/frontend/src/hooks/use-system.ts Normal file
View file

@ -0,0 +1,130 @@
// 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 { useEffect, useState } from "react";
export interface GpuDevice {
index?: number;
index_kind?: string;
visible_ordinal?: number;
name?: string;
memory_total_gb?: number;
vram_used_gb?: number;
vram_free_gb?: number;
vram_utilization_pct?: number | null;
}
export interface SystemInfoResponse {
platform: string;
python_version: string;
device_backend: "cuda" | "rocm" | "cpu" | "mlx" | "xpu";
uptime_seconds: number | null;
cpu: {
logical_count: number;
physical_count: number;
usage_percent: number;
frequency_mhz: number | null;
};
memory: {
total_gb: number;
available_gb: number;
percent_used: number;
process_used_mb: number;
};
disk: {
total_gb: number;
free_gb: number;
percent_used: number;
};
gpu: {
available: boolean;
backend?: string;
backend_cuda_visible_devices?: string | null;
parent_visible_gpu_ids?: number[];
index_kind?: string;
devices: GpuDevice[];
};
ml_packages: {
torch?: string;
transformers?: string;
};
}
let cachedSystem: SystemInfoResponse | null = null;
let systemFetchPromise: Promise<SystemInfoResponse> | null = null;
const DEFAULT_SYSTEM: SystemInfoResponse = {
platform: "Unknown",
python_version: "Unknown",
device_backend: "cpu",
uptime_seconds: 0,
cpu: { logical_count: 0, physical_count: 0, usage_percent: 0, frequency_mhz: null },
memory: { total_gb: 0, available_gb: 0, percent_used: 0, process_used_mb: 0 },
disk: { total_gb: 0, free_gb: 0, percent_used: 0 },
gpu: { available: false, devices: [] },
ml_packages: {}
};
async function fetchSystemOnce({
force = false,
}: { force?: boolean } = {}): Promise<SystemInfoResponse> {
if (systemFetchPromise) return systemFetchPromise;
if (!force && cachedSystem) return cachedSystem;
systemFetchPromise = (async () => {
try {
const res = await authFetch("/api/system");
if (!res.ok) throw new Error(`HTTP ${res.status}`);
const data = await res.json();
cachedSystem = data as SystemInfoResponse;
return cachedSystem;
} catch {
cachedSystem = null;
return DEFAULT_SYSTEM;
} finally {
systemFetchPromise = null;
}
})();
return systemFetchPromise;
}
interface UseSystemInfoOptions {
pollMs?: number;
enabled?: boolean;
}
export function useSystemInfo({
pollMs,
enabled = true,
}: UseSystemInfoOptions = {}): SystemInfoResponse {
const [systemInfo, setSystemInfo] = useState<SystemInfoResponse>(cachedSystem ?? DEFAULT_SYSTEM);
useEffect(() => {
if (!enabled) return;
let cancelled = false;
let timeoutId: number | null = null;
const update = (force: boolean) => {
void fetchSystemOnce({ force })
.then((info) => {
if (!cancelled) setSystemInfo(info);
})
.finally(() => {
if (cancelled || !pollMs) return;
timeoutId = window.setTimeout(() => update(true), pollMs);
});
};
update(Boolean(pollMs));
return () => {
cancelled = true;
if (timeoutId !== null) window.clearTimeout(timeoutId);
};
}, [enabled, pollMs]);
return systemInfo;
}

3
studio/frontend/src/i18n/AGENTS.md
View file

@ -2,10 +2,11 @@
- `locales/en.ts` is the complete baseline message file.
- Non-English locale files may be partial. Missing keys must fall back to English at runtime.
- Use BCP 47 locale tags for new languages, for example `zh-CN`, `ja-JP`, and `ko-KR`.
- Use BCP 47 locale tags for new languages, for example `zh-CN`, `pt-BR`, `ja-JP`, and `ko-KR`.
- Do not change fallback logic to hide missing translations.
- Do not add automatic DOM translation, MutationObserver text replacement, or runtime guess-based translation.
- Preserve interpolation variables exactly, for example `{count}`, `{model}`, and `{provider}`.
- Keep product and technical names unchanged unless there is an established localized name, for example `Unsloth Studio`, `LoRA`, `GGUF`, and `Hugging Face`.
- Keep translation changes small and reviewable. Prefer separate commits for runtime changes, UI migration, and locale text.
- When adding user-facing Studio UI text, add the English message key first and add non-English overrides only when the translation is clear.
- Run `npx tsx src/i18n/check-parity.ts` before committing to ensure there are no shape mismatches or placeholder discrepancies in the non-English overlays.

6
studio/frontend/src/i18n/check-parity.ts
View file

@ -3,13 +3,14 @@
// Parity check between en.ts and every non-English locale.
// - Locale files may be partial; missing keys must fall back to English.
// - All zh-CN keys must exist in en (no extras).
// - All non-English keys must exist in en (no extras).
// - Placeholder set must match per leaf between en and the overlay.
//
// Run: npx tsx src/i18n/check-parity.ts
import { en } from "./locales/en.ts";
import { zhCN } from "./locales/zh-CN.ts";
import { ptBR } from "./locales/pt-br.ts";
import { ja } from "./locales/ja.ts";
type Tree = { readonly [k: string]: string | Tree };
@ -90,6 +91,7 @@ function checkExtras(
const overlays: Record<string, Tree> = {
"zh-CN": zhCN as unknown as Tree,
"pt-BR": ptBR as unknown as Tree,
"ja": ja as unknown as Tree,
};
let anyError = false;
@ -112,4 +114,4 @@ for (const [locale, overlay] of Object.entries(overlays)) {
}
if (anyError) process.exit(1);
console.log("\nAll locale overlays pass parity.");
console.log("\nAll locale overlays pass parity.");

73
studio/frontend/src/i18n/locales/en.ts
View file

@ -93,6 +93,7 @@ export const en = {
general: "General",
profile: "Profile",
appearance: "Appearance",
resources: "System",
chat: "Chat",
connections: "Connections",
apiKeys: "API",
@ -194,6 +195,20 @@ export const en = {
maxUploadSize: "Training dataset upload cap",
maxUploadSizeDescription: "Default is {defaultSize} MB.",
},
rag: {
sectionTitle: "Documents & RAG",
embeddingModel: "Embedding model",
embeddingModelDescription:
"Hugging Face model or local path used to index and search your documents. Default is {defaultModel}.",
reindexWarning:
"Only affects newly indexed documents. Re-upload existing ones after changing the model.",
emptyError: "Enter a Hugging Face model id or local path.",
loadError: "Failed to load the embedding model setting.",
saveError: "Failed to save the embedding model.",
saved: "Embedding model saved.",
saveAnyway: "Save anyway",
resetAction: "Reset to default",
},
storage: {
sectionTitle: "Storage",
modelsFolder: "Models folder",
@ -262,6 +277,58 @@ export const en = {
"Keep the sidebar expanded instead of collapsing to icons.",
},
},
resources: {
title: "System",
description: "Monitor this Studio server's hardware and storage.",
liveUpdates: "Live updates",
floatingWindow: "Floating window",
disableOverlay: "Disable overlay",
liveMonitor: {
title: "Live monitor",
cpu: "CPU",
ram: "RAM",
disk: "Disk",
vram: "VRAM",
cpuCores: "{logical} logical / {physical} physical cores",
currentLoad: "Current load",
free: "{value} free",
noGpu: "No visible GPU",
},
gpu: {
title: "GPU devices",
noGpu: "No visible GPU detected. CPU-only resources are shown above.",
unknownDevice: "Unknown GPU",
deviceWithIndex: "GPU {index}",
vramUtilization: "VRAM",
used: "{value} used",
free: "{value} free",
total: "{value} total",
},
storage: {
title: "Storage",
systemDisk: "System disk",
diskUsage: "{used} used / {total}",
diskFree: "{free} free",
modelsFolder: "Models folder",
modelsFolderDescription: "Where downloaded models are stored.",
openAction: "Open",
copyAction: "Copy path",
copied: "Path copied",
openError: "Couldn't open the folder",
copyError: "Couldn't copy the path",
},
environment: {
title: "Environment",
backend: "Backend",
python: "Python",
torch: "Torch",
transformers: "Transformers",
uptime: "Uptime",
processMemory: "Process memory",
notInstalled: "Not installed",
unknown: "Unknown",
},
},
chat: {
title: "Chat",
description: "Manage chat history stored on this device.",
@ -360,8 +427,10 @@ export const en = {
usageTools: "Tools",
exampleCurlTools: "curl + tools",
examplePythonTools: "Python + tools",
exampleJavaScriptTools: "JavaScript + tools",
exampleCurlAdvanced: "curl + advanced",
examplePythonAdvanced: "Python + advanced",
exampleJavaScriptAdvanced: "JavaScript + advanced",
osUnix: "Linux / macOS / WSL",
osWindows: "Windows",
secureHttps: "Secure HTTPS",
@ -372,6 +441,10 @@ export const en = {
copy: "Copy",
copied: "Copied",
setupDocs: "Setup docs:",
codingAgents: "Coding agents",
codingAgentsHint:
"Launch a coding agent against this server. It uses the loaded model; a local server mints an API key automatically, a remote one includes it in the command.",
codingAgentsSwap: "Swap claude for codex, openclaw, opencode, hermes, or pi.",
relativeNever: "never",
relativeJustNow: "just now",
relativeHoursAgo: "{count}h ago",

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