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

Merge main and resolve installer conflicts

Browse source
This commit is contained in:
Unsloth 2026-07-18 22:53:06 -07:00
commit b9f50ed00c
406 changed files with 65192 additions and 6001 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -527,6 +527,154 @@ case "$MODE" in
echo "[claude] attribution A/B OK (suppressed HIT, header=1 MISS)"
;;
# ── resume: does a launched agent's session survive exit and resume? ────
# Unlike the other modes, this drives the real LAUNCH path (`unsloth start
# <agent> ...`, the interactive default), not the --no-launch recipe. That
# path relocates each agent's home to a throwaway temp dir wiped on exit, so
# a session cannot be resumed -- unless --persist routes it to the stable
# Unsloth agents dir instead. We run one headless turn per pass and check
# whether the turn left a session in a persistent store (deterministic, no
# reliance on the model recalling anything), for a baseline pass and a
# --persist pass, and assert the expected split for this agent.
resume)
CODEWORD="PLATYPUS7"
T1="Remember this codeword for later: ${CODEWORD}. Reply with just the word OK."
T2="What codeword did I ask you to remember? Reply with just that word."
WORK="$WORKDIR_BASE/${AGENT}-resume"
# STABLE_HOME: the stable dir that --no-launch (and --persist) relocate to.
# Read it from a --no-launch probe (which also writes the agent's config
# there). codex/pi relocate their whole home/HOME here; opencode/claude keep
# their session data in a fixed user dir, so STABLE_HOME stays empty for them.
parse_connect
case "$AGENT" in
codex) STABLE_HOME="$(raw_env CODEX_HOME)" ;;
pi) STABLE_HOME="$(raw_env HOME)" ;;
*) STABLE_HOME="" ;;
esac
# The persistent stores a session would land in if it were NOT wiped. We
# count files here before/after each turn; a positive delta means the
# session persisted (is resumable), zero means it went to a wiped temp dir.
resume_tracked_dirs() {
case "$AGENT" in
codex) printf '%s\n' "$HOME/.codex" ;;
opencode) printf '%s\n' "$HOME/.local/share/opencode" "$HOME/.config/opencode" ;;
claude) printf '%s\n' "$HOME/.claude" ;;
pi) printf '%s\n' "$HOME/.pi" ;;
*) : ;;
esac
[ -n "$STABLE_HOME" ] && printf '%s\n' "$STABLE_HOME"
}
count_session_files() {
local total=0 d n
while IFS= read -r d; do
[ -n "$d" ] && [ -d "$d" ] || continue
n="$(find "$d" -type f 2>/dev/null | wc -l)"; total=$((total + n))
done < <(resume_tracked_dirs)
echo "$total"
}
# The headless first-turn subcommand per agent (mirrors file-edit's map),
# forwarded verbatim through the launch path as passthrough args.
set_t1_cmd() {
case "$AGENT" in
claude) T1_CMD=("${CLAUDE_CONNECT_FLAGS[@]}" -p "$T1") ;;
codex) T1_CMD=(exec "$T1") ;;
opencode) T1_CMD=(run "$T1") ;;
pi) T1_CMD=(-p "$T1") ;;
*) guide_fail "resume mode does not cover agent '$AGENT'" ;;
esac
}
# Run one headless turn through the launch path. $1=outfile, $2="" or
# "--persist", rest = the agent subcommand. --yolo auto-approves so no tool
# prompt can hang; --api-key attaches to the already-served CI model.
launch_turn() {
local out="$1" rflag="$2"; shift 2
local flag=(); [ -n "$rflag" ] && flag=("$rflag")
run_timed "$out" unsloth start "$AGENT" "${flag[@]}" --yolo \
--api-key "$UNSLOTH_API_KEY" "$@"
local rc=$?
redact "$out"
return "$rc"
}
# One pass: fresh work dir, one planting turn, set RESULT to PERSISTED/WIPED
# from the session-store delta. Runs in the main shell (not a command
# substitution) so a hang's guide_fail actually fails the job and the
# progress lines reach the CI log. $1 = "" (baseline) or "--persist".
RESULT=""
run_pass() {
local rflag="$1" label="baseline"
[ -n "$rflag" ] && label="resume"
rm -rf "$WORK"; mkdir -p "$WORK"
set_t1_cmd
local out="$LOGS_DIR/${AGENT}-resume-${label}.txt"
local before after rc
before="$(count_session_files)"
pushd "$WORK" >/dev/null || guide_fail "could not enter work dir $WORK"
launch_turn "$out" "$rflag" "${T1_CMD[@]}"; rc=$?
popd >/dev/null || true
after="$(count_session_files)"
echo "[$AGENT] ${label}: session files ${before} -> ${after} (rc=${rc})"
# The turn must succeed for the delta to mean anything: an agent that writes a
# session file then errors would otherwise be misread as PERSISTED. Mirror the
# file-edit mode and fail the pass on a non-zero launch (the flagship codex recall
# below stays WARN-only, driven by its own launch_turn calls).
[ "$rc" -eq 0 ] || { echo "[$AGENT] ${label} transcript (tail):"; tail -30 "$out" 2>/dev/null || true; \
guide_fail "resume ${label} turn for ${AGENT} exited non-zero (rc=${rc})"; }
if [ "$after" -gt "$before" ]; then RESULT="PERSISTED"; else RESULT="WIPED"; fi
}
run_pass ""; BASELINE="$RESULT"
# Only the temp-dir agents (codex/pi) need the --persist pass to prove the fix.
# opencode/claude persist either way, so the baseline already proves it and a
# second full CPU turn only risks a timeout; skip it for them.
case "$AGENT" in
codex|pi) run_pass "--persist"; RESUME="$RESULT" ;;
*) RESUME="n/a (persists either way)" ;;
esac
# Expected: codex/pi relocate their whole home to the temp dir, so a plain
# launch is WIPED and only --persist PERSISTS. opencode/claude keep their
# session data in a fixed user dir, so the baseline already PERSISTS.
case "$AGENT" in
codex|pi) EXPECT_BASELINE="WIPED" ;;
opencode|claude) EXPECT_BASELINE="PERSISTED" ;;
esac
echo "──────────────────────────────────────────────"
echo "[$AGENT] RESUME EXPERIMENT"
echo " baseline (unsloth start ${AGENT}): ${BASELINE} (expected ${EXPECT_BASELINE})"
echo " with --persist (unsloth start ${AGENT} --persist): ${RESUME}"
echo "──────────────────────────────────────────────"
[ "$BASELINE" = "$EXPECT_BASELINE" ] \
|| guide_fail "baseline resume behavior for ${AGENT} was ${BASELINE}, expected ${EXPECT_BASELINE}"
case "$AGENT" in
codex|pi)
[ "$RESUME" = "PERSISTED" ] \
|| guide_fail "--persist did not persist ${AGENT}'s session (got ${RESUME}); the session dir is still not stable" ;;
esac
# Flagship behavioral proof (codex only, WARN-only): after a --persist plant,
# resume the session and check the model actually recalls the codeword. A
# miss is not a failure (the CI model is small); the mechanism gate above is
# the real assertion.
if [ "$AGENT" = "codex" ]; then
rm -rf "$WORK"; mkdir -p "$WORK"
( cd "$WORK" && launch_turn "$LOGS_DIR/codex-resume-plant.txt" "--persist" exec "$T1" ) || true
( cd "$WORK" && launch_turn "$LOGS_DIR/codex-resume-recall.txt" "--persist" exec resume --last "$T2" ) || true
if grep -q "$CODEWORD" "$LOGS_DIR/codex-resume-recall.txt" 2>/dev/null; then
echo "[codex] behavioral recall HIT: resumed session remembered ${CODEWORD}"
else
echo "::warning::[codex] behavioral recall MISS (small CI model); mechanism gate still passed"
fi
fi
echo "[$AGENT] resume OK"
;;
*)
echo "agent-guides-drive.sh: unknown mode '$MODE'" >&2
exit 2

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

@ -268,10 +268,12 @@ jobs:
tests/saving/test_save_shell_injection.py \
tests/saving/test_patch_saving_none_tokenizer.py \
tests/saving/test_fix_sentencepiece_gguf_robustness.py \
tests/saving/test_fix_sentencepiece_tokenizer_guard.py \
tests/saving/test_compressed_export_schemes.py \
tests/saving/test_export_api_surface.py \
tests/saving/test_export_dispatch.py \
tests/saving/test_imatrix_export.py \
tests/saving/test_gguf_single_pass_export.py \
tests/utils/test_attention_masks.py \
tests/utils/test_trunc_normal_patch.py \
tests/python/test_fast_language_model_text_only.py
@ -357,10 +359,12 @@ jobs:
tests/saving/test_save_shell_injection.py \
tests/saving/test_patch_saving_none_tokenizer.py \
tests/saving/test_fix_sentencepiece_gguf_robustness.py \
tests/saving/test_fix_sentencepiece_tokenizer_guard.py \
tests/saving/test_compressed_export_schemes.py \
tests/saving/test_export_api_surface.py \
tests/saving/test_export_dispatch.py \
tests/saving/test_imatrix_export.py \
tests/saving/test_gguf_single_pass_export.py \
tests/utils/test_attention_masks.py \
tests/utils/test_trunc_normal_patch.py \
tests/python/test_fast_language_model_text_only.py \

14
.github/workflows/cross-platform-parity-ci.yml vendored
View file

@ -1,7 +1,7 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
# Runs tests/python/test_cross_platform_parity.py on Windows and macOS.
# Runs installer parity and autostart opt-out tests on Windows and macOS.
#
# Why: that test is the guard that install.sh and install.ps1 stay in
# sync, but today it only runs on ubuntu-latest (auto-discovered by
@ -21,6 +21,7 @@ on:
paths:
- 'install.sh'
- 'install.ps1'
- 'tests/test_installer_skip_autostart.py'
- 'tests/python/test_cross_platform_parity.py'
- '.github/workflows/cross-platform-parity-ci.yml'
push:
@ -28,6 +29,7 @@ on:
paths:
- 'install.sh'
- 'install.ps1'
- 'tests/test_installer_skip_autostart.py'
- 'tests/python/test_cross_platform_parity.py'
- '.github/workflows/cross-platform-parity-ci.yml'
workflow_dispatch:
@ -57,5 +59,11 @@ jobs:
python-version: '3.12'
cache: 'pip'
- run: python -m pip install -U pip pytest
- name: Cross-platform parity test
run: python -m pytest tests/python/test_cross_platform_parity.py -q
- name: Cross-platform parity tests
env:
UNSLOTH_NO_TORCH: '1'
run: >-
python -m pytest
tests/python/test_cross_platform_parity.py
tests/test_installer_skip_autostart.py
-q

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

@ -471,6 +471,176 @@ jobs:
redacted-configs/
retention-days: 7
# ═════════════════════════════════════════════════════════════════════
# Job: resume
# Does a conversation started with `unsloth start <agent>` survive exit
# and resume? This drives the REAL launch path (not the --no-launch
# recipe the other jobs use). A plain launch relocates the agent home to
# a temp dir wiped on exit, so codex/pi cannot resume; --persist routes the
# session to the stable Unsloth agents dir so it persists. opencode/claude
# keep their session data in a fixed user dir, so they persist either way.
# Dispatch-only: it is an end-to-end experiment, not a PR gate.
# ═════════════════════════════════════════════════════════════════════
resume:
name: resume (${{ matrix.agent }})
if: github.event_name == 'workflow_dispatch'
runs-on: ubuntu-latest
timeout-minutes: 60
strategy:
fail-fast: false
matrix:
# codex/pi relocate their whole home (resume broken without --persist);
# opencode/claude keep session data in a fixed dir (resume already works).
# One agent from each class proves the split end to end; openclaw/hermes
# share codex's relocation mechanism and are covered by the unit tests.
agent: [codex, opencode, claude, pi]
env:
GGUF_REPO: unsloth/gemma-4-E4B-it-GGUF
GGUF_FILE: gemma-4-E4B-it-UD-Q4_K_XL.gguf
STUDIO_PORT: '18904'
steps:
- uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
with:
persist-credentials: false
- name: Linux deps for llama.cpp prebuilt
run: |
sudo apt-get update
sudo apt-get install -y --no-install-recommends \
libcurl4-openssl-dev libssl-dev jq
- uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0
with:
node-version: '22'
- uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0
with:
python-version: '3.12'
cache: 'pip'
- name: Restore GGUF model file
id: cache-gguf
uses: actions/cache/restore@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5
continue-on-error: true
with:
path: gguf-cache
key: ${{ runner.os }}-gguf-${{ env.GGUF_REPO }}-${{ env.GGUF_FILE }}-v1
- name: Download GGUF if cache miss
id: download-gguf
if: steps.cache-gguf.outputs.cache-hit != 'true' || steps.cache-gguf.outcome != 'success'
env:
HF_TOKEN: ${{ secrets.HF_TOKEN }}
run: |
python -m pip install --upgrade huggingface_hub
mkdir -p gguf-cache
bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE" gguf-cache
- name: Save GGUF model file
if: always() && steps.download-gguf.outcome == 'success'
uses: actions/cache/save@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5
with:
path: gguf-cache
key: ${{ runner.os }}-gguf-${{ env.GGUF_REPO }}-${{ env.GGUF_FILE }}-v1
- name: Install Studio (--local, --no-torch)
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
HF_TOKEN: ${{ secrets.HF_TOKEN }}
run: |
mkdir -p logs
set -o pipefail
bash install.sh --local --no-torch 2>&1 | tee logs/install.log
- name: Serve unsloth run --disable-tools (gemma-4-E4B)
run: |
unsloth studio reset-password
bash .github/scripts/serve-unsloth-run.sh \
--gguf-file "$GITHUB_WORKSPACE/gguf-cache/${GGUF_FILE}" \
--port "$STUDIO_PORT" --log-dir logs \
--extra "--seed $UNSLOTH_SEED --temp 0" \
--health-timeout 900
- name: Preflight the agent's API dialect (class-a isolation)
env:
AGENT: ${{ matrix.agent }}
run: |
set -uo pipefail
B="$UNSLOTH_BASE_URL"; K="$UNSLOTH_API_KEY"
preflight_fail() {
echo "::error::[server/API regression] agent=$AGENT: $* (preflight failed BEFORE install/connect). Endpoint contract lives in studio/backend/routes/**.";
exit 1
}
code=$(curl -s -o /tmp/pf.json -w '%{http_code}' "$B/v1/models" \
-H "Authorization: Bearer $K") || true
[ "$code" = "200" ] || preflight_fail "/v1/models returned HTTP $code"
case "$AGENT" in
claude)
code=$(curl -s -o /tmp/pf.json -w '%{http_code}' "$B/v1/messages" \
-H "Authorization: Bearer $K" -H 'content-type: application/json' \
--max-time 120 \
-d "{\"model\":\"$UNSLOTH_MODEL_ID\",\"max_tokens\":16,\"messages\":[{\"role\":\"user\",\"content\":\"Hi\"}]}") || true
[ "$code" = "200" ] || preflight_fail "/v1/messages returned HTTP $code"
;;
codex)
code=$(curl -s -o /tmp/pf.json -w '%{http_code}' "$B/v1/responses" \
-H "Authorization: Bearer $K" -H 'content-type: application/json' \
--max-time 120 \
-d "{\"model\":\"$UNSLOTH_MODEL_ID\",\"input\":\"Hi\",\"max_output_tokens\":16,\"stream\":true}") || true
[ "$code" = "200" ] || preflight_fail "/v1/responses returned HTTP $code"
;;
*)
code=$(curl -s -o /tmp/pf.json -w '%{http_code}' "$B/v1/chat/completions" \
-H "Authorization: Bearer $K" -H 'content-type: application/json' \
--max-time 120 \
-d "{\"model\":\"$UNSLOTH_MODEL_ID\",\"max_tokens\":16,\"messages\":[{\"role\":\"user\",\"content\":\"Hi\"}]}") || true
[ "$code" = "200" ] || preflight_fail "/v1/chat/completions returned HTTP $code"
;;
esac
echo "preflight OK for $AGENT"
- name: Install agent CLI (class-b isolation)
env:
AGENT: ${{ matrix.agent }}
run: bash .github/scripts/agent-guides-install.sh "$AGENT"
- name: Resume experiment (launch path)
env:
AGENT: ${{ matrix.agent }}
run: bash .github/scripts/agent-guides-drive.sh resume "$AGENT"
- name: Collect server logs (debug)
if: always()
run: |
mkdir -p logs/studio-logs
cp -r "$HOME/.unsloth/studio/logs/." logs/studio-logs/ 2>/dev/null || true
if [ -n "${UNSLOTH_API_KEY:-}" ]; then
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
- name: Stop Studio
if: always()
run: |
if [ -n "${UNSLOTH_SERVER_PID:-}" ] && [ "${UNSLOTH_SERVER_PID}" != "0" ]; then
kill "${UNSLOTH_SERVER_PID}" 2>/dev/null || true
fi
sleep 2
ss -tln 2>/dev/null | grep ":${STUDIO_PORT}" || true
- name: Upload logs
if: always()
continue-on-error: true
uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
with:
name: resume-${{ matrix.agent }}-log
path: |
logs/
agent-workdir/
redacted-configs/
retention-days: 7
# ═════════════════════════════════════════════════════════════════════
# Job 3: prompt-cache
# (a) curl 2-turn /v1/chat/completions: assert turn-2 cached_tokens > 0

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

@ -19,6 +19,19 @@ on:
permissions:
contents: read
env:
DESKTOP_RELEASE_NOTES: |
Desktop app for Unsloth Studio.
**macOS**: Download the Apple Silicon `.dmg`.
**Windows**: Download the `-setup.exe` installer.
**Linux**: Download `.deb` for Ubuntu/Debian. `.AppImage` is experimental.
> Linux in-app updates are AppImage-oriented. Package installs should update by downloading a new package.
> Linux AppImage can show a blank window on some Tauri/WebKitGTK + Wayland/Mesa stacks; use `.deb` when available.
> Linux AppImage on Ubuntu 24.04+ may require: `sudo apt install libfuse2t64`
> First-run system dependency elevation is supported on Ubuntu/Debian. Other Linux distributions should install system packages manually.
concurrency:
group: release-desktop-${{ github.repository }}
cancel-in-progress: false
@ -295,14 +308,6 @@ jobs:
PY
build:
# TODO: split into a "build (no secrets)" + "publish (secrets)" job pair
# with actions/upload-artifact handoff so the matrix build cannot
# publish a Release on its own. The current matrix runs across
# Linux/macOS/Windows in a single job, so the split needs artefact
# collection across the OS matrix and is out of scope for this
# hardening pass.
permissions:
contents: write # tauri-apps/tauri-action creates / uploads a GitHub Release
strategy:
fail-fast: false
max-parallel: 1
@ -311,15 +316,21 @@ jobs:
- platform: macos-latest
args: '--target aarch64-apple-darwin'
label: macOS (Apple Silicon)
artifact: macos-aarch64
release_arch: aarch64
# - platform: macos-latest
# args: '--target x86_64-apple-darwin'
# label: macOS (Intel)
- platform: ubuntu-22.04
args: ''
label: Linux (x64)
artifact: linux-x64
release_arch: x64
- platform: windows-latest
args: ''
label: Windows (x64)
artifact: windows-x64
release_arch: x64
name: Build ${{ matrix.label }}
needs: prepare-version
@ -465,41 +476,18 @@ jobs:
if (chmodIdx !== -1 && sha256Idx > chmodIdx) {
throw new Error('Desktop Linux release must verify the linuxdeploy digest before chmod +x');
}
const releaseBodies = [];
for (let i = 0; i < lines.length; i += 1) {
const match = lines[i].match(/^(\s*)releaseBody:\s*\|\s*$/);
if (!match) continue;
const baseIndent = match[1].length;
const bodyLines = [];
i += 1;
for (; i < lines.length; i += 1) {
const line = lines[i];
if (line.trim() === '') {
bodyLines.push('');
continue;
}
const indent = line.match(/^\s*/)[0].length;
if (indent <= baseIndent) {
i -= 1;
break;
}
bodyLines.push(line.slice(baseIndent + 2));
}
releaseBodies.push(bodyLines.join('\n'));
const releaseBody = process.env.DESKTOP_RELEASE_NOTES;
if (!releaseBody) {
throw new Error('DESKTOP_RELEASE_NOTES must not be empty');
}
if (releaseBodies.length === 0) {
throw new Error('Expected at least one desktop release body');
if (/\brpm\b|\.rpm/i.test(releaseBody)) {
throw new Error('Desktop release body must not advertise RPM packages');
}
for (const body of releaseBodies) {
if (/\brpm\b|\.rpm/i.test(body)) {
throw new Error('Desktop release body must not advertise RPM packages');
}
if (/AppImage.*universal|universal.*AppImage/i.test(body)) {
throw new Error('Desktop release body must not advertise AppImage as universal');
}
if (!/AppImage.*experimental/i.test(body)) {
throw new Error('Desktop release body must mark AppImage as experimental');
}
if (/AppImage.*universal|universal.*AppImage/i.test(releaseBody)) {
throw new Error('Desktop release body must not advertise AppImage as universal');
}
if (!/AppImage.*experimental/i.test(releaseBody)) {
throw new Error('Desktop release body must mark AppImage as experimental');
}
JS
@ -644,48 +632,33 @@ jobs:
dest="$tools_dir/linuxdeploy-x86_64.AppImage"
curl -fsSL "$LINUXDEPLOY_URL" -o "$dest"
# Verify the digest BEFORE the binary is ever marked executable. The
# next step builds the AppImage with the Tauri signing key and a
# contents:write GITHUB_TOKEN in scope, so a substituted linuxdeploy
# that ran here could exfiltrate signing material or tamper with
# published release artifacts. Fail closed on any mismatch.
# next step builds the AppImage with the Tauri signing key, so a
# substituted linuxdeploy that ran here could exfiltrate signing
# material or tamper with release artifacts. Fail closed on any
# mismatch.
echo "${LINUXDEPLOY_SHA256} ${dest}" | sha256sum -c -
chmod +x "$dest"
# ── Linux: build + sign + upload ──
# ── Linux: build + sign ──
- name: Build Linux app
id: build_linux
if: matrix.platform == 'ubuntu-22.04'
uses: tauri-apps/tauri-action@84b9d35b5fc46c1e45415bdb6144030364f7ebc5
env:
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
TAURI_SIGNING_PRIVATE_KEY: ${{ secrets.TAURI_SIGNING_PRIVATE_KEY }}
TAURI_SIGNING_PRIVATE_KEY_PASSWORD: ${{ secrets.TAURI_SIGNING_PRIVATE_KEY_PASSWORD }}
XDG_CACHE_HOME: ${{ runner.temp }}/tauri-tools-cache
with:
projectPath: studio
tauriScript: npx --prefix . tauri
tagName: ${{ needs.prepare-version.outputs.desktop_release_tag }}
releaseName: 'Unsloth Studio (Desktop) ${{ needs.prepare-version.outputs.studio_version }}'
releaseBody: |
Desktop app for Unsloth Studio.
**macOS**: Download the Apple Silicon `.dmg`.
**Windows**: Download the `-setup.exe` installer.
**Linux**: Download `.deb` for Ubuntu/Debian. `.AppImage` is experimental.
> Linux in-app updates are AppImage-oriented. Package installs should update by downloading a new package.
> Linux AppImage can show a blank window on some Tauri/WebKitGTK + Wayland/Mesa stacks; use `.deb` when available.
> Linux AppImage on Ubuntu 24.04+ may require: `sudo apt install libfuse2t64`
> First-run system dependency elevation is supported on Ubuntu/Debian. Other Linux distributions should install system packages manually.
releaseDraft: ${{ inputs.draft }}
prerelease: ${{ needs.prepare-version.outputs.prerelease }}
args: -v ${{ matrix.args }}
# ── macOS: build + sign + notarize + upload ──
# ── macOS: build + sign + notarize ──
- name: Build macOS app
id: build_macos
if: matrix.platform == 'macos-latest'
uses: tauri-apps/tauri-action@84b9d35b5fc46c1e45415bdb6144030364f7ebc5
env:
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
TAURI_SIGNING_PRIVATE_KEY: ${{ secrets.TAURI_SIGNING_PRIVATE_KEY }}
TAURI_SIGNING_PRIVATE_KEY_PASSWORD: ${{ secrets.TAURI_SIGNING_PRIVATE_KEY_PASSWORD }}
APPLE_SIGNING_IDENTITY: ${{ secrets.APPLE_SIGNING_IDENTITY }}
@ -695,29 +668,14 @@ jobs:
with:
projectPath: studio
tauriScript: npx --prefix . tauri
tagName: ${{ needs.prepare-version.outputs.desktop_release_tag }}
releaseName: 'Unsloth Studio (Desktop) ${{ needs.prepare-version.outputs.studio_version }}'
releaseBody: |
Desktop app for Unsloth Studio.
**macOS**: Download the Apple Silicon `.dmg`.
**Windows**: Download the `-setup.exe` installer.
**Linux**: Download `.deb` for Ubuntu/Debian. `.AppImage` is experimental.
> Linux in-app updates are AppImage-oriented. Package installs should update by downloading a new package.
> Linux AppImage can show a blank window on some Tauri/WebKitGTK + Wayland/Mesa stacks; use `.deb` when available.
> Linux AppImage on Ubuntu 24.04+ may require: `sudo apt install libfuse2t64`
> First-run system dependency elevation is supported on Ubuntu/Debian. Other Linux distributions should install system packages manually.
releaseDraft: ${{ inputs.draft }}
prerelease: ${{ needs.prepare-version.outputs.prerelease }}
args: -v ${{ matrix.args }}
# ── Windows: build + sign + upload ──
# ── Windows: build + sign ──
- name: Build Windows app
id: build_windows
if: matrix.platform == 'windows-latest'
uses: tauri-apps/tauri-action@84b9d35b5fc46c1e45415bdb6144030364f7ebc5
env:
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
TAURI_SIGNING_PRIVATE_KEY: ${{ secrets.TAURI_SIGNING_PRIVATE_KEY }}
TAURI_SIGNING_PRIVATE_KEY_PASSWORD: ${{ secrets.TAURI_SIGNING_PRIVATE_KEY_PASSWORD }}
AZURE_CLIENT_ID: ${{ secrets.AZURE_CLIENT_ID }}
@ -728,35 +686,83 @@ jobs:
with:
projectPath: studio
tauriScript: npx --prefix . tauri
tagName: ${{ needs.prepare-version.outputs.desktop_release_tag }}
releaseName: 'Unsloth Studio (Desktop) ${{ needs.prepare-version.outputs.studio_version }}'
releaseBody: |
Desktop app for Unsloth Studio.
**macOS**: Download the Apple Silicon `.dmg`.
**Windows**: Download the `-setup.exe` installer.
**Linux**: Download `.deb` for Ubuntu/Debian. `.AppImage` is experimental.
> Linux in-app updates are AppImage-oriented. Package installs should update by downloading a new package.
> Linux AppImage can show a blank window on some Tauri/WebKitGTK + Wayland/Mesa stacks; use `.deb` when available.
> Linux AppImage on Ubuntu 24.04+ may require: `sudo apt install libfuse2t64`
> First-run system dependency elevation is supported on Ubuntu/Debian. Other Linux distributions should install system packages manually.
releaseDraft: ${{ inputs.draft }}
prerelease: ${{ needs.prepare-version.outputs.prerelease }}
args: -v ${{ matrix.args }}
# Release process note: only non-draft workflow runs advance the public
# desktop-latest updater channel. Draft builds are for private review; if a
# draft is manually published later, this channel intentionally remains
# unchanged until a narrow manual channel-publish flow is added or a public
# desktop release is created by running this workflow with draft=false.
publish-updater-channel:
name: Publish desktop updater channel
- name: Stage release assets
shell: bash
env:
ARTIFACT_PATHS: ${{ steps.build_linux.outputs.artifactPaths || steps.build_macos.outputs.artifactPaths || steps.build_windows.outputs.artifactPaths }}
RELEASE_ARCH: ${{ matrix.release_arch }}
run: |
set -euo pipefail
if command -v python3 >/dev/null 2>&1; then
PYTHON=python3
else
PYTHON=python
fi
"$PYTHON" <<'PY'
import json
import os
import pathlib
import re
import shutil
import sys
import unicodedata
raw_paths = os.environ.get('ARTIFACT_PATHS', '')
try:
artifact_paths = json.loads(raw_paths)
except json.JSONDecodeError as error:
sys.exit(f'Invalid tauri-action artifactPaths output: {error}')
if not isinstance(artifact_paths, list) or not artifact_paths:
sys.exit('tauri-action did not return any release artifacts')
destination = pathlib.Path(os.environ['RUNNER_TEMP'], 'desktop-release-assets')
destination.mkdir(parents=True, exist_ok=True)
staged = []
for raw_path in artifact_paths:
source = pathlib.Path(raw_path)
if not source.is_file():
continue
name = source.name
for extension in ('.app.tar.gz.sig', '.app.tar.gz'):
if name.endswith(extension):
name = f'{name[:-len(extension)]}_{os.environ["RELEASE_ARCH"]}{extension}'
break
name = unicodedata.normalize('NFD', name)
name = ''.join(character for character in name if not unicodedata.combining(character))
name = re.sub(r'[ ()\[\]{}]', '.', name)
while '..' in name:
name = name.replace('..', '.')
target = destination / name
if target.exists():
sys.exit(f'Duplicate staged release asset name: {name}')
shutil.copy2(source, target)
staged.append(name)
if not staged:
sys.exit('No release files were staged')
print('Staged release assets:')
print('\n'.join(sorted(staged)))
PY
- name: Upload signed release assets
uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
with:
name: desktop-release-${{ matrix.artifact }}
path: ${{ runner.temp }}/desktop-release-assets/*
if-no-files-found: error
compression-level: 0
retention-days: 1
# Only this job gets write access; builds hand off signed files via artifacts.
# Draft runs do not advance the public desktop-latest channel.
publish-release:
name: Publish desktop release
needs: [prepare-version, build]
if: ${{ !inputs.draft }}
runs-on: ubuntu-latest
permissions:
contents: write
contents: write # create the versioned Release and replace updater-channel metadata
env:
GH_REPO: ${{ github.repository }}
APP_VERSION: ${{ needs.prepare-version.outputs.app_version }}
@ -765,7 +771,164 @@ jobs:
DESKTOP_PRERELEASE: ${{ needs.prepare-version.outputs.prerelease }}
steps:
- name: Harden runner (audit)
uses: step-security/harden-runner@a5ad31d6a139d249332a2605b85202e8c0b78450 # v2.19.1
with:
egress-policy: audit
- name: Download signed release assets
uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8.0.1
with:
pattern: desktop-release-*
path: ${{ runner.temp }}/desktop-release-assets
merge-multiple: true
- name: Validate release asset set
shell: bash
run: |
set -euo pipefail
python3 <<'PY'
import pathlib
import os
import sys
asset_dir = pathlib.Path(os.environ['RUNNER_TEMP'], 'desktop-release-assets')
files = [path for path in asset_dir.iterdir() if path.is_file()]
required_suffixes = (
'.dmg',
'.app.tar.gz',
'.app.tar.gz.sig',
'.deb',
'.AppImage',
'.AppImage.sig',
'-setup.exe',
'-setup.exe.sig',
)
for suffix in required_suffixes:
matches = [path for path in files if path.name.endswith(suffix)]
if len(matches) != 1:
sys.exit(f'Expected exactly one {suffix} release asset, found {len(matches)}')
if any(path.name == 'latest.json' for path in files):
sys.exit('Build artifacts must not supply latest.json')
print('\n'.join(sorted(path.name for path in files)))
PY
- name: Create or validate versioned release
shell: bash
env:
GH_TOKEN: ${{ github.token }}
RELEASE_DRAFT: ${{ inputs.draft }}
run: |
set -euo pipefail
notes_file="$RUNNER_TEMP/desktop-release-notes.md"
printf '%s\n' "$DESKTOP_RELEASE_NOTES" > "$notes_file"
release_json="$RUNNER_TEMP/versioned-release.json"
# REST tag lookup omits drafts; `gh release view` also checks pending tags.
if gh release view "$DESKTOP_RELEASE_TAG" \
--json tagName,isDraft,isPrerelease > "$release_json" 2>/dev/null; then
python3 <<'PY'
import json
import os
import pathlib
import sys
release = json.loads(pathlib.Path(os.environ['RUNNER_TEMP'], 'versioned-release.json').read_text())
expected_draft = os.environ['RELEASE_DRAFT'].lower() == 'true'
expected_prerelease = os.environ['DESKTOP_PRERELEASE'].lower() == 'true'
if release.get('tagName') != os.environ['DESKTOP_RELEASE_TAG']:
sys.exit('Existing desktop release tag does not match the requested tag')
if bool(release.get('isDraft')) != expected_draft:
sys.exit('Existing desktop release draft state does not match the workflow input')
if bool(release.get('isPrerelease')) != expected_prerelease:
sys.exit('Existing desktop release prerelease state does not match the requested version')
PY
else
release_flags=(
--title "Unsloth Studio (Desktop) ${STUDIO_VERSION}"
--notes-file "$notes_file"
--target "$GITHUB_SHA"
)
if [ "$RELEASE_DRAFT" = "true" ]; then
release_flags+=(--draft)
fi
if [ "$DESKTOP_PRERELEASE" = "true" ]; then
release_flags+=(--prerelease)
fi
gh release create "$DESKTOP_RELEASE_TAG" "${release_flags[@]}"
fi
- name: Publish versioned release assets
shell: bash
env:
GH_TOKEN: ${{ github.token }}
run: |
set -euo pipefail
gh release upload "$DESKTOP_RELEASE_TAG" "$RUNNER_TEMP/desktop-release-assets"/* --clobber
- name: Generate and publish versioned updater metadata
shell: bash
env:
GH_TOKEN: ${{ github.token }}
run: |
set -euo pipefail
python3 <<'PY'
import datetime
import json
import os
import pathlib
import sys
import urllib.parse
asset_dir = pathlib.Path(os.environ['RUNNER_TEMP'], 'desktop-release-assets')
files = [path for path in asset_dir.iterdir() if path.is_file()]
def exactly_one(suffix: str) -> pathlib.Path:
matches = [path for path in files if path.name.endswith(suffix)]
if len(matches) != 1:
sys.exit(f'Expected exactly one {suffix} updater asset, found {len(matches)}')
return matches[0]
def entry(signature_suffix: str) -> dict[str, str]:
signature_path = exactly_one(signature_suffix)
bundle_name = signature_path.name.removesuffix('.sig')
bundle_path = asset_dir / bundle_name
if not bundle_path.is_file():
sys.exit(f'Missing updater bundle for {signature_path.name}: {bundle_name}')
encoded_tag = urllib.parse.quote(os.environ['DESKTOP_RELEASE_TAG'], safe='')
encoded_name = urllib.parse.quote(bundle_name, safe='')
return {
'signature': signature_path.read_text(),
'url': (
f'https://github.com/{os.environ["GITHUB_REPOSITORY"]}/releases/download/'
f'{encoded_tag}/{encoded_name}'
),
}
darwin = entry('.app.tar.gz.sig')
linux = entry('.AppImage.sig')
windows = entry('.exe.sig')
notes = pathlib.Path(os.environ['RUNNER_TEMP'], 'desktop-release-notes.md').read_text()
metadata = {
'version': os.environ['APP_VERSION'],
'notes': notes,
'pub_date': datetime.datetime.now(datetime.timezone.utc).isoformat(timespec='milliseconds').replace('+00:00', 'Z'),
'platforms': {
'darwin-aarch64': darwin,
'darwin-aarch64-app': darwin,
'linux-x86_64': linux,
'linux-x86_64-appimage': linux,
'windows-x86_64': windows,
'windows-x86_64-nsis': windows,
},
}
output = pathlib.Path(os.environ['RUNNER_TEMP'], 'latest.json')
output.write_text(json.dumps(metadata, indent=2) + '\n')
PY
gh release upload "$DESKTOP_RELEASE_TAG" "$RUNNER_TEMP/latest.json" --clobber
- name: Download versioned updater metadata
if: ${{ !inputs.draft }}
shell: bash
env:
GH_TOKEN: ${{ github.token }}
@ -790,6 +953,7 @@ jobs:
test -s "$RUNNER_TEMP/desktop-updater/latest.json"
- name: Validate versioned updater metadata
if: ${{ !inputs.draft }}
shell: bash
run: |
python3 <<'PY'
@ -849,6 +1013,7 @@ jobs:
PY
- name: Ensure desktop updater channel release
if: ${{ !inputs.draft }}
shell: bash
env:
GH_TOKEN: ${{ github.token }}
@ -881,6 +1046,7 @@ jobs:
PY
- name: Prevent updater channel downgrade
if: ${{ !inputs.draft }}
shell: bash
env:
GH_TOKEN: ${{ github.token }}
@ -971,6 +1137,7 @@ jobs:
PY
- name: Publish desktop updater channel metadata
if: ${{ !inputs.draft }}
shell: bash
env:
GH_TOKEN: ${{ github.token }}

6
.github/workflows/security-audit.yml vendored
View file

@ -2,8 +2,8 @@
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
# Multi-language supply-chain audit. Triggers:
# - PRs touching any dependency manifest (Python / npm / Cargo) or
# this workflow file,
# - PRs touching any dependency manifest (Python / npm / Cargo), a
# scanner or its allowlist baseline, or this workflow file,
# - push to main / pip,
# - nightly @ 04:13 UTC so newly-published advisories surface even
# when no PR opens,
@ -57,7 +57,9 @@ on:
- 'studio/src-tauri/Cargo.lock'
- 'pyproject.toml'
- 'scripts/scan_packages.py'
- 'scripts/scan_packages_baseline.json'
- 'scripts/scan_npm_packages.py'
- 'scripts/scan_npm_packages_baseline.json'
- '.github/workflows/security-audit.yml'
push:
branches: [main, pip]

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

@ -444,6 +444,8 @@ jobs:
python - <<'PY'
import json
import os
import time
import urllib.error
import urllib.request
BASE = os.environ["BASE_URL"]
@ -464,10 +466,26 @@ jobs:
"Content-Type": "application/json",
},
)
with urllib.request.urlopen(req, timeout = timeout) as resp:
return resp.status, json.loads(resp.read().decode())
# Shared CI runners stall sporadically, so retry transport-level
# failures only; HTTP status errors surface immediately. Bounded
# to fit the job's timeout-minutes: short probes get 3 full
# attempts, long probes one retry capped at 300s (a healthy
# server answers a retry quickly; a stalled one never does).
attempts = 3 if timeout <= 300 else 2
for attempt in range(attempts):
try:
t = timeout if attempt == 0 else min(timeout, 300)
with urllib.request.urlopen(req, timeout = t) as resp:
return resp.status, json.loads(resp.read().decode())
except urllib.error.HTTPError:
raise
except (TimeoutError, ConnectionError, urllib.error.URLError) as exc:
if attempt == attempts - 1:
raise
print(f"[retry] {path}: {exc!r}", flush = True)
time.sleep(15)
def post_sse(path, body, *, timeout = 600):
def post_sse(path, body, *, timeout = 600, retries = 1, complete_on = None):
"""POST a streaming request and accumulate the assistant
text deltas. The server-side agentic loop ALWAYS returns
SSE regardless of the request's `stream` field, so any
@ -483,6 +501,22 @@ jobs:
invocation markers / tool output, since
`delta.content` alone is not evidence
that the tool path executed.
A shared CI runner can stall the stream transport (the
connection opening, or a mid-stream read) even when Studio
is healthy, so retry a stall once with a fresh request
capped at 300s. A stall means the stream did NOT complete,
so partial events are normally NOT returned (an early
tool_start with no tool_end is not proof the tool loop
finished). The one exception is `complete_on`: an optional
predicate over the events collected so far -- when a stall
happens after it is already satisfied (the tool ran and
produced its result before the trailing read timed out),
those events are returned rather than discarded, so the
stall-after-answer case still counts. HTTP status errors
surface immediately; a stall that yields no completed result
across all attempts re-raises so the caller can rotate to
the next seed.
"""
body = {**body, "stream": True}
data = json.dumps(body).encode()
@ -495,26 +529,45 @@ jobs:
"Content-Type": "application/json",
},
)
parts = []
events = []
with urllib.request.urlopen(req, timeout = timeout) as resp:
for raw in resp:
line = raw.decode().strip()
if not line.startswith("data: "):
continue
payload = line[6:]
if payload == "[DONE]":
break
events.append(payload)
try:
chunk = json.loads(payload)
except json.JSONDecodeError:
continue
for choice in chunk.get("choices", []):
delta = choice.get("delta", {}) or {}
if delta.get("content"):
parts.append(delta["content"])
return "".join(parts), events
for attempt in range(retries + 1):
parts = []
events = []
t = timeout if attempt == 0 else min(timeout, 300)
try:
with urllib.request.urlopen(req, timeout = t) as resp:
for raw in resp:
line = raw.decode().strip()
if not line.startswith("data: "):
continue
payload = line[6:]
if payload == "[DONE]":
break
events.append(payload)
try:
chunk = json.loads(payload)
except json.JSONDecodeError:
continue
for choice in chunk.get("choices", []):
delta = choice.get("delta", {}) or {}
if delta.get("content"):
parts.append(delta["content"])
return "".join(parts), events
except urllib.error.HTTPError:
raise
except (TimeoutError, ConnectionError, urllib.error.URLError) as exc:
# A stall after the tool already produced its result is
# the case this probe exists to tolerate: keep those
# events. But a stall with only an early tool_start (no
# completed output) is not proof the tool loop finished,
# so it must not pass -- retry once, then raise so
# _run_tool_probe rotates to the next seed.
if complete_on is not None and complete_on(events):
print(f"[retry-sse] {path}: {exc!r}; keeping {len(events)} completed events", flush = True)
return "".join(parts), events
if attempt == retries:
raise
print(f"[retry-sse] {path}: {exc!r}", flush = True)
time.sleep(15)
_STUDIO_TOOL_TYPES = {
"tool_start", "tool_end", "tool_use", "tool_result",
@ -651,17 +704,55 @@ jobs:
"""
attempts_log = []
best = None
# Cap the wall-clock spent rotating through stalled seeds so a
# persistent no-data wedge fails fast (clean assertion) instead
# of being killed by the job's timeout-minutes. A healthy or
# merely degenerate round answers in seconds, so all seeds still
# run in the normal case; only stalls consume the budget.
probe_deadline = time.monotonic() + 300
for attempt_i in range(max_attempts):
# Cap each read by the budget still remaining (not just a flat
# 180s) and skip an attempt too small to finish, so the whole
# rotation stays within ~300s -- two probes then fit the job's
# timeout-minutes even if every seed stalls.
remaining = int(probe_deadline - time.monotonic())
if attempt_i and remaining < 30:
print(f"[tools] {label}: seed-rotation budget spent after {attempt_i} attempts", flush = True)
break
attempt_seed = SEED + attempt_i
content, events = post_sse("/v1/chat/completions", {
"messages": [{"role": "user", "content": prompt}],
"enable_tools": True,
"enabled_tools": enabled,
"session_id": f"{session}-att{attempt_i}",
"temperature": TOOL_PROBE_TEMP,
"seed": attempt_seed,
"max_tokens": 600,
})
try:
# Bounded per-attempt timeout, no inner retry -- the seed
# loop IS the retry, so a stall raises quickly and rotates
# rather than spending post_sse's full 600+300s. complete_on
# keeps a stall that already produced the tool result (only
# the trailing read timed out) instead of discarding it.
content, events = post_sse("/v1/chat/completions", {
"messages": [{"role": "user", "content": prompt}],
"enable_tools": True,
"permission_mode": "full",
"enabled_tools": enabled,
"session_id": f"{session}-att{attempt_i}",
"temperature": TOOL_PROBE_TEMP,
"seed": attempt_seed,
"max_tokens": 600,
}, timeout = min(180, remaining), retries = 0,
complete_on = lambda ev: _tool_invoked(ev) and _tool_output_contains(ev, *needles))
except urllib.error.HTTPError:
# HTTPError subclasses URLError, so re-raise a real 4xx/5xx
# here instead of letting the transport-stall handler below
# swallow it and rotate seeds -- an endpoint status failure
# must surface, not be masked as missing tool evidence.
raise
except (TimeoutError, ConnectionError, urllib.error.URLError) as exc:
# A transport stall that outlived post_sse's own retry:
# log it as a failed attempt and rotate to the next seed
# rather than sinking the whole probe on one bad stream.
attempts_log.append({
"attempt": attempt_i, "seed": attempt_seed,
"transport_error": repr(exc),
})
print(f"[tools] retry {label} attempt {attempt_i}: transport {exc!r}", flush = True)
continue
invoked = _tool_invoked(events)
produced = _tool_output_contains(events, *needles)
attempts_log.append({
@ -722,15 +813,19 @@ jobs:
# enough that requiring a tool_call marker would create
# red-herring failures from infra rather than from Studio.
try:
# Best-effort and bounded: a single 180s attempt keeps a stall
# from eating the job's timeout-minutes (it already WARNs, so a
# retry buys nothing).
content, events = post_sse("/v1/chat/completions", {
"messages": [{"role": "user", "content": "Search the web for 'unsloth ai github' and summarise."}],
"enable_tools": True,
"permission_mode": "full",
"enabled_tools": ["web_search"],
"session_id": "ci-tool-calling-web",
"temperature": 0.0,
"seed": SEED,
"max_tokens": 400,
})
}, timeout = 180, retries = 0)
print(
f"[tools] PASS web_search stream ({len(content)} chars in content, "
f"{len(events)} raw events)"
@ -938,6 +1033,8 @@ jobs:
import base64
import json
import os
import time
import urllib.error
import urllib.request
from openai import OpenAI
from anthropic import Anthropic
@ -956,8 +1053,24 @@ jobs:
"Content-Type": "application/json",
},
)
with urllib.request.urlopen(req, timeout = timeout) as resp:
return resp.status, json.loads(resp.read().decode())
# Shared CI runners stall sporadically, so retry transport-level
# failures only; HTTP status errors surface immediately. Bounded
# to fit the job's timeout-minutes: short probes get 3 full
# attempts, long probes one retry capped at 300s (a healthy
# server answers a retry quickly; a stalled one never does).
attempts = 3 if timeout <= 300 else 2
for attempt in range(attempts):
try:
t = timeout if attempt == 0 else min(timeout, 300)
with urllib.request.urlopen(req, timeout = t) as resp:
return resp.status, json.loads(resp.read().decode())
except urllib.error.HTTPError:
raise
except (TimeoutError, ConnectionError, urllib.error.URLError) as exc:
if attempt == attempts - 1:
raise
print(f"[retry] {path}: {exc!r}", flush = True)
time.sleep(15)
# ── 1. response_format = json_object (JSON mode) ─────────────
# llama.cpp's HTTP server supports OpenAI-compatible JSON

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

@ -430,6 +430,8 @@ jobs:
python - <<'PY'
import json
import os
import time
import urllib.error
import urllib.request
BASE = os.environ["BASE_URL"]
@ -450,14 +452,41 @@ jobs:
"Content-Type": "application/json",
},
)
with urllib.request.urlopen(req, timeout = timeout) as resp:
return resp.status, json.loads(resp.read().decode())
# Shared CI runners stall sporadically, so retry transport-level
# failures only; HTTP status errors surface immediately. Bounded
# to fit the job's timeout-minutes: short probes get 3 full
# attempts, long probes one retry capped at 300s (a healthy
# server answers a retry quickly; a stalled one never does).
attempts = 3 if timeout <= 300 else 2
for attempt in range(attempts):
try:
t = timeout if attempt == 0 else min(timeout, 300)
with urllib.request.urlopen(req, timeout = t) as resp:
return resp.status, json.loads(resp.read().decode())
except urllib.error.HTTPError:
raise
except (TimeoutError, ConnectionError, urllib.error.URLError) as exc:
if attempt == attempts - 1:
raise
print(f"[retry] {path}: {exc!r}", flush = True)
time.sleep(15)
def post_sse(path, body, *, timeout = 600):
def post_sse(path, body, *, timeout = 600, retries = 1, soft = False):
"""POST a streaming request and accumulate the assistant
text deltas. The server-side agentic loop ALWAYS returns
SSE regardless of the request's `stream` field, so any
call with enable_tools=true must use this helper."""
call with enable_tools=true must use this helper.
A shared CI runner can stall the stream transport (the
connection opening, or a mid-stream read) even when Studio
is healthy, so harden the read three ways: retry a stall
once with a fresh request capped at 300s; return any text
already streamed before a stall (a stall on the trailing
tokens, after the answer arrived, still counts); and when
every attempt yields nothing, a hard call re-raises while a
soft call (the best-effort server-side tool probes) returns
None so the caller can WARN instead of sinking the whole
job. HTTP status errors always surface immediately."""
body = {**body, "stream": True}
data = json.dumps(body).encode()
req = urllib.request.Request(
@ -469,24 +498,43 @@ jobs:
"Content-Type": "application/json",
},
)
parts = []
with urllib.request.urlopen(req, timeout = timeout) as resp:
for raw in resp:
line = raw.decode().strip()
if not line.startswith("data: "):
continue
payload = line[6:]
if payload == "[DONE]":
break
try:
chunk = json.loads(payload)
except json.JSONDecodeError:
continue
for choice in chunk.get("choices", []):
delta = choice.get("delta", {}) or {}
if delta.get("content"):
parts.append(delta["content"])
return "".join(parts)
for attempt in range(retries + 1):
parts = []
t = timeout if attempt == 0 else min(timeout, 300)
try:
with urllib.request.urlopen(req, timeout = t) as resp:
for raw in resp:
line = raw.decode().strip()
if not line.startswith("data: "):
continue
payload = line[6:]
if payload == "[DONE]":
break
try:
chunk = json.loads(payload)
except json.JSONDecodeError:
continue
for choice in chunk.get("choices", []):
delta = choice.get("delta", {}) or {}
if delta.get("content"):
parts.append(delta["content"])
return "".join(parts)
except urllib.error.HTTPError:
raise
except (TimeoutError, ConnectionError, urllib.error.URLError) as exc:
# Text already streamed is a valid signal -- keep it
# rather than re-running a heavy generation.
if parts:
joined = "".join(parts)
print(f"[retry-sse] {path}: {exc!r}; keeping {len(joined)} partial chars", flush = True)
return joined
if attempt == retries:
if soft:
print(f"[tools] WARN {path}: SSE transport stalled with no data ({exc!r}) -- non-blocking", flush = True)
return None
raise
print(f"[retry-sse] {path}: {exc!r}", flush = True)
time.sleep(15)
# ── 1. Standard OpenAI function calling ──────────────────────
weather_tool = {
@ -557,16 +605,23 @@ jobs:
# macos-14 free runner is ~10 tok/s on Qwen3.5-2B Q4_K_XL;
# cap max_tokens tightly so each SSE round stays under ~30s
# even when the model stalls in a degenerate output state.
# retries=0 on the best-effort probes: this job's 25-minute cap
# allows a 10-minute model load, so a no-data stall must be a
# single 180s attempt (not 180+15+180s) to leave room for the
# thinking checks. A soft/best-effort probe only WARNs anyway.
content = post_sse("/v1/chat/completions", {
"messages": [{"role": "user", "content": "What is 123 * 456? Use the python tool to compute it and tell me the number."}],
"enable_tools": True,
"permission_mode": "full",
"enabled_tools": ["python"],
"session_id": "ci-tool-calling-py",
"temperature": TEMP,
"seed": SEED,
"max_tokens": 128,
}, timeout = 180)
if "56088" in content or "56,088" in content:
}, timeout = 180, retries = 0, soft = True)
if content is None:
print("[tools] WARN python tool: SSE transport stalled after retries -- non-blocking")
elif "56088" in content or "56,088" in content:
print(f"[tools] PASS python tool ({len(content)} chars, found 56088)")
else:
# Empty stream is a known Mac-quant degeneracy too; log
@ -593,12 +648,13 @@ jobs:
content = post_sse("/v1/chat/completions", {
"messages": [{"role": "user", "content": "Search the web for 'unsloth ai github' and summarise."}],
"enable_tools": True,
"permission_mode": "full",
"enabled_tools": ["web_search"],
"session_id": "ci-tool-calling-web",
"temperature": TEMP,
"seed": SEED,
"max_tokens": 96,
}, timeout = 180)
}, timeout = 180, retries = 0)
print(f"[tools] PASS web_search stream ({len(content)} chars)")
except Exception as exc:
print(f"[tools] WARN web_search probe failed (non-blocking): {exc}")
@ -825,6 +881,8 @@ jobs:
import base64
import json
import os
import time
import urllib.error
import urllib.request
from openai import OpenAI
from anthropic import Anthropic
@ -848,8 +906,24 @@ jobs:
"Content-Type": "application/json",
},
)
with urllib.request.urlopen(req, timeout = timeout) as resp:
return resp.status, json.loads(resp.read().decode())
# Shared CI runners stall sporadically, so retry transport-level
# failures only; HTTP status errors surface immediately. Bounded
# to fit the job's timeout-minutes: short probes get 3 full
# attempts, long probes one retry capped at 300s (a healthy
# server answers a retry quickly; a stalled one never does).
attempts = 3 if timeout <= 300 else 2
for attempt in range(attempts):
try:
t = timeout if attempt == 0 else min(timeout, 300)
with urllib.request.urlopen(req, timeout = t) as resp:
return resp.status, json.loads(resp.read().decode())
except urllib.error.HTTPError:
raise
except (TimeoutError, ConnectionError, urllib.error.URLError) as exc:
if attempt == attempts - 1:
raise
print(f"[retry] {path}: {exc!r}", flush = True)
time.sleep(15)
# ── 1. response_format = json_object (JSON mode) ─────────────
# llama.cpp's HTTP server supports OpenAI-compatible JSON

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

@ -634,6 +634,8 @@ jobs:
python - <<'PY'
import json
import os
import time
import urllib.error
import urllib.request
BASE = os.environ["BASE_URL"]
@ -656,10 +658,41 @@ jobs:
"Content-Type": "application/json",
},
)
with urllib.request.urlopen(req, timeout = timeout) as resp:
return resp.status, json.loads(resp.read().decode())
# Shared CI runners stall sporadically, so retry transport-level
# failures only; HTTP status errors surface immediately. Bounded
# to fit the job's timeout-minutes: short probes get 3 full
# attempts, long probes one retry capped at 300s (a healthy
# server answers a retry quickly; a stalled one never does).
attempts = 3 if timeout <= 300 else 2
for attempt in range(attempts):
try:
t = timeout if attempt == 0 else min(timeout, 300)
with urllib.request.urlopen(req, timeout = t) as resp:
return resp.status, json.loads(resp.read().decode())
except urllib.error.HTTPError:
raise
except (TimeoutError, ConnectionError, urllib.error.URLError) as exc:
if attempt == attempts - 1:
raise
print(f"[retry] {path}: {exc!r}", flush = True)
time.sleep(15)
def post_sse(path, body, *, timeout = 600):
def post_sse(path, body, *, timeout = 600, retries = 1, soft = False):
# The server-side agentic loop always answers over SSE. A
# shared CI runner can stall the stream transport (the
# connection opening, or a mid-stream read) even when Studio
# is healthy, so harden the read three ways:
# * retry a transport stall once with a fresh request,
# capped at 300s (a healthy server answers a retry
# quickly, a wedged one never does);
# * return any text already streamed before a stall, so a
# stall on the trailing tokens -- after the answer
# arrived -- still counts;
# * when every attempt yields nothing, a hard call
# re-raises while a soft call (the best-effort
# server-side tool probes) returns None so the caller
# can WARN instead of sinking the whole job.
# HTTP status errors always surface immediately.
body = {**body, "stream": True}
data = json.dumps(body).encode()
req = urllib.request.Request(
@ -671,24 +704,43 @@ jobs:
"Content-Type": "application/json",
},
)
parts = []
with urllib.request.urlopen(req, timeout = timeout) as resp:
for raw in resp:
line = raw.decode().strip()
if not line.startswith("data: "):
continue
payload = line[6:]
if payload == "[DONE]":
break
try:
chunk = json.loads(payload)
except json.JSONDecodeError:
continue
for choice in chunk.get("choices", []):
delta = choice.get("delta", {}) or {}
if delta.get("content"):
parts.append(delta["content"])
return "".join(parts)
for attempt in range(retries + 1):
parts = []
t = timeout if attempt == 0 else min(timeout, 300)
try:
with urllib.request.urlopen(req, timeout = t) as resp:
for raw in resp:
line = raw.decode().strip()
if not line.startswith("data: "):
continue
payload = line[6:]
if payload == "[DONE]":
break
try:
chunk = json.loads(payload)
except json.JSONDecodeError:
continue
for choice in chunk.get("choices", []):
delta = choice.get("delta", {}) or {}
if delta.get("content"):
parts.append(delta["content"])
return "".join(parts)
except urllib.error.HTTPError:
raise
except (TimeoutError, ConnectionError, urllib.error.URLError) as exc:
# Text already streamed is a valid signal -- keep it
# rather than re-running a heavy generation.
if parts:
joined = "".join(parts)
print(f"[retry-sse] {path}: {exc!r}; keeping {len(joined)} partial chars", flush = True)
return joined
if attempt == retries:
if soft:
print(f"[tools] WARN {path}: SSE transport stalled with no data ({exc!r}) -- non-blocking", flush = True)
return None
raise
print(f"[retry-sse] {path}: {exc!r}", flush = True)
time.sleep(15)
# ── 1. Standard OpenAI function calling ──────────────────────
weather_tool = {
@ -731,16 +783,24 @@ jobs:
)
# ── 2. Server-side python tool ───────────────────────────────
# Bound each soft probe to a single 180s attempt (timeout=180,
# retries=0): this job runs two of them back-to-back under a
# 30-minute cap, so the default 600+15+300s per stall could hit
# the workflow timeout before the thinking checks run. A soft
# probe only WARNs anyway, so a retry buys nothing.
content = post_sse("/v1/chat/completions", {
"messages": [{"role": "user", "content": "What is 123 * 456? Use the python tool to compute it and tell me the number."}],
"enable_tools": True,
"permission_mode": "full",
"enabled_tools": ["python"],
"session_id": "ci-tool-calling-py",
"temperature": TEMP,
"seed": SEED,
"max_tokens": 600,
})
if "56088" in content or "56,088" in content:
}, timeout = 180, retries = 0, soft = True)
if content is None:
print("[tools] WARN python tool: SSE transport stalled after retries -- non-blocking")
elif "56088" in content or "56,088" in content:
print(f"[tools] PASS python tool ({len(content)} chars, found 56088)")
else:
assert content, "python tool: SSE stream empty"
@ -757,13 +817,16 @@ jobs:
content = post_sse("/v1/chat/completions", {
"messages": [{"role": "user", "content": "Use the terminal tool to run `echo hello-bash-tool` and tell me the exact output."}],
"enable_tools": True,
"permission_mode": "full",
"enabled_tools": ["terminal"],
"session_id": "ci-tool-calling-bash",
"temperature": TEMP,
"seed": SEED,
"max_tokens": 600,
})
if "hello-bash-tool" in content:
}, timeout = 180, retries = 0, soft = True)
if content is None:
print("[tools] WARN terminal tool: SSE transport stalled after retries -- non-blocking")
elif "hello-bash-tool" in content:
print(f"[tools] PASS terminal tool ({len(content)} chars)")
else:
assert content, "terminal tool: SSE stream empty"
@ -779,12 +842,13 @@ jobs:
content = post_sse("/v1/chat/completions", {
"messages": [{"role": "user", "content": "Search the web for 'unsloth ai github' and summarise."}],
"enable_tools": True,
"permission_mode": "full",
"enabled_tools": ["web_search"],
"session_id": "ci-tool-calling-web",
"temperature": TEMP,
"seed": SEED,
"max_tokens": 400,
})
}, timeout = 180, retries = 0)
print(f"[tools] PASS web_search stream ({len(content)} chars)")
except Exception as exc:
print(f"[tools] WARN web_search probe failed (non-blocking): {exc}")
@ -1063,6 +1127,8 @@ jobs:
import base64
import json
import os
import time
import urllib.error
import urllib.request
from openai import OpenAI
from anthropic import Anthropic
@ -1082,8 +1148,24 @@ jobs:
"Content-Type": "application/json",
},
)
with urllib.request.urlopen(req, timeout = timeout) as resp:
return resp.status, json.loads(resp.read().decode())
# Shared CI runners stall sporadically, so retry transport-level
# failures only; HTTP status errors surface immediately. Bounded
# to fit the job's timeout-minutes: short probes get 3 full
# attempts, long probes one retry capped at 300s (a healthy
# server answers a retry quickly; a stalled one never does).
attempts = 3 if timeout <= 300 else 2
for attempt in range(attempts):
try:
t = timeout if attempt == 0 else min(timeout, 300)
with urllib.request.urlopen(req, timeout = t) as resp:
return resp.status, json.loads(resp.read().decode())
except urllib.error.HTTPError:
raise
except (TimeoutError, ConnectionError, urllib.error.URLError) as exc:
if attempt == attempts - 1:
raise
print(f"[retry] {path}: {exc!r}", flush = True)
time.sleep(15)
# ── 1. response_format = json_object (JSON mode) ─────────────
status, data = post("/v1/chat/completions", {
@ -1334,42 +1416,75 @@ jobs:
try { Add-MpPreference -ExclusionPath $p -ErrorAction Stop } catch { }
}
- name: Hide Visual Studio + CMake (simulate a host with no build tools)
- name: Prepare no-build-tools simulation
shell: pwsh
run: |
$ErrorActionPreference = 'Stop'
# A Program Files dir can hold a transient handle (Defender / MSBuild node)
# so Rename-Item intermittently fails with "Access is denied"; retry to ride it out.
function Rename-WithRetry($Path, $NewName) {
for ($i = 1; $i -le 6; $i++) {
try { Rename-Item -LiteralPath $Path -NewName $NewName -ErrorAction Stop; return }
catch { if ($i -eq 6) { throw }; Start-Sleep -Seconds 3 }
$root = Join-Path $env:GITHUB_WORKSPACE 'no-build-tools'
$pf = Join-Path $root 'ProgramFiles'
$pfx86 = Join-Path $root 'ProgramFilesx86'
New-Item -ItemType Directory -Force -Path $pf, $pfx86 | Out-Null
$blocked = [System.Collections.Generic.HashSet[string]]::new([StringComparer]::OrdinalIgnoreCase)
foreach ($tool in @('cmake', 'cl.exe')) {
foreach ($cmd in (Get-Command $tool -All -ErrorAction SilentlyContinue)) {
if ($cmd.Source) {
$dir = Split-Path -Parent $cmd.Source
if ($dir) {
[void] $blocked.Add(
[Environment]::ExpandEnvironmentVariables($dir).Trim().Trim('"').TrimEnd('\'))
}
}
}
}
# Rename the Visual Studio install roots (incl. the Installer that holds
# vswhere.exe) so Find-VsBuildTools' vswhere + filesystem scan both miss.
foreach ($d in @("$env:ProgramFiles\Microsoft Visual Studio", "${env:ProgramFiles(x86)}\Microsoft Visual Studio")) {
if (Test-Path -LiteralPath $d) {
Rename-WithRetry $d ((Split-Path $d -Leaf) + '.vsoff')
Write-Host "Hid VS: $d"
}
# Normalized comparison so registry spellings (trailing slash,
# unexpanded %VAR%) still match.
function Test-Blocked([string]$p) {
$n = [Environment]::ExpandEnvironmentVariables($p).Trim().Trim('"').TrimEnd('\')
return $blocked.Contains($n)
}
# Surgically rename each cmake executable on PATH (not its parent dir --
# cmake can share a dir with other shims) so Get-Command cmake fails.
$hidden = @()
foreach ($c in (Get-Command cmake -All -ErrorAction SilentlyContinue)) {
if ($c.Source -and (Test-Path -LiteralPath $c.Source)) {
Rename-WithRetry $c.Source ((Split-Path $c.Source -Leaf) + '.off')
$hidden += $c.Source
Write-Host "Hid cmake: $($c.Source)"
}
$pathParts = $env:Path -split [IO.Path]::PathSeparator |
Where-Object { $_ -and -not (Test-Blocked $_) }
$noBuildToolsPath = $pathParts -join [IO.Path]::PathSeparator
# install.ps1's Refresh-SessionPath and setup.ps1's Refresh-Environment
# rebuild the session Path from these scopes mid-install, so filter
# them too. Originals are saved for the cleanup step.
foreach ($scope in @('Machine', 'User')) {
$orig = [Environment]::GetEnvironmentVariable('Path', $scope)
if (-not $orig) { continue }
Set-Content -LiteralPath (Join-Path $root "orig-path-$scope.txt") -Value $orig -NoNewline
$kept = ($orig -split ';' | Where-Object { $_ -and -not (Test-Blocked $_) }) -join ';'
[Environment]::SetEnvironmentVariable('Path', $kept, $scope)
Write-Host "Filtered $scope Path scope."
}
"NO_BUILD_TOOLS_PROGRAMFILES=$pf" | Out-File -FilePath $env:GITHUB_ENV -Append -Encoding utf8
"NO_BUILD_TOOLS_PROGRAMFILES_X86=$pfx86" | Out-File -FilePath $env:GITHUB_ENV -Append -Encoding utf8
"NO_BUILD_TOOLS_PATH<<NO_BUILD_TOOLS_PATH_EOF" | Out-File -FilePath $env:GITHUB_ENV -Append -Encoding utf8
$noBuildToolsPath | Out-File -FilePath $env:GITHUB_ENV -Append -Encoding utf8
"NO_BUILD_TOOLS_PATH_EOF" | Out-File -FilePath $env:GITHUB_ENV -Append -Encoding utf8
Write-Host "ProgramFiles simulation root: $pf"
Write-Host "ProgramFiles(x86) simulation root: $pfx86"
if ($blocked.Count -gt 0) {
Write-Host "Removed build-tool PATH dirs:"
$blocked | Sort-Object | ForEach-Object { Write-Host " $_" }
} else {
Write-Host "No cmake or cl.exe PATH dirs found to remove."
}
("HIDDEN_CMAKE=" + ($hidden -join '|')) | Out-File -FilePath $env:GITHUB_ENV -Append -Encoding utf8
- name: Assert Visual Studio + CMake are genuinely undetectable
shell: pwsh
run: |
$ErrorActionPreference = 'Stop'
# Set in-script: the runner does not apply step-level env keys with
# parentheses (`ProgramFiles(x86)`), so vswhere still found VS.
if (-not $env:NO_BUILD_TOOLS_PROGRAMFILES) { Write-Error "NO_BUILD_TOOLS_* env missing (Prepare step did not run?)"; exit 1 }
$env:ProgramFiles = $env:NO_BUILD_TOOLS_PROGRAMFILES
${env:ProgramFiles(x86)} = $env:NO_BUILD_TOOLS_PROGRAMFILES_X86
$env:Path = $env:NO_BUILD_TOOLS_PATH
. (Join-Path $env:GITHUB_WORKSPACE 'tests/studio_setup_ps1/Get-FunctionSource.ps1')
$setup = Join-Path $env:GITHUB_WORKSPACE 'studio/setup.ps1'
foreach ($fn in @('Resolve-VsGeneratorFromLabel', 'Find-VsBuildTools')) {
@ -1394,6 +1509,10 @@ jobs:
# Withheld on PR: this step runs checked-out PR code; public GGUF still downloads.
HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }}
run: |
# Set in-script (see the assert step); child processes inherit these.
$env:ProgramFiles = $env:NO_BUILD_TOOLS_PROGRAMFILES
${env:ProgramFiles(x86)} = $env:NO_BUILD_TOOLS_PROGRAMFILES_X86
$env:Path = $env:NO_BUILD_TOOLS_PATH
New-Item -ItemType Directory -Force -Path logs | Out-Null
$ProgressPreference = 'SilentlyContinue'
& ./install.ps1 --local --no-torch *>&1 | Tee-Object -FilePath logs/install.log
@ -1480,19 +1599,19 @@ jobs:
[ -n "$CONTENT" ] && [ "$CONTENT" != "null" ] || { echo "::error::empty completion"; exit 1; }
echo "Inference OK without Visual Studio: $CONTENT"
- name: Restore Visual Studio + CMake
- name: Clean no-build-tools simulation
if: always()
shell: pwsh
run: |
foreach ($d in @("$env:ProgramFiles\Microsoft Visual Studio", "${env:ProgramFiles(x86)}\Microsoft Visual Studio")) {
$off = "$d.vsoff"
if (Test-Path -LiteralPath $off) { Rename-Item -LiteralPath $off -NewName (Split-Path $d -Leaf); Write-Host "Restored $d" }
}
if ($env:HIDDEN_CMAKE) {
foreach ($src in ($env:HIDDEN_CMAKE -split '\|')) {
if ($src -and (Test-Path -LiteralPath "$src.off")) { Rename-Item -LiteralPath "$src.off" -NewName (Split-Path $src -Leaf) }
$root = Join-Path $env:GITHUB_WORKSPACE 'no-build-tools'
foreach ($scope in @('Machine', 'User')) {
$saved = Join-Path $root "orig-path-$scope.txt"
if (Test-Path -LiteralPath $saved) {
[Environment]::SetEnvironmentVariable('Path', (Get-Content -LiteralPath $saved -Raw), $scope)
Write-Host "Restored $scope Path scope."
}
}
Remove-Item -LiteralPath $root -Recurse -Force -ErrorAction SilentlyContinue
- name: Stop Studio
if: always()
@ -1540,21 +1659,34 @@ jobs:
with:
python-version: '3.12'
- name: Hide Visual Studio
- name: Prepare no-build-tools simulation
shell: pwsh
run: |
$ErrorActionPreference = 'Stop'
# Retry the rename: a Program Files dir can hold a transient handle that
# makes Rename-Item intermittently fail with "Access is denied".
function Rename-WithRetry($Path, $NewName) {
for ($i = 1; $i -le 6; $i++) {
try { Rename-Item -LiteralPath $Path -NewName $NewName -ErrorAction Stop; return }
catch { if ($i -eq 6) { throw }; Start-Sleep -Seconds 3 }
$root = Join-Path $env:GITHUB_WORKSPACE 'no-build-tools'
$pf = Join-Path $root 'ProgramFiles'
$pfx86 = Join-Path $root 'ProgramFilesx86'
New-Item -ItemType Directory -Force -Path $pf, $pfx86 | Out-Null
$blocked = [System.Collections.Generic.HashSet[string]]::new([StringComparer]::OrdinalIgnoreCase)
foreach ($tool in @('cmake', 'cl.exe')) {
foreach ($cmd in (Get-Command $tool -All -ErrorAction SilentlyContinue)) {
if ($cmd.Source) {
$dir = Split-Path -Parent $cmd.Source
if ($dir) { [void] $blocked.Add($dir) }
}
}
}
foreach ($d in @("$env:ProgramFiles\Microsoft Visual Studio", "${env:ProgramFiles(x86)}\Microsoft Visual Studio")) {
if (Test-Path -LiteralPath $d) { Rename-WithRetry $d ((Split-Path $d -Leaf) + '.vsoff'); Write-Host "Hid VS: $d" }
}
$pathParts = $env:Path -split [IO.Path]::PathSeparator |
Where-Object { $_ -and -not $blocked.Contains($_) }
$noBuildToolsPath = $pathParts -join [IO.Path]::PathSeparator
"NO_BUILD_TOOLS_PROGRAMFILES=$pf" | Out-File -FilePath $env:GITHUB_ENV -Append -Encoding utf8
"NO_BUILD_TOOLS_PROGRAMFILES_X86=$pfx86" | Out-File -FilePath $env:GITHUB_ENV -Append -Encoding utf8
"NO_BUILD_TOOLS_PATH<<NO_BUILD_TOOLS_PATH_EOF" | Out-File -FilePath $env:GITHUB_ENV -Append -Encoding utf8
$noBuildToolsPath | Out-File -FilePath $env:GITHUB_ENV -Append -Encoding utf8
"NO_BUILD_TOOLS_PATH_EOF" | Out-File -FilePath $env:GITHUB_ENV -Append -Encoding utf8
- name: Windows CUDA and ROCm prebuilts exist in unslothai/llama.cpp (what GPU users download, no VS)
env:
@ -1577,25 +1709,34 @@ jobs:
echo "Windows CUDA and ROCm prebuilts are available -- GPU users get them without compiling."
- name: The prebuilt resolver runs without Visual Studio
shell: pwsh
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
run: |
$ErrorActionPreference = 'Stop'
# pwsh: bash cannot export `ProgramFiles(x86)`; set in-script so the
# python child inherits the overrides.
$env:ProgramFiles = $env:NO_BUILD_TOOLS_PROGRAMFILES
${env:ProgramFiles(x86)} = $env:NO_BUILD_TOOLS_PROGRAMFILES_X86
$env:Path = $env:NO_BUILD_TOOLS_PATH
# Resolver-only (no GPU on hosted runners, so the host resolves to the
# CPU bundle). The point is that resolution needs no compiler/VS.
python -m pip install --upgrade huggingface_hub
python studio/install_llama_prebuilt.py --resolve-prebuilt latest --output-format json > /tmp/resolve.json || {
echo "::error::resolver exited non-zero"; cat /tmp/resolve.json || true; exit 1; }
cat /tmp/resolve.json
echo "Prebuilt resolver ran with no Visual Studio present."
if ($LASTEXITCODE -ne 0) { Write-Host "::error::pip install huggingface_hub failed"; exit 1 }
python studio/install_llama_prebuilt.py --resolve-prebuilt latest --output-format json > resolve.json
if ($LASTEXITCODE -ne 0) {
Write-Host "::error::resolver exited non-zero"
if (Test-Path resolve.json) { Get-Content resolve.json }
exit 1
}
Get-Content resolve.json
Write-Host "Prebuilt resolver ran with no Visual Studio present."
- name: Restore Visual Studio
- name: Clean no-build-tools simulation
if: always()
shell: pwsh
run: |
foreach ($d in @("$env:ProgramFiles\Microsoft Visual Studio", "${env:ProgramFiles(x86)}\Microsoft Visual Studio")) {
$off = "$d.vsoff"
if (Test-Path -LiteralPath $off) { Rename-Item -LiteralPath $off -NewName (Split-Path $d -Leaf); Write-Host "Restored $d" }
}
Remove-Item -LiteralPath (Join-Path $env:GITHUB_WORKSPACE 'no-build-tools') -Recurse -Force -ErrorAction SilentlyContinue
# ── folded from studio-setup-ps1-vs2026.yml: setup.ps1 unit tests + real-VS detection + vcredist ──
pester:

25
README.md
View file

@ -84,7 +84,7 @@ Use the same command to update.
```bash
unsloth studio -p 8888
```
For cloud or global access, add `-H 0.0.0.0`. By default, Unsloth is accessible only locally.
For LAN or cloud access, add `-H 0.0.0.0` (raw port only; add `--cloudflare` for a public URL). By default, Unsloth is accessible only locally.
Launching from the terminal never opens a browser -- open the printed URL yourself. The desktop shortcut opens your default browser once the server is up; to launch the server without that (e.g. when running Studio as a browser PWA / app window), answer "n" at the installer's browser prompt, pass `--no-browser` to the launcher, or set `UNSLOTH_STUDIO_NO_BROWSER=1`.
@ -214,10 +214,23 @@ By default `unsloth studio` binds to `127.0.0.1` (this machine only). To reach i
```bash
unsloth studio --secure -p 8888
```
- `-H 0.0.0.0`: bind the raw port on all network interfaces, reachable from anywhere on the network. Only use this on a trusted network.
- `-H 0.0.0.0`: bind the raw port on all network interfaces, reachable from anywhere on the network (subject to your firewall). It does not create a public internet URL; add `--cloudflare` to also publish an internet-reachable `https://*.trycloudflare.com` link even behind a firewall. Only use this on a network you trust.
```bash
unsloth studio -H 0.0.0.0 -p 8888
```
The Cloudflare tunnel is **off by default**: `-H 0.0.0.0` exposes the raw port only, not a public internet URL. Pair the wildcard bind with `--cloudflare` (`unsloth studio -H 0.0.0.0 --cloudflare`) to also publish a public `https://*.trycloudflare.com` link, or prefer `--secure` (above), which keeps the raw port private. `--cloudflare` has no effect on a loopback bind.
The first time Studio is published on a public URL (`--secure` or `--cloudflare`) with the auto-generated admin password still in place, it asks for a new admin password in the terminal (masked input with confirmation) before the public link goes up. Without an attached terminal it warns instead and keeps the bootstrap deadline: Studio shuts down after `UNSLOTH_STUDIO_BOOTSTRAP_TIMEOUT` (default 1 hour) unless the password is changed in the web UI.
For headless setups that cannot answer that prompt, set the initial admin password non-interactively with `--password` (only takes effect when no password is set yet; if one already exists it is a hard error, so rotate later with `unsloth studio reset-password`):
```bash
unsloth studio --secure --password 'your-strong-password' # visible in `ps`/history
UNSLOTH_STUDIO_PASSWORD='your-strong-password' unsloth studio --secure # via env var
printf '%s\n' 'your-strong-password' | unsloth studio --secure --password - # via stdin
```
A literal `--password VALUE` is visible in the process list and shell history, so prefer the `UNSLOTH_STUDIO_PASSWORD` env var or `--password -` (stdin) for automation. This applies to any launch (public or a headless `-H 0.0.0.0` bind), and the password is set in the parent before the server binds, so it never reaches a re-executed child process.
Server-side tools (web search, Python and terminal code execution) run as your user and are on by default. Anyone who can reach the server with the API key can run code on this machine, so keep your API key private and pass `--disable-tools` when exposing Studio.
@ -232,6 +245,14 @@ curl -fsSL https://unsloth.ai/install.sh | UNSLOTH_NO_TORCH=1 sh
$env:UNSLOTH_NO_TORCH=1; irm https://unsloth.ai/install.ps1 | iex
```
Skip the post-install prompt that starts Studio (useful for automated installs):
```bash
curl -fsSL https://unsloth.ai/install.sh | UNSLOTH_SKIP_AUTOSTART=1 sh
```
```powershell
$env:UNSLOTH_SKIP_AUTOSTART=1; irm https://unsloth.ai/install.ps1 | iex
```
Pin the Python version:
```bash
curl -fsSL https://unsloth.ai/install.sh | UNSLOTH_PYTHON=3.12 sh

51
install.ps1
View file

@ -6,6 +6,7 @@
# irm | iex cannot forward arguments, so web installs take options as env vars set
# before the pipe (flags still work via .\install.ps1):
# $env:UNSLOTH_NO_TORCH=1; irm https://unsloth.ai/install.ps1 | iex # skip PyTorch (GGUF-only)
# $env:UNSLOTH_SKIP_AUTOSTART=1; irm https://unsloth.ai/install.ps1 | iex # do not prompt to launch
# $env:UNSLOTH_PYTHON='3.12'; irm https://unsloth.ai/install.ps1 | iex # pin Python version
# $env:UNSLOTH_STUDIO_HOME='C:\path'; irm https://unsloth.ai/install.ps1 | iex
# .\install.ps1 --no-torch # equivalent flag
@ -90,6 +91,7 @@ function Install-UnslothStudio {
if ($TauriMode) {
exit $Code
}
throw $Message
}
# ── Parse flags ──
@ -98,6 +100,7 @@ function Install-UnslothStudio {
$RepoRoot = ""
$TauriMode = $false
$SkipTorch = $false
$SkipAutostart = $false
$ShortcutsOnly = $false
# Launcher browser auto-open: "" = undecided (prompt, else keep existing, else on).
$OpenBrowserPref = ""
@ -134,6 +137,7 @@ function Install-UnslothStudio {
# Env-var equivalent for web installs; an explicit flag still wins.
if ($env:UNSLOTH_NO_TORCH -in @('1', 'true', 'yes', 'on')) { $SkipTorch = $true }
if ($env:UNSLOTH_SKIP_AUTOSTART -in @('1', 'true', 'yes', 'on')) { $SkipAutostart = $true }
# Propagate to child processes so they also respect verbose mode.
# Process-scoped -- does not persist.
@ -473,6 +477,17 @@ function Install-UnslothStudio {
param(
[Parameter(Mandatory = $true)][ScriptBlock]$Command
)
# Installer-pinned index installs (torch) must beat an inherited uv mirror
# (#6898): when the command pins an index, clear every uv index env var so
# it wins, then restore in finally. Other installs keep the user's mirror.
$savedUvIndex = $null
if ($Command.ToString() -match '--default-index') {
$savedUvIndex = @{}
foreach ($n in 'UV_DEFAULT_INDEX', 'UV_INDEX_URL', 'UV_INDEX', 'UV_EXTRA_INDEX_URL') {
$savedUvIndex[$n] = [Environment]::GetEnvironmentVariable($n)
Remove-Item "Env:$n" -ErrorAction SilentlyContinue
}
}
$prevEap = $ErrorActionPreference
$ErrorActionPreference = "Continue"
try {
@ -492,6 +507,7 @@ function Install-UnslothStudio {
return [int]$LASTEXITCODE
} finally {
$ErrorActionPreference = $prevEap
if ($savedUvIndex) { foreach ($n in $savedUvIndex.Keys) { if ($null -ne $savedUvIndex[$n]) { Set-Item "Env:$n" $savedUvIndex[$n] } } }
}
}
@ -2197,7 +2213,7 @@ exit 0
if ($SkipTorch) {
# No-torch: install unsloth + unsloth-zoo with --no-deps, then
# runtime deps (typer, safetensors, transformers, etc.) with --no-deps.
$baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated no-torch)" { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.7.2" "unsloth-zoo>=2026.7.2" }
$baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated no-torch)" { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.7.3" "unsloth-zoo>=2026.7.3" }
if ($baseInstallExit -eq 0) {
# Resolve pydantic WITH deps so pip pins pydantic-core
# to the matching version (no-torch-runtime.txt below
@ -2211,7 +2227,7 @@ exit 0
}
}
} else {
$baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated)" { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.7.2" "unsloth-zoo>=2026.7.2" }
$baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated)" { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.7.3" "unsloth-zoo>=2026.7.3" }
}
if ($baseInstallExit -ne 0) {
Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red
@ -2242,7 +2258,7 @@ exit 0
# ABI-incompatible torchvision/torchaudio on AMD's per-arch index.
$visionSpec = if ($ROCmGfxArch -and $torchvisionFloorMap.ContainsKey($ROCmGfxArch)) { $torchvisionFloorMap[$ROCmGfxArch] } else { "torchvision" }
$audioSpec = if ($ROCmGfxArch -and $torchaudioFloorMap.ContainsKey($ROCmGfxArch)) { $torchaudioFloorMap[$ROCmGfxArch] } else { "torchaudio" }
$torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch (AMD ROCm)" { uv pip install --python $VenvPython --force-reinstall --index-url $ROCmIndexUrl $torchSpec $visionSpec $audioSpec }
$torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch (AMD ROCm)" { uv pip install --python $VenvPython --force-reinstall --default-index $ROCmIndexUrl $torchSpec $visionSpec $audioSpec }
if ($torchInstallExit -ne 0) {
# Transient AMD-index failure: fall back to a CPU base so the install
# still completes; Studio setup retries ROCm afterwards.
@ -2251,7 +2267,7 @@ exit 0
# torch (e.g. 2.10.0+rocm on gfx110X/gfx90a) that still satisfies the CPU
# torch>= range, so without it uv would keep the ROCm build and only swap
# the companions -- a mismatched venv the flavor-repair block won't fix.
$torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch (CPU fallback)" { uv pip install --python $VenvPython --force-reinstall "torch>=2.4,<2.11.0" torchvision torchaudio --index-url $TorchIndexUrl }
$torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch (CPU fallback)" { uv pip install --python $VenvPython --force-reinstall "torch>=2.4,<2.11.0" torchvision torchaudio --default-index $TorchIndexUrl }
if ($torchInstallExit -ne 0) {
Write-Host "[ERROR] Failed to install PyTorch (ROCm and CPU base both failed, exit code $torchInstallExit)" -ForegroundColor Red
return (Exit-InstallFailure "Failed to install PyTorch (exit code $torchInstallExit)" $torchInstallExit)
@ -2265,7 +2281,7 @@ exit 0
} else {
Write-TauriLog "STEP" "Installing PyTorch"
substep "installing PyTorch ($TorchIndexUrl)..."
$torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch" { uv pip install --python $VenvPython "torch>=2.4,<2.11.0" torchvision torchaudio --index-url $TorchIndexUrl }
$torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch" { uv pip install --python $VenvPython "torch>=2.4,<2.11.0" torchvision torchaudio --default-index $TorchIndexUrl }
if ($torchInstallExit -ne 0) {
Write-Host "[ERROR] Failed to install PyTorch (exit code $torchInstallExit)" -ForegroundColor Red
return (Exit-InstallFailure "Failed to install PyTorch (exit code $torchInstallExit)" $torchInstallExit)
@ -2277,7 +2293,7 @@ exit 0
if ($SkipTorch) {
# No-torch: install unsloth + unsloth-zoo with --no-deps, then
# runtime deps (typer, safetensors, transformers, etc.) with --no-deps.
$baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (no-torch)" { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.7.2" "unsloth-zoo>=2026.7.2" }
$baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (no-torch)" { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.7.3" "unsloth-zoo>=2026.7.3" }
if ($baseInstallExit -eq 0) {
# Same pydantic-with-deps trick as the migrated branch.
$baseInstallExit = Invoke-InstallCommandRetry -Label "install pydantic" { uv pip install --python $VenvPython pydantic }
@ -2289,7 +2305,7 @@ exit 0
}
}
} elseif ($StudioLocalInstall) {
$baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (local)" { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.7.2" "unsloth-zoo>=2026.7.2" }
$baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (local)" { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.7.3" "unsloth-zoo>=2026.7.3" }
} else {
$baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth" { uv pip install --python $VenvPython --upgrade-package unsloth -- "$PackageName" }
}
@ -2317,7 +2333,7 @@ exit 0
Write-TauriLog "STEP" "Installing unsloth"
substep "installing unsloth (this may take a few minutes)..."
if ($StudioLocalInstall) {
$baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (auto torch backend)" { uv pip install --python $VenvPython "unsloth-zoo>=2026.7.2" "unsloth>=2026.7.2" --torch-backend=auto }
$baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (auto torch backend)" { uv pip install --python $VenvPython "unsloth-zoo>=2026.7.3" "unsloth>=2026.7.3" --torch-backend=auto }
if ($baseInstallExit -ne 0) {
Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red
return (Exit-InstallFailure "Failed to install unsloth (exit code $baseInstallExit)" $baseInstallExit)
@ -2348,7 +2364,7 @@ exit 0
# keeps a stale torch==X+cpu against a CUDA index and setup.ps1 then loops on
# "torch cpu != required cuXXX". Reinstall the right triplet when a GPU build is
# expected: CUDA from $TorchIndexUrl, ROCm from $ROCmIndexUrl (repo.amd.com gfx*
# is a PEP 503 index uv resolves via --index-url, same URL the fresh ROCm install
# is a PEP 503 index uv resolves via --default-index, same URL the fresh ROCm install
# above uses). --no-torch / CPU-only hosts (expected cpu) are no-ops.
if (-not $SkipTorch) {
$expectedTorchTag = Get-ExpectedTorchFlavorTag -TorchIndexUrl $TorchIndexUrl -ROCmIndexUrl $ROCmIndexUrl
@ -2364,7 +2380,7 @@ exit 0
$visionSpec = if ($ROCmGfxArch -and $torchvisionFloorMap.ContainsKey($ROCmGfxArch)) { $torchvisionFloorMap[$ROCmGfxArch] } else { "torchvision" }
$audioSpec = if ($ROCmGfxArch -and $torchaudioFloorMap.ContainsKey($ROCmGfxArch)) { $torchaudioFloorMap[$ROCmGfxArch] } else { "torchaudio" }
substep "PyTorch flavor mismatch (installed $installedTorchTag, need ROCm) -- reinstalling correct build..." "Yellow"
$torchFixExit = Invoke-InstallCommand { uv pip install --python $VenvPython --force-reinstall --index-url $ROCmIndexUrl $rocmSpec $visionSpec $audioSpec }
$torchFixExit = Invoke-InstallCommand { uv pip install --python $VenvPython --force-reinstall --default-index $ROCmIndexUrl $rocmSpec $visionSpec $audioSpec }
if ($torchFixExit -ne 0) {
Write-Host "[ERROR] Failed to reinstall PyTorch with the correct ROCm build (exit code $torchFixExit)" -ForegroundColor Red
return (Exit-InstallFailure "Failed to reinstall PyTorch (ROCm) (exit code $torchFixExit)" $torchFixExit)
@ -2373,7 +2389,7 @@ exit 0
} elseif ($expectedTorchTag -ne 'rocm') {
# CUDA: stale +cpu (or wrong cuXXX) against a CUDA index -> reinstall triplet.
substep "PyTorch flavor mismatch (installed $installedTorchTag, need $expectedTorchTag) -- reinstalling correct build..." "Yellow"
$torchFixExit = Invoke-InstallCommand { uv pip install --python $VenvPython "torch>=2.4,<2.11.0" torchvision torchaudio --index-url $TorchIndexUrl --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio }
$torchFixExit = Invoke-InstallCommand { uv pip install --python $VenvPython "torch>=2.4,<2.11.0" torchvision torchaudio --default-index $TorchIndexUrl --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio }
if ($torchFixExit -ne 0) {
Write-Host "[ERROR] Failed to reinstall PyTorch with the correct CUDA build (exit code $torchFixExit)" -ForegroundColor Red
return (Exit-InstallFailure "Failed to reinstall PyTorch ($expectedTorchTag) (exit code $torchFixExit)" $torchFixExit)
@ -2668,9 +2684,10 @@ exit 0
# Diagnostic only; never block install on a probe failure.
}
# In interactive terminals, ask the user before starting Studio.
# In interactive terminals, ask the user before starting Studio unless the
# caller explicitly disabled the post-install prompt.
# In non-interactive environments (CI, Docker) just print instructions.
$IsInteractive = [Environment]::UserInteractive -and (-not [Console]::IsInputRedirected)
$IsInteractive = (-not $SkipAutostart) -and [Environment]::UserInteractive -and (-not [Console]::IsInputRedirected)
# Background watcher for the foreground launch below: once the server is
# healthy, open the browser per the persisted preference (mirrors the
# desktop launcher). Guarded by the per-install root id so a different
@ -2710,8 +2727,8 @@ exit 0
} else {
step "launch" "to start later, run:"
substep "unsloth studio -p 8888"
substep "(add -H 0.0.0.0 to allow network / cloud access)"
substep "(add --secure for a public Cloudflare HTTPS link; anyone with the API key can run code)"
substep "(add -H 0.0.0.0 for LAN / cloud access; exposes the raw port only, not a public URL)"
substep "(add -H 0.0.0.0 --cloudflare for a public Cloudflare HTTPS link, or --secure to keep the raw port private; anyone with the API key can run code)"
Write-Host ""
}
} else {
@ -2731,8 +2748,8 @@ exit 0
substep "& $_actLiteral"
substep "unsloth studio -p 8888"
}
substep "(add -H 0.0.0.0 to allow network / cloud access)"
substep "(add --secure for a public Cloudflare HTTPS link; anyone with the API key can run code)"
substep "(add -H 0.0.0.0 for LAN / cloud access; exposes the raw port only, not a public URL)"
substep "(add -H 0.0.0.0 --cloudflare for a public Cloudflare HTTPS link, or --secure to keep the raw port private; anyone with the API key can run code)"
Write-Host ""
}
}

59
install.sh
View file

@ -8,8 +8,9 @@
#
# Piped installs take options as env vars after the pipe (a bare `| sh --no-torch`
# makes sh reject --no-torch as its own option). Flags still work via ./install.sh:
# curl -fsSL https://unsloth.ai/install.sh | UNSLOTH_NO_TORCH=1 sh # skip PyTorch (GGUF-only)
# curl -fsSL https://unsloth.ai/install.sh | UNSLOTH_PYTHON=3.12 sh # pin Python version
# curl -fsSL https://unsloth.ai/install.sh | UNSLOTH_NO_TORCH=1 sh # skip PyTorch (GGUF-only)
# curl -fsSL https://unsloth.ai/install.sh | UNSLOTH_SKIP_AUTOSTART=1 sh # do not prompt to launch
# curl -fsSL https://unsloth.ai/install.sh | UNSLOTH_PYTHON=3.12 sh # pin Python version
# curl -fsSL https://unsloth.ai/install.sh | UNSLOTH_STUDIO_HOME=/abs/path sh
# Equivalent flags: ./install.sh --no-torch --python 3.12 (or pipe them: sh -s -- --no-torch)
# ./install.sh --no-browser: launcher starts the server without opening the browser.
@ -50,6 +51,7 @@ PACKAGE_NAME="unsloth"
TAURI_MODE=false
_USER_PYTHON=""
_NO_TORCH_FLAG=false
_SKIP_AUTOSTART=false
_VERBOSE=false
_SHORTCUTS_ONLY=false
# Launcher browser auto-open: "" = undecided (prompt, else keep existing, else on).
@ -93,6 +95,7 @@ done
# Env-var equivalents for piped installs; an explicit flag still wins.
case "${UNSLOTH_NO_TORCH:-}" in 1|true|TRUE|yes|YES|on|ON) _NO_TORCH_FLAG=true ;; esac
case "${UNSLOTH_SKIP_AUTOSTART:-}" in 1|true|TRUE|yes|YES|on|ON) _SKIP_AUTOSTART=true ;; esac
[ -z "$_USER_PYTHON" ] && [ -n "${UNSLOTH_PYTHON:-}" ] && _USER_PYTHON="$UNSLOTH_PYTHON"
if [ "$_VERBOSE" = true ]; then
@ -164,6 +167,12 @@ run_maybe_quiet() {
run_install_cmd() {
_label="$1"
shift
# Installer-pinned index installs (torch) must beat an inherited uv mirror
# (#6898): when we pass --default-index, neutralize every uv index env var so
# the pinned index wins. Other installs keep the user's mirror.
case " $* " in
*" --default-index "*) set -- env -u UV_DEFAULT_INDEX -u UV_INDEX_URL -u UV_INDEX -u UV_EXTRA_INDEX_URL "$@" ;;
esac
if _is_verbose; then
"$@" && return 0
_rc=$?
@ -1664,6 +1673,7 @@ _maybe_reroute_strixhalo_to_2404() {
# Forward explicit ROCm-bootstrap consent (e.g. Tauri) so the child auto-enables the
# GPU instead of falling back to the desktop-app prompt path.
[ "${UNSLOTH_ROCM_WSL_AUTO:-0}" = "1" ] && _rr_exports="$_rr_exports; export UNSLOTH_ROCM_WSL_AUTO=1"
[ "$_SKIP_AUTOSTART" = true ] && _rr_exports="$_rr_exports; export UNSLOTH_SKIP_AUTOSTART=1"
_rr_args=""
[ "$PACKAGE_NAME" != "unsloth" ] && _rr_args="$_rr_args --package $(_rr_q "$PACKAGE_NAME")"
[ -n "$_USER_PYTHON" ] && _rr_args="$_rr_args --python $(_rr_q "$_USER_PYTHON")"
@ -2233,9 +2243,9 @@ _expected_torch_flavor_tag() {
esac
}
# Whether index ($1) supports a plain --index-url reinstall. pytorch.org cuXXX /
# Whether index ($1) supports a plain --default-index reinstall. pytorch.org cuXXX /
# rocmX.Y AND the repo.amd.com gfx* indexes are all PEP 503 simple indexes that uv
# resolves (torch + every transitive dep) via --index-url -- the same URLs the
# resolves (torch + every transitive dep) via --default-index -- the same URLs the
# fresh-install paths above already use -- so a stale wheel is auto-repairable.
# Unknown/odd-mirror leaves -> no, so we warn rather than risk a wrong reinstall.
_torch_index_repairable() {
@ -2749,7 +2759,7 @@ if [ "$_MIGRATED" = true ]; then
# to prevent transitive torch resolution.
run_install_cmd_retry "install unsloth (migrated no-torch)" uv pip install --python "$_VENV_PY" --no-deps \
--reinstall-package unsloth --reinstall-package unsloth-zoo \
"unsloth>=2026.7.2" "unsloth-zoo>=2026.7.2"
"unsloth>=2026.7.3" "unsloth-zoo>=2026.7.3"
# Resolve pydantic WITH deps so pip pins pydantic-core to the
# matching version (no-torch-runtime.txt below is --no-deps).
# All transitive deps are torch-free.
@ -2764,7 +2774,7 @@ if [ "$_MIGRATED" = true ]; then
# overrides file, so UV_OVERRIDE is unset and this positional is the only cover.
run_install_cmd_retry "install unsloth (migrated)" uv pip install --python "$_VENV_PY" \
--reinstall-package unsloth --reinstall-package unsloth-zoo \
"unsloth>=2026.7.2" "unsloth-zoo>=2026.7.2" ${_MLX_LM_EXCLUDE_ARG:-}
"unsloth>=2026.7.3" "unsloth-zoo>=2026.7.3" ${_MLX_LM_EXCLUDE_ARG:-}
fi
if [ "$STUDIO_LOCAL_INSTALL" = true ]; then
substep "overlaying local repo (editable)..."
@ -2787,7 +2797,7 @@ if [ "$_MIGRATED" = true ]; then
substep "repairing ROCm torch (overwritten by dependency resolution)..."
run_install_cmd_retry "repair ROCm torch" uv pip install --python "$_VENV_PY" \
"$TORCH_CONSTRAINT" torchvision torchaudio \
--index-url "$TORCH_INDEX_URL" \
--default-index "$TORCH_INDEX_URL" \
--force-reinstall
fi
;;
@ -2913,7 +2923,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
substep "[WARN] Radeon repo lacks a compatible wheel set for this Python; falling back to ROCm index ($TORCH_INDEX_URL)" "$C_WARN"
run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" \
"$TORCH_CONSTRAINT" torchvision torchaudio \
--index-url "$TORCH_INDEX_URL"
--default-index "$TORCH_INDEX_URL"
else
substep "installing PyTorch from Radeon repo (${_RADEON_BASE_URL})..."
# Pass explicit wheel URLs so the matched trio is
@ -2936,18 +2946,18 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
substep "[WARN] Radeon repo unavailable; falling back to ROCm index ($TORCH_INDEX_URL)" "$C_WARN"
run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" \
"$TORCH_CONSTRAINT" torchvision torchaudio \
--index-url "$TORCH_INDEX_URL"
--default-index "$TORCH_INDEX_URL"
fi
else
substep "[WARN] Radeon GPU detected but could not detect full ROCm version; falling back to ROCm index" "$C_WARN"
run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" \
"$TORCH_CONSTRAINT" torchvision torchaudio \
--index-url "$TORCH_INDEX_URL"
--default-index "$TORCH_INDEX_URL"
fi
else
substep "installing PyTorch ($TORCH_INDEX_URL)..."
run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" "$TORCH_CONSTRAINT" torchvision torchaudio \
--index-url "$TORCH_INDEX_URL"
--default-index "$TORCH_INDEX_URL"
fi
# AMD ROCm: install bitsandbytes (once, after torch, for all ROCm paths).
# Gate on SKIP_TORCH=false so a user running with --no-torch on a ROCm
@ -2968,7 +2978,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
# runtime deps (typer, safetensors, transformers, etc.) with --no-deps.
run_install_cmd_retry "install unsloth (no-torch)" uv pip install --python "$_VENV_PY" --no-deps \
--upgrade-package unsloth --upgrade-package unsloth-zoo \
"unsloth>=2026.7.2" "unsloth-zoo>=2026.7.2"
"unsloth>=2026.7.3" "unsloth-zoo>=2026.7.3"
# Same pydantic-with-deps trick as the migrated branch.
run_install_cmd_retry "install pydantic (with deps for compatible core)" \
uv pip install --python "$_VENV_PY" pydantic
@ -2986,7 +2996,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
fi
elif [ "$STUDIO_LOCAL_INSTALL" = true ]; then
run_install_cmd_retry "install unsloth (local)" uv pip install --python "$_VENV_PY" \
--upgrade-package unsloth "unsloth>=2026.7.2" "unsloth-zoo>=2026.7.2"
--upgrade-package unsloth "unsloth>=2026.7.3" "unsloth-zoo>=2026.7.3"
substep "overlaying local repo (editable)..."
run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps
substep "overlaying unsloth-zoo from git main..."
@ -3007,7 +3017,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
substep "repairing ROCm torch (overwritten by dependency resolution)..."
run_install_cmd_retry "repair ROCm torch" uv pip install --python "$_VENV_PY" \
"$TORCH_CONSTRAINT" torchvision torchaudio \
--index-url "$TORCH_INDEX_URL" \
--default-index "$TORCH_INDEX_URL" \
--force-reinstall
fi
;;
@ -3018,7 +3028,7 @@ else
tauri_log "STEP" "Installing Unsloth"
substep "installing unsloth (this may take a few minutes)..."
if [ "$STUDIO_LOCAL_INSTALL" = true ]; then
run_install_cmd_retry "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" "unsloth-zoo>=2026.7.2" "unsloth>=2026.7.2" --torch-backend=auto
run_install_cmd_retry "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" "unsloth-zoo>=2026.7.3" "unsloth>=2026.7.3" --torch-backend=auto
substep "overlaying local repo (editable)..."
run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps
substep "overlaying unsloth-zoo from git main..."
@ -3042,14 +3052,14 @@ if [ "$SKIP_TORCH" = false ] && [ -n "${TORCH_INDEX_URL:-}" ]; then
_installed_torch_ver=$("$_VENV_PY" -c "import torch; print(torch.__version__)" 2>/dev/null || true)
_installed_torch_tag=""
[ -n "$_installed_torch_ver" ] && _installed_torch_tag=$(_torch_flavor_tag "$_installed_torch_ver")
# Repair when flavor is wrong AND the index is plain --index-url reinstallable
# Repair when flavor is wrong AND the index is plain --default-index reinstallable
# (cuXXX / rocmX.Y / repo.amd.com gfx*); an unknown mirror leaf -> warn only.
if [ -n "$_installed_torch_tag" ] && [ "$_installed_torch_tag" != "$_expected_torch_tag" ] \
&& [ "$(_torch_index_repairable "$TORCH_INDEX_URL")" = "yes" ]; then
substep "PyTorch flavor mismatch (installed $_installed_torch_tag, need $_expected_torch_tag) -- reinstalling correct build..."
run_install_cmd "reinstall PyTorch ($_expected_torch_tag)" uv pip install --python "$_VENV_PY" \
"$TORCH_CONSTRAINT" torchvision torchaudio \
--index-url "$TORCH_INDEX_URL" \
--default-index "$TORCH_INDEX_URL" \
--reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio
_installed_torch_ver=$("$_VENV_PY" -c "import torch; print(torch.__version__)" 2>/dev/null || true)
_installed_torch_tag=""
@ -3060,7 +3070,7 @@ if [ "$SKIP_TORCH" = false ] && [ -n "${TORCH_INDEX_URL:-}" ]; then
substep "[WARN] PyTorch is CPU-only but a $_expected_torch_tag GPU build was expected for this machine." "$C_WARN"
substep "[WARN] Training and GPU inference will run on CPU until this is fixed." "$C_WARN"
substep "[WARN] Re-run this installer, or reinstall the GPU build manually:" "$C_WARN"
substep "[WARN] uv pip install --python \"$_VENV_PY\" \"$TORCH_CONSTRAINT\" torchvision torchaudio --index-url $TORCH_INDEX_URL --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio" "$C_WARN"
substep "[WARN] uv pip install --python \"$_VENV_PY\" \"$TORCH_CONSTRAINT\" torchvision torchaudio --default-index $TORCH_INDEX_URL --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio" "$C_WARN"
fi
fi
fi
@ -3328,9 +3338,10 @@ _post_install_browser_watch() {
) &
}
# In interactive terminals, ask the user before starting Studio.
# In interactive terminals, ask the user before starting Studio unless the
# caller explicitly disabled the post-install prompt.
# In non-interactive environments (Docker, CI, cloud-init) just print instructions.
if [ -t 1 ]; then
if [ "$_SKIP_AUTOSTART" != true ] && [ -t 1 ]; then
echo ""
printf " Start Unsloth Studio now? [Y/n] "
# No readable answer (closed/EOF tty) defaults to no; Enter is still yes.
@ -3371,8 +3382,8 @@ if [ -t 1 ]; then
*)
step "launch" "to start later, run:"
substep "unsloth studio -p 8888"
substep "(add -H 0.0.0.0 to allow network / cloud access)"
substep "(add --secure for a public Cloudflare HTTPS link; anyone with the API key can run code)"
substep "(add -H 0.0.0.0 for LAN / cloud access; exposes the raw port only, not a public URL)"
substep "(add -H 0.0.0.0 --cloudflare for a public Cloudflare HTTPS link, or --secure to keep the raw port private; anyone with the API key can run code)"
echo ""
;;
esac
@ -3393,7 +3404,7 @@ else
substep "source $_li_act_q"
substep "unsloth studio -p 8888"
fi
substep "(add -H 0.0.0.0 to allow network / cloud access)"
substep "(add --secure for a public Cloudflare HTTPS link; anyone with the API key can run code)"
substep "(add -H 0.0.0.0 for LAN / cloud access; exposes the raw port only, not a public URL)"
substep "(add -H 0.0.0.0 --cloudflare for a public Cloudflare HTTPS link, or --secure to keep the raw port private; anyone with the API key can run code)"
echo ""
fi

7
pyproject.toml
View file

@ -42,6 +42,7 @@ version = {attr = "unsloth.models._utils.__version__"}
include-package-data = true
[tool.setuptools.package-data]
unsloth_cli = ["codex_fallback_prompt.md"]
studio = [
"*.sh",
"*.ps1",
@ -73,7 +74,7 @@ triton = [
]
huggingfacenotorch = [
"unsloth_zoo>=2026.7.2",
"unsloth_zoo>=2026.7.3",
"wheel>=0.42.0",
"packaging",
"numpy",
@ -94,7 +95,7 @@ huggingfacenotorch = [
]
huggingface = [
"unsloth[huggingfacenotorch]",
"unsloth_zoo>=2026.7.2",
"unsloth_zoo>=2026.7.3",
"torchvision",
"unsloth[triton]",
]
@ -579,7 +580,7 @@ colab-ampere-torch220 = [
"flash-attn>=2.6.3 ; ('linux' in sys_platform)",
]
colab-new = [
"unsloth_zoo>=2026.7.2",
"unsloth_zoo>=2026.7.3",
"packaging",
"tyro",
"transformers>=4.51.3,!=4.52.0,!=4.52.1,!=4.52.2,!=4.52.3,!=4.53.0,!=4.54.0,!=4.55.0,!=4.55.1,!=4.57.0,!=4.57.4,!=4.57.5,!=5.0.0,!=5.1.0,<=5.5.0",

326
scripts/scan_packages_baseline.json
View file

File diff suppressed because one or more lines are too long

34
studio/MCP.md Normal file
View file

@ -0,0 +1,34 @@
# Unsloth Studio MCP server
Studio can expose a local MCP server so an MCP client can inspect models and
GPU state, validate recipes, start or stop training, inspect recipe output, and
export a loaded model.
The server is disabled by default. Enable it for a local Studio process with:
```bash
UNSLOTH_STUDIO_ENABLE_MCP=1 \
UNSLOTH_STUDIO_MCP_TOKEN='use-a-local-secret' \
unsloth studio
```
The endpoint is `http://127.0.0.1:8888/mcp/` when Studio uses its default port
(a request to `/mcp` redirects to the canonical `/mcp/`). Use the actual Studio
port when it is configured differently.
The high-impact tools are:
- `studio_status` and `list_local_models` for discovery
- `get_training_status`, `start_training`, `stop_training`, and `list_training_runs`
- `validate_recipe`, `get_recipe_job_status`, and `get_recipe_job_dataset`
- `load_checkpoint` and `export_gguf`
`start_training` accepts the same fields as the Studio `TrainingStartRequest`.
The request is validated by the existing Pydantic model before a subprocess is
started. Export paths use the existing Studio validation as well.
The endpoint always requires `UNSLOTH_STUDIO_MCP_TOKEN` and checks an exact
Bearer token for both HTTP and WebSocket connections. Keep it on localhost
unless the deployment has an authenticated reverse proxy. The MCP endpoint is
intentionally opt-in because tools can consume GPU memory, write model
artifacts, and stop active work.

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

@ -18,6 +18,10 @@ from utils.paths import auth_db_path, ensure_dir
DB_PATH = auth_db_path()
DEFAULT_ADMIN_USERNAME = "unsloth"
# Single source for the password policy; models/auth.py ChangePasswordRequest
# and the terminal prompt both enforce it. Keep the unsloth_cli mirror in sync.
MIN_PASSWORD_LENGTH = 8
# Plaintext bootstrap password file beside auth.db, deleted on first password
# change so the credential never lingers on disk.
_BOOTSTRAP_PW_PATH = DB_PATH.parent / ".bootstrap_password"
@ -79,11 +83,42 @@ def _load_bootstrap_password() -> Optional[str]:
def clear_bootstrap_password() -> None:
"""Delete the persisted bootstrap password file (called after password change)."""
"""Delete the persisted bootstrap password file (after a password change).
Best-effort: the new hash is already committed, so a locked/undeletable file
(Windows AV, read-only auth dir) must not fail the change.
"""
global _bootstrap_password
_bootstrap_password = None
if _BOOTSTRAP_PW_PATH.is_file():
_BOOTSTRAP_PW_PATH.unlink(missing_ok = True)
try:
_BOOTSTRAP_PW_PATH.unlink(missing_ok = True)
except OSError as e:
# Removal failed (Windows AV, read-only auth dir). The hash is already
# committed, so don't fail the change -- but truncate the file so its
# stale plaintext can't be re-seeded by generate_bootstrap_password()
# if a later reset-password deletes auth.db and re-validates it.
try:
_BOOTSTRAP_PW_PATH.write_text("")
cleared = True
except OSError:
cleared = False
import sys
if cleared:
message = (
f"Warning: could not delete {_BOOTSTRAP_PW_PATH.name} ({e}); "
"cleared its contents so the old bootstrap password cannot be reused."
)
else:
# Neither removed nor truncated: stale plaintext is still on disk
# and would be reused if auth.db is reset. Don't claim otherwise.
message = (
f"Warning: could not delete or clear {_BOOTSTRAP_PW_PATH.name} ({e}); "
"its old bootstrap password is still on disk. Remove it manually to "
"prevent reuse after a reset."
)
print(message, file = sys.stderr, flush = True)
def _hash_token(token: str) -> str:
@ -547,8 +582,18 @@ def ensure_default_admin() -> bool:
return False
def update_password(username: str, new_password: str) -> bool:
"""Update password, clear first-login requirement, rotate JWT secret."""
def update_password(
username: str,
new_password: str,
*,
revoke_refresh_tokens: bool = False,
) -> bool:
"""Update password, clear first-login requirement, rotate JWT secret.
``revoke_refresh_tokens`` deletes the user's refresh tokens in the SAME
transaction: a separate delete could fail after the password commit and
leave a pre-change token still able to mint access tokens.
"""
from .hashing import hash_password
salt, pwd_hash = hash_password(new_password)
@ -563,6 +608,8 @@ def update_password(username: str, new_password: str) -> bool:
""",
(salt, pwd_hash, jwt_secret, username),
)
if revoke_refresh_tokens and cursor.rowcount > 0:
conn.execute("DELETE FROM refresh_tokens WHERE username = ?", (username,))
conn.commit()
if cursor.rowcount > 0:
clear_bootstrap_password()

282
studio/backend/auth/terminal_prompt.py Normal file
View file

@ -0,0 +1,282 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Interactive terminal prompt that forces a bootstrap password change before
Studio is exposed on a public Cloudflare URL (``--secure`` / ``--cloudflare``).
Masked input echoes one ``*`` per keystroke (unlike ``getpass``). Works on
Windows (``msvcrt``) and Linux/macOS (``termios``). All output goes to stderr so
redirected stdout never swallows the prompt.
Mirrored for the CLI at ``unsloth_cli/commands/_password_prompt.py`` (the CLI
cannot import the Studio backend package); keep the two in sync.
"""
from __future__ import annotations
import os
import sys
from typing import Callable, TextIO
_CTRL_C = "\x03"
_CTRL_D = "\x04"
_CTRL_Z = "\x1a"
_BACKSPACES = ("\x7f", "\x08")
_SUBMITS = ("\r", "\n")
# Env var that supplies the initial admin password non-interactively (mirror in
# unsloth_cli/commands/_password_prompt.py). Keep the name in sync.
SUPPLIED_PASSWORD_ENV = "UNSLOTH_STUDIO_PASSWORD"
def _getch_windows() -> str: # pragma: no cover - exercised via fake on Linux CI
import msvcrt
ch = msvcrt.getwch()
# Function/arrow keys arrive as a two-wchar \x00/\xe0 sequence; consume the
# second half and report a no-op control char.
if ch in ("\x00", "\xe0"):
msvcrt.getwch()
return "\x00"
return ch
class _RestoreTtyOnSignals:
"""Restore terminal attrs if SIGTERM/SIGHUP kills the prompt mid-read.
A finally block can't run when a signal terminates the process, leaving the
shared terminal in cbreak/no-echo. Best-effort: no-op off the main thread or
where the signals are absent.
"""
def __init__(self, fd: int, old_attrs) -> None:
self._fd = fd
self._old_attrs = old_attrs
self._previous: list = []
def __enter__(self) -> "_RestoreTtyOnSignals":
import signal
import termios
def _restore_and_reraise(signum, frame):
termios.tcsetattr(self._fd, termios.TCSADRAIN, self._old_attrs)
signal.signal(signum, signal.SIG_DFL)
signal.raise_signal(signum)
for name in ("SIGTERM", "SIGHUP"):
sig = getattr(signal, name, None)
if sig is None:
continue
try:
self._previous.append((sig, signal.signal(sig, _restore_and_reraise)))
except (ValueError, OSError): # non-main thread / unsupported
pass
return self
def __exit__(self, *exc) -> None:
import signal
for sig, previous in self._previous:
try:
signal.signal(sig, previous)
except (ValueError, OSError):
pass
class _prompt_raw_mode:
"""Hold cbreak + cleared ISIG (no echo) on stdin for the WHOLE prompt line,
restoring when the line finishes (and on SIGTERM/SIGHUP).
Echo must never re-enable mid-line: cbreak echoes on receipt, so a keystroke
arriving while echo is on would appear in cleartext. One cbreak block for the
whole line closes that window. No-op when stdin is not a real terminal, so
the _getch seam can be faked in tests.
"""
def __enter__(self) -> "_prompt_raw_mode":
self._fd = None
self._old_attrs = None
self._signals = None
try:
import termios
import tty
except ImportError: # non-POSIX (Windows uses msvcrt, no mode to hold)
return self
try:
fd = sys.stdin.fileno()
old_attrs = termios.tcgetattr(fd)
except (AttributeError, ValueError, OSError, termios.error):
return self # redirected / captured stdin (tests): nothing to hold
self._fd = fd
self._old_attrs = old_attrs
self._signals = _RestoreTtyOnSignals(fd, old_attrs)
self._signals.__enter__()
# cbreak (not raw) keeps output post-processing while disabling echo/line
# buffering. It leaves ISIG on, so clear it and surface Ctrl-C as \x03 to
# the caller loop, which restores the tty itself.
tty.setcbreak(fd, termios.TCSADRAIN)
new_attrs = termios.tcgetattr(fd)
new_attrs[3] &= ~termios.ISIG
termios.tcsetattr(fd, termios.TCSADRAIN, new_attrs)
return self
def __exit__(self, *exc) -> None:
if self._old_attrs is None:
return
import termios
try:
termios.tcsetattr(self._fd, termios.TCSADRAIN, self._old_attrs)
finally:
if self._signals is not None:
self._signals.__exit__(*exc)
def _getch_posix() -> str: # pragma: no cover - needs a real tty
# Terminal already in cbreak+no-echo for the whole line (_prompt_raw_mode),
# so just read. Byte-at-a-time incremental decode so a multi-byte UTF-8 char
# straddling a read boundary isn't dropped.
import codecs
fd = sys.stdin.fileno()
decoder = codecs.getincrementaldecoder(sys.stdin.encoding or "utf-8")("replace")
while True:
b = os.read(fd, 1)
if not b:
return "" # stream EOF; caller raises EOFError
ch = decoder.decode(b)
if ch:
return ch
_getch: Callable[[], str] = _getch_windows if os.name == "nt" else _getch_posix
def _read_password(prompt: str, *, out: "TextIO | None" = None) -> str:
"""Read one masked line: echo ``*`` per char, support backspace editing.
Raises KeyboardInterrupt on Ctrl-C and EOFError on Ctrl-D/Ctrl-Z with an
empty buffer; the terminal is restored on every exit path.
"""
if out is None:
out = sys.stderr
out.write(prompt)
out.flush()
chars: list[str] = []
with _prompt_raw_mode():
while True:
key = _getch()
if key == "": # stream ended mid-line: abort, don't submit a partial
out.write("\n")
out.flush()
raise EOFError
for ch in key: # a paste can deliver several chars per read
if ch in _SUBMITS:
out.write("\n")
out.flush()
return "".join(chars)
if ch == _CTRL_C:
out.write("\n")
out.flush()
raise KeyboardInterrupt
if ch in (_CTRL_D, _CTRL_Z):
if not chars:
out.write("\n")
out.flush()
raise EOFError
continue # ignore mid-input
if ch in _BACKSPACES:
if chars:
chars.pop()
out.write("\b \b")
out.flush()
continue
if ch < " ": # other control characters (tab, escape, ...)
continue
chars.append(ch)
out.write("*")
out.flush()
def should_prompt_password_change(
*, tunnel_will_start: bool, requires_change: bool, stdin_isatty: bool, stderr_isatty: bool
) -> bool:
"""Whether to block startup on an interactive terminal password change.
True only when the tunnel is actually about to start, the admin still has
the seeded password, and both stdin and stderr are real terminals (headless
launches keep the bootstrap-timeout protection instead of hanging).
"""
return tunnel_will_start and requires_change and stdin_isatty and stderr_isatty
def prompt_for_password_change(
*,
min_length: int,
is_current_password: Callable[[str], bool],
apply_change: Callable[[str], None],
username: str = "unsloth",
out: "TextIO | None" = None,
) -> bool:
"""Force a new admin password before public exposure; True on success.
Loops until a valid, confirmed password is committed via ``apply_change``.
Ctrl-C / EOF returns False; the caller must then abort the launch.
"""
if out is None:
out = sys.stderr
out.write(
"\n"
"Unsloth Studio will be exposed on the public internet, so set a\n"
"password now. Ctrl+C to abort.\n\n"
)
out.flush()
try:
while True:
new_password = _read_password("New password: ", out = out)
if len(new_password) < min_length:
out.write(f"Password must be at least {min_length} characters; try again.\n")
out.flush()
continue
if is_current_password(new_password):
out.write(
"New password must differ from the current bootstrap password; try again.\n"
)
out.flush()
continue
confirmation = _read_password("Confirm new password: ", out = out)
if confirmation != new_password:
out.write("Passwords do not match; try again.\n")
out.flush()
continue
apply_change(new_password)
out.write(f"Password updated for '{username}'.\n")
out.flush()
return True
except (KeyboardInterrupt, EOFError):
out.write("Password change aborted; not exposing Studio.\n")
out.flush()
return False
def resolve_supplied_password(cli_value: "str | None", out: "TextIO | None" = None) -> "str | None":
"""Resolve a non-interactive initial admin password, or None if unset.
Precedence: an explicit ``--password`` (literal ``-`` reads a line from
stdin), then the ``UNSLOTH_STUDIO_PASSWORD`` env var; empty/omitted means off.
A literal argv value is visible in the process list, so a note points at the
env var or stdin instead. Mirror of the CLI helper -- keep the two in sync.
"""
if out is None:
out = sys.stderr
if cli_value == "-":
line = sys.stdin.readline()
if not line:
return None
return line.rstrip("\r\n") or None
if cli_value:
out.write(
"Note: --password is visible in the process list and shell history; "
f"prefer {SUPPLIED_PASSWORD_ENV} or --password - (stdin).\n"
)
out.flush()
return cli_value
return os.environ.get(SUPPLIED_PASSWORD_ENV) or None

4
studio/backend/colab.py
View file

@ -323,8 +323,8 @@ def start(port: int = 8888, *, cloudflare: bool = False):
logger.info(" Starting server...")
try:
# cloudflare=False: this helper owns the tunnel. run_server's default True
# would tunnel this 0.0.0.0 bind if Colab detection fails, breaking the opt-out.
# cloudflare=False: this helper owns the tunnel (Colab's own
# start(cloudflare=...) drives it), so pin it off explicitly.
app = run_server(
host = "0.0.0.0",
port = port,

7
studio/backend/core/data_recipe/service.py
View file

@ -9,6 +9,8 @@ import os
from pathlib import Path
from typing import Any
from utils.paths import recipe_datasets_root
from .jsonable import to_jsonable
from .local_callable_validators import (
register_oxc_local_callable_validators,
@ -277,6 +279,11 @@ def create_data_designer(recipe: dict[str, Any], *, artifact_path: str | None =
_apply_data_designer_image_context_patch()
from data_designer.interface.data_designer import DataDesigner # pyright: ignore[reportMissingImports]
if artifact_path is None:
# DataDesigner defaults to cwd/artifacts; packaged Studio can run with
# cwd=/, so keep default callers on Studio's writable recipe artifact root.
artifact_path = str(recipe_datasets_root())
recipe = _strip_frontend_model_config_metadata(recipe)
model_providers = build_model_providers(recipe)
_validate_recipe_runtime_support(recipe, model_providers)

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

@ -132,6 +132,11 @@ class ExportOrchestrator:
"""True while an export / load / cleanup command is running."""
return self._export_active
def is_worker_alive(self) -> bool:
"""True while the persistent export subprocess is running (op or idle)."""
proc = self._proc
return proc is not None and proc.is_alive()
def was_cancelled(self) -> bool:
"""True if the in-flight (or most recent) run was cancelled by the user."""
return self._cancel_requested
@ -204,6 +209,23 @@ class ExportOrchestrator:
def _spawn_subprocess(self, config: dict) -> None:
"""Spawn a new export subprocess."""
# Last-resort recheck for spawns outside an active op. Inside an op, _export_active is set and
# load_checkpoint already rechecked, so a reservation here is an install about to observe
# is_export_active() and abort; raising would kill this export for an install that never proceeds.
from utils.transformers_version import sidecar_swap_in_progress
from utils.transformers_version import sidecar_swap_kind
_swap_kind = sidecar_swap_kind()
# Inside an active op an INSTALL reservation is about to abort on the
# is_export_active check, but a lazy REPAIR has no such check and can be
# rebuilding the sidecar right now, so it must always refuse the spawn.
if _swap_kind == "repair" or (_swap_kind is not None and not self._export_active):
from utils.transformers_version import SidecarSwapInProgress
raise SidecarSwapInProgress(
"A transformers installation is replacing the latest sidecar; "
"retry when it completes."
)
from utils.native_path_leases import (
native_path_secret_removed_for_child_start,
run_without_native_path_secret,
@ -231,11 +253,17 @@ class ExportOrchestrator:
adopt_pid(self._proc.pid) # bind to parent lifetime (Windows job / sweep)
logger.info("Export subprocess started (pid=%s)", self._proc.pid)
def _shutdown_subprocess(self, timeout: float = 10.0) -> None:
"""Gracefully shut down the export subprocess."""
def _shutdown_subprocess(self, timeout: float = 10.0) -> bool:
"""Gracefully shut down the export subprocess.
Returns True only once the worker is confirmed dead. If it survives
terminate/kill (e.g. wedged in an uninterruptible CUDA syscall that outlives
SIGKILL) the live handle is KEPT, not nulled, so is_worker_alive() and the
pre-swap liveness guard can still observe the survivor instead of a cleared
handle and refuse the destructive sidecar swap."""
if self._proc is None or not self._proc.is_alive():
self._proc = None
return
return True
self._drain_queue()
@ -265,10 +293,20 @@ class ExportOrchestrator:
except Exception:
pass
if self._proc is not None and self._proc.is_alive():
# Survived SIGKILL (uninterruptible syscall): keep the handle so callers
# and the pre-swap guard see a live worker rather than a nulled one.
logger.error(
"Export subprocess still alive after terminate/kill; "
"preserving its handle for the pre-swap liveness check"
)
return False
self._proc = None
self._cmd_queue = None
self._resp_queue = None
logger.info("Export subprocess shut down")
return True
def _cleanup(self):
"""atexit handler."""
@ -339,9 +377,10 @@ class ExportOrchestrator:
if rtype == "status":
message = resp.get("message", "")
logger.info("Export subprocess status: %s", message)
# Surface status in the live log panel for high-level progress.
# One structured export_progress line per phase (consolidated in the
# server log, like training/download progress); also shown live.
if message:
logger.info("export_progress", phase = message)
self._append_log(
{
"stream": "status",
@ -409,14 +448,44 @@ class ExportOrchestrator:
self._export_active = True
op_success, op_message = False, ""
try:
# Handshake with the sidecar install route: _export_active is set above, so either this
# recheck refuses BEFORE tearing down the old worker (keeping the loaded checkpoint), or
# the install sees is_export_active() and 409s. The spawn-time recheck stays as a last resort.
from utils.transformers_version import sidecar_swap_in_progress
if sidecar_swap_in_progress():
from utils.transformers_version import SidecarSwapInProgress
op_message = (
"A transformers installation is replacing the latest "
"sidecar; retry when it completes."
)
raise SidecarSwapInProgress(op_message)
# Always kill any existing subprocess and spawn fresh.
if self._ensure_subprocess_alive():
self._shutdown_subprocess()
if self._shutdown_subprocess() is False:
# Survivor still holds GPU memory (a wedged CUDA syscall outliving
# SIGKILL); its handle is kept so is_worker_alive() and the pre-swap
# guard still see it. Do not spawn a second worker over it -- fail so
# the load can retry once it exits.
op_message = (
"The current export worker did not exit and still holds GPU "
"memory; not starting a new checkpoint load over it. Retry shortly."
)
return False, op_message
elif self._proc is not None:
self._shutdown_subprocess(timeout = 2)
logger.info("Spawning fresh export subprocess for '%s'", checkpoint_path)
self._spawn_subprocess(sub_config)
try:
self._spawn_subprocess(sub_config)
except Exception:
# The old worker is already gone; a stale current_checkpoint
# would make the Export page claim a loaded checkpoint that
# the next op then fails on with "no subprocess running".
self.current_checkpoint = None
self.is_vision = False
self.is_peft = False
raise
try:
resp = self._wait_response("loaded")
@ -560,6 +629,18 @@ class ExportOrchestrator:
self._export_active = True
op_success, op_message, op_output_path = False, "", None
try:
# Handshake with the sidecar install route (see load_checkpoint): _export_active is set
# above, so this recheck refuses before the command is sent, or the install sees the active
# op and 409s. Without it, an install would block in cleanup_memory behind a long export op.
from utils.transformers_version import sidecar_swap_in_progress
if sidecar_swap_in_progress():
from utils.transformers_version import SidecarSwapInProgress
op_message = (
"A transformers installation is replacing the latest "
"sidecar; retry when it completes."
)
raise SidecarSwapInProgress(op_message)
cmd = {"type": "export", "export_type": export_type, **params}
try:
self._send_cmd(cmd)

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

@ -236,6 +236,17 @@ def _handle_load(backend, cmd: dict, resp_queue: Any) -> None:
checkpoint_path = cmd["checkpoint_path"]
max_seq_length = cmd.get("max_seq_length", 2048)
load_in_4bit = cmd.get("load_in_4bit", True)
# Latest-sidecar checkpoints load 16-bit here too: bnb 4-bit feeds quantized
# expert weights into unvalidated paths (same flip as the chat worker).
if load_in_4bit:
from utils.transformers_version import latest_tier_active_for
if latest_tier_active_for(checkpoint_path, cmd.get("hf_token")):
load_in_4bit = False
logger.info(
"Latest-transformers sidecar active for %s - forcing a 16-bit "
"export load (4-bit is disabled for brand-new architectures)",
checkpoint_path,
)
trust_remote_code = cmd.get("trust_remote_code", False)
# Auto-enable trust_remote_code for NemotronH/Nano models.
@ -387,6 +398,19 @@ def _handle_export(backend, cmd: dict, resp_queue: Any) -> None:
# orchestrator spawns a fresh subprocess per checkpoint load, resetting it.
_log_forward_gate.set()
# Phase milestone so the heavy export step shows in the server log; the
# merge/save/convert itself only forwards stdout to the live panel.
_phase = {
"merged": f"Exporting merged model ({cmd.get('format_type', '16-bit (FP16)')})...",
"gguf": f"Exporting GGUF ({cmd.get('quantization_method', 'Q4_K_M')})...",
"lora": "Exporting LoRA adapter...",
"base": "Exporting base model...",
}.get(export_type, f"Exporting ({export_type})...")
_send_response(
resp_queue,
{"type": "status", "message": _phase, "ts": time.time()},
)
output_path: Any = None
try:
if export_type == "merged":

387
studio/backend/core/inference/_html_to_md.py
View file

@ -7,6 +7,11 @@ Minimal HTML-to-Markdown converter using only the standard library.
Replaces the external ``html2text`` (GPL-3.0) dependency with a ~250-line
``html.parser.HTMLParser`` subclass. Covers headings, links, bold/italic,
lists, tables, blockquotes, code blocks, and entity decoding.
``main_content=True`` also applies a readability-style heuristic: scope
conversion to the page's ``<article>`` (else ``<main>``) subtree when it
carries substantial text, and strip known boilerplate fragments (skip-links,
error placeholders, session banners, cookie prompts) from the result.
"""
from __future__ import annotations
@ -27,8 +32,138 @@ _SKIP_TAGS = frozenset(
"math",
"nav",
"footer",
# Never-rendered / form-chrome elements, not page content.
"template",
"dialog",
"button",
"select",
"datalist",
}
)
# <aside> is NOT skipped: docs use it for admonition callouts (real content);
# page-furniture asides are excluded by the main-content scoping pass instead.
# Void elements never produce an end tag, so they must not join the
# open-element stack used to bound hidden subtrees.
_VOID_TAGS = frozenset(
{
"area",
"base",
"br",
"col",
"embed",
"hr",
"img",
"input",
"link",
"meta",
"param",
"source",
"track",
"wbr",
}
)
def _style_hides_element(style: str) -> bool:
"""True when an inline ``style`` sets ``display:none`` / ``visibility:hidden``.
Parsed per property so an unrelated value that merely contains ``none`` is
not misread as hidden."""
lowered = style.lower()
if "none" not in lowered and "hidden" not in lowered:
return False
for declaration in style.split(";"):
prop, sep, value = declaration.partition(":")
if not sep:
continue
prop = prop.strip().lower()
# Drop any !important flag and keep the first token of the value.
value = value.split("!", 1)[0].strip().lower()
if prop == "display" and value == "none":
return True
if prop == "visibility" and value == "hidden":
return True
return False
def _is_hidden_element(attr_dict: dict) -> bool:
"""True when the element is not rendered: ``hidden`` attribute,
``aria-hidden="true"``, or an inline ``style`` hiding it. Such JS-only
placeholders ship in the HTML but must not reach the output. ``hidden`` is
enumerated: any present value (even ``hidden="false"``) means not rendered."""
if "hidden" in attr_dict:
return True
if (attr_dict.get("aria-hidden") or "").strip().lower() == "true":
return True
return _style_hides_element(attr_dict.get("style") or "")
# HTML5 optional end tags: a listed start tag implicitly closes an open element
# of the key type (as browsers do), else an unclosed ``<p hidden>``/``<li hidden>``
# swallows every following sibling. Keys: closable elements; values: closers.
_P_CLOSING_TAGS = frozenset(
{
"address",
"article",
"aside",
"blockquote",
"details",
"div",
"dl",
"fieldset",
"figcaption",
"figure",
"footer",
"form",
"h1",
"h2",
"h3",
"h4",
"h5",
"h6",
"header",
"hgroup",
"hr",
"main",
"menu",
"nav",
"ol",
"p",
"pre",
"section",
"table",
"ul",
}
)
_IMPLICIT_CLOSERS: dict = {
"p": _P_CLOSING_TAGS,
"li": frozenset({"li"}),
"dt": frozenset({"dt", "dd"}),
"dd": frozenset({"dt", "dd"}),
"tr": frozenset({"tr"}),
"td": frozenset({"td", "th", "tr"}),
"th": frozenset({"td", "th", "tr"}),
"option": frozenset({"option", "optgroup"}),
"optgroup": frozenset({"optgroup"}),
}
# Item tag -> container tags that re-scope it: a nested container makes an inner
# item a descendant, not an optional-close sibling, so recovery must stop there
# rather than close (and un-hide) the outer item and leak its nested content.
_CLOSE_BARRIERS: dict = {
"li": frozenset({"ul", "ol", "menu"}),
"dt": frozenset({"dl"}),
"dd": frozenset({"dl"}),
"tr": frozenset({"table"}),
"td": frozenset({"table"}),
"th": frozenset({"table"}),
"option": frozenset({"select", "datalist"}),
"optgroup": frozenset({"select", "datalist"}),
}
_BLOCK_TAGS = frozenset(
{
"p",
@ -51,13 +186,33 @@ _INLINE_EMPHASIS = {"strong": "**", "b": "**", "em": "*", "i": "*"}
class _MarkdownRenderer(HTMLParser):
"""HTMLParser subclass that emits Markdown tokens into a list."""
"""HTMLParser subclass that emits Markdown tokens into a list.
def __init__(self):
``scope_tags`` restricts emission to the subtree(s) of the given tags
(e.g. ``{"article"}``): outside them every handler is a no-op, which is
how the readability-style main-content pass drops page furniture.
"""
def __init__(self, scope_tags: frozenset[str] | None = None):
super().__init__(convert_charrefs = False)
self._out: list[str] = []
self._skip_depth: int = 0
# Main-content scoping: emit only while inside a scope tag.
self._scope_tags = scope_tags
self._scope_depth: int = 0
# Output boundaries per top-level scope element, so a caller can size each
# candidate alone and a swarm of tiny sibling cards can't clear the threshold.
self.scope_segments: list[str] = []
self._scope_seg_start: int | None = None
# Hidden-subtree tracking: stack of open non-void tags plus the indices
# where a hidden element started. End tags pop to the matching tag, so
# an omitted </p>/<li> close cannot leave the renderer stuck hidden.
self._open_tags: list[str] = []
self._hidden_marks: list[int] = []
# Link state
self._link_href: str | None = None
self._link_text_parts: list[str] = []
@ -150,16 +305,95 @@ class _MarkdownRenderer(HTMLParser):
# ------------------------------------------------------------------
# Tag handlers
# ------------------------------------------------------------------
# Structural bookkeeping shared by every start tag (skip/hidden/scope).
def _close_implicit(self, tag: str) -> None:
"""HTML5 optional-end-tag recovery for a start tag about to open.
Pops each implicitly-closed ancestor (and its hidden marks), scanning the
whole stack so an open ``<p>``/``<li>`` still closes under an unclosed inline
``<span>``. Stops at a ``_CLOSE_BARRIERS`` container so recovery never crosses
a nested list/table/dl and leaks the outer item's hidden content. Runs even
for skipped ``<nav>``/``<footer>``, which also close ``<p>``."""
barriers = _CLOSE_BARRIERS.get(tag, ())
while True:
close_at = None
for i in range(len(self._open_tags) - 1, -1, -1):
name = self._open_tags[i]
if tag in _IMPLICIT_CLOSERS.get(name, ()):
close_at = i
break
# A barrier container re-scopes the item; stop before it.
if name in barriers:
break
if close_at is None:
break
del self._open_tags[close_at:]
while self._hidden_marks and self._hidden_marks[-1] >= close_at:
self._hidden_marks.pop()
def _enter_tag(self, tag: str, attr_dict: dict) -> bool:
"""Track open/hidden/scope state; return True when the tag's content
should be rendered (False = suppressed). Caller runs ``_close_implicit``
first so recovery also fires for skipped tags."""
if tag not in _VOID_TAGS:
self._open_tags.append(tag)
if _is_hidden_element(attr_dict):
self._hidden_marks.append(len(self._open_tags) - 1)
elif _is_hidden_element(attr_dict):
# Void elements never join the stack, so suppress a hidden one inline.
return False
if self._scope_tags is not None and tag in self._scope_tags:
if self._scope_depth == 0:
self._scope_seg_start = len(self._out)
self._scope_depth += 1
if self._hidden_marks:
return False
if self._scope_tags is not None and self._scope_depth == 0:
return False
return True
def _exit_tag(self, tag: str) -> bool:
"""Pop to the matching open tag; return True when the end tag should
be rendered (False = it closed inside a hidden / out-of-scope region)."""
suppressed = bool(self._hidden_marks) or (
self._scope_tags is not None and self._scope_depth == 0
)
if tag not in _VOID_TAGS:
# Pop to the innermost matching open tag (recovers omitted closes).
for i in range(len(self._open_tags) - 1, -1, -1):
if self._open_tags[i] == tag:
del self._open_tags[i:]
while self._hidden_marks and self._hidden_marks[-1] >= i:
self._hidden_marks.pop()
break
if self._scope_tags is not None and tag in self._scope_tags and self._scope_depth > 0:
self._scope_depth -= 1
if self._scope_depth == 0 and self._scope_seg_start is not None:
self.scope_segments.append("".join(self._out[self._scope_seg_start :]))
self._scope_seg_start = None
return not suppressed
def handle_starttag(self, tag: str, attrs: list[tuple[str, str | None]]) -> None:
tag = tag.lower()
if self._skip_depth:
# Inside a skipped subtree: only track nested skip depth.
if tag in _SKIP_TAGS:
self._skip_depth += 1
return
# Recover optional end tags before the skip decision: a skipped
# <nav>/<footer> still implicitly closes an open <p>, releasing its
# hidden mark so following siblings render.
self._close_implicit(tag)
if tag in _SKIP_TAGS:
self._skip_depth += 1
return
if self._skip_depth:
return
attr_dict = dict(attrs)
if not self._enter_tag(tag, attr_dict):
return
if tag in _HEADING_TAGS:
level = int(tag[1])
@ -250,6 +484,9 @@ class _MarkdownRenderer(HTMLParser):
if self._skip_depth:
return
if not self._exit_tag(tag):
return
if tag in _HEADING_TAGS:
self._emit("\n\n")
@ -308,8 +545,13 @@ class _MarkdownRenderer(HTMLParser):
# ------------------------------------------------------------------
# Text / entity handlers
# ------------------------------------------------------------------
def _text_suppressed(self) -> bool:
if self._skip_depth or self._hidden_marks:
return True
return self._scope_tags is not None and self._scope_depth == 0
def handle_data(self, data: str) -> None:
if self._skip_depth:
if self._text_suppressed():
return
if self._in_pre:
self._pre_parts.append(data)
@ -326,12 +568,12 @@ class _MarkdownRenderer(HTMLParser):
self._emit(text)
def handle_entityref(self, name: str) -> None:
if self._skip_depth:
if self._text_suppressed():
return
self._emit(html.unescape(f"&{name};"))
def handle_charref(self, name: str) -> None:
if self._skip_depth:
if self._text_suppressed():
return
self._emit(html.unescape(f"&#{name};"))
@ -366,6 +608,14 @@ class _MarkdownRenderer(HTMLParser):
else:
self._out.append("\n\n" + prefixed + "\n\n")
# A scope left open by truncated HTML never reached _exit_tag, so its output
# never joined scope_segments and would score 0. Flush the still-open segment
# here (after the side-buffers) so a truncated main-content page is scored.
if self._scope_seg_start is not None:
self.scope_segments.append("".join(self._out[self._scope_seg_start :]))
self._scope_seg_start = None
self._scope_depth = 0
# Post-processing
def _cleanup(text: str) -> str:
@ -399,17 +649,124 @@ def _cleanup(text: str) -> str:
return "\n".join(out).strip()
# Public API
def html_to_markdown(source_html: str) -> str:
"""Convert HTML to Markdown (headings, links, emphasis, lists, tables, blockquotes, code, entities).
# Known boilerplate fragments stripped from main-content conversions, matched
# only against short lines. Sources: GitHub page furniture / client-side error
# placeholders, skip-links, cookie banners.
_BOILERPLATE_FRAGMENTS = (
"skip to content",
"skip to main content",
"there was an error while loading",
"please reload this page",
"you can't perform that action at this time",
"you signed in with another tab or window",
"you signed out in another tab or window",
"you switched accounts on another tab or window",
"reload to refresh your session",
"you must be signed in to change notification settings",
"uh oh!",
"{{ message }}",
"this website uses cookies",
"we use cookies",
"accept all cookies",
"manage cookie preferences",
)
# Only shorter lines are eligible for boilerplate dropping; real content
# sentences quoting a fragment run longer.
_BOILERPLATE_MAX_LINE_CHARS = 300
``<script>``, ``<style>``, and ``<head>`` are stripped entirely.
"""
# Normalize line endings before parsing.
source_html = source_html.replace("\r\n", "\n").replace("\r", "\n")
renderer = _MarkdownRenderer()
# Normalized furniture phrases for whole-segment matching. See _line_is_boilerplate.
_BOILERPLATE_NORMALIZED = frozenset(
re.sub(r"\s+", " ", fragment).strip().casefold().rstrip(".!:")
for fragment in _BOILERPLATE_FRAGMENTS
)
def _line_is_boilerplate(line: str) -> bool:
"""True only when a whole line is composed of known furniture phrases.
Splits on sentence terminators and requires every segment to be furniture, so a
line stacking several phrases is dropped while prose that merely quotes one is
kept (its other words leave a non-furniture segment)."""
normalized = re.sub(r"\s+", " ", line).strip().casefold()
if not normalized:
return False
segments = [segment.strip().rstrip(".!:") for segment in re.split(r"[.!]", normalized)]
segments = [segment for segment in segments if segment]
return bool(segments) and all(segment in _BOILERPLATE_NORMALIZED for segment in segments)
def _strip_boilerplate_lines(text: str) -> str:
"""Drop short lines that consist entirely of known page-furniture phrases.
Fenced code blocks are preserved verbatim: boilerplate never renders
inside ``<pre>``, while READMEs legitimately quote error strings."""
out: list[str] = []
in_fence = False
for line in text.split("\n"):
if line.lstrip().startswith("```"):
in_fence = not in_fence
out.append(line)
continue
if not in_fence and len(line) <= _BOILERPLATE_MAX_LINE_CHARS and _line_is_boilerplate(line):
continue
out.append(line)
# Collapse blank runs the dropped lines may have left behind.
return re.sub(r"\n{3,}", "\n\n", "\n".join(out)).strip()
def _render(source_html: str, scope_tags: frozenset[str] | None) -> str:
renderer = _MarkdownRenderer(scope_tags = scope_tags)
renderer.feed(source_html)
renderer.close()
renderer.flush_pending()
raw = "".join(renderer._out)
return _cleanup(raw)
def _select_main_scope_render(source_html: str, tag: str) -> tuple[int, str]:
"""Length and boilerplate-stripped render of the largest single ``<tag>``
subtree. Sizing candidates one at a time stops many tiny sibling cards from
clearing the threshold together, and returning that one subtree keeps
unrelated siblings (related cards, comment threads) out of the output."""
renderer = _MarkdownRenderer(scope_tags = frozenset({tag}))
renderer.feed(source_html)
renderer.close()
renderer.flush_pending()
best_len = 0
best_render = ""
for seg in renderer.scope_segments:
rendered = _strip_boilerplate_lines(_cleanup(seg))
if len(rendered) > best_len:
best_len = len(rendered)
best_render = rendered
return best_len, best_render
# A scoped conversion below this size is judged not to be the page's main
# content (e.g. an empty <article> stub) and the next candidate is tried.
_MIN_MAIN_CONTENT_CHARS = 200
# Public API
def html_to_markdown(source_html: str, *, main_content: bool = False) -> str:
"""Convert HTML to Markdown (headings, links, emphasis, lists, tables, blockquotes, code, entities).
``<script>``, ``<style>``, and ``<head>`` are stripped entirely, as are
subtrees hidden from rendering (``hidden`` / ``aria-hidden="true"``).
``main_content=True`` applies a readability-style heuristic for page
fetches: prefer the ``<article>`` subtree (GitHub renders READMEs there),
then ``<main>``, falling back to the whole document, and strip known
boilerplate fragments from the result.
"""
# Normalize line endings before parsing.
source_html = source_html.replace("\r\n", "\n").replace("\r", "\n")
if main_content:
for scope_tag in ("article", "main"):
# Render only the chosen subtree so sibling <article>/<main>
# elements do not leak in once the largest passes the size gate.
length, rendered = _select_main_scope_render(source_html, scope_tag)
if length >= _MIN_MAIN_CONTENT_CHARS:
return rendered
return _strip_boilerplate_lines(_render(source_html, None))
return _render(source_html, None)

110
studio/backend/core/inference/_vulkan_probe.py Normal file
View file

@ -0,0 +1,110 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Standalone free-VRAM probe for the bundled ggml Vulkan backend.
Run in a short-lived subprocess (``python _vulkan_probe.py <bindir>``) so the
Vulkan instance never lives in the long-running backend process. Loads the
bundled ggml Vulkan backend from ``<bindir>`` and prints one
``<idx>\\t<free_bytes>\\t<is_igpu>\\t<total_bytes>`` line per device to stdout.
Indices are ggml's own Vulkan device ordinals, which need not match nvidia-smi
order. ``is_igpu`` (from ggml's device type) is ``1`` for an integrated GPU
sharing system RAM. ``total_bytes`` is the device-local heap; the reader uses
it to reserve absolute headroom on a discrete card (parity with the CUDA/ROCm
fit) and ignores it for an iGPU, whose "VRAM" is shared system RAM.
Uses only the standard library so it stays runnable as a bare script.
"""
import ctypes
import os
import sys
# ggml_backend_dev_type enum (ggml-backend.h): CPU=0, GPU=1, IGPU=2, ...
_GGML_BACKEND_DEVICE_TYPE_IGPU = 2
def _igpu_flags(base, lib, count: int) -> list[bool]:
"""Per-device integrated-GPU flags via ggml's backend registry.
The Vulkan reg enumerates devices in the same order as
``ggml_backend_vk_get_device_memory`` (each context uses ``ctx->device =
i``), so reg index == device ordinal. Returns all-False on any failure so
the reader never over-caps a discrete card.
"""
flags = [False] * count
try:
lib.ggml_backend_vk_reg.restype = ctypes.c_void_p
lib.ggml_backend_vk_reg.argtypes = []
base.ggml_backend_reg_dev_count.restype = ctypes.c_size_t
base.ggml_backend_reg_dev_count.argtypes = [ctypes.c_void_p]
base.ggml_backend_reg_dev_get.restype = ctypes.c_void_p
base.ggml_backend_reg_dev_get.argtypes = [ctypes.c_void_p, ctypes.c_size_t]
base.ggml_backend_dev_type.restype = ctypes.c_int
base.ggml_backend_dev_type.argtypes = [ctypes.c_void_p]
reg = lib.ggml_backend_vk_reg()
if not reg:
return flags
dev_count = base.ggml_backend_reg_dev_count(reg)
for i in range(min(count, dev_count)):
dev = base.ggml_backend_reg_dev_get(reg, i)
if dev:
flags[i] = base.ggml_backend_dev_type(dev) == _GGML_BACKEND_DEVICE_TYPE_IGPU
except Exception:
# Best-effort: any failure degrades to "discrete" so the memory
# readings still get through instead of crashing the probe.
pass
return flags
def main() -> int:
if len(sys.argv) < 2:
return 0
bindir = sys.argv[1]
# Hold add_dll_directory's handle for the rest of main() (the documented
# idiom) so bindir stays on the search path while the sibling ggml DLLs
# resolve below.
_dll_dir = None
if sys.platform == "win32":
base_name, vk_name = "ggml-base.dll", "ggml-vulkan.dll"
try:
_dll_dir = os.add_dll_directory(bindir)
except Exception:
pass
else:
base_name, vk_name = "libggml-base.so", "libggml-vulkan.so"
# RTLD_GLOBAL exposes ggml-base's symbols to ggml-vulkan on POSIX. getattr
# falls back to 0 where the flag doesn't exist (Windows CDLL ignores mode).
_rtld_global = getattr(ctypes, "RTLD_GLOBAL", 0)
try:
base = ctypes.CDLL(os.path.join(bindir, base_name), mode = _rtld_global)
lib = ctypes.CDLL(os.path.join(bindir, vk_name), mode = _rtld_global)
except OSError as e:
print(f"ggml-vulkan load failed: {e}", file = sys.stderr)
return 1
lib.ggml_backend_vk_get_device_count.restype = ctypes.c_int
lib.ggml_backend_vk_get_device_count.argtypes = []
lib.ggml_backend_vk_get_device_memory.restype = None
lib.ggml_backend_vk_get_device_memory.argtypes = [
ctypes.c_int,
ctypes.POINTER(ctypes.c_size_t),
ctypes.POINTER(ctypes.c_size_t),
]
count = lib.ggml_backend_vk_get_device_count()
igpu = _igpu_flags(base, lib, count)
rows = []
for i in range(count):
free, total = ctypes.c_size_t(0), ctypes.c_size_t(0)
lib.ggml_backend_vk_get_device_memory(i, ctypes.byref(free), ctypes.byref(total))
rows.append("%d\t%d\t%d\t%d" % (i, free.value, int(igpu[i]), total.value))
sys.stdout.write("\n".join(rows))
return 0
if __name__ == "__main__":
raise SystemExit(main())

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

@ -10,10 +10,242 @@ native-chat-template fallback used by the transformers and MLX backends.
import copy
import json
import logging
from dataclasses import dataclass
from typing import Optional
_THINK_OPEN = "<think>"
_THINK_CLOSE = "</think>"
_GEMMA_CHANNEL_START = "<|channel>"
_GEMMA_THOUGHT_OPEN = "<|channel>thought"
_GEMMA_THOUGHT_CLOSE = "<channel|>"
_GEMMA_TEMPLATE_OPENERS = (
_GEMMA_THOUGHT_OPEN + "\n",
_GEMMA_THOUGHT_OPEN + "\\n",
_GEMMA_THOUGHT_OPEN + _GEMMA_THOUGHT_CLOSE,
)
def _tokenizer_objects(tokenizer) -> tuple:
"""Return a processor/tokenizer and its distinct nested tokenizer."""
if tokenizer is None:
return ()
nested = getattr(tokenizer, "tokenizer", None)
return (tokenizer,) if nested is None or nested is tokenizer else (tokenizer, nested)
def _selected_template_strings_from_value(
template,
tools = None,
*,
prefer_tool_use: bool = True,
) -> tuple[str, ...]:
"""Return the named chat template matching HF's default selection rules."""
tools = tools or None
if isinstance(template, str):
return (template,)
if not isinstance(template, dict):
return ()
if prefer_tool_use and tools and isinstance(template.get("tool_use"), str):
return (template["tool_use"],)
if isinstance(template.get("default"), str):
return (template["default"],)
values = tuple(value for value in template.values() if isinstance(value, str))
return values if len(values) == 1 else ()
def _selected_chat_template_strings(tokenizer, tools = None) -> tuple[str, ...]:
"""Return the active chat template selected for this request."""
tools = tools or None
getter = getattr(tokenizer, "get_chat_template", None)
if callable(getter):
for kwargs in ({"chat_template": None, "tools": tools}, {"tools": tools}, {}):
try:
selected = getter(**kwargs)
except Exception:
continue
if isinstance(selected, str):
return (selected,)
# ProcessorMixin.apply_chat_template does not switch to "tool_use" implicitly;
# it uses "default" unless chat_template= names another template.
is_processor = getattr(tokenizer, "tokenizer", None) is not None and callable(
getattr(tokenizer, "apply_chat_template", None)
)
return _selected_template_strings_from_value(
getattr(tokenizer, "chat_template", None),
tools,
prefer_tool_use = not is_processor,
)
def _detect_reasoning_channel_markers_from_templates(
templates: tuple[str, ...],
) -> Optional[tuple[str, str]]:
"""Return Gemma native reasoning markers only when a template emits them."""
if any(opener in template for template in templates for opener in _GEMMA_TEMPLATE_OPENERS):
return _GEMMA_THOUGHT_OPEN, _GEMMA_THOUGHT_CLOSE
return None
def detect_reasoning_channel_markers(tokenizer, tools = None) -> Optional[tuple[str, str]]:
"""Return native Gemma thought-channel markers supported by a tokenizer.
Detection uses the active chat template rather than model names or vocabulary
membership. Some models expose Gemma control tokens without using the native
thought-channel response protocol, and those must keep normal
``skip_special_tokens`` streaming.
"""
for obj in _tokenizer_objects(tokenizer):
templates = _selected_chat_template_strings(obj, tools)
if templates:
return _detect_reasoning_channel_markers_from_templates(templates)
return None
def detect_reasoning_channel_markers_from_template(
template, tools = None
) -> Optional[tuple[str, str]]:
"""Return native Gemma thought-channel markers from a raw template value."""
return _detect_reasoning_channel_markers_from_templates(
_selected_template_strings_from_value(template, tools)
)
def detect_reasoning_channel_markers_from_model_info(
tokenizer,
model_info: Optional[dict] = None,
tools = None,
) -> Optional[tuple[str, str]]:
"""Return reasoning markers from the active or cached native template."""
markers = detect_reasoning_channel_markers(tokenizer, tools = tools)
if markers is not None or not isinstance(model_info, dict):
return markers
native_templates = (
model_info.get("native_chat_template"),
(model_info.get("chat_template_info") or {}).get("template"),
)
for template in native_templates:
markers = detect_reasoning_channel_markers_from_template(template, tools)
if markers is not None:
return markers
return None
@dataclass(frozen = True)
class ChatTemplateRenderResult:
"""Prompt plus response-protocol metadata selected by the renderer."""
prompt: str
reasoning_channel_markers: Optional[tuple[str, str]] = None
def _split_partial_marker(text: str, marker: str) -> tuple[str, str]:
"""Hold the longest suffix that may become ``marker`` in the next chunk."""
for length in range(min(len(text), len(marker) - 1), 0, -1):
if text.endswith(marker[:length]):
return text[:-length], text[-length:]
return text, ""
class ReasoningChannelNormalizer:
"""Incrementally convert one native reasoning channel to ``<think>``.
The parser follows mlx-vlm's streaming boundary behavior but emits Studio's
established canonical text contract. Only the configured opening and
closing markers are consumed; tool-call and other control markers remain
available to downstream parsers.
"""
def __init__(self, opening_marker: str, closing_marker: str):
self._opening_marker = opening_marker
self._closing_marker = closing_marker
self._buffer = ""
self._in_reasoning = False
self._reasoning_done = False
self._skip_opening_newline = False
def feed(self, text: str) -> str:
"""Consume a raw text delta and return the stable canonical delta."""
self._buffer += text or ""
output: list[str] = []
while self._buffer:
if self._reasoning_done:
output.append(self._buffer)
self._buffer = ""
break
if self._in_reasoning and self._skip_opening_newline:
if self._buffer.startswith("\n"):
self._buffer = self._buffer[1:]
self._skip_opening_newline = False
if not self._buffer:
break
marker = self._closing_marker if self._in_reasoning else self._opening_marker
index = self._buffer.find(marker)
if index < 0:
stable, self._buffer = _split_partial_marker(self._buffer, marker)
output.append(stable)
break
output.append(self._buffer[:index])
self._buffer = self._buffer[index + len(marker) :]
if self._in_reasoning:
output.append(_THINK_CLOSE)
self._in_reasoning = False
self._reasoning_done = True
else:
output.append(_THINK_OPEN)
self._in_reasoning = True
self._skip_opening_newline = True
return "".join(output)
def finish(self) -> str:
"""Flush a naturally completed stream and close an open think block."""
output = self.drain()
if self._in_reasoning:
output += _THINK_CLOSE
self._in_reasoning = False
self._reasoning_done = True
return output
def drain(self) -> str:
"""Flush buffered literal text without synthesizing a closing tag."""
output = self._buffer
self._buffer = ""
return output
def normalize_reasoning_snapshots(
stream,
tokenizer = None,
cancel_event = None,
markers: Optional[tuple[str, str]] = None,
tools = None,
):
"""Normalize a prefix-monotonic cumulative text stream when supported."""
markers = markers or detect_reasoning_channel_markers(tokenizer, tools = tools)
if markers is None:
yield from stream
return
normalizer = ReasoningChannelNormalizer(*markers)
raw_output = ""
normalized_output = ""
for snapshot in stream:
if not snapshot.startswith(raw_output):
raise RuntimeError("Reasoning normalization requires cumulative text snapshots")
delta = normalizer.feed(snapshot[len(raw_output) :])
raw_output = snapshot
if delta:
normalized_output += delta
yield normalized_output
cancelled = cancel_event is not None and cancel_event.is_set()
tail = normalizer.drain() if cancelled else normalizer.finish()
if tail:
normalized_output += tail
yield normalized_output
def detect_think_prefill(prompt: Optional[str], special_tokens = None) -> str:
@ -166,7 +398,8 @@ def render_native_template(
preserve_thinking: Optional[bool] = None,
apply_fn = None,
hf_token: Optional[str] = None,
) -> Optional[str]:
return_metadata: bool = False,
):
"""Render ``messages`` + ``tools`` with the model's NATIVE chat template.
Some Unsloth override templates (e.g. ``mistral``, ``gemma-4``) do not emit
@ -175,7 +408,9 @@ def render_native_template(
tool-calling syntax. It is loaded straight from the repo (bypassing any
override on the live tokenizer) and cached on ``model_info``. Returns the
rendered prompt only if the native template actually emits the tools (render
differs with vs without tools); otherwise ``None``.
differs with vs without tools); otherwise ``None``. With ``return_metadata``,
returns ``ChatTemplateRenderResult`` so callers can stream with the response
protocol selected by this request's template.
``hf_token`` is the token the model was loaded with -- passed to the repo load
so a gated/private model's native template can still be fetched (otherwise the
@ -261,7 +496,16 @@ def render_native_template(
exc,
)
return None
return with_tools if with_tools != no_tools else None
if with_tools == no_tools:
return None
if return_metadata:
return ChatTemplateRenderResult(
with_tools,
_detect_reasoning_channel_markers_from_templates(
_selected_template_strings_from_value(native_tpl, tools)
),
)
return with_tools
def render_with_native_template_fallback(
@ -277,7 +521,8 @@ def render_with_native_template_fallback(
preserve_thinking: Optional[bool] = None,
apply_fn = None,
hf_token: Optional[str] = None,
) -> str:
return_metadata: bool = False,
):
"""Return ``formatted_prompt``, swapping in a native-template render when an
override template dropped the ``tools`` schema.
@ -285,9 +530,27 @@ def render_with_native_template_fallback(
them (detected by comparison, robust against tool names in the system prompt),
re-render with the model's native template. Shared by the transformers and MLX
backends so both advertise tools consistently. ``hf_token`` is forwarded so a
gated/private model's native template can still be fetched."""
gated/private model's native template can still be fetched. With
``return_metadata``, returns the selected prompt plus reasoning-channel markers
for the exact template used by this request."""
live_markers = detect_reasoning_channel_markers(tokenizer, tools = tools)
def _result(prompt: str, markers = live_markers):
if return_metadata:
return ChatTemplateRenderResult(prompt, markers)
return prompt
if not tools:
return formatted_prompt
# Gemma 4 can emit its native reasoning protocol even when a generation-time
# Unsloth override rendered a marker-free prompt. Preserve the live-verified
# no-tools thinking behavior without letting cached native metadata describe
# unrelated tool prompts that kept the active override.
markers = live_markers
if markers is None:
markers = detect_reasoning_channel_markers_from_model_info(
tokenizer, model_info, tools = None
)
return _result(formatted_prompt, markers)
if apply_fn is None:
apply_fn = apply_chat_template_for_generation
# Probe whether the live template dropped the schema. A tools-requiring template
@ -307,9 +570,9 @@ def render_with_native_template_fallback(
active_model_name,
exc,
)
return formatted_prompt
return _result(formatted_prompt)
if formatted_prompt != probe_no_tools:
return formatted_prompt # template already emits the tools schema
return _result(formatted_prompt) # template already emits the tools schema
native_prompt = render_native_template(
model_info = model_info,
active_model_name = active_model_name,
@ -320,6 +583,7 @@ def render_with_native_template_fallback(
preserve_thinking = preserve_thinking,
apply_fn = apply_fn,
hf_token = hf_token,
return_metadata = return_metadata,
)
if native_prompt:
logger.info(
@ -328,4 +592,4 @@ def render_with_native_template_fallback(
active_model_name,
)
return native_prompt
return formatted_prompt
return _result(formatted_prompt)

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

@ -5,7 +5,7 @@
from unsloth import FastLanguageModel, FastVisionModel
from unsloth.chat_templates import get_chat_template
from transformers import TextStreamer
from transformers import TextIteratorStreamer, TextStreamer
from peft import PeftModel, PeftModelForCausalLM
import json
@ -15,6 +15,7 @@ from pathlib import Path
from typing import Optional, Union, Generator, Tuple
from utils.models import ModelConfig, get_base_model_from_lora
from utils.paths import is_model_cached
from utils.transformers_dtype import dtype_kwargs
from utils.utils import format_error_message
from utils.hardware import (
get_device,
@ -31,6 +32,11 @@ from core.inference.chat_eos import (
chat_eos_repair,
resolve_chat_turn_end_eos_ids_using,
)
from core.inference.chat_template_helpers import (
ReasoningChannelNormalizer,
detect_reasoning_channel_markers,
detect_think_prefill,
)
from core.inference.presence_penalty import _make_presence_penalty_processor
from io import StringIO
import structlog
@ -186,6 +192,53 @@ class HarmonyTextStreamer:
self._queue.put(new_content)
class ReasoningTextIteratorStreamer(TextIteratorStreamer):
"""TextIteratorStreamer that preserves native channel tokens until parsed."""
def __init__(
self,
tokenizer,
*,
markers: tuple[str, str],
skip_prompt: bool = True,
timeout: float = 0.2,
cancel_event = None,
**decode_kwargs,
):
decode_kwargs["skip_special_tokens"] = False
super().__init__(tokenizer, skip_prompt = skip_prompt, timeout = timeout, **decode_kwargs)
self._normalizer = ReasoningChannelNormalizer(*markers)
self._cancel_event = cancel_event
self._aborted = False
def abort(self):
"""Mark generation as failed so ``end`` drains without closing."""
self._aborted = True
def on_finalized_text(
self,
text: str,
stream_end: bool = False,
):
"""Queue canonical deltas, closing only on natural stream completion."""
delta = self._normalizer.feed(text)
if delta:
self.text_queue.put(delta, timeout = self.timeout)
if stream_end:
cancelled = self._aborted or (
self._cancel_event is not None and self._cancel_event.is_set()
)
tail = self._normalizer.drain() if cancelled else self._normalizer.finish()
if tail:
self.text_queue.put(tail, timeout = self.timeout)
self.text_queue.put(self.stop_signal, timeout = self.timeout)
class _GenerationThreadError(RuntimeError):
"""Generation worker failures that should propagate through stream routes."""
class InferenceBackend:
"""Unified inference backend supporting text, vision, and LoRA models"""
@ -440,7 +493,7 @@ class InferenceBackend:
feature_extractor = tokenizer.feature_extractor,
processor = tokenizer,
return_language = True,
torch_dtype = torch.float16,
**dtype_kwargs(torch.float16),
)
self.models[model_name]["model"] = model
self.models[model_name]["tokenizer"] = tokenizer
@ -832,8 +885,10 @@ class InferenceBackend:
nudge_tool_calls: Optional[bool] = None,
tool_call_timeout: int = 300,
session_id: Optional[str] = None,
thread_id: Optional[str] = None,
rag_scope: Optional[dict] = None,
presence_penalty: float = 0.0,
reasoning_prefilled: bool = False,
):
"""Run an agentic tool loop on top of ``generate_chat_response``.
@ -885,7 +940,9 @@ class InferenceBackend:
max_tool_iterations = max_tool_iterations,
tool_call_timeout = tool_call_timeout,
session_id = session_id,
thread_id = thread_id,
rag_scope = rag_scope,
reasoning_prefilled = reasoning_prefilled,
)
def generate_chat_response(
@ -957,8 +1014,7 @@ class InferenceBackend:
thread can toggle adapters under the generation lock.
"""
if not self.active_model_name:
yield "Error: No active model"
return
raise RuntimeError("No active model")
model_info = self.models[self.active_model_name]
is_vision = model_info.get("is_vision", False)
@ -1046,6 +1102,7 @@ class InferenceBackend:
template_messages = [{"role": "system", "content": system_prompt}] + messages
else:
template_messages = messages
reasoning_channel_markers_resolved = False
try:
if not (hasattr(tokenizer, "chat_template") and tokenizer.chat_template):
raise ValueError(
@ -1055,6 +1112,7 @@ class InferenceBackend:
f"Please use a model that includes a chat template, or manually set "
f"one via tokenizer.chat_template before inference."
)
reasoning_channel_markers = None
formatted_prompt = self._apply_chat_template_for_generation(
tokenizer,
template_messages,
@ -1070,7 +1128,7 @@ class InferenceBackend:
render_with_native_template_fallback,
)
formatted_prompt = render_with_native_template_fallback(
render_result = render_with_native_template_fallback(
formatted_prompt = formatted_prompt,
tokenizer = tokenizer,
model_info = model_info,
@ -1082,13 +1140,19 @@ class InferenceBackend:
preserve_thinking = preserve_thinking,
apply_fn = self._apply_chat_template_for_generation,
hf_token = model_info.get("hf_token"),
return_metadata = True,
)
formatted_prompt = render_result.prompt
reasoning_channel_markers = render_result.reasoning_channel_markers
reasoning_channel_markers_resolved = True
logger.debug(f"Formatted prompt: {formatted_prompt[:200]}...")
except Exception as e:
logger.error(f"Error applying chat template: {e}")
# Fall back to manual formatting
formatted_prompt = self.format_chat_prompt(messages, system_prompt)
reasoning_channel_markers = None
reasoning_channel_markers_resolved = True
# Step 3: generate
yield from self.generate_stream(
@ -1102,6 +1166,8 @@ class InferenceBackend:
cancel_event = cancel_event,
_adapter_state = _adapter_state,
presence_penalty = presence_penalty,
reasoning_channel_markers = reasoning_channel_markers,
reasoning_channel_markers_resolved = reasoning_channel_markers_resolved,
)
def _generate_vision_response(
@ -1187,21 +1253,27 @@ class InferenceBackend:
# Stream with TextIteratorStreamer + background thread
try:
from core.inference.chat_template_helpers import detect_think_prefill
# Re-emit an open <think> prefill swallowed by skip_prompt (see
# generate_stream).
think_prefix = detect_think_prefill(
prompt_text, getattr(raw_tokenizer, "all_special_tokens", None)
)
from transformers import TextIteratorStreamer
import threading
streamer = TextIteratorStreamer(
streamer = self._make_text_streamer(
raw_tokenizer,
protocol_source = processor,
# The text-only VLM fallback above did not render with the
# processor template, so its native markers do not describe
# this request's response protocol.
reasoning_channel_markers = detect_reasoning_channel_markers(processor)
if image
else None,
reasoning_channel_markers_resolved = True,
skip_prompt = True,
skip_special_tokens = True,
timeout = 0.2,
cancel_event = cancel_event,
use_harmony = self._is_gpt_oss_model(),
)
generation_kwargs = dict(
@ -1223,6 +1295,10 @@ class InferenceBackend:
)
if _pp is not None:
generation_kwargs["logits_processor"] = _pp
stopping_criteria = self._cancel_stopping_criteria(cancel_event)
if stopping_criteria is not None:
generation_kwargs["stopping_criteria"] = stopping_criteria
active_stop_token_ids = self._generation_stop_token_ids(model, generation_kwargs)
err: dict[str, str] = {}
@ -1232,6 +1308,8 @@ class InferenceBackend:
model.generate(**generation_kwargs)
except Exception as e:
err["msg"] = str(e)
if hasattr(streamer, "abort"):
streamer.abort()
logger.error(f"Vision generation error in thread: {e}")
finally:
try:
@ -1248,12 +1326,17 @@ class InferenceBackend:
if think_prefix:
yield think_prefix
from queue import Empty
import time
generation_complete = False
cancel_deadline = None
try:
while True:
if cancel_event is not None and cancel_event.is_set():
break
if cancel_deadline is None:
cancel_deadline = time.monotonic() + 10
elif time.monotonic() >= cancel_deadline:
break
try:
new_token = next(streamer)
except StopIteration:
@ -1262,27 +1345,48 @@ class InferenceBackend:
except Empty:
if not thread.is_alive():
generation_complete = True
output = yield from self._drain_streamer_tail(
streamer, output, active_stop_token_ids
)
break
if cancel_deadline is not None:
remaining = cancel_deadline - time.monotonic()
if remaining <= 0:
break
thread.join(timeout = remaining)
if thread.is_alive():
break
generation_complete = True
output = yield from self._drain_streamer_tail(
streamer, output, active_stop_token_ids
)
break
continue
if new_token:
output += new_token
cleaned = self._clean_generated_text(output)
output, cleaned = self._append_stream_delta(
output, new_token, active_stop_token_ids
)
yield cleaned
finally:
if cancel_event is not None and not generation_complete:
cancel_event.set()
thread.join(timeout = 10)
join_timeout = 10
if cancel_deadline is not None:
join_timeout = max(0, cancel_deadline - time.monotonic())
thread.join(timeout = join_timeout)
if thread.is_alive():
logger.warning(
"Vision generation thread did not exit after cancel/join timeout"
)
if err.get("msg"):
yield f"Error: {err['msg']}"
raise _GenerationThreadError(err["msg"])
except _GenerationThreadError:
raise
except Exception as e:
logger.error(f"Vision generation error: {e}")
yield f"Error: {str(e)}"
raise
def generate_audio_input_response(
self,
@ -1407,11 +1511,13 @@ class InferenceBackend:
)
if err.get("msg"):
yield f"Error: {err['msg']}"
raise _GenerationThreadError(err["msg"])
except _GenerationThreadError:
raise
except Exception as e:
logger.error(f"Audio input generation error: {e}")
yield f"Error: {str(e)}"
raise
def generate_whisper_response(
self,
@ -1444,6 +1550,86 @@ class InferenceBackend:
from utils.datasets import is_gpt_oss_model_name
return is_gpt_oss_model_name(model_name or self.active_model_name or "")
def _make_text_streamer(
self,
tokenizer,
*,
protocol_source = None,
reasoning_channel_markers = None,
reasoning_channel_markers_resolved: bool = False,
skip_prompt: bool = True,
timeout: float = 0.2,
cancel_event = None,
use_harmony: bool = False,
):
"""Create the streamer matching this model's native response protocol."""
if use_harmony:
try:
return HarmonyTextStreamer(
tokenizer,
skip_prompt = skip_prompt,
timeout = timeout,
)
except Exception as e:
logger.warning(f"HarmonyTextStreamer init failed, falling back: {e}")
return TextIteratorStreamer(
tokenizer,
skip_prompt = skip_prompt,
skip_special_tokens = True,
timeout = timeout,
)
markers = (
reasoning_channel_markers
if reasoning_channel_markers_resolved
else reasoning_channel_markers
or detect_reasoning_channel_markers(protocol_source or tokenizer)
)
if markers is not None:
return ReasoningTextIteratorStreamer(
tokenizer,
markers = markers,
skip_prompt = skip_prompt,
timeout = timeout,
cancel_event = cancel_event,
)
return TextIteratorStreamer(
tokenizer,
skip_prompt = skip_prompt,
skip_special_tokens = True,
timeout = timeout,
)
def _append_stream_delta(
self,
output: str,
new_token: str,
stop_token_ids = None,
):
"""Append a streamer delta and apply response-boundary cleanup."""
output += new_token
return output, self._clean_generated_text(output, stop_token_ids = stop_token_ids)
def _drain_streamer_tail(
self,
streamer,
output: str,
stop_token_ids = None,
):
"""Drain queued streamer text after the producer exits."""
while True:
try:
new_token = next(streamer)
except StopIteration:
return output
except Exception:
return output
if new_token:
output, cleaned = self._append_stream_delta(
output, new_token, stop_token_ids = stop_token_ids
)
yield cleaned
def generate_stream(
self,
prompt: str,
@ -1456,6 +1642,8 @@ class InferenceBackend:
cancel_event = None,
_adapter_state = None,
presence_penalty: float = 0.0,
reasoning_channel_markers = None,
reasoning_channel_markers_resolved: bool = False,
) -> Generator[str, None, None]:
"""Generate a streaming text response (text models only).
@ -1464,8 +1652,7 @@ class InferenceBackend:
``presence_penalty`` matches the GGUF sampling path via a logits processor (0 disables it).
"""
if not self.active_model_name:
yield "Error: No active model"
return
raise RuntimeError("No active model")
model_info = self.models[self.active_model_name]
model = model_info["model"]
@ -1478,9 +1665,7 @@ class InferenceBackend:
try:
inputs = tokenizer(prompt, return_tensors = "pt").to(model.device)
from transformers import TextIteratorStreamer
import threading
from core.inference.chat_template_helpers import detect_think_prefill
# skip_prompt swallows an open <think> prefilled by the template;
# re-emit it so the frontend can render the thinking block.
@ -1491,30 +1676,16 @@ class InferenceBackend:
else detect_think_prefill(prompt, getattr(tokenizer, "all_special_tokens", None))
)
# gpt-oss models: HarmonyTextStreamer parses the multi-channel
# harmony protocol into <think> tags
if self._is_gpt_oss_model():
try:
streamer = HarmonyTextStreamer(
tokenizer,
skip_prompt = True,
timeout = 0.2,
)
except Exception as e:
logger.warning(f"HarmonyTextStreamer init failed, falling back: {e}")
streamer = TextIteratorStreamer(
tokenizer,
skip_prompt = True,
skip_special_tokens = True,
timeout = 0.2,
)
else:
streamer = TextIteratorStreamer(
tokenizer,
skip_prompt = True,
skip_special_tokens = True,
timeout = 0.2,
)
streamer = self._make_text_streamer(
tokenizer,
protocol_source = model_info.get("tokenizer"),
reasoning_channel_markers = reasoning_channel_markers,
reasoning_channel_markers_resolved = reasoning_channel_markers_resolved,
skip_prompt = True,
timeout = 0.2,
cancel_event = cancel_event,
use_harmony = self._is_gpt_oss_model(),
)
generation_kwargs = dict(
**inputs,
@ -1532,27 +1703,16 @@ class InferenceBackend:
if tokenizer.pad_token_id is None
else tokenizer.pad_token_id,
)
active_stop_token_ids = self._generation_stop_token_ids(model, generation_kwargs)
# Presence penalty (GGUF parity); prompt_len excludes prompt tokens.
_pp = _make_presence_penalty_processor(
presence_penalty, int(inputs["input_ids"].shape[1])
)
if _pp is not None:
generation_kwargs["logits_processor"] = _pp
if cancel_event is not None:
from transformers.generation.stopping_criteria import (
StoppingCriteria,
StoppingCriteriaList,
)
class _CancelCriteria(StoppingCriteria):
def __init__(self, ev):
self.ev = ev
def __call__(self, input_ids, scores, **kwargs):
return self.ev.is_set()
generation_kwargs["stopping_criteria"] = StoppingCriteriaList(
[_CancelCriteria(cancel_event)]
)
stopping_criteria = self._cancel_stopping_criteria(cancel_event)
if stopping_criteria is not None:
generation_kwargs["stopping_criteria"] = stopping_criteria
def generate_fn():
with self._generation_lock:
@ -1562,6 +1722,8 @@ class InferenceBackend:
model.generate(**generation_kwargs)
except Exception as e:
err["msg"] = str(e)
if hasattr(streamer, "abort"):
streamer.abort()
logger.error(f"Generation error: {e}")
finally:
try:
@ -1579,12 +1741,17 @@ class InferenceBackend:
if think_prefix:
yield think_prefix
from queue import Empty
import time
generation_complete = False
cancel_deadline = None
try:
while True:
if cancel_event is not None and cancel_event.is_set():
break
if cancel_deadline is None:
cancel_deadline = time.monotonic() + 10
elif time.monotonic() >= cancel_deadline:
break
try:
new_token = next(streamer)
except StopIteration:
@ -1593,11 +1760,27 @@ class InferenceBackend:
except Empty:
if not thread.is_alive():
generation_complete = True
output = yield from self._drain_streamer_tail(
streamer, output, active_stop_token_ids
)
break
if cancel_deadline is not None:
remaining = cancel_deadline - time.monotonic()
if remaining <= 0:
break
thread.join(timeout = remaining)
if thread.is_alive():
break
generation_complete = True
output = yield from self._drain_streamer_tail(
streamer, output, active_stop_token_ids
)
break
continue
if new_token:
output += new_token
cleaned = self._clean_generated_text(output)
output, cleaned = self._append_stream_delta(
output, new_token, active_stop_token_ids
)
yield cleaned
finally:
# Set cancel_event only on early exit (user cancel), NOT on
@ -1606,16 +1789,21 @@ class InferenceBackend:
# disrupt the next serialized request (e.g. compare mode).
if cancel_event is not None and not generation_complete:
cancel_event.set()
thread.join(timeout = 10)
join_timeout = 10
if cancel_deadline is not None:
join_timeout = max(0, cancel_deadline - time.monotonic())
thread.join(timeout = join_timeout)
if thread.is_alive():
logger.warning("Generation thread did not exit after cancel/join timeout")
if err.get("msg"):
yield f"Error: {err['msg']}"
raise _GenerationThreadError(err["msg"])
except _GenerationThreadError:
raise
except Exception as e:
logger.error(f"Error during generation: {e}")
yield f"Error: {str(e)}"
raise
# ── Audio (TTS) Generation ────────────────────────────────────
@ -2104,8 +2292,42 @@ class InferenceBackend:
return img.resize(new_size, Image.Resampling.LANCZOS)
return img
def _clean_generated_text(self, text: str) -> str:
"""Strip leaked special tokens using the tokenizer's own token list."""
def _generation_stop_token_ids(self, model, generation_kwargs: dict):
"""Return the stop-token ids active for a ``generate`` call."""
if "eos_token_id" in generation_kwargs:
return generation_kwargs.get("eos_token_id")
generation_config = getattr(model, "generation_config", None)
eos_token_id = getattr(generation_config, "eos_token_id", None)
if eos_token_id is not None:
return eos_token_id
config = getattr(model, "config", None)
return getattr(config, "eos_token_id", None)
def _cancel_stopping_criteria(self, cancel_event):
"""Build a Transformers stopping criteria list for user cancellation."""
if cancel_event is None:
return None
from transformers.generation.stopping_criteria import (
StoppingCriteria,
StoppingCriteriaList,
)
class _CancelCriteria(StoppingCriteria):
def __init__(self, ev):
self.ev = ev
def __call__(self, input_ids, scores, **kwargs):
return self.ev.is_set()
return StoppingCriteriaList([_CancelCriteria(cancel_event)])
def _clean_generated_text(
self,
text: str,
*,
stop_token_ids = None,
) -> str:
"""Strip leaked response-boundary tokens after streaming."""
if self._is_gpt_oss_model():
# HarmonyTextStreamer emits clean <think>...</think>. Strip any
# harmony protocol tokens and other gpt-oss tokens (e.g.
@ -2115,10 +2337,28 @@ class InferenceBackend:
return text.strip()
tokenizer = self.models.get(self.active_model_name, {}).get("tokenizer")
tokenizer = getattr(tokenizer, "tokenizer", tokenizer)
if tokenizer:
for token in getattr(tokenizer, "all_special_tokens", []):
if token in text:
text = text.replace(token, "")
if stop_token_ids is None:
stop_token_ids = self.models.get(self.active_model_name, {}).get(
"chat_turn_end_eos_ids"
)
if isinstance(stop_token_ids, int):
stop_token_ids = (stop_token_ids,)
for token_id in stop_token_ids or ():
try:
token = tokenizer.convert_ids_to_tokens(int(token_id))
except Exception:
token = None
if isinstance(token, str) and token and text.endswith(token):
text = text[: -len(token)]
elif (
isinstance(token, str)
and token
and text.endswith("</think>")
and text[: -len("</think>")].endswith(token)
):
text = text[: -len("</think>") - len(token)] + "</think>"
return text.strip()
def _load_chat_template_info(self, model_name: str):

368
studio/backend/core/inference/llama_admission.py Normal file
View file

@ -0,0 +1,368 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
"""Admission control for local llama-server generation requests.
The helpers in this module deliberately know nothing about FastAPI, SSE, or the
OpenAI-compatible route shape. They only coordinate how many upstream generation
requests may be active for one llama-server backend and provide a cancellable
FIFO queue for excess requests.
"""
from __future__ import annotations
import asyncio
import os
import threading
from collections import deque
from dataclasses import dataclass
from typing import Deque, Optional
ADMISSION_CONTROL_ENV = "UNSLOTH_OPENAI_COMPAT_ADMISSION_CONTROL"
ADMISSION_QUEUE_TIMEOUT_ENV = "UNSLOTH_OPENAI_COMPAT_ADMISSION_QUEUE_TIMEOUT"
ADMISSION_KEEPALIVE_INTERVAL_ENV = "UNSLOTH_OPENAI_COMPAT_ADMISSION_KEEPALIVE_INTERVAL"
ADMISSION_MAX_QUEUE_ENV = "UNSLOTH_OPENAI_COMPAT_ADMISSION_MAX_QUEUE"
DEFAULT_ADMISSION_ENABLED = True
DEFAULT_ADMISSION_QUEUE_TIMEOUT_S = None
DEFAULT_ADMISSION_KEEPALIVE_INTERVAL_S = 5.0
DEFAULT_ADMISSION_MAX_QUEUE = 64
@dataclass(frozen = True)
class LlamaAdmissionConfig:
enabled: bool = DEFAULT_ADMISSION_ENABLED
queue_timeout_s: Optional[float] = DEFAULT_ADMISSION_QUEUE_TIMEOUT_S
keepalive_interval_s: float = DEFAULT_ADMISSION_KEEPALIVE_INTERVAL_S
max_queue: Optional[int] = DEFAULT_ADMISSION_MAX_QUEUE
@dataclass(frozen = True)
class LlamaAdmissionSnapshot:
key: str
capacity: int
active: int
queued: int
class LlamaAdmissionError(Exception):
def __init__(
self,
message: str,
*,
snapshot: Optional[LlamaAdmissionSnapshot] = None,
):
super().__init__(message)
self.snapshot = snapshot
class LlamaAdmissionQueueFull(LlamaAdmissionError):
pass
class LlamaAdmissionTimeout(LlamaAdmissionError):
pass
class LlamaAdmissionCancelled(LlamaAdmissionError):
pass
def _bool_env(name: str, default: bool) -> bool:
value = os.environ.get(name)
if value is None or not value.strip():
return default
value = value.strip().lower()
if value in {"1", "true", "yes", "on"}:
return True
if value in {"0", "false", "no", "off"}:
return False
return default
def _optional_positive_float_env(name: str, default: Optional[float]) -> Optional[float]:
value = os.environ.get(name)
if value is None or not value.strip():
return default
try:
parsed = float(value.strip())
except ValueError:
return default
return parsed if parsed > 0 else None
def _positive_float_env(name: str, default: float) -> float:
value = os.environ.get(name)
if value is None or not value.strip():
return default
try:
parsed = float(value.strip())
except ValueError:
return default
return parsed if parsed > 0 else default
def _optional_positive_int_env(name: str, default: Optional[int]) -> Optional[int]:
value = os.environ.get(name)
if value is None or not value.strip():
return default
try:
parsed = int(value.strip())
except ValueError:
return default
return parsed if parsed > 0 else None
def llama_admission_config_from_env() -> LlamaAdmissionConfig:
return LlamaAdmissionConfig(
enabled = _bool_env(ADMISSION_CONTROL_ENV, DEFAULT_ADMISSION_ENABLED),
queue_timeout_s = _optional_positive_float_env(
ADMISSION_QUEUE_TIMEOUT_ENV,
DEFAULT_ADMISSION_QUEUE_TIMEOUT_S,
),
keepalive_interval_s = _positive_float_env(
ADMISSION_KEEPALIVE_INTERVAL_ENV,
DEFAULT_ADMISSION_KEEPALIVE_INTERVAL_S,
),
max_queue = _optional_positive_int_env(
ADMISSION_MAX_QUEUE_ENV,
DEFAULT_ADMISSION_MAX_QUEUE,
),
)
@dataclass
class _Waiter:
loop: asyncio.AbstractEventLoop
future: asyncio.Future
cancelled: bool = False
granted_lease: Optional["LlamaAdmissionLease"] = None
class LlamaAdmissionLease:
def __init__(self, queue: Optional["LlamaAdmissionQueue"]):
self._queue = queue
self._released = False
self._release_lock = threading.Lock()
def release(self) -> None:
queue = None
with self._release_lock:
if self._released:
return
self._released = True
queue = self._queue
if queue is not None:
queue.release()
async def __aenter__(self) -> "LlamaAdmissionLease":
return self
async def __aexit__(self, *_args) -> None:
self.release()
class LlamaAdmissionReservation:
def __init__(
self,
*,
queue: Optional["LlamaAdmissionQueue"],
lease: Optional[LlamaAdmissionLease] = None,
waiter: Optional[_Waiter] = None,
snapshot: Optional[LlamaAdmissionSnapshot] = None,
):
self._queue = queue
self._lease = lease
self._waiter = waiter
self.snapshot = snapshot
@property
def is_cancelled(self) -> bool:
return self._lease is None and self._waiter is None
def lease_nowait(self) -> Optional[LlamaAdmissionLease]:
if self._lease is not None:
return self._lease
if self._waiter is None or not self._waiter.future.done():
return None
if self._waiter.future.cancelled():
self._waiter.cancelled = True
self._waiter = None
return None
self._lease = self._waiter.future.result()
self._waiter = None
return self._lease
async def wait(self, timeout_s: float) -> Optional[LlamaAdmissionLease]:
lease = self.lease_nowait()
if lease is not None:
return lease
if self._waiter is None:
return None
waiter = self._waiter
try:
await asyncio.wait_for(asyncio.shield(waiter.future), timeout = timeout_s)
except asyncio.CancelledError:
if waiter.future.cancelled():
waiter.cancelled = True
if self._waiter is waiter:
self._waiter = None
return None
raise
return self.lease_nowait()
def cancel(self) -> None:
lease = self.lease_nowait()
if lease is not None:
lease.release()
self._lease = None
return
if self._queue is not None and self._waiter is not None:
self._queue.cancel(self._waiter)
self._waiter = None
def snapshot_now(self) -> Optional[LlamaAdmissionSnapshot]:
if self._queue is None:
return self.snapshot
return self._queue.snapshot()
class LlamaAdmissionQueue:
def __init__(self, key: str):
self.key = key
self._lock = threading.Lock()
self._active = 0
self._capacity = 1
self._waiters: Deque[_Waiter] = deque()
def reserve(self, *, capacity: int, config: LlamaAdmissionConfig) -> LlamaAdmissionReservation:
capacity = max(1, int(capacity or 1))
if not config.enabled:
return LlamaAdmissionReservation(
queue = None,
lease = LlamaAdmissionLease(None),
snapshot = LlamaAdmissionSnapshot(self.key, capacity, 0, 0),
)
loop = asyncio.get_running_loop()
with self._lock:
self._capacity = capacity
self._prune_waiters_locked()
self._grant_waiters_locked()
if self._active < self._capacity and not self._waiters:
self._active += 1
return LlamaAdmissionReservation(
queue = self,
lease = LlamaAdmissionLease(self),
snapshot = self._snapshot_locked(),
)
if config.max_queue is not None and len(self._waiters) >= config.max_queue:
raise LlamaAdmissionQueueFull(
"llama-server generation queue is full",
snapshot = self._snapshot_locked(),
)
waiter = _Waiter(
loop = loop,
future = loop.create_future(),
)
self._waiters.append(waiter)
return LlamaAdmissionReservation(
queue = self,
waiter = waiter,
snapshot = self._snapshot_locked(),
)
def release(self) -> None:
with self._lock:
if self._active > 0:
self._active -= 1
self._grant_waiters_locked()
def cancel(self, waiter: _Waiter) -> None:
lease_to_release = None
with self._lock:
waiter.cancelled = True
try:
self._waiters.remove(waiter)
except ValueError:
pass
if waiter.granted_lease is not None:
lease_to_release = waiter.granted_lease
waiter.granted_lease = None
if not waiter.future.done():
waiter.loop.call_soon_threadsafe(waiter.future.cancel)
if lease_to_release is not None:
lease_to_release.release()
def snapshot(self) -> LlamaAdmissionSnapshot:
with self._lock:
self._prune_waiters_locked()
return self._snapshot_locked()
def is_idle(self) -> bool:
with self._lock:
self._prune_waiters_locked()
return self._active == 0 and not self._waiters
def _grant_waiters_locked(self) -> None:
self._prune_waiters_locked()
while self._waiters and self._active < self._capacity:
waiter = self._waiters.popleft()
if waiter.cancelled or waiter.future.done():
continue
self._active += 1
lease = LlamaAdmissionLease(self)
waiter.granted_lease = lease
waiter.loop.call_soon_threadsafe(self._deliver_lease, waiter, lease)
def _deliver_lease(self, waiter: _Waiter, lease: LlamaAdmissionLease) -> None:
if waiter.cancelled or waiter.future.done():
waiter.granted_lease = None
if not waiter.future.done():
waiter.future.cancel()
lease.release()
return
try:
waiter.future.set_result(lease)
waiter.granted_lease = None
except asyncio.InvalidStateError:
waiter.granted_lease = None
lease.release()
def _prune_waiters_locked(self) -> None:
self._waiters = deque(
waiter for waiter in self._waiters if not waiter.cancelled and not waiter.future.done()
)
def _snapshot_locked(self) -> LlamaAdmissionSnapshot:
return LlamaAdmissionSnapshot(
key = self.key,
capacity = self._capacity,
active = self._active,
queued = len(self._waiters),
)
_QUEUES_LOCK = threading.Lock()
_QUEUES: dict[str, LlamaAdmissionQueue] = {}
def get_llama_admission_queue(key: str) -> LlamaAdmissionQueue:
with _QUEUES_LOCK:
queue = _QUEUES.get(key)
if queue is None:
queue = LlamaAdmissionQueue(key)
_QUEUES[key] = queue
# base_url carries a fresh ephemeral port on every model load, so
# each load registers a new key. Drop the now-idle queues from prior
# loads so the registry can't grow without bound on a long-running
# server. Queues with in-flight requests are kept until they drain.
for stale_key in [k for k in _QUEUES if k != key and _QUEUES[k].is_idle()]:
del _QUEUES[stale_key]
return queue
def reset_llama_admission_queues() -> None:
with _QUEUES_LOCK:
_QUEUES.clear()

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

File diff suppressed because it is too large Load diff

774
studio/backend/core/inference/mcp_client.py
View file

@ -4,11 +4,16 @@
from __future__ import annotations
import asyncio
import atexit
import concurrent.futures
import hashlib
import json
import os
import shlex
import sys
import threading
import time
import uuid
from typing import Any, Optional
from loggers import get_logger
@ -115,6 +120,18 @@ def join_stdio_command(parts: list[str]) -> str:
return shlex.join(parts)
def _stdio_log_id(url: str) -> str:
"""A non-secret label for logs. stdio commands can embed credentials in argv
(e.g. ``npx server --token sk-...``), so never log the raw command; use the
executable basename plus a short digest of the full command instead."""
try:
parts = parse_stdio_command(url)
exe = os.path.basename(parts[0]) if parts else "<empty>"
except Exception: # noqa: BLE001
exe = "<invalid>"
return f"{exe}#{hashlib.sha256(url.encode()).hexdigest()[:12]}"
def stdio_mcp_enabled() -> bool:
"""stdio MCP servers spawn local processes as the backend user (bypassing the
sandbox), so allowed only when the host is the user's own machine. On startup
@ -192,7 +209,6 @@ async def clear_oauth_tokens_async(url: str) -> None:
auth = OAuth(mcp_url = url, token_storage = _oauth_store())
await auth.token_storage_adapter.clear()
except Exception as exc: # noqa: BLE001
# Cleanup is best-effort; the row delete still wins.
logger.warning("Failed to clear OAuth tokens for %s: %s", url, exc)
@ -237,6 +253,502 @@ def _client(
return Client(transport_cls(url = url, headers = headers or None, auth = auth))
# Persistent stdio sessions: a stdio MCP server owns live state (a browser, a
# DB handle), so keep one connected client per (command, env, chat session) on
# a dedicated event-loop thread instead of respawning per call.
_STDIO_SESSION_IDLE_TTL = 300.0
_STDIO_SESSION_REAP_INTERVAL = 30.0
_STDIO_CONNECT_TIMEOUT = 60.0 # allows first-run `npx -y ...` package download
_STDIO_CLOSE_TIMEOUT = 10.0
_STDIO_WEDGE_MARGIN = 15.0
# Cap concurrent persistent sessions: each owns a subprocess + loop thread, and
# the scope includes a caller-supplied thread_id, so an unbounded cache is a
# resource-exhaustion surface. Overridable via env for large deployments.
try:
_STDIO_MAX_SESSIONS = max(1, int(os.environ.get("UNSLOTH_STUDIO_MAX_STDIO_MCP_SESSIONS", "32")))
except ValueError:
_STDIO_MAX_SESSIONS = 32
def _is_tool_error(exc: BaseException) -> bool:
"""A tool-level failure (the tool ran and errored) leaves the transport alive,
so the session is kept; fastmcp raises ToolError for these. Anything else from
call_tool is transport-level. Version-safe (fastmcp 3.0.2 has no dead probe)."""
try:
from fastmcp.exceptions import ToolError
except Exception: # noqa: BLE001
return False
return isinstance(exc, ToolError)
def _transport_dead(session) -> bool:
"""Best-effort, version-adaptive liveness probe for a cached stdio client.
``Client.is_connected()`` only checks a session object exists, not that the
subprocess is alive, so it is never used here. Returns True only when the
transport is positively gone; unknown returns False (the call surfaces it)."""
client = getattr(session, "client", None)
if client is None:
return True
transport = getattr(client, "transport", None)
probe = getattr(transport, "_is_session_dead", None)
if callable(probe):
try:
if probe():
return True
except Exception: # noqa: BLE001
pass
connect_task = getattr(transport, "_connect_task", None)
if connect_task is not None:
try:
if connect_task.done():
return True
except Exception: # noqa: BLE001
pass
return False
class _SessionWedged(Exception):
pass
class _SessionClosed(Exception):
"""The session was closed (server update/delete/shutdown) mid-call."""
def _abort_future(future) -> None:
# Let the cancelled coroutine unwind before its loop is stopped.
future.cancel()
try:
future.result(1.0)
except BaseException: # noqa: BLE001
pass
class _StdioSession:
def __init__(self, url: str, headers: Optional[dict]):
self.url = url
self.headers = headers
self.client = None
self.closed = threading.Event()
self.defunct = False # discarded; close once in_flight drains (see _retire)
self._close_lock = threading.Lock()
self.call_lock = threading.Lock() # serializes tool calls on this session
self.last_used = time.monotonic()
self.in_flight = 0 # guarded by _stdio_sessions_lock
# On Windows a bare new_event_loop() can be a SelectorEventLoop (if any
# component set that policy), which cannot spawn subprocesses natively;
# force a ProactorEventLoop so the stdio transport always works.
if sys.platform == "win32":
self.loop = asyncio.ProactorEventLoop()
else:
self.loop = asyncio.new_event_loop()
self._thread = threading.Thread(
target = self._run_loop, name = "mcp-stdio-session", daemon = True
)
self._thread.start()
def _run_loop(self) -> None:
asyncio.set_event_loop(self.loop)
try:
self.loop.run_forever()
finally:
self.loop.close()
def connect(self, timeout: Optional[float], cancel_event) -> None:
async def _open():
client = _client(self.url, self.headers)
await client.__aenter__()
# Publish on the loop thread with no await in between: if an abort
# races a just-completed connect, close() still sees the client and
# __aexit__s it instead of orphaning the subprocess.
self.client = client
return client
future = asyncio.run_coroutine_threadsafe(_open(), self.loop)
# timeout=None means unlimited (no connect deadline); a finite caller
# timeout still bounds connect by min(timeout, _STDIO_CONNECT_TIMEOUT).
window = None if timeout is None else min(timeout, _STDIO_CONNECT_TIMEOUT)
deadline = None if window is None else time.monotonic() + window
while True:
if cancel_event is not None and cancel_event.is_set():
_abort_future(future)
raise _MCPCancelled
try:
future.result(0.05)
return
except (concurrent.futures.TimeoutError, asyncio.TimeoutError):
if future.done():
raise # the connect itself failed fast; don't wait out the window
if deadline is not None and time.monotonic() >= deadline:
_abort_future(future)
raise asyncio.TimeoutError
def is_connected(self) -> bool:
client = self.client
if client is None:
return False
probe = getattr(client, "is_connected", None)
try:
return bool(probe()) if callable(probe) else True
except Exception:
return False
def run(self, coro, timeout: Optional[float]):
self.last_used = time.monotonic()
future = asyncio.run_coroutine_threadsafe(coro, self.loop)
# The coroutine enforces the tool timeout; the margin only catches a
# wedged loop. No deadline at all when the caller set none -- but poll
# so a session closed under us (server update/delete) can't hang the
# request thread forever on a stopped loop.
deadline = None if timeout is None else time.monotonic() + timeout + _STDIO_WEDGE_MARGIN
try:
while True:
try:
return future.result(0.25)
except concurrent.futures.CancelledError:
# Only close() cancels in-flight tasks (in _shutdown).
raise _SessionClosed
except (concurrent.futures.TimeoutError, asyncio.TimeoutError):
if future.done():
raise # the call's own timeout; the session stays usable
if self.closed.is_set():
future.cancel()
raise _SessionClosed
if deadline is not None and time.monotonic() >= deadline:
future.cancel()
raise _SessionWedged
finally:
self.last_used = time.monotonic()
def close(self) -> None:
# Idempotent: a discard racing close_stdio_sessions() may close twice.
# Setting `closed` first also unblocks run() waiters (they poll it).
with self._close_lock:
if self.closed.is_set():
return
self.closed.set()
loop = getattr(self, "loop", None)
loop_alive = loop is not None and not loop.is_closed()
if loop_alive:
async def _shutdown() -> None:
# Runs on the loop thread, so it serializes with an aborted
# connect() that finished anyway and just published its client.
client, self.client = self.client, None
if client is not None:
await client.__aexit__(None, None, None)
# Cancel in-flight calls so they unwind before loop.stop
# (their run() waiters have already been released via `closed`).
for task in asyncio.all_tasks():
if task is not asyncio.current_task():
task.cancel()
try:
asyncio.run_coroutine_threadsafe(_shutdown(), loop).result(_STDIO_CLOSE_TIMEOUT)
except Exception as exc: # noqa: BLE001
logger.warning(
"MCP stdio session close failed for %s: %s",
_stdio_log_id(getattr(self, "url", "")),
exc,
)
try:
loop.call_soon_threadsafe(loop.stop)
except RuntimeError:
pass
else:
self.client = None
thread = getattr(self, "_thread", None)
if thread is not None:
thread.join(timeout = 5.0)
_stdio_sessions: dict[tuple, _StdioSession] = {}
# Per-key locks so a slow connect/close never blocks unrelated servers; the
# global lock only guards the dicts.
class _StdioKeyLock:
"""A per-key lock that can be removed once nobody references it."""
def __init__(self) -> None:
self.lock = threading.Lock()
self.users = 0 # guarded by _stdio_sessions_lock
_stdio_key_locks: dict[tuple, _StdioKeyLock] = {}
_stdio_sessions_lock = threading.Lock()
_stdio_reaper_started = False
# close_stdio_sessions() can only close sessions already published in
# _stdio_sessions; one still inside connect() would be missed and cached
# stale. Bump a generation on every close so that connect discards its
# session instead of publishing it. Guarded by _stdio_sessions_lock.
_stdio_close_all_gen = 0
_stdio_url_close_gen: dict[str, int] = {}
_stdio_cfg_close_gen: dict[tuple, int] = {}
# close_stdio_sessions(url): match any env for that command.
_ANY_HEADERS = object()
def _headers_key(headers: Optional[dict]) -> tuple:
return tuple(sorted((headers or {}).items()))
def _url_close_key(url: str) -> str:
# Commands/URLs (token args, embedded credentials) and env values can hold
# secrets and these maps are never pruned; key by digest so closed/edited
# configs don't retain them in memory forever.
return hashlib.sha256(url.encode()).hexdigest()
def _cfg_close_key(url: str, headers: Optional[dict]) -> str:
return hashlib.sha256(repr((url, _headers_key(headers))).encode()).hexdigest()
def _stdio_close_generation(url: str, headers: Optional[dict]) -> tuple[int, int, int]:
return (
_stdio_close_all_gen,
_stdio_url_close_gen.get(_url_close_key(url), 0),
_stdio_cfg_close_gen.get(_cfg_close_key(url, headers), 0),
)
def _session_key(url: str, headers: Optional[dict], scope: Optional[str]) -> tuple:
return (url, _headers_key(headers), scope or "")
def _checkout_stdio_session(key: tuple) -> Optional[_StdioSession]:
session = _stdio_sessions.get(key)
if session is not None and session.is_connected():
session.last_used = time.monotonic()
session.in_flight += 1
return session
return None
def _borrow_stdio_key_lock(key: tuple) -> _StdioKeyLock:
"""Return a stable per-key lock while a caller waits for/connects it."""
key_lock = _stdio_key_locks.setdefault(key, _StdioKeyLock())
key_lock.users += 1
return key_lock
def _discard_stdio_key_lock(key: tuple) -> None:
key_lock = _stdio_key_locks.get(key)
if key_lock is not None and key_lock.users == 0 and key not in _stdio_sessions:
_stdio_key_locks.pop(key, None)
def _return_stdio_key_lock(key: tuple, key_lock: _StdioKeyLock) -> None:
with _stdio_sessions_lock:
key_lock.users -= 1
_discard_stdio_key_lock(key)
def _get_stdio_session(
url: str, headers: Optional[dict], scope: Optional[str], deadline, cancel_event, config_check
) -> _StdioSession:
"""``deadline`` is the caller's absolute monotonic budget (None = no limit):
the key-lock wait and the connect share it, so a slow startup can't stack
full timeout windows (see _call_stdio_tool)."""
global _stdio_reaper_started
key = _session_key(url, headers, scope)
with _stdio_sessions_lock:
session = _checkout_stdio_session(key)
if session is not None:
return session
key_lock = _borrow_stdio_key_lock(key)
try:
# Poll the acquire with connect()'s deadline/cancel semantics: a second
# same-scope call must not block uncancellably behind another caller's
# slow startup (e.g. a first-run npx download).
remaining = None if deadline is None else max(0.0, deadline - time.monotonic())
# timeout=None means no key-lock deadline (only cancel unblocks it).
window = None if remaining is None else min(remaining, _STDIO_CONNECT_TIMEOUT)
lock_deadline = None if window is None else time.monotonic() + window
while not key_lock.lock.acquire(timeout = 0.05):
if cancel_event is not None and cancel_event.is_set():
raise _MCPCancelled
if lock_deadline is not None and time.monotonic() >= lock_deadline:
raise asyncio.TimeoutError
try:
stale = None
with _stdio_sessions_lock:
session = _checkout_stdio_session(key)
if session is not None:
return session
if key in _stdio_sessions:
stale = _stdio_sessions.pop(key)
generation = _stdio_close_generation(url, headers)
if stale is not None:
_retire_stdio_session(stale)
session = _StdioSession(url, headers)
try:
session.connect(
None if deadline is None else max(0.0, deadline - time.monotonic()),
cancel_event,
)
except Exception:
session.close()
raise
# A caller can read the server row, then lose to an update/delete whose close ran
# before our generation snapshot. Re-verify the row after connect; the generation check
# below covers a close landing between this check and publish.
if config_check is not None:
try:
current = bool(config_check())
except Exception: # noqa: BLE001
current = False
if not current:
session.close()
raise RuntimeError("MCP server was updated or removed while connecting")
evicted: list = []
with _stdio_sessions_lock:
closed_while_connecting = _stdio_close_generation(url, headers) != generation
if not closed_while_connecting:
session.in_flight = 1
evicted = _evict_stdio_lru_locked() # bound the cache (LRU idle)
_stdio_sessions[key] = session
if not _stdio_reaper_started:
_stdio_reaper_started = True
threading.Thread(
target = _stdio_session_reaper, name = "mcp-stdio-reaper", daemon = True
).start()
atexit.register(close_stdio_sessions)
for victim in evicted:
logger.info("Evicting LRU idle stdio MCP session: %s", _stdio_log_id(victim.url))
victim.close()
if closed_while_connecting:
session.close()
raise RuntimeError("MCP server was updated or removed while connecting")
return session
finally:
key_lock.lock.release()
finally:
_return_stdio_key_lock(key, key_lock)
def _release_stdio_session(session: _StdioSession) -> None:
victims: list = []
with _stdio_sessions_lock:
session.in_flight = max(0, session.in_flight - 1)
session.last_used = time.monotonic()
close_now = session.defunct and session.in_flight == 0
# Re-enforce the cap once a burst's sessions go idle. Insert-time eviction
# only trims idle sessions, so it can overshoot while every cached session
# is busy; reclaim that overshoot here instead of waiting for the idle
# reaper. Never evict the session we just used (its last_used is newest).
while len(_stdio_sessions) > _STDIO_MAX_SESSIONS:
idle = [
(s.last_used, k)
for k, s in _stdio_sessions.items()
if s.in_flight == 0 and s is not session
]
if not idle:
break
_, oldest = min(idle, key = lambda item: item[0])
victims.append(_stdio_sessions.pop(oldest))
_discard_stdio_key_lock(oldest)
if close_now:
session.close()
for victim in victims:
victim.close()
def _retire_stdio_session(session: _StdioSession) -> None:
"""Close a discarded session, but only once no other borrower is mid-call
on it -- overlapping same-scope calls share one client, and one call's
timeout must not kill another's in-flight request. The last borrower's
_release_stdio_session() performs the deferred close."""
with _stdio_sessions_lock:
session.defunct = True
busy = session.in_flight > 0
if not busy:
session.close()
def _drop_stdio_session(key: tuple, session: _StdioSession) -> None:
with _stdio_sessions_lock:
if _stdio_sessions.get(key) is session:
_stdio_sessions.pop(key)
_discard_stdio_key_lock(key)
_retire_stdio_session(session)
def _evict_stdio_lru_locked() -> list:
"""Caller holds _stdio_sessions_lock. Evict least-recently-used *idle*
sessions until the cache is under the cap. Returns the evicted sessions so
the caller can close them OUTSIDE the lock. If every session is busy the
cache may transiently overshoot rather than kill an in-flight call."""
victims: list = []
while len(_stdio_sessions) >= _STDIO_MAX_SESSIONS:
idle = [(s.last_used, k) for k, s in _stdio_sessions.items() if s.in_flight == 0]
if not idle:
break
_, oldest = min(idle, key = lambda item: item[0])
victims.append(_stdio_sessions.pop(oldest))
_discard_stdio_key_lock(oldest)
return victims
def close_stdio_sessions(url: Optional[str] = None, headers = _ANY_HEADERS) -> None:
"""Close persistent stdio sessions: all of them (``url`` None), every env
for one command (``headers`` omitted), or one server config (url + headers).
Two server rows can share a command with different envs; editing one must
not kill the other's live state, so the routes pass the edited row's env."""
global _stdio_close_all_gen
# HTTP/SSE servers are never cached as stdio sessions, so a specific non-stdio
# url has nothing to close and must not accrue a close-generation entry.
if url is not None and not is_stdio(url):
return
hk = None if headers is _ANY_HEADERS else _headers_key(headers)
with _stdio_sessions_lock:
if url is None:
_stdio_close_all_gen += 1
elif hk is None:
uk = _url_close_key(url)
_stdio_url_close_gen[uk] = _stdio_url_close_gen.get(uk, 0) + 1
else:
cfg = _cfg_close_key(url, headers)
_stdio_cfg_close_gen[cfg] = _stdio_cfg_close_gen.get(cfg, 0) + 1
keys = [
k
for k in _stdio_sessions
if (url is None or k[0] == url) and (hk is None or k[1] == hk)
]
sessions = [_stdio_sessions.pop(k) for k in keys]
for key in keys:
_discard_stdio_key_lock(key)
for session in sessions:
session.close()
def _reap_idle_stdio_sessions(now: Optional[float] = None) -> None:
now = time.monotonic() if now is None else now
with _stdio_sessions_lock:
expired = [
key
for key, session in _stdio_sessions.items()
if session.in_flight == 0 and now - session.last_used >= _STDIO_SESSION_IDLE_TTL
]
sessions = [_stdio_sessions.pop(key) for key in expired]
for key in expired:
_discard_stdio_key_lock(key)
for session in sessions:
logger.info("Closing idle stdio MCP session: %s", _stdio_log_id(session.url))
session.close()
def _stdio_session_reaper() -> None:
while True:
time.sleep(_STDIO_SESSION_REAP_INTERVAL)
try:
_reap_idle_stdio_sessions()
except Exception as exc: # noqa: BLE001
logger.debug("stdio session reaper iteration failed: %s", exc)
async def list_tools_async(
url: str,
headers: Optional[dict] = None,
@ -251,11 +763,10 @@ async def list_tools_async(
return await asyncio.wait_for(_fetch(), timeout = timeout)
# Discovered-tool cache, keyed by MCP server id. get_enabled_mcp_tools()
# probes a server only on a cache miss, keeping MCP discovery off the chat
# send's critical path -- tool schemas are stable within a session. The
# /refresh route warms it; a URL/header/OAuth change or a delete evicts it.
# Successful probes are cached indefinitely.
# Discovered-tool cache, keyed by MCP server id. get_enabled_mcp_tools() probes a server only
# on a cache miss, keeping MCP discovery off the chat send's critical path -- tool schemas are
# stable within a session. The /refresh route warms it; a URL/header/OAuth change or a delete
# evicts it. Successful probes are cached indefinitely.
_tool_cache: dict[str, list[dict]] = {}
# server_id -> monotonic time before which a failed server must not be
@ -298,23 +809,210 @@ def invalidate_tool_cache(server_id: Optional[str] = None) -> None:
_probe_cooloff_until.pop(server_id, None)
MCP_IMAGES_SENTINEL = "__MCP_IMAGES__:"
MAX_IMAGE_PAYLOAD_CHARS = 12_000_000
def _flatten_result(result: Any) -> str:
parts = []
images = []
omitted = 0
budget = MAX_IMAGE_PAYLOAD_CHARS
for block in getattr(result, "content", None) or []:
text = getattr(block, "text", None)
if text:
parts.append(str(text))
continue
data = getattr(block, "data", None)
mime = getattr(block, "mimeType", None)
if data and isinstance(mime, str) and mime.startswith("image/"):
data = str(data)
if len(data) > budget:
omitted += 1
continue
budget -= len(data)
images.append({"data": data, "mimeType": mime})
body = "\n".join(parts)
if not body:
structured = getattr(result, "structured_content", None)
body = str(structured) if structured is not None else ""
if images or omitted:
notes = []
if images:
n = len(images)
notes.append(f"{n} image{'s' if n > 1 else ''} attached; displayed to the user")
if omitted:
notes.append(f"{omitted} image{'s' if omitted > 1 else ''} omitted (too large)")
note = f"[{'; '.join(notes)}]"
body = f"{body}\n{note}" if body else note
if getattr(result, "is_error", False):
# "Error: " prefix triggers tool_call_parser's TOOL_ERROR_PREFIXES nudge.
return f"Error: {body}" if body else "Error: tool returned no content"
body = f"Error: {body}" if body else "Error: tool returned no content"
if images:
body += "\n" + MCP_IMAGES_SENTINEL + json.dumps(images)
return body
async def _race_tool_call(call_coro, timeout: Optional[float], cancel_event) -> Any:
"""Await ``call_coro`` under ``timeout``, polling ``cancel_event`` so a
/cancel POST interrupts even mid-network-read."""
async def _watch_cancel() -> None:
while cancel_event is not None and not cancel_event.is_set():
await asyncio.sleep(0.05)
if cancel_event is not None and cancel_event.is_set():
call_coro.close()
raise _MCPCancelled
call_task = asyncio.create_task(call_coro)
if cancel_event is None:
return await asyncio.wait_for(call_task, timeout = timeout)
watch_task = asyncio.create_task(_watch_cancel())
try:
done, pending = await asyncio.wait(
{call_task, watch_task},
timeout = timeout,
return_when = asyncio.FIRST_COMPLETED,
)
finally:
for t in (call_task, watch_task):
if not t.done():
t.cancel()
if not done:
raise asyncio.TimeoutError
if call_task in done:
return call_task.result()
raise _MCPCancelled
def _call_stdio_tool(
url: str,
headers: Optional[dict],
name: str,
args: dict,
timeout,
cancel_event,
scope: Optional[str],
config_check,
) -> Any:
if cancel_event is not None and cancel_event.is_set():
raise _MCPCancelled
# One deadline covers the key-lock wait, connect, call-lock wait, and the
# call itself, matching the one-shot/HTTP paths where the timeout wrapped
# connect plus call in a single window.
deadline = None if timeout is None else time.monotonic() + timeout
def _remaining() -> Optional[float]:
return None if deadline is None else max(0.0, deadline - time.monotonic())
# Callers without a Studio session id must retain the former one-shot
# behavior: no browser/cookie/tool state can leak into another request.
# Use an ephemeral key (and close it below) rather than the shared empty
# scope that the persistent-session cache used previously.
def _config_ok() -> bool:
if config_check is None:
return True
try:
return bool(config_check())
except Exception: # noqa: BLE001
return False
ephemeral = not scope
if ephemeral:
scope = f"request-{uuid.uuid4().hex}"
key = _session_key(url, headers, scope)
# attempt 0 may find the cached session stale/dead *before* dispatch and
# reconnect once (safe); attempt 1 is a freshly connected session.
for attempt in (0, 1):
session = _get_stdio_session(url, headers, scope, deadline, cancel_event, config_check)
try:
# Serialize calls per session: overlapping same-scope calls must
# not interleave operations on one stateful server (browser, REPL).
while not session.call_lock.acquire(timeout = 0.05):
if cancel_event is not None and cancel_event.is_set():
raise _MCPCancelled
rem = _remaining()
if rem is not None and rem <= 0:
raise asyncio.TimeoutError
except BaseException:
# Never touched the transport: keep the session for its borrower.
_release_stdio_session(session)
if ephemeral:
_drop_stdio_session(key, session)
raise
discard_session = ephemeral
retry = False
try:
# We may have waited on the call lock while another caller's timeout retired this
# session, a server update/delete invalidated it, or a reused subprocess died. Re-check
# all three before dispatch so we never run on a retired/dead client or a stale config.
if session.closed.is_set():
# Intentional close (server update/delete/shutdown): don't retry on stale config.
discard_session = True
raise RuntimeError("MCP server was updated or removed during the call")
elif session.defunct:
# A concurrent same-scope caller's timeout retired this session;
# move to a fresh one instead of reusing the retired client.
discard_session = True
if attempt == 0:
retry = True
else:
raise RuntimeError("MCP server session was retired during the call")
elif not _config_ok():
discard_session = True
raise RuntimeError("MCP server was updated or removed during the call")
elif _transport_dead(session):
# Dead BEFORE dispatch: no request was sent, so reconnect + retry.
discard_session = True
if attempt == 0:
retry = True
else:
raise RuntimeError("MCP server connection is not available")
else:
rem = _remaining()
coro = _race_tool_call(session.client.call_tool(name, args), rem, cancel_event)
return session.run(coro, rem)
except (_MCPCancelled, asyncio.TimeoutError):
# _race_tool_call cancels the pending call but cancellation is
# cooperative. Never return this client to the cache while the
# timed-out/cancelled operation might still run on its transport.
discard_session = True
raise
except _SessionWedged:
discard_session = True
raise asyncio.TimeoutError
except _SessionClosed:
# close_stdio_sessions() shut this session mid-call (server
# update/delete/shutdown); don't retry on the stale config.
discard_session = True
raise RuntimeError("MCP server was updated or removed during the call")
except Exception as exc:
if session.closed.is_set():
# An intentional close (server update/delete) can surface as a plain transport
# error or AttributeError instead of _SessionClosed; don't mistake it for a crash.
discard_session = True
raise RuntimeError("MCP server was updated or removed during the call")
# ToolError leaves the transport alive -> keep the session so its state
# survives. Any other exception is transport-level (dead subprocess,
# broken pipe): evict so it can't poison the scope, but DO NOT replay
# (the tool may already have run); the next call opens a fresh session.
if not _is_tool_error(exc):
discard_session = True
raise
finally:
# Set defunct + remove from the cache BEFORE releasing the call lock,
# so a queued same-scope borrower observes the retirement and opens a
# fresh session instead of reusing this one.
_release_stdio_session(session)
if discard_session:
_drop_stdio_session(key, session)
session.call_lock.release()
if not retry:
break
raise RuntimeError("unreachable")
def call_tool_sync(
url: str,
headers: Optional[dict],
@ -323,55 +1021,35 @@ def call_tool_sync(
timeout: Optional[float] = 300.0,
use_oauth: bool = False,
cancel_event = None,
scope: Optional[str] = None,
config_check = None,
) -> str:
"""Synchronously call an MCP tool.
"""Synchronously call an MCP tool. stdio servers reuse a persistent session
keyed by (command, env, scope) only when ``scope`` is provided; calls
without one stay one-shot. HTTP servers always stay one-shot.
``cancel_event`` (threading.Event) cancels the in-flight call when set.
``config_check`` (callable -> bool) re-validates the caller's server config
before a fresh stdio session is cached; False fails the call."""
``cancel_event``: optional ``threading.Event``. When set, the in-flight call is
cancelled and a cancellation Error returned. Polled alongside the tool call via
``asyncio.wait`` so a /cancel POST interrupts even mid-network-read.
"""
async def _call() -> Any:
async def _one_shot() -> Any:
async with _client(url, headers, use_oauth) as client:
return await client.call_tool(name, args)
async def _watch_cancel() -> None:
# 50 ms cadence keeps cancellation responsive without busy-looping;
# matches routes/inference.py's cancel watcher cadence.
while cancel_event is not None and not cancel_event.is_set():
await asyncio.sleep(0.05)
async def _race() -> Any:
# Check cancellation before spawning the call task so a pre-set event
# short-circuits before opening the transport / HTTP connection.
if cancel_event is not None and cancel_event.is_set():
raise _MCPCancelled
call_task = asyncio.create_task(_call())
if cancel_event is None:
return await asyncio.wait_for(call_task, timeout = timeout)
watch_task = asyncio.create_task(_watch_cancel())
try:
done, pending = await asyncio.wait(
{call_task, watch_task},
timeout = timeout,
return_when = asyncio.FIRST_COMPLETED,
)
finally:
for t in (call_task, watch_task):
if not t.done():
t.cancel()
if not done:
raise asyncio.TimeoutError
if call_task in done:
return call_task.result()
raise _MCPCancelled
# raise_on_error=False lets an is_error result (which may still carry
# image content) reach _flatten_result instead of FastMCP raising ToolError
# and dropping the images. Transport failures still raise (handled below).
return await client.call_tool(name, args, raise_on_error = False)
try:
result = asyncio.run(_race())
if is_stdio(url):
result = _call_stdio_tool(
url, headers, name, args, timeout, cancel_event, scope, config_check
)
else:
result = asyncio.run(_race_tool_call(_one_shot(), timeout, cancel_event))
except _MCPCancelled:
return f"Error: MCP tool '{name}' cancelled"
except asyncio.TimeoutError:
return f"Error: MCP tool '{name}' timed out after {timeout:g}s"
suffix = f" after {timeout:g}s" if timeout is not None else ""
return f"Error: MCP tool '{name}' timed out{suffix}"
except Exception as exc:
logger.exception("MCP call_tool failed for %s: %s", name, exc)
return f"Error: MCP tool '{name}' failed: {exc}"

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

@ -5,15 +5,124 @@ Drop-in replacement for InferenceBackend — same interface, uses mlx-lm/mlx-vlm
instead of torch/transformers for model loading and generation.
"""
import json
import os
import threading
from typing import Optional, Generator
from core.inference.message_content import content_to_text
from core.inference.runtime_context import runtime_context_length
from core.inference.chat_template_helpers import (
ReasoningChannelNormalizer,
normalize_reasoning_snapshots,
)
from loggers import get_logger
logger = get_logger(__name__)
def _mlx_vlm_model_config(model):
"""Return the loaded MLX model config and its type, preferring whichever of
config / _config actually carries a model_type."""
def _model_type(cfg):
return cfg.get("model_type") if isinstance(cfg, dict) else getattr(cfg, "model_type", None)
configs = [
cfg
for cfg in (getattr(model, "config", None), getattr(model, "_config", None))
if cfg is not None
]
for cfg in configs:
model_type = _model_type(cfg)
if model_type is not None:
return cfg, model_type
return (configs[0] if configs else None), None
def _render_registered_vlm_prompt(processor, model, messages, num_images):
"""Render through mlx-vlm when it declares a formatter for this model."""
from mlx_vlm import prompt_utils
config, model_type = _mlx_vlm_model_config(model)
if config is None:
return None
if model_type not in getattr(prompt_utils, "MODEL_CONFIG", {}):
return None
rendered = prompt_utils.apply_chat_template(
processor,
config,
messages,
add_generation_prompt = True,
num_images = num_images,
)
if isinstance(rendered, str) and rendered.strip():
return rendered
raise RuntimeError("mlx-vlm's registered renderer returned an empty prompt.")
def _count_vlm_images(content):
if isinstance(content, list):
return sum(_count_vlm_images(item) for item in content)
if not isinstance(content, dict):
return 0
if str(content.get("type", "")).lower() in ("image", "image_url", "input_image"):
return 1
return _count_vlm_images(content.get("content"))
def _vlm_media_reprs(content):
if isinstance(content, list):
values = (
{str(content), json.dumps(content, ensure_ascii = False)}
if _count_vlm_images(content)
else set()
)
for item in content:
values.update(_vlm_media_reprs(item))
return values
if not isinstance(content, dict):
return set()
if str(content.get("type", "")).lower() in ("image", "image_url", "input_image"):
return {str(content), json.dumps(content, ensure_ascii = False)}
return _vlm_media_reprs(content.get("content"))
def _prompt_serializes_vlm_media(prompt, messages):
"""Detect templates that embed the exact structured media object repr."""
media_reprs = set()
for message in messages:
if isinstance(message, dict):
media_reprs.update(_vlm_media_reprs(message.get("content")))
text_content = [
content_to_text(message.get("content")) for message in messages if isinstance(message, dict)
]
return any(
prompt.count(media_repr) > sum(content.count(media_repr) for content in text_content)
for media_repr in media_reprs
)
def _vlm_prompt_issue(prompt, messages):
if not isinstance(prompt, str) or not prompt.strip():
return "an empty prompt"
if _prompt_serializes_vlm_media(prompt, messages):
return "serialized structured image content"
return None
def _vlm_messages_have_tool_history(messages):
return any(
isinstance(message, dict)
and (
message.get("role") == "tool"
or message.get("tool_calls")
or message.get("tool_call_id")
)
for message in messages
)
def _build_generation_stats(prompt_n, prompt_tps, gen_n, gen_tps):
"""Map mlx stream stats onto the usage/timings shape llama-server emits."""
prompt_n = int(prompt_n or 0)
@ -422,15 +531,13 @@ class MLXInferenceBackend:
{"type": "text", "text": content},
]
elif isinstance(content, list):
has_image = any(
p.get("type") == "image" for p in content if isinstance(p, dict)
)
has_image = _count_vlm_images(content) > 0
if not has_image:
content.insert(0, {"type": "image"})
break
if self._is_vlm:
yield from self._generate_vlm(
stream = self._generate_vlm(
full_messages,
image,
temperature,
@ -447,7 +554,7 @@ class MLXInferenceBackend:
presence_penalty = presence_penalty,
)
else:
yield from self._generate_text(
stream = self._generate_text(
full_messages,
temperature,
top_p,
@ -462,6 +569,7 @@ class MLXInferenceBackend:
preserve_thinking = preserve_thinking,
presence_penalty = presence_penalty,
)
yield from stream
def _generate_text(
self,
@ -506,7 +614,7 @@ class MLXInferenceBackend:
# probe and native render share a renderer. (VLM renders via the
# processor for image tokens and is not wired here.)
model_info = self.models.get(self.active_model_name, {})
prompt = render_with_native_template_fallback(
render_result = render_with_native_template_fallback(
formatted_prompt = prompt,
tokenizer = self._tokenizer,
model_info = model_info,
@ -517,7 +625,10 @@ class MLXInferenceBackend:
reasoning_effort = reasoning_effort,
preserve_thinking = preserve_thinking,
hf_token = model_info.get("hf_token"),
return_metadata = True,
)
prompt = render_result.prompt
reasoning_channel_markers = render_result.reasoning_channel_markers
# An open <think> prefilled by the template lives in the prompt, not
# the generated tokens; re-emit it so the frontend renders the block.
@ -551,7 +662,17 @@ class MLXInferenceBackend:
if not logits_processors:
logits_processors = None
preserve_native_channels = reasoning_channel_markers is not None
token_ids = []
normalizer = (
ReasoningChannelNormalizer(*reasoning_channel_markers)
if reasoning_channel_markers is not None
else None
)
# MLX consumers diff cumulative snapshots. Keep a prompt-prefilled
# <think> prefix on every native-protocol snapshot just as the normal
# decoding path does below.
normalized_output = think_prefix
logger.info(
"Generating: prompt_len=%d, max_tokens=%d, model=%s, tokenizer=%s",
len(prompt),
@ -575,12 +696,19 @@ class MLXInferenceBackend:
**gen_kwargs,
):
final_response = response
token_ids.append(response.token)
cumulative = self._tokenizer.decode(
token_ids,
skip_special_tokens = True,
)
yield think_prefix + cumulative
if preserve_native_channels:
piece = getattr(response, "text", None) or ""
delta = normalizer.feed(piece)
if delta:
normalized_output += delta
yield normalized_output
else:
token_ids.append(response.token)
cumulative = self._tokenizer.decode(
token_ids,
skip_special_tokens = True,
)
yield think_prefix + cumulative
if cancel_event and cancel_event.is_set():
break
@ -597,6 +725,12 @@ class MLXInferenceBackend:
getattr(final_response, "generation_tokens", 0),
getattr(final_response, "generation_tps", 0.0),
)
if normalizer is not None:
cancelled = cancel_event is not None and cancel_event.is_set()
tail = normalizer.drain() if cancelled else normalizer.finish()
if tail:
normalized_output += tail
yield normalized_output
def _generate_vlm(
self,
@ -632,17 +766,87 @@ class MLXInferenceBackend:
):
chat_target = getattr(self._processor, "tokenizer", self._processor)
prompt = apply_chat_template_for_generation(
chat_target,
messages,
tools = tools,
enable_thinking = enable_thinking,
reasoning_effort = reasoning_effort,
preserve_thinking = preserve_thinking,
)
# mlx_vlm's stream_generate handles pixel_values (None for text-only)
images = [image] if image is not None else None
attached_images = 0 if images is None else len(images)
structured_images = sum(
_count_vlm_images(message.get("content"))
for message in messages
if isinstance(message, dict)
)
if structured_images != attached_images:
raise RuntimeError(
f"VLM conversation contains {structured_images} structured image "
f"item(s) for {attached_images} attached image(s)."
)
prompt = None
has_tool_history = _vlm_messages_have_tool_history(messages)
prompt_error = None
try:
prompt = apply_chat_template_for_generation(
chat_target,
messages,
tools = tools,
enable_thinking = enable_thinking,
reasoning_effort = reasoning_effort,
preserve_thinking = preserve_thinking,
)
except Exception as exc:
if images is None or has_tool_history:
raise
prompt_error = exc
prompt_issue = (
_vlm_prompt_issue(prompt, messages) if prompt_error is None else "a rendering error"
)
if prompt_issue and has_tool_history:
raise RuntimeError(
f"VLM chat template returned {prompt_issue} and cannot be recovered "
"without dropping tool-call history."
) from prompt_error
if images is not None and prompt_issue:
if tools or any(
value is not None
for value in (enable_thinking, reasoning_effort, preserve_thinking)
):
if prompt_error is not None:
raise prompt_error
raise RuntimeError(
f"VLM chat template returned {prompt_issue} and cannot be recovered "
"without dropping requested tools or reasoning controls."
)
try:
recovered_prompt = _render_registered_vlm_prompt(
self._processor,
self._model,
messages,
len(images),
)
except Exception as recovery_error:
if prompt_error is not None:
raise prompt_error
raise RuntimeError(
f"VLM chat template returned {prompt_issue}; model-aware "
f"recovery failed: {recovery_error}"
) from recovery_error
if recovered_prompt is None:
if prompt_error is not None:
raise prompt_error
raise RuntimeError(
f"VLM chat template returned {prompt_issue}, and no registered "
"MLX VLM renderer was available for this model."
)
recovered_issue = _vlm_prompt_issue(recovered_prompt, messages)
if recovered_issue:
if prompt_error is not None:
raise prompt_error
raise RuntimeError(
f"Model-aware VLM rendering returned {recovered_issue} for "
f"{attached_images} attached image(s)."
)
prompt = recovered_prompt
elif prompt_issue:
raise RuntimeError(f"VLM chat template returned {prompt_issue}.") from prompt_error
from core.inference.chat_template_helpers import detect_think_prefill
@ -685,31 +889,37 @@ class MLXInferenceBackend:
elif _rep_active:
vlm_kwargs["repetition_penalty"] = float(repetition_penalty)
with self._generation_lock:
final_response = None
try:
for response in vlm_stream(
self._model,
self._processor,
prompt,
images,
**vlm_kwargs,
):
final_response = response
token_text = response.text if hasattr(response, "text") else str(response)
cumulative += token_text
yield cumulative
if cancel_event and cancel_event.is_set():
break
finally:
# mlx_vlm exposes the same stats fields as mlx_lm.
if final_response is not None:
self.last_generation_stats = _build_generation_stats(
getattr(final_response, "prompt_tokens", 0),
getattr(final_response, "prompt_tps", 0.0),
getattr(final_response, "generation_tokens", 0),
getattr(final_response, "generation_tps", 0.0),
)
def _stream_vlm_snapshots():
nonlocal cumulative
with self._generation_lock:
final_response = None
try:
for response in vlm_stream(
self._model,
self._processor,
prompt,
images,
**vlm_kwargs,
):
final_response = response
token_text = response.text if hasattr(response, "text") else str(response)
cumulative += token_text
yield cumulative
if cancel_event and cancel_event.is_set():
break
finally:
# mlx_vlm exposes the same stats fields as mlx_lm.
if final_response is not None:
self.last_generation_stats = _build_generation_stats(
getattr(final_response, "prompt_tokens", 0),
getattr(final_response, "prompt_tps", 0.0),
getattr(final_response, "generation_tokens", 0),
getattr(final_response, "generation_tps", 0.0),
)
yield from normalize_reasoning_snapshots(
_stream_vlm_snapshots(), chat_target, cancel_event, tools = tools
)
def generate_with_adapter_control(
self,

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

@ -59,7 +59,32 @@ class GenStreamError(str):
"Error:" by checking isinstance(chunk, GenStreamError).
"""
__slots__ = ()
__slots__ = ("public",)
def __new__(
cls,
value,
*,
public: bool = False,
):
obj = str.__new__(cls, value)
obj.public = bool(public)
return obj
class GenStreamErrorRaised(RuntimeError):
"""Internal exception form of ``GenStreamError`` for generator boundaries."""
__slots__ = ("public",)
def __init__(
self,
value,
*,
public: bool = False,
):
super().__init__(value)
self.public = bool(public)
class InferenceOrchestrator:
@ -174,6 +199,21 @@ class InferenceOrchestrator:
def _spawn_subprocess(self, config: dict) -> None:
"""Spawn a new inference subprocess."""
# Same recheck as the training/export spawns, REPAIR reservations only: a
# repair swaps without holding the lifecycle gate this load's caller owns,
# while an install cannot swap until this gate is released (and then its
# queued-load snapshot aborts it), so tolerating installs here lets the
# load win instead of failing both sides. Also covers the OpenAI
# auto-switch path, which enters _load_model_impl without route guards.
from utils.transformers_version import (
SidecarSwapInProgress,
sidecar_swap_kind,
)
if sidecar_swap_kind() == "repair":
raise SidecarSwapInProgress(
"A transformers repair is replacing the latest sidecar; retry when it completes."
)
from utils.native_path_leases import (
native_path_secret_removed_for_child_start,
run_without_native_path_secret,
@ -210,12 +250,24 @@ class InferenceOrchestrator:
if self._cancel_event is not None:
self._cancel_event.set()
def _shutdown_subprocess(self, timeout: float = 10.0) -> None:
"""Gracefully shut down the inference subprocess."""
def is_worker_alive(self) -> bool:
"""True while the inference subprocess is running, even with no model
active (a failed load can leave a live worker holding sidecar modules)."""
proc = self._proc
return proc is not None and proc.is_alive()
def _shutdown_subprocess(self, timeout: float = 10.0) -> bool:
"""Gracefully shut down the inference subprocess.
Returns True only once the worker is confirmed dead. If it survives
terminate/kill (e.g. wedged in an uninterruptible CUDA syscall that outlives
SIGKILL) the live handle is KEPT, not nulled, so is_worker_alive() and the
pre-swap liveness guard can still observe the survivor instead of a cleared
handle and refuse the destructive sidecar swap."""
self._stop_dispatcher() # before killing subprocess
if self._proc is None or not self._proc.is_alive():
self._proc = None
return
return True
# 1. Cancel any ongoing generation first (instant via mp.Event)
self._cancel_generation()
@ -252,12 +304,22 @@ class InferenceOrchestrator:
except Exception:
pass
if self._proc is not None and self._proc.is_alive():
# Survived SIGKILL (uninterruptible syscall): keep the handle so callers
# and the pre-swap guard see a live worker rather than a nulled one.
logger.error(
"Inference subprocess still alive after terminate/kill; "
"preserving its handle for the pre-swap liveness check"
)
return False
self._proc = None
self._cmd_queue = None
self._resp_queue = None
self._cancel_event = None
self._drain_event = None
logger.info("Inference subprocess shut down")
return True
def _cleanup(self):
"""atexit handler."""
@ -494,13 +556,19 @@ class InferenceOrchestrator:
initial_resp_queue = self._resp_queue
while True:
if self._proc is not initial_proc or self._resp_queue is not initial_resp_queue:
yield GenStreamError(f"Error: {self._subprocess_crash_message(crash_context)}")
yield GenStreamError(
f"Error: {self._subprocess_crash_message(crash_context)}",
public = True,
)
return
resp = read_one(read_timeout)
if resp is None:
# Check subprocess health
if not self._ensure_subprocess_alive():
yield GenStreamError(f"Error: {self._subprocess_crash_message(crash_context)}")
yield GenStreamError(
f"Error: {self._subprocess_crash_message(crash_context)}",
public = True,
)
return
continue
@ -652,11 +720,11 @@ class InferenceOrchestrator:
GPU work stays serialized; this only avoids orchestrator lock contention.
"""
if not self._ensure_subprocess_alive():
yield GenStreamError("Error: Inference subprocess is not running")
yield GenStreamError("Error: Inference subprocess is not running", public = True)
return
if not self.active_model_name:
yield GenStreamError("Error: No active model")
yield GenStreamError("Error: No active model", public = True)
return
# Latch the target model so the recheck below can detect a switch that completed
# between _start_dispatcher and mailbox registration (mirrors the locked path's
@ -667,7 +735,7 @@ class InferenceOrchestrator:
# so without this early-out a compare request would enqueue a generate on the
# outgoing model and delay the switch.
if self._unload_pending:
yield GenStreamError("Error: model is being unloaded")
yield GenStreamError("Error: model is being unloaded", public = True)
return
# Ensure the dispatcher runs. _start_dispatcher serializes concurrent starters under
@ -739,7 +807,7 @@ class InferenceOrchestrator:
# _stop_dispatcher joins the dispatcher, which itself takes that lock.
if orphaned_dispatcher:
self._stop_dispatcher()
yield GenStreamError("Error: model is being unloaded")
yield GenStreamError("Error: model is being unloaded", public = True)
return
try:
@ -882,6 +950,13 @@ class InferenceOrchestrator:
# Public API — same interface as InferenceBackend
# ------------------------------------------------------------------
# Monotonic count of PUBLISHED loads; lets the install route detect a load
# (including a same-model reload) that completed while it waited on the gate.
# Bumped when the load result is published, not at load start: a start-time
# bump is already visible when the installer snapshots mid-load, so the
# completed reload would look unchanged and get unloaded by the swap.
load_generation: int = 0
def load_model(
self,
config, # ModelConfig
@ -935,13 +1010,36 @@ class InferenceOrchestrator:
sub_config["resolved_gpu_ids"] = resolved_gpu_ids
sub_config["gpu_selection"] = gpu_selection
# Recheck the sidecar reservation BEFORE tearing the old worker down,
# for REPAIRS only: an install holds this same lifecycle gate, so it
# cannot swap while this load runs, and its queued-load snapshot
# aborts it after this load publishes -- the load wins cleanly.
# Raising here (repair) keeps the current model loaded.
from utils.transformers_version import (
SidecarSwapInProgress,
sidecar_swap_kind,
)
if sidecar_swap_kind() == "repair":
raise SidecarSwapInProgress(
"A transformers repair is replacing the latest sidecar; "
"retry when it completes."
)
# Always kill the existing subprocess and spawn fresh: reusing one
# after unsloth patches torch internals breaks getsource on reload.
if self._ensure_subprocess_alive():
self._cancel_generation()
time.sleep(0.3)
self._shutdown_subprocess()
if self._shutdown_subprocess() is False:
# The worker survived terminate/kill (e.g. a wedged CUDA syscall that
# outlives SIGKILL). Its handle is kept, so is_worker_alive() and the
# pre-swap guard still see it; do not spawn a second worker over one
# still holding GPU memory. Fail so the load can retry once it exits.
raise RuntimeError(
"The current inference worker did not exit and still holds GPU "
"memory; not starting a new model over it. Retry shortly."
)
elif self._proc is not None:
self._shutdown_subprocess(timeout = 2)
@ -1030,6 +1128,7 @@ class InferenceOrchestrator:
return False
model_info = resp.get("model_info", {})
self.active_model_name = model_info.get("identifier", model_name)
self.load_generation += 1
# A load always spawns a fresh subprocess holding only this model, so
# mirror that. A lingering stale name would pass unload_model's "not in
# self.models" guard, and the worker's absent-name fallback would unload
@ -1061,8 +1160,15 @@ class InferenceOrchestrator:
self.models.clear()
raise Exception(error)
except Exception:
except Exception as exc:
self.loading_models.discard(model_name)
from utils.transformers_version import SidecarSwapInProgress
if isinstance(exc, SidecarSwapInProgress) and self._ensure_subprocess_alive():
# Raised before the old worker was torn down: the previous model
# is still live, so keep the mirrors (clearing them would let the
# installer treat the worker as inactive and kill it unreported).
raise
self.active_model_name = None
self.models.clear()
raise
@ -1293,12 +1399,15 @@ class InferenceOrchestrator:
nudge_tool_calls: Optional[bool] = None,
tool_call_timeout: int = 300,
session_id: Optional[str] = None,
thread_id: Optional[str] = None,
rag_scope: Optional[dict] = None,
confirm_tool_calls: bool = False,
bypass_permissions: bool = False,
permission_mode: Optional[str] = None,
use_adapter: Optional[Union[bool, str]] = None,
stats_holder: Optional[dict] = None,
presence_penalty: float = 0.0,
reasoning_prefilled: bool = False,
**_unused,
):
"""Run the safetensors agentic tool loop in the parent process,
@ -1337,12 +1446,27 @@ class InferenceOrchestrator:
presence_penalty = presence_penalty,
)
if use_adapter is not None:
yield from self.generate_with_adapter_control(
stream = self.generate_with_adapter_control(
use_adapter = use_adapter,
**common_kwargs,
)
else:
yield from self.generate_chat_response(**common_kwargs)
stream = self.generate_chat_response(**common_kwargs)
close_stream = False
try:
for chunk in stream:
if isinstance(chunk, GenStreamError):
close_stream = True
raise GenStreamErrorRaised(str(chunk), public = chunk.public)
yield chunk
finally:
if close_stream:
close = getattr(stream, "close", None)
if callable(close):
try:
close()
except Exception:
logger.debug("failed to close errored generation stream", exc_info = True)
initial = list(messages)
if system_prompt:
@ -1359,9 +1483,12 @@ class InferenceOrchestrator:
max_tool_iterations = max_tool_iterations,
tool_call_timeout = tool_call_timeout,
session_id = session_id,
thread_id = thread_id,
rag_scope = rag_scope,
confirm_tool_calls = confirm_tool_calls,
bypass_permissions = bypass_permissions,
permission_mode = permission_mode,
reasoning_prefilled = reasoning_prefilled,
)
def generate_with_adapter_control(
@ -1410,11 +1537,11 @@ class InferenceOrchestrator:
readers don't consume each other's tokens off the shared resp_queue.
"""
if not self._ensure_subprocess_alive():
yield GenStreamError("Error: Inference subprocess is not running")
yield GenStreamError("Error: Inference subprocess is not running", public = True)
return
if not self.active_model_name:
yield GenStreamError("Error: No active model")
yield GenStreamError("Error: No active model", public = True)
return
expected_model = self.active_model_name
@ -1431,7 +1558,7 @@ class InferenceOrchestrator:
# so we never generate on the wrong one.
if self._unload_pending or self.active_model_name != expected_model:
# Won the lock handoff during a switch; don't start on the outgoing model.
yield GenStreamError("Error: model is being unloaded")
yield GenStreamError("Error: model is being unloaded", public = True)
return
request_id = str(uuid.uuid4())
image_b64 = self._pil_to_base64(image) if image is not None else None
@ -1616,10 +1743,10 @@ class InferenceOrchestrator:
) -> Generator[str, None, None]:
"""Shared inner logic for audio input generation (Whisper + ASR)."""
if not self._ensure_subprocess_alive():
yield GenStreamError("Error: Inference subprocess is not running")
yield GenStreamError("Error: Inference subprocess is not running", public = True)
return
if not self.active_model_name:
yield GenStreamError("Error: No active model")
yield GenStreamError("Error: No active model", public = True)
return
expected_model = self.active_model_name
@ -1628,7 +1755,7 @@ class InferenceOrchestrator:
# cleared or swapped the model while we waited.
if self._unload_pending or self.active_model_name != expected_model:
# Won the lock handoff during a switch; don't start on the outgoing model.
yield GenStreamError("Error: model is being unloaded")
yield GenStreamError("Error: model is being unloaded", public = True)
return
request_id = str(uuid.uuid4())

3
studio/backend/core/inference/passthrough_healing.py
View file

@ -41,6 +41,9 @@ _HEAL_SIGNALS = (
"<|tool_call>",
"<function=",
"[TOOL_CALLS]",
# TML Inkling native call marker (leaks as text when the server-side
# parser misses a narration-then-call turn).
"<|content_invoke_tool_json|>",
)

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

@ -276,8 +276,9 @@ PROVIDER_REGISTRY: dict[str, dict[str, Any]] = {
"auth_header": "Authorization",
"auth_prefix": "Bearer ",
"notes": (
"Local Ollama server. OpenAI-compatible /v1/chat/completions; "
"no API key. Surfaced via CUSTOM_PROVIDER_PRESETS in the frontend."
"Ollama server (local or cloud). OpenAI-compatible "
"/v1/chat/completions; API key optional (required by Ollama "
"cloud). Surfaced via CUSTOM_PROVIDER_PRESETS in the frontend."
),
"hidden": True,
},

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

@ -15,6 +15,7 @@ parses tool calls from the cumulative text and dispatches via
"""
import bisect
import inspect
import re
import threading
from typing import Callable, Generator, Optional
@ -49,6 +50,7 @@ from core.inference.tool_call_parser import (
# pattern lists, so the safetensors streaming strip stays aligned with the parser.
from core.tool_healing import (
_REHEARSAL_TAIL_STRIP_RE,
_THINK_CLOSE_RE,
_strip_bracket_tag_calls,
_think_spans_outside_tool_markup,
apply_tool_strip_patterns,
@ -56,10 +58,12 @@ from core.tool_healing import (
)
from core.inference.tool_loop_controller import (
ToolLoopController,
append_deferred_nudges,
coerce_tool_arguments,
status_for_tool,
tool_event_provenance,
)
from core.inference.tool_stream_exec import stream_tool_execution
from state.tool_approvals import (
TOOL_REJECTED_MESSAGE,
abort_tool_decision,
@ -301,6 +305,45 @@ def _status_for_tool(tool_name: str, arguments: dict) -> str:
return status_for_tool(tool_name, arguments)
def _reprompt_intent_text(text: str, *, reasoning_prefilled: bool = False) -> str:
"""Return visible answer text for the plan-without-action classifier.
Safetensors reasoning shares the cumulative text channel with the answer.
Forward-looking phrases inside ``<think>`` / ``[THINK]`` are private
planning, not a user-visible promise to call a tool. Match GGUF's behavior:
classify visible content when present and fall back to reasoning only for a
reasoning-only stall.
"""
prefilled_reasoning = ""
if reasoning_prefilled:
close = _THINK_CLOSE_RE.search(text)
if close is None:
return text.strip()
prefilled_reasoning = text[: close.end()].strip()
text = text[close.end() :].strip()
if not text:
return prefilled_reasoning
spans = _think_spans_outside_tool_markup(text)
if not spans:
return text.strip()
visible: list[str] = []
reasoning: list[str] = []
cursor = 0
for start, end in spans:
visible.append(text[cursor:start])
reasoning.append(text[start:end])
cursor = end
visible.append(text[cursor:])
visible_text = "".join(visible).strip()
reasoning_text = "".join(reasoning).strip()
if visible_text:
return visible_text
return "\n".join(part for part in (prefilled_reasoning, reasoning_text) if part).strip()
def _looks_like_enabled_bare_json(text: str, enabled_tool_names: Optional[set]) -> bool:
"""True when ``text`` opens with an ENABLED markerless bare-JSON call; an ordinary JSON answer returns False."""
probe = strip_llama3_leading_sentinels(text.lstrip())
@ -402,6 +445,22 @@ def _tool_event_provenance(**flags: object) -> dict[str, object]:
return tool_event_provenance(**flags)
def _accepts_output_callback(func: Callable[..., str]) -> bool:
"""Whether an injectable ``execute_tool`` supports ``output_callback``.
The loop's ``execute_tool`` is a parameter (tests inject fakes), so forward
the live-output kwarg only when the callable declares it or takes ``**kwargs``.
"""
try:
sig = inspect.signature(func)
except (TypeError, ValueError):
return False
params = sig.parameters
if "output_callback" in params:
return True
return any(p.kind is inspect.Parameter.VAR_KEYWORD for p in params.values())
def _call_single_turn(single_turn, conversation: list, active_tools: list[dict]):
"""Call a single-turn generator with active tool schemas when supported."""
try:
@ -424,9 +483,12 @@ def run_safetensors_tool_loop(
max_tool_iterations: int = 25,
tool_call_timeout: int = 300,
session_id: Optional[str] = None,
thread_id: Optional[str] = None,
rag_scope: Optional[dict] = None,
confirm_tool_calls: bool = False,
bypass_permissions: bool = False,
permission_mode: Optional[str] = None,
reasoning_prefilled: bool = False,
) -> Generator[dict, None, None]:
"""Drive an agentic tool loop on top of a cumulative-text generator.
@ -452,10 +514,27 @@ def run_safetensors_tool_loop(
"""
conversation = list(messages)
# Forced first-pass RAG (mirrors the GGUF loop) so doc Qs don't lose to web_search.
# Normalize the mode (mirrors the GGUF loop): "full" and
# bypass_permissions are the same switch; unset/unknown behaves as "ask".
# "off" keeps the sandbox but never prompts.
if permission_mode == "full":
bypass_permissions = True
elif bypass_permissions:
permission_mode = "full"
elif permission_mode not in ("ask", "auto", "off"):
permission_mode = "ask"
# Forced first-pass RAG (mirrors the GGUF loop) so doc Qs don't lose to
# web_search. Skip only when a retrieval call would actually prompt (ask
# mode); auto never gates the safe search_knowledge_base tool.
from core.inference.tools import build_rag_autoinject
_auto = None if confirm_tool_calls else build_rag_autoinject(conversation, rag_scope)
# off never prompts, so (like auto) it must not lose first-pass retrieval
# even if a direct caller passes a stale confirm_tool_calls flag.
_skip_autoinject = (
confirm_tool_calls and not bypass_permissions and permission_mode not in ("auto", "off")
)
_auto = None if _skip_autoinject else build_rag_autoinject(conversation, rag_scope)
if _auto:
for _ev in _auto["events"]:
yield _ev
@ -533,12 +612,24 @@ def run_safetensors_tool_loop(
provisional_render_html_started = False
provisional_resolved = False
provisional_render_html_id = f"call_{next_call_id}"
# Live-args offset for the provisional render_html card: the drained call
# text streams as tool_args so the canvas shows the HTML being written.
_live_args_streamed_upto = -1
# When a human confirmation gate is active the real tool_start is keyed
# by an approval id and carries awaiting_confirmation, so an early
# provisional card (keyed by tool_call_id, no approval) would show the
# tool as "running" before the user has approved it. Suppress the early
# card in that case and let the gated tool_start be the first signal.
_provisional_confirm_gated = bool(confirm_tool_calls) and not bypass_permissions
# In auto mode render_html is always safe and never prompts, so keep its
# early canvas card (the frontend sends confirm_tool_calls=true alongside
# auto); mirrors the GGUF path's _confirm_gated exemption.
from core.inference.tools import is_always_safe_tool
_provisional_confirm_gated = (
bool(confirm_tool_calls)
and not bypass_permissions
and not (permission_mode == "auto" and is_always_safe_tool("render_html"))
)
gen = _call_single_turn(single_turn, conversation, active_tools)
prev_cumulative = ""
@ -595,6 +686,30 @@ def run_safetensors_tool_loop(
"arguments": {},
"provenance": _tool_event_provenance(provisional = True),
}
# Backlog first: everything drained so far.
yield {
"type": "tool_args",
"tool_call_id": provisional_render_html_id,
"tool_name": "render_html",
"text": content_accum,
}
_live_args_streamed_upto = len(content_accum)
elif (
provisional_render_html_started
and not provisional_resolved
and _live_args_streamed_upto >= 0
and len(content_accum) > _live_args_streamed_upto
):
# Still writing the call: stream the fragment so the canvas
# renders live. Display only; content_accum still feeds the
# stream-end parser verbatim.
yield {
"type": "tool_args",
"tool_call_id": provisional_render_html_id,
"tool_name": "render_html",
"text": content_accum[_live_args_streamed_upto:],
}
_live_args_streamed_upto = len(content_accum)
continue
if detect_state == _state_streaming:
@ -635,6 +750,13 @@ def run_safetensors_tool_loop(
"arguments": {},
"provenance": _tool_event_provenance(provisional = True),
}
yield {
"type": "tool_args",
"tool_call_id": provisional_render_html_id,
"tool_name": "render_html",
"text": content_accum,
}
_live_args_streamed_upto = len(content_accum)
continue
cumulative_display = candidate
cleaned = strip_tool_markup_streaming(
@ -791,6 +913,13 @@ def run_safetensors_tool_loop(
"arguments": {},
"provenance": _tool_event_provenance(provisional = True),
}
yield {
"type": "tool_args",
"tool_call_id": provisional_render_html_id,
"tool_name": "render_html",
"text": content_accum,
}
_live_args_streamed_upto = len(content_accum)
elif is_prefix and (is_rehearsal_prefix or len(stripped) < _MAX_BUFFER_CHARS):
# A rehearsal prefix is self-bounded; the buffer cap must not cut long MCP names short.
continue
@ -868,7 +997,10 @@ def run_safetensors_tool_loop(
# (GGUF loop parity). The retry is gated on nudge_tool_calls so
# Studio callers (which send True) always nudge, while API callers
# who omit the flag keep today's no-reprompt behavior (opt-in).
stripped_answer = content_accum.strip()
intent_text = _reprompt_intent_text(
content_accum,
reasoning_prefilled = reasoning_prefilled,
)
if (
auto_heal_tool_calls
and nudge_tool_calls
@ -877,7 +1009,7 @@ def run_safetensors_tool_loop(
and not rag_autoinjected
and not tool_denied
and not any(record.executed for record in tool_controller.history)
and is_short_intent_without_action(stripped_answer)
and is_short_intent_without_action(intent_text)
):
reprompt_count += 1
logger.info(
@ -885,9 +1017,9 @@ def run_safetensors_tool_loop(
"calling tools (%d chars)",
reprompt_count,
MAX_ACT_REPROMPTS,
len(stripped_answer),
len(intent_text),
)
conversation.append({"role": "assistant", "content": stripped_answer})
conversation.append({"role": "assistant", "content": intent_text})
tool_hint = " or ".join(_active_tool_names(active_tools)) or "an available tool"
conversation.append(
{
@ -1012,6 +1144,9 @@ def run_safetensors_tool_loop(
assistant_msg: dict = {"role": "assistant", "content": content_text}
assistant_appended = False
# Collect no-op nudges and flush them after the batch, so a no-op doesn't
# abort it and drop the parallel calls that follow.
deferred_noop_msgs: list = []
for tc in tool_calls or []:
func = tc.get("function", {}) or {}
@ -1040,12 +1175,12 @@ def run_safetensors_tool_loop(
"provenance": decision.provenance,
}
completion = tool_controller.record_noop(decision)
conversation.append(completion.model_message())
deferred_noop_msgs.append(completion.model_message())
logger.info(
"Suppressed local safetensors tool call as internal no-op: "
f"action={decision.action} tool={decision.tool_name}"
)
break
continue
if not assistant_appended:
assistant_msg["tool_calls"] = [decision.as_assistant_tool_call()]
@ -1055,8 +1190,17 @@ def run_safetensors_tool_loop(
assistant_msg.setdefault("tool_calls", []).append(decision.as_assistant_tool_call())
# Bypass wins over the confirm gate at the loop level too, so a
# direct internal caller passing both flags never prompts.
needs_confirm = bool(confirm_tool_calls) and not bypass_permissions
# direct internal caller passing both flags never prompts. In
# "auto" mode only calls detected as potentially unsafe pause.
# "off" never prompts (sandbox stays on).
needs_confirm = (
bool(confirm_tool_calls) and not bypass_permissions and permission_mode != "off"
)
if needs_confirm and permission_mode == "auto":
from core.inference.tools import is_potentially_unsafe_tool_call
needs_confirm = is_potentially_unsafe_tool_call(
decision.tool_name, decision.arguments
)
approval_id = new_approval_id() if needs_confirm else ""
decision_slot = begin_tool_decision(session_id, approval_id) if needs_confirm else None
start_event = decision.tool_start_event()
@ -1109,16 +1253,30 @@ def run_safetensors_tool_loop(
):
result = RAG_SEARCH_CAP_NUDGE
else:
try:
result = execute_tool(
decision.tool_name,
decision.arguments,
# Execute in a worker thread so live stdout chunks and heartbeats
# stream while the tool blocks (the SSE route turns heartbeats into
# keepalives). execute_tool is injectable; pass output_callback
# only when it accepts it.
def _invoke_tool(_output_callback, _decision = decision):
kwargs = dict(
cancel_event = cancel_event,
timeout = eff_timeout,
session_id = session_id,
thread_id = thread_id,
rag_scope = rag_scope,
disable_sandbox = bypass_permissions,
)
if _accepts_output_callback(execute_tool):
kwargs["output_callback"] = _output_callback
return execute_tool(_decision.tool_name, _decision.arguments, **kwargs)
try:
result = yield from stream_tool_execution(
_invoke_tool,
tool_name = decision.tool_name,
tool_call_id = decision.tool_call_id,
cancel_event = cancel_event,
)
except Exception as exc:
logger.exception("Tool %s raised: %s", decision.tool_name, exc)
result = f"Error: tool raised an exception: {exc}"
@ -1133,6 +1291,8 @@ def run_safetensors_tool_loop(
yield completion.tool_end_event()
conversation.append(completion.tool_message())
append_deferred_nudges(conversation, deferred_noop_msgs)
# Clear the status badge before the next turn.
yield {"type": "status", "text": ""}

6
studio/backend/core/inference/sandbox_site/__init__.py Normal file
View file

@ -0,0 +1,6 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
# Package marker only, so wheel builds ship this directory. It goes on the
# sandbox PYTHONPATH so site machinery imports the sibling ``sitecustomize`` at
# startup; nothing in the backend imports it directly.

313
studio/backend/core/inference/sandbox_site/sitecustomize.py Normal file
View file

@ -0,0 +1,313 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Sandbox-side compatibility shim for ChatGPT code-interpreter paths.
Models habitually write to /mnt/data (or /mnt/outputs, /home/sandbox,
/workspace), none of which exist in the Studio sandbox. This module sits on the
sandbox subprocess PYTHONPATH (see ``tools._build_safe_env``), so it loads at
interpreter startup in every sandboxed ``python`` run and any Python the
``terminal`` tool launches.
It remaps those prefixes onto the CWD in ``open`` / ``io.open``, ``os.open``,
``os.makedirs`` / ``os.mkdir`` and ``pathlib.Path.mkdir``. A write/create to a
convention prefix always heals onto the CWD; a READ heals only when the mapped
target already exists (re-reading an earlier write), so a genuinely missing
input stays truthful on the path the model used instead of silently reading a
same-basename workdir file. Since prefix lists cannot cover every invented path,
``open`` / ``io.open`` also get a create-mode fallback: an absolute path outside
the CWD whose parent is missing is redirected to the basename in the CWD. Reads
and mkdir never use the fallback (an arbitrary absolute directory can legitimately
succeed). It is collision-safe: it refuses to redirect onto an existing CWD file
(letting open raise). The patch set (io.open, os.open, os.mkdir, Path.mkdir, and
the <3.11 ``_NormalAccessor.open``) covers the low-level entry points pathlib
routes through. A one-line stderr notice fires on the first remap, and everything
is wrapped in try/except so a failure never breaks the interpreter.
Identical with and without output streaming because the child env is.
"""
import builtins
import io
import json
import os
import sys
# Code-interpreter convention prefixes. Remapping is gated on the prefix being
# ABSENT (see _remap) so a genuine host mount / user dir is never shadowed.
_PREFIXES = ("/mnt/data", "/mnt/outputs", "/home/sandbox", "/workspace")
# /tmp exists on the host; separate only to note that. The absence gate applies alike.
_CONDITIONAL_PREFIXES = ("/tmp/outputs",)
_notified = False
# Invented absolute write path -> healed CWD target, so re-writing the same
# artifact re-serves it instead of tripping the anti-clobber guard.
_remapped_writes: dict = {}
# Each tool call is a fresh subprocess (in-process map starts empty), so this
# on-disk sidecar carries the map across runs. It records only sources the
# fallback healed, so an unrelated same-basename file is never adopted.
_REMAP_SIDECAR = ".unsloth_sandbox_remap.json"
def _note(subject, original, mapped):
"""Print the one-shot stderr notice so the model learns the real location.
``subject`` is what "does not exist" (the prefix, or the whole invented
path); ``original`` is echoed in the ``(original -> mapped)`` tail.
"""
global _notified
if _notified:
return
_notified = True
print(
f"note: {subject} does not exist in this sandbox; "
f"using the working directory instead ({original} -> {mapped})",
file = sys.stderr,
)
def _contained_join(cwd, rel):
"""Join ``rel`` onto ``cwd`` so the result can never escape ``cwd``.
A habit path can carry ``..`` segments; joining verbatim would let the target
climb above the sandbox. ``..`` components are dropped and empty / ``.`` ones
ignored, keeping the result under ``cwd``.
"""
parts = []
for part in rel.split("/"):
if part == "" or part == ".":
continue
if part == "..":
if parts:
parts.pop()
continue
parts.append(part)
return os.path.join(cwd, *parts) if parts else cwd
def _map_onto_cwd(
prefix,
text,
notify = True,
):
"""Map ``<prefix>/rest`` onto ``./rest`` in the CWD, noting it once.
The suffix is contained under the CWD (see ``_contained_join``) so a path
like ``/mnt/data/../other_session/file`` cannot escape the workdir.
``notify`` is False when the caller may keep the original path (a read), so
the one-shot notice is not spent on a remap that never happens.
"""
rel = text[len(prefix) :].lstrip("/")
mapped = _contained_join(os.getcwd(), rel)
if notify:
_note(prefix, text, mapped)
return mapped
def _sidecar_path(cwd):
return os.path.join(cwd, _REMAP_SIDECAR)
def _load_sidecar(cwd):
"""Return the persisted ``source -> healed target`` map, or {} on any error
(missing/corrupt/foreign sidecar degrades to in-process-only behaviour)."""
try:
with open(_sidecar_path(cwd)) as fh:
data = json.load(fh)
except Exception: # noqa: BLE001 - a bad sidecar must never break user code
return {}
return data if isinstance(data, dict) else {}
def _record_sidecar(cwd, source, target):
"""Persist ``source -> target`` so the next run re-serves it.
Written atomically (temp + ``os.replace``) and wrapped so a read-only/full
filesystem never breaks the interpreter. The path is inside the CWD, so the
patched ``open`` leaves it untouched (no remap, no recursion).
"""
try:
data = _load_sidecar(cwd)
if data.get(source) == target:
return
data[source] = target
tmp = _sidecar_path(cwd) + ".tmp"
with open(tmp, "w") as fh:
json.dump(data, fh)
os.replace(tmp, _sidecar_path(cwd))
except Exception: # noqa: BLE001 - persistence is best effort only
pass
def _is_creating_mode(mode):
"""True only when an ``open()`` mode string can CREATE a missing file.
Only ``w`` / ``a`` / ``x`` create. ``r+`` / ``rb+`` require the path to exist,
so they must not trip the write fallback (which would corrupt an unrelated
same-basename file); ``w+`` / ``a+`` / ``x+`` still match.
"""
return isinstance(mode, str) and any(c in mode for c in ("w", "a", "x"))
def _remap_open(file, mode):
"""Remap for ``open()`` / ``io.open()``.
A prefix remap runs first: a write/create heals onto the CWD; a READ heals
only when the mapped target already exists (re-reading an earlier write),
else the original path is kept so a genuine missing input fails truthfully
instead of silently reading a same-basename workdir file. Only if no prefix
matched and the call creates does the fallback kick in: an absolute target
outside the CWD whose parent is missing is redirected to the basename in the
CWD, unless ``CWD/<basename>`` already exists (an unrelated file), in which
case the original path is kept so open raises.
"""
creating = _is_creating_mode(mode)
# notify=False: emit the notice only once we commit to the mapping below.
mapped = _remap(file, notify = False)
if mapped is not file:
# Write always heals; a read only when the mapped target exists (else keep
# the original path so a missing input stays truthful).
if creating or os.path.exists(mapped):
# Commit: emit the notice now (the notify=False peek above deferred it).
_remap(file, notify = True)
return mapped
return file
if not creating:
return file
try:
text = os.fspath(file)
except TypeError:
return file
# bytes paths left untouched (str-only, matching the prefix remaps).
if not isinstance(text, str) or not os.path.isabs(text):
return file
cwd = os.getcwd()
# Already inside the CWD: a real target the model meant; leave it alone.
if text == cwd or text.startswith(cwd + os.sep):
return file
parent = os.path.dirname(text)
# Redirect only when the parent is missing; an existing external directory is
# a deliberate target and stays truthful (os.path.exists follows symlinks).
if parent and os.path.exists(parent):
return file
base = os.path.basename(text)
# A trailing sep or '.'/'..' basename would redirect onto the CWD or its
# parent; refuse and let open raise.
if base in ("", ".", ".."):
return file
remapped = os.path.join(cwd, base)
# Never clobber an unrelated file sharing this basename (lexists catches
# dangling symlinks). But a target this fallback already healed for the same
# invented path (in-process map or cross-run sidecar) is the artifact being
# re-written, so re-serve it instead of raising on every overwrite.
if os.path.lexists(remapped) and remapped not in (
_remapped_writes.get(text),
_load_sidecar(cwd).get(text),
):
return file
_remapped_writes[text] = remapped
_record_sidecar(cwd, text, remapped)
_note(text, text, remapped)
return remapped
def _remap(path, notify = True):
"""Map ``<prefix>/rest`` onto ``./rest`` in the CWD; other paths pass through.
``notify`` is forwarded to ``_map_onto_cwd``; ``_remap_open`` passes False so
a read that keeps its original path emits no false notice.
"""
try:
text = os.fspath(path)
except TypeError:
return path
if not isinstance(text, str):
return path
for prefix in _PREFIXES + _CONDITIONAL_PREFIXES:
# Heal only while the real prefix directory is absent, so a genuine host
# mount / user directory at that prefix is never shadowed.
if (text == prefix or text.startswith(prefix + "/")) and not os.path.exists(prefix):
return _map_onto_cwd(prefix, text, notify = notify)
return path
def _install():
import pathlib
original_open = builtins.open
original_io_open = io.open
original_os_open = os.open
original_makedirs = os.makedirs
original_mkdir = os.mkdir
original_path_mkdir = pathlib.Path.mkdir
def _open(
file,
mode = "r",
*args,
**kwargs,
):
return original_open(_remap_open(file, mode), mode, *args, **kwargs)
def _io_open(
file,
mode = "r",
*args,
**kwargs,
):
return original_io_open(_remap_open(file, mode), mode, *args, **kwargs)
# mkdir/makedirs get only the prefix remap, never the write-mode fallback:
# an arbitrary absolute directory can legitimately succeed on the host.
def _makedirs(name, *args, **kwargs):
return original_makedirs(_remap(name), *args, **kwargs)
def _mkdir(path, *args, **kwargs):
return original_mkdir(_remap(path), *args, **kwargs)
def _os_open(
path,
flags,
mode = 0o777,
*,
dir_fd = None,
):
# Path.touch() etc. go through os.open, not builtins.open. Only O_CREAT
# can create, so only it maps to "creating" mode; O_TRUNC / O_APPEND
# without O_CREAT still require the file to exist, so behave as a read.
logical_mode = "w" if (flags & os.O_CREAT) else "r"
mapped = _remap_open(path, logical_mode)
if dir_fd is None:
return original_os_open(mapped, flags, mode)
return original_os_open(mapped, flags, mode, dir_fd = dir_fd)
def _path_mkdir(self, *args, **kwargs):
# pathlib probes Path.is_dir()/os.stat (unpatched) on FileExistsError, so
# a bare os.mkdir remap would still raise when the target exists. Remap
# the receiver up front so parents/exist_ok stays idempotent.
mapped = _remap(self)
target = self if mapped is self else self.__class__(mapped)
return original_path_mkdir(target, *args, **kwargs)
builtins.open = _open
# pathlib.Path.open / write_text / read_text call io.open directly, so patch both.
io.open = _io_open
# Python < 3.11 only: pathlib's accessor captured the ORIGINAL io.open at
# import (``_NormalAccessor.open = io.open``), so the io.open patch misses it.
# Repoint it at the same wrapper (staticmethod to stay unbound); 3.11+ dropped
# the accessor, so this is a no-op there.
accessor = getattr(pathlib, "_NormalAccessor", None)
if accessor is not None and hasattr(accessor, "open"):
accessor.open = staticmethod(_io_open)
# Path.touch() and other low-level opens call os.open directly, so patch it too.
os.open = _os_open
os.makedirs = _makedirs
# Path.mkdir(parents=True) calls os.mkdir per component, so patch os.mkdir;
# patch Path.mkdir itself too so exist_ok/parents land on the mapped path.
os.mkdir = _mkdir
pathlib.Path.mkdir = _path_mkdir
try:
_install()
except Exception: # noqa: BLE001 - a broken shim must never break user code
pass

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

@ -54,6 +54,8 @@ TOOL_XML_SIGNALS = (
# Kimi K2 / Moonshot.
"<|tool_calls_section_begin|>",
"<|tool_call_begin|>",
# TML Inkling native call marker.
"<|content_invoke_tool_json|>",
)

44
studio/backend/core/inference/tool_loop_controller.py
View file

@ -233,15 +233,50 @@ def is_tool_error(result: str) -> bool:
return isinstance(result, str) and result.lstrip().startswith(TOOL_ERROR_PREFIXES)
def _strip_mcp_image_suffix(result: str) -> str:
"""Drop a trailing __MCP_IMAGES__ envelope only when it is the valid JSON
image array appended by _flatten_result, so legit tool text that merely
mentions the marker is not truncated."""
head, sep, payload = result.rpartition("\n__MCP_IMAGES__:")
if not sep:
return result
try:
images = json.loads(payload)
except (ValueError, RecursionError):
return result
if not isinstance(images, list) or not images:
return result
if not all(
isinstance(img, dict)
and isinstance(img.get("data"), str)
and isinstance(img.get("mimeType"), str)
for img in images
):
return result
return head.rstrip()
def strip_result_for_model(result: str) -> str:
"""Remove frontend-only sentinels (image paths, RAG source map) before
feeding the result back to the model."""
result = _strip_mcp_image_suffix(result)
for sentinel in ("__IMAGES__:", "__RAG_SOURCES__:"):
if sentinel in result:
result = result.split(sentinel, 1)[0].rstrip()
return result
def append_deferred_nudges(conversation: list, msgs: Sequence[dict]) -> None:
"""Append a batch's no-op nudges as one deduped ``role=user`` message.
Deferred to after the batch's tool results so a no-op never splits an
assistant's ``tool_calls`` from their ``role=tool`` results.
"""
contents = list(dict.fromkeys(msg["content"] for msg in msgs))
if contents:
conversation.append({"role": "user", "content": "\n\n".join(contents)})
def _tool_name_from_schema(tool: Mapping[str, Any]) -> str:
function = tool.get("function")
if not isinstance(function, Mapping):
@ -253,8 +288,9 @@ def _tool_name_from_schema(tool: Mapping[str, Any]) -> str:
def _noop_result(reason: NoopReason, tool_name: str) -> str:
if reason == "duplicate":
return (
"The previous tool request was not executed because this exact "
"tool call already completed successfully. Do not repeat the same "
f"One earlier request to call tool '{tool_name}' in this batch was "
"not executed because an identical call had already completed "
"successfully. Do not repeat the same "
"tool call. Continue with a different enabled tool if that would "
"materially help, or provide the final answer if you have enough "
"information."
@ -267,8 +303,8 @@ def _noop_result(reason: NoopReason, tool_name: str) -> str:
"the requested final note or answer."
)
return (
f"The previous tool request was not executed because tool "
f"'{tool_name}' is not enabled for this request. Provide the "
f"One earlier request to call tool '{tool_name}' in this batch was "
"not executed because that tool is not enabled for this request. Provide the "
"final answer now without calling more tools."
)

284
studio/backend/core/inference/tool_stream_exec.py Normal file
View file

@ -0,0 +1,284 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Streaming wrapper around blocking server-side tool execution.
``stream_tool_execution`` runs a blocking tool call in a worker thread and
turns it into a generator that yields:
* ``{"type": "tool_output", "tool_name", "tool_call_id", "text"}`` -- an
incremental stdout/stderr chunk (python/terminal tools) for live UI output;
* ``{"type": "heartbeat"}`` -- emitted whenever nothing else has been yielded
for ``heartbeat_interval_s`` seconds, so the SSE route can write a
keepalive and reverse proxies (Cloudflare tunnels cap idle streams at
~100 s) never see a silent connection while a tool runs;
and *returns* the tool's final result string via ``StopIteration.value``
(``result = yield from stream_tool_execution(...)``). The returned result is
byte-identical to calling the tool directly, so tool-result parsing, nudging,
and healing downstream are untouched.
"""
from __future__ import annotations
import inspect
import queue
import threading
import time
from typing import Any, Callable, Generator
from loggers import get_logger
logger = get_logger(__name__)
def accepts_output_callback(func: Callable[..., str]) -> bool:
"""Whether an injectable ``execute_tool`` supports ``output_callback``.
``execute_tool`` is replaceable (tests inject fakes / the pre-PR signature),
so forward the kwarg only when the callable declares it or takes ``**kwargs``
(passing it unconditionally would ``TypeError`` on an old signature).
"""
try:
params = inspect.signature(func).parameters
except (TypeError, ValueError):
return False
if "output_callback" in params:
return True
return any(p.kind is inspect.Parameter.VAR_KEYWORD for p in params.values())
# Cadence of heartbeat events while a tool blocks with no output. Well under
# common proxy idle caps (Cloudflare ~100 s, nginx default 60 s).
TOOL_HEARTBEAT_INTERVAL_S = 10.0
# How often the wrapper wakes to poll for output / completion / cancellation.
_POLL_INTERVAL_S = 0.25
# Upper bound on how long teardown waits for the worker once the stream is
# closed or errors. A cancel-observing tool returns within this after
# ``cancel_event`` is set; a cancel-ignoring one is a daemon left to finish on
# its own rather than blocking teardown for the tool's full timeout.
_WORKER_JOIN_TIMEOUT_S = 5.0
# Cap on total streamed live-output characters per tool call, bounding the
# transient UI stream so a tight print loop cannot flood the SSE channel. Much
# higher than the model-visible result cap (tools._MAX_OUTPUT_CHARS) since the
# UI keeps the live stream as the displayed output when the result is truncated.
TOOL_OUTPUT_STREAM_MAX_CHARS = 400_000
_STREAM_CAPPED_NOTICE = "\n... (further live output not streamed)\n"
def _drain_queue(q: "queue.Queue", sentinel: object, max_chars: int | None) -> tuple[str, bool]:
"""Pull every currently-queued item, joining chunks in FIFO order.
With ``max_chars`` set, stop concatenating at the budget and discard the
remaining chunks in place, bounding peak allocation when a chatty tool queues
far more than the cap before the consumer wakes. The crossing chunk is sliced
to one char past the budget, enough for the caller's truncation to stay
byte-identical. Returns ``(joined_text, hit_sentinel)``; the surplus is still
scanned so completion is detected promptly.
"""
parts: list[str] = []
total = 0
dropping = False
hit_sentinel = False
while True:
try:
item = q.get_nowait()
except queue.Empty:
break
if item is sentinel:
hit_sentinel = True
break
if dropping:
continue
if max_chars is not None and total + len(item) > max_chars:
# Keep one char past the budget as the overflow signal; drop the rest.
parts.append(item[: max(0, max_chars - total) + 1])
dropping = True
continue
parts.append(item)
total += len(item)
return "".join(parts), hit_sentinel
def stream_tool_execution(
invoke: Callable[[Callable[[str], None]], str],
*,
tool_name: str,
tool_call_id: str = "",
cancel_event: Any = None,
heartbeat_interval_s: float = TOOL_HEARTBEAT_INTERVAL_S,
poll_interval_s: float = _POLL_INTERVAL_S,
) -> Generator[dict, None, str]:
"""Run ``invoke(output_callback)`` in a thread; yield live events; return the result.
``invoke`` receives a thread-safe ``callable(str)`` it may call with
incremental output chunks (or ignore entirely). Exceptions raised by the
tool propagate to the caller unchanged after the worker thread finishes.
``cancel_event`` is the request-level cancellation signal already handed to
the tool. If the consumer closes this generator early (an SSE disconnect
calls ``gen.close()``, raising ``GeneratorExit`` at a ``yield``), the wrapper
sets it so a cancel-observing tool stops, then joins the worker with a bounded
timeout. Set ONLY on that abnormal-exit path, never on a clean finish, because
the event is shared across a turn's tool calls and setting it early would
abort the next tool.
"""
output_queue: queue.Queue[Any] = queue.Queue()
done_sentinel = object()
outcome: dict[str, Any] = {}
# Bound accepted output at the PRODUCER boundary: the consumer-side cap alone
# wouldn't stop a fast worker enqueuing unboundedly while a slow SSE client
# backpressures. Accept at most one char past the cap (so the consumer still
# emits the capped notice) and drop the rest. The final result is captured
# independently, so this never changes the byte-identical result.
accepted_output_chars = 0
accepted_output_lock = threading.Lock()
def _on_output(text: str) -> None:
nonlocal accepted_output_chars
if not text:
return
with accepted_output_lock:
remaining = TOOL_OUTPUT_STREAM_MAX_CHARS + 1 - accepted_output_chars
if remaining <= 0:
return
accepted = text[:remaining]
accepted_output_chars += len(accepted)
output_queue.put(accepted)
def _run() -> None:
try:
outcome["result"] = invoke(_on_output)
except BaseException as exc: # noqa: BLE001 - re-raised on the caller side
outcome["error"] = exc
finally:
# Posted after the result/error is recorded; wakes the consumer
# immediately so fast tools pay no poll-interval latency.
output_queue.put(done_sentinel)
worker = threading.Thread(
target = _run,
daemon = True,
name = f"tool-exec-{tool_name or 'unknown'}",
)
worker.start()
# Heartbeats are paced by counting idle queue polls rather than a wall clock
# (tests patch ``time.monotonic`` globally, so the wrapper must not read it).
idle_polls_per_heartbeat = max(1, int(round(heartbeat_interval_s / poll_interval_s)))
idle_polls = 0
streamed_chars = 0
stream_capped = False
finished = False
def _drain_pending(max_chars: int | None = None) -> str:
nonlocal finished
text, hit_sentinel = _drain_queue(output_queue, done_sentinel, max_chars)
if hit_sentinel:
finished = True
return text
def _drain_and_drop() -> None:
"""Discard the current and every queued chunk without concatenating.
Past the cap every chunk is dropped, so don't pay to build a combined
string only to drop it. Still detect completion so the loop can exit.
"""
nonlocal finished
while True:
try:
item = output_queue.get_nowait()
except queue.Empty:
return
if item is done_sentinel:
finished = True
return
abnormal_exit = False
try:
while not finished:
try:
item = output_queue.get(timeout = poll_interval_s)
except queue.Empty:
# A disconnect sets cancel_event while the worker is silent;
# surface a heartbeat this poll so the route regains control and
# tears down at once, not after a full heartbeat interval.
if cancel_event is not None and cancel_event.is_set():
yield {"type": "heartbeat"}
continue
idle_polls += 1
if idle_polls >= idle_polls_per_heartbeat:
idle_polls = 0
yield {"type": "heartbeat"}
continue
if item is done_sentinel:
break
if stream_capped:
# Past the cap: drop this chunk and every queued sibling (see
# _drain_and_drop). Pace with one time.sleep per poll (not
# time.monotonic -- tests patch the clock), counted as an idle
# poll so heartbeats keep flowing while the queue stays non-empty.
_drain_and_drop()
if finished:
break
time.sleep(poll_interval_s)
idle_polls += 1
if idle_polls >= idle_polls_per_heartbeat:
idle_polls = 0
yield {"type": "heartbeat"}
continue
# Bound the join to the remaining budget so the crossing batch can't
# allocate far past the cap (surplus is truncated below anyway); the
# prefix is long enough that truncation stays byte-identical.
budget = TOOL_OUTPUT_STREAM_MAX_CHARS - streamed_chars
chunk = item + _drain_pending(max_chars = budget - len(item))
idle_polls = 0
if streamed_chars + len(chunk) > TOOL_OUTPUT_STREAM_MAX_CHARS:
chunk = chunk[: max(0, TOOL_OUTPUT_STREAM_MAX_CHARS - streamed_chars)]
chunk += _STREAM_CAPPED_NOTICE
stream_capped = True
streamed_chars += len(chunk)
if chunk:
yield {
"type": "tool_output",
"tool_name": tool_name,
"tool_call_id": tool_call_id,
"text": chunk,
}
except BaseException:
# The loop only raises when the consumer closes us early: an SSE
# disconnect calls gen.close() (GeneratorExit at the yield) or the route
# throws in. Signal cancellation so a cancel-observing tool returns; the
# daemon worker is then abandoned (see finally). Re-raise so the caller
# sees the real cause (GeneratorExit must not be swallowed). Runs ONLY on
# abnormal exit, so the shared cancel_event is never set out from under
# the next tool in a clean multi-tool turn.
abnormal_exit = True
if cancel_event is not None:
try:
cancel_event.set()
except Exception:
pass
raise
finally:
# Clean finish: the worker already recorded its result and queued the
# sentinel we consumed, so this join returns at once. Abnormal exit:
# cancel_event is set and the daemon worker abandoned, so join with a zero
# timeout -- teardown never blocks the caller (the route may close this
# generator on the event loop), and the daemon cannot outlive the process.
worker.join(timeout = 0 if abnormal_exit else _WORKER_JOIN_TIMEOUT_S)
error = outcome.get("error")
if error is not None:
raise error
# Returned verbatim (the loop's record_result handles non-str), so the
# final tool result is byte-identical to a direct execute_tool call.
return outcome.get("result")

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

File diff suppressed because it is too large Load diff

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

@ -291,6 +291,18 @@ def _handle_load(backend, config: dict, resp_queue: Any) -> None:
hf_token = _clean_token(config.get("hf_token"))
load_in_4bit = _resolve_lora_4bit(mc, config.get("load_in_4bit", True))
# Latest-transformers sidecar models load 16-bit: bnb 4-bit feeds quantized
# expert weights into unvalidated paths (e.g. grouped-MoE torch._grouped_mm).
if load_in_4bit:
from utils.transformers_version import latest_tier_active_for
if latest_tier_active_for(config["model_name"], hf_token):
load_in_4bit = False
logger.info(
"Latest-transformers sidecar active for %s - forcing a 16-bit "
"load (4-bit is disabled for brand-new architectures)",
config["model_name"],
)
trust_remote_code = config.get("trust_remote_code", False)
if not trust_remote_code and _needs_nemotron_trust(config["model_name"], hf_token = hf_token):
trust_remote_code = True
@ -906,7 +918,31 @@ def run_inference_process(
)
return
# ── Resolve the effective base once, before activation/gates/install (no ML import) ──
# ── Windows: check Triton availability ──
# Placed ahead of the torchao stub below (which imports torch on win32 to detect ROCm),
# matching the training and export workers' gate-then-stub ordering.
if sys.platform == "win32":
try:
import triton # noqa: F401
logger.info("Triton available — torch.compile enabled")
except ImportError:
os.environ["TORCHDYNAMO_DISABLE"] = "1"
logger.warning(
"Triton not found on Windows — torch.compile disabled. "
'Install for better performance: pip install "triton-windows<3.7"'
)
# ── Stub torchao on Windows ROCm before ANY transformers import ──
# Must precede every path that pulls transformers, not just the ML imports in section 2:
# a local LoRA adapter with no recorded base reaches transformers here via
# _resolve_base_model -> utils.models. See core/_torchao_stub.py; no-op off Windows ROCm.
from core._torchao_stub import install_torchao_windows_rocm_stub
install_torchao_windows_rocm_stub()
# ── Resolve the effective base once, before activation/gates/install ──
# No ML import on the common path; a local adapter with no recorded base pulls
# transformers via utils.models, which is why the stub above precedes this.
# A remote LoRA's base is in its Hub adapter_config.json (else surfaced only by ModelConfig
# after import). _lora_base is set only for a genuine adapter, never a full fine-tune's base.
import json as _json
@ -944,19 +980,7 @@ def run_inference_process(
)
return
# ── 1b. Windows: check Triton availability (must precede import torch) ──
if sys.platform == "win32":
try:
import triton # noqa: F401
logger.info("Triton available — torch.compile enabled")
except ImportError:
os.environ["TORCHDYNAMO_DISABLE"] = "1"
logger.warning(
"Triton not found on Windows — torch.compile disabled. "
'Install for better performance: pip install "triton-windows<3.7"'
)
# ── 1c. Security gates, then SSM/Mamba kernels, BEFORE importing transformers ──
# ── 1b. Security gates, then SSM/Mamba kernels, BEFORE importing transformers ──
# transformers snapshots its optional-backend gates at import, so a hybrid model's kernels
# must be installed before the import below ("mamba-ssm is required" otherwise). The gates
# are metadata-only, so run them first and refuse a blocked model before any native build.

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

@ -21,6 +21,7 @@ from functools import lru_cache
from typing import Callable
from utils.hardware.hardware import DeviceType, get_device
from utils.transformers_dtype import dtype_kwargs
from . import config
@ -157,9 +158,7 @@ def _get(model_name: str | None = None):
device = _device()
logger.info("loading embedding model %s on %s", name, device)
_guard_model_security(name)
_model = SentenceTransformer(
name, device = device, model_kwargs = {"torch_dtype": "float16"}
)
_model = SentenceTransformer(name, device = device, model_kwargs = dtype_kwargs("float16"))
_name = name
return _model

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

@ -103,7 +103,7 @@ def _markdown_incomplete(markdown: str, plain: str) -> bool:
return markdown_letters < _PDF_INCOMPLETE_RATIO * plain_letters
def _pdf_markdown(doc) -> list[str] | None:
def _pdf_markdown(doc, pages: range | None = None) -> list[str] | None:
"""Per-page layout-aware Markdown (tables, headings, lists) via pymupdf4llm; index
i maps to page i+1. Returns None when the lib is missing, extraction fails, or the
page count does not line up, so the caller falls back to plain PyMuPDF text."""
@ -112,28 +112,44 @@ def _pdf_markdown(doc) -> list[str] | None:
except Exception:
return None
try:
chunks = pymupdf4llm.to_markdown(
doc,
page_chunks = True,
show_progress = False,
)
kwargs = {"page_chunks": True, "show_progress": False}
if pages is not None:
kwargs["pages"] = list(pages)
chunks = pymupdf4llm.to_markdown(doc, **kwargs)
except Exception: # noqa: BLE001 - never let Markdown extraction break ingestion
logger.warning("pymupdf4llm extraction failed; using plain text", exc_info = True)
return None
if not isinstance(chunks, list) or len(chunks) != doc.page_count:
expected_pages = doc.page_count if pages is None else len(pages)
if not isinstance(chunks, list) or len(chunks) != expected_pages:
return None
return [str(c.get("text") or "") for c in chunks]
def _pdf(path: str, want_images: bool) -> tuple[list[Page], list[ParsedImage]]:
def _pdf(
source: str | bytes,
want_images: bool,
max_pages: int | None = None,
) -> tuple[list[Page], list[ParsedImage], int]:
import fitz # PyMuPDF
pages: list[Page] = []
images: list[ParsedImage] = []
doc = fitz.open(path)
doc = (
fitz.open(stream = source, filetype = "pdf") if isinstance(source, bytes) else fitz.open(source)
)
try:
md = _pdf_markdown(doc) if config.PDF_MARKDOWN else None
for i, page in enumerate(doc):
if doc.needs_pass:
raise ValueError("encrypted PDF requires a password")
total_pages = doc.page_count
page_numbers = range(total_pages if max_pages is None else min(total_pages, max_pages))
if not config.PDF_MARKDOWN:
md = None
elif max_pages is None:
md = _pdf_markdown(doc)
else:
md = _pdf_markdown(doc, page_numbers)
for i, page_number in enumerate(page_numbers):
page = doc[page_number]
plain = page.get_text("text") or ""
candidate = md[i] if md else ""
# Prefer layout-aware Markdown (keeps tables/headings legible for retrieval),
@ -147,7 +163,7 @@ def _pdf(path: str, want_images: bool) -> tuple[list[Page], list[ParsedImage]]:
text = candidate
else:
text = plain
pages.append(_page(text, i + 1))
pages.append(_page(text, page_number + 1))
if want_images:
for img in page.get_images(full = True):
xref = img[0]
@ -161,13 +177,22 @@ def _pdf(path: str, want_images: bool) -> tuple[list[Page], list[ParsedImage]]:
images.append(
ParsedImage(
image_bytes = image_bytes,
page_number = i + 1,
page_number = page_number + 1,
xref = xref,
)
)
finally:
doc.close()
return pages, images
return pages, images, total_pages
def parse_pdf_bytes(data: bytes, *, max_pages: int | None = None) -> tuple[list[Page], int]:
"""Extract PDF pages from an in-memory download using the ingestion parser.
Returns the (capped) pages plus the document's full page count, so a caller
that set ``max_pages`` can tell a fully-read short PDF from a truncated one."""
pages, _images, total_pages = _pdf(data, want_images = False, max_pages = max_pages)
return pages, total_pages
def _merge_rects(boxes: list) -> list:
@ -416,7 +441,7 @@ def parse(path: str, *, want_images: bool = False):
ext = os.path.splitext(path)[1].lower()
if ext == ".pdf":
pages, images = _pdf(path, want_images)
pages, images, _total = _pdf(path, want_images)
return (pages, images) if want_images else pages
if ext == ".docx":

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

@ -124,10 +124,12 @@ def apply_tool_strip_patterns(
# Pre-compiled patterns for tool-call XML parsing.
_TC_JSON_START_RE = re.compile(r"<tool_call>\s*\{")
# <|content_invoke_tool_json|> is TML Inkling's native call marker; its JSON uses
# an ``args`` key and the block closes with <|end_message|>.
_TC_JSON_START_RE = re.compile(r"(?:<tool_call>|<\|content_invoke_tool_json\|>)\s*\{")
_TC_GEMMA_START_RE = re.compile(r"<\|tool_call>\s*call\s*:\s*([\w.\-]+)\s*\{")
_TC_FUNC_START_RE = re.compile(r"<function=([\w-]+)>\s*")
_TC_END_TAG_RE = re.compile(r"</tool_call>")
_TC_END_TAG_RE = re.compile(r"</tool_call>|<\|end_message\|>")
_TC_GEMMA_END_TAG_RE = re.compile(r"<tool_call\|>")
_TC_FUNC_CLOSE_RE = re.compile(r"\s*</function>\s*$")
# Horizontal-whitespace trailing class keeps the wrapping newline; _trim_param_value trims it.
@ -686,12 +688,24 @@ def parse_tool_calls_from_text(
if kind == "json":
obj = json.loads(content[m.end() - 1 : brace_end + 1])
name = obj.get("name", "")
# Accept ``parameters`` alias for ``arguments`` (Llama-3.2 drift inside Hermes).
# Accept ``parameters`` alias for ``arguments`` (Llama-3.2 drift inside
# Hermes) and ``args`` (TML Inkling native calls).
arguments = obj.get("arguments")
if arguments is None:
arguments = obj.get("parameters", {})
arguments = obj.get("parameters")
if arguments is None:
arguments = obj.get("args", {})
if isinstance(arguments, dict):
arguments = json.dumps(arguments)
# Inkling echoes the bare tool name (and a role opener) before the
# marker: <|message_model|>NAME<|content_invoke_tool_json|>{...}.
# Fold that echo into the markup span so promotion removes it too.
if name and content.startswith("<|content_invoke_tool_json|>", start):
pre = content[:start]
if pre.endswith(name):
start -= len(name)
if content[:start].endswith("<|message_model|>"):
start -= len("<|message_model|>")
else:
name = m.group(1)
arguments = json.dumps(_gemma_arguments_to_json(content[m.end() : brace_end]))

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

@ -11,8 +11,10 @@ import os
import sys
import types
# Prevent tokenizer parallelism deadlocks when datasets forks.
os.environ["TOKENIZERS_PARALLELISM"] = "false"
# Off on Linux so datasets' forked map() workers can't deadlock. On spawn platforms
# (Windows/macOS) map() runs in-process, so keep the fast tokenizer's Rust threads on
# (the only parallelism single-process tokenize gets; off makes prep run serially).
os.environ["TOKENIZERS_PARALLELISM"] = "true" if sys.platform in ("win32", "darwin") else "false"
# Make compiled cache modules importable by any subprocess. On spawn platforms
# (Windows/macOS) spawned dataset.map() workers re-import top-level modules, and
@ -70,7 +72,7 @@ from core.inference.llama_cpp import _hf_offline_if_dns_dead
from utils.models import is_vision_model, detect_audio_type
from utils.models.model_config import _env_offline
from utils.datasets import format_and_template_dataset
from utils.datasets import MODEL_TO_TEMPLATE_MAPPER, TEMPLATE_TO_RESPONSES_MAPPER
from utils.datasets.completion_masking import apply_completion_masking
from utils.datasets.iterable import is_streaming_dataset as detect_streaming_dataset
from utils.datasets.raw_text import prepare_raw_text_dataset, resolve_column_names
from utils.paths import (
@ -931,7 +933,7 @@ class UnslothTrainer:
use_gradient_checkpointing = "unsloth"
elif use_gradient_checkpointing in ("true", "1", "yes"):
use_gradient_checkpointing = True
elif use_gradient_checkpointing in ("false", "0", "no"):
elif use_gradient_checkpointing in ("false", "0", "no", "none", "off"):
use_gradient_checkpointing = False
else:
# Invalid value -> "unsloth"
@ -3455,8 +3457,6 @@ class UnslothTrainer:
# ========== TRAIN ON RESPONSES ONLY ==========
# Raw-text datasets always train on all tokens.
instruction_part = None
response_part = None
is_cpt = training_args.get("is_cpt", False)
train_on_responses_enabled = (
False
@ -3473,113 +3473,93 @@ class UnslothTrainer:
# DeepSeek OCR handles this internally in its collator, so skip
# Audio VLM handles label masking in its collator, so skip
# Markers auto-detected from the chat template first, manual table
# as fallback; gpt-oss stays on its manual markers. See
# apply_completion_masking.
if (
train_on_responses_enabled
and not self.is_audio_vlm
and not self.is_audio
and not (is_deepseek_ocr or dataset_final_format == "alpaca")
):
try:
logger.info("Configuring train on responses only...\n")
from unsloth.chat_templates import train_on_responses_only
# Template mapping for this model
model_name_lower = self.model_name.lower()
logger.info("Configuring train on responses only...\n")
if model_name_lower in MODEL_TO_TEMPLATE_MAPPER:
template_name = MODEL_TO_TEMPLATE_MAPPER[model_name_lower]
logger.info(f"Detected template: {template_name}\n")
def _notify(level, message):
if level == "warning":
logger.warning(message)
else:
logger.info(f"{message}\n")
if template_name in TEMPLATE_TO_RESPONSES_MAPPER:
instruction_part = TEMPLATE_TO_RESPONSES_MAPPER[template_name][
"instruction"
]
response_part = TEMPLATE_TO_RESPONSES_MAPPER[template_name]["response"]
# No try/except: the helper handles detection failures and
# double misses itself, so an exception here is a real masking
# failure that must fail the run, not silently train on full
# sequences.
self.trainer, masking_applied = apply_completion_masking(
self.trainer,
self.model_name,
train_on_responses_only,
num_proc = config_args["dataset_num_proc"],
notify = _notify,
)
logger.info(f"Instruction marker: {instruction_part[:50]}...\n")
logger.info(f"Response marker: {response_part[:50]}...\n")
if not masking_applied:
train_on_responses_enabled = False
if masking_applied:
try:
# ── Safety net: check if all samples were filtered out ──
# train_on_responses_only masks non-response tokens with -100; a
# row becomes all -100 (Unsloth drops it) when the response
# template is not found in the formatted text. Usually a
# dataset/template mismatch (already-formatted data, or 'Train on
# completions' on data that doesn't match the model's chat
# template); only sometimes max_seq_length truncating the response
# away. Skip this len()-based check for streaming.
if detect_streaming_dataset(self.trainer.train_dataset):
logger.info("Skipping post-filter length check for streaming dataset\n")
else:
logger.info(
f"No response mapping found for template: {template_name}\n"
filtered_len = len(self.trainer.train_dataset)
original_dataset_obj = (
dataset["dataset"] if isinstance(dataset, dict) else dataset
)
train_on_responses_enabled = False
else:
logger.info(f"No template mapping found for model: {self.model_name}\n")
train_on_responses_enabled = False
except Exception as e:
logger.warning(f"Could not configure train on responses: {e}")
train_on_responses_enabled = False
# Apply train on responses only if we have valid parts
if (
train_on_responses_enabled
and instruction_part
and response_part
and not self.is_audio_vlm
and not self.is_audio
and not (is_deepseek_ocr or dataset_final_format == "alpaca")
):
try:
from unsloth.chat_templates import train_on_responses_only
self.trainer = train_on_responses_only(
self.trainer,
instruction_part = instruction_part,
response_part = response_part,
num_proc = config_args["dataset_num_proc"],
)
logger.info("Train on responses only configured successfully\n")
# ── Safety net: check if all samples were filtered out ──
# train_on_responses_only masks non-response tokens with -100;
# a row becomes all -100 (and Unsloth drops it) when the response
# template is not found in the formatted text. That is usually a
# dataset/template mismatch (already-formatted data, or 'Train on
# completions' applied to data that doesn't match the model's chat
# template), and only sometimes max_seq_length truncating the
# response away. Skip this len()-based check for streaming.
if detect_streaming_dataset(self.trainer.train_dataset):
logger.info("Skipping post-filter length check for streaming dataset\n")
else:
filtered_len = len(self.trainer.train_dataset)
original_dataset_obj = (
dataset["dataset"] if isinstance(dataset, dict) else dataset
)
original_len = len(original_dataset_obj)
dropped = original_len - filtered_len
drop_pct = round(100 * dropped / original_len, 1) if original_len > 0 else 0
if filtered_len == 0 or drop_pct > 30:
max_seq = training_args.get("max_seq_length", 2048)
error_msg = (
f"{dropped}/{original_len} samples ({drop_pct}%) were "
f"dropped after applying 'Train on completions': after "
f"masking, those rows had no trainable response tokens "
f"left. The usual cause is that this model's response "
f"template was not found in the formatted samples, so "
f"every token was masked out. That typically means the "
f"dataset is already formatted, or its structure does "
f"not match the model's chat template, so 'Train on "
f"completions' should be turned off for this dataset. "
f"Less commonly, a max_seq_length ({max_seq}) shorter "
f"than the prompt can truncate the response away; only "
f"raise it if your samples are actually longer than that."
original_len = len(original_dataset_obj)
dropped = original_len - filtered_len
drop_pct = (
round(100 * dropped / original_len, 1) if original_len > 0 else 0
)
logger.error(error_msg)
self._update_progress(error = error_msg, is_training = False)
return
if dropped > 0:
logger.info(
f"⚠️ {dropped}/{original_len} samples "
f"({drop_pct}%) were dropped (all labels "
f"masked). {filtered_len} samples remain.\n"
)
logger.info(f"Post-filter dataset size: {filtered_len} samples\n")
if filtered_len == 0 or drop_pct > 30:
max_seq = training_args.get("max_seq_length", 2048)
error_msg = (
f"{dropped}/{original_len} samples ({drop_pct}%) were "
f"dropped after applying 'Train on completions': after "
f"masking, those rows had no trainable response tokens "
f"left. The usual cause is that this model's response "
f"template was not found in the formatted samples, so "
f"every token was masked out. That typically means the "
f"dataset is already formatted, or its structure does "
f"not match the model's chat template, so 'Train on "
f"completions' should be turned off for this dataset. "
f"Less commonly, a max_seq_length ({max_seq}) shorter "
f"than the prompt can truncate the response away; only "
f"raise it if your samples are actually longer than that."
)
logger.error(error_msg)
self._update_progress(error = error_msg, is_training = False)
return
except Exception as e:
logger.warning(f"Failed to apply train on responses only: {e}")
train_on_responses_enabled = False
if dropped > 0:
logger.info(
f"⚠️ {dropped}/{original_len} samples "
f"({drop_pct}%) were dropped (all labels "
f"masked). {filtered_len} samples remain.\n"
)
logger.info(f"Post-filter dataset size: {filtered_len} samples\n")
except Exception as e:
logger.warning(f"Post-masking dataset size check failed: {e}")
else:
if train_on_responses_enabled and is_deepseek_ocr:
logger.info("Train on responses handled by DeepSeek OCR collator\n")

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

@ -39,6 +39,27 @@ from utils.paths import outputs_root
logger = get_logger(__name__)
def _env_int(name: str, default: int) -> int:
try:
raw = (os.environ.get(name) or "").strip()
return int(raw) if raw else default
except ValueError:
return default
# Stop-watchdog escalation timeouts. Primary trigger: a short grace once "complete"
# (save done). Absolute cap is a backstop: long for save=True so a slow save is never
# killed mid-write, shorter for a cancel that has nothing to save.
_STOP_GRACE_S = _env_int("UNSLOTH_STUDIO_TRAINING_STOP_GRACE_S", 15)
_STOP_TIMEOUT_S = _env_int("UNSLOTH_STUDIO_TRAINING_STOP_TIMEOUT_S", 600)
_CANCEL_TIMEOUT_S = _env_int("UNSLOTH_STUDIO_TRAINING_CANCEL_TIMEOUT_S", 120)
# Watchdog DB finalize: a few short retries so a transient SQLite lock doesn't lose the
# terminal state, since the watchdog is the sole finalizer once _proc is dropped.
_DB_FINALIZE_RETRIES = 3
_DB_FINALIZE_RETRY_S = 0.5
_pyplot = None
_pyplot_failed = False
@ -741,11 +762,22 @@ class TrainingBackend:
self._pump_running: bool = False
self._lock = threading.Lock()
# Stop watchdog: after a stop is requested, escalates to force_terminate()
# if the worker does not exit on its own within a bounded time. The watched
# proc is tracked so a new run always gets its own watcher.
self._stop_watchdog: Optional[threading.Thread] = None
self._stop_watchdog_proc: Optional[mp.Process] = None
self._complete_seen = threading.Event()
# Progress state (updated by pump thread from subprocess events)
self._progress = TrainingProgress()
self._should_stop = False
self._cancel_requested = False # True only for stop(save=False)
# Throttled training-status logging to the server log (not one line/step).
self._last_progress_log_ts: float = 0.0
self._last_progress_log_step: int = -1
# Training metrics (consumed by routes for SSE and /metrics)
self.loss_history: list = []
self.lr_history: list = []
@ -765,6 +797,7 @@ class TrainingBackend:
self._metric_buffer: list[dict] = []
self._run_finalized: bool = False
self._db_run_created: bool = False
self._db_create_in_progress: bool = False
self._db_total_steps_set: bool = False
self._db_config: Optional[dict] = None
self._db_started_at: Optional[str] = None
@ -852,91 +885,129 @@ class TrainingBackend:
else:
defer_auto_selection = True
# Synchronous validation passed -> free VRAM (export + chat) now, before
# auto-selection and the spawn, so placement sees the freed memory.
if before_spawn is not None:
try:
before_spawn()
except Exception:
logger.warning("before_spawn hook failed; continuing", exc_info = True)
# Handshake with the sidecar install route: mark the spawn in progress BEFORE rechecking
# the reservation, so either this recheck aborts, or the install's is_training_active()
# sees this flag (or the recorded proc) and refuses.
from utils.transformers_version import sidecar_swap_in_progress
if defer_auto_selection:
resolved_gpu_ids, gpu_selection = prepare_gpu_selection(None, **gpu_selection_kwargs)
config["resolved_gpu_ids"] = resolved_gpu_ids
config["gpu_selection"] = gpu_selection
from .worker import run_training_process
self._spawn_in_progress = True
if sidecar_swap_in_progress():
self._spawn_in_progress = False
from utils.transformers_version import SidecarSwapInProgress
raise SidecarSwapInProgress(
"A transformers installation is replacing the latest sidecar; "
"retry when it completes."
)
# Any exception between the handshake above and the flag reset below would
# otherwise leave _spawn_in_progress latched, wedging is_training_active
# (and the install route) until restart.
try:
with native_path_secret_removed_for_child_start():
event_queue = _CTX.Queue()
stop_queue = _CTX.Queue()
# Synchronous validation passed -> free VRAM (export + chat) now, before
# auto-selection and the spawn, so placement sees the freed memory. Runs AFTER the handshake
# so a lost race to an install can't tear down chat/export for a training run that never spawns.
if before_spawn is not None:
try:
before_spawn()
except Exception:
logger.warning("before_spawn hook failed; continuing", exc_info = True)
proc = _CTX.Process(
target = run_without_native_path_secret,
args = (run_training_process,),
kwargs = {
"event_queue": event_queue,
"stop_queue": stop_queue,
"config": config,
},
daemon = True,
)
proc.start()
from utils.process_lifetime import adopt_pid
if defer_auto_selection:
try:
resolved_gpu_ids, gpu_selection = prepare_gpu_selection(
None, **gpu_selection_kwargs
)
except Exception:
# Flag is already set; a failed GPU selection must not leave is_training_active stuck True.
self._spawn_in_progress = False
raise
config["resolved_gpu_ids"] = resolved_gpu_ids
config["gpu_selection"] = gpu_selection
from .worker import run_training_process
try:
with native_path_secret_removed_for_child_start():
event_queue = _CTX.Queue()
stop_queue = _CTX.Queue()
proc = _CTX.Process(
target = run_without_native_path_secret,
args = (run_training_process,),
kwargs = {
"event_queue": event_queue,
"stop_queue": stop_queue,
"config": config,
},
daemon = True,
)
proc.start()
from utils.process_lifetime import adopt_pid
adopt_pid(proc.pid) # bind to parent lifetime (Windows job / sweep)
except Exception:
logger.error("Failed to start training subprocess", exc_info = True)
self._spawn_in_progress = False
return False
logger.info("Training subprocess started (pid=%s)", proc.pid)
# Reset state (old pump thread dead, proc.start() succeeded).
self.current_job_id = job_id
self._should_stop = False
self._cancel_requested = False
self._complete_seen.clear()
self._progress = TrainingProgress(
is_training = True, status_message = "Initializing training..."
)
# Reset the progress-log throttle so the new run always logs its first step,
# even if it starts within 30s of a prior run whose last logged step matches.
self._last_progress_log_ts = 0.0
self._last_progress_log_step = -1
self.loss_history.clear()
self.lr_history.clear()
self.step_history.clear()
self.grad_norm_history.clear()
self.grad_norm_step_history.clear()
self.eval_loss_history.clear()
self.eval_step_history.clear()
self.eval_enabled = False
self._output_dir = None
self._metric_buffer.clear()
self._run_finalized = False
self._db_run_created = False
self._db_create_in_progress = False # a stale watchdog create can't block this run
self._db_total_steps_set = False
self._db_config = _sanitize_db_config(config)
self._db_started_at = datetime.now(timezone.utc).isoformat()
# Start each job Xet-first; keep config so a stall can respawn over HTTP.
self._last_full_config = config
self._in_model_load = False
self._xet_fallback_used = False
self._needs_xet_respawn = False
# Create the DB run row before the pump can consume events, so it appears
# in history during model loading and a fast terminal worker can't race the
# pump into a duplicate create/finalize. From here the pump only finalizes.
self._ensure_db_run_created()
# Assign handles and start the pump together under the lock so a concurrent
# poll can't see a live _proc with no pump and spawn a duplicate.
new_pump = threading.Thread(target = self._pump_loop, daemon = True)
with self._lock:
self._pump_running = False
self._event_queue = event_queue
self._stop_queue = stop_queue
self._proc = proc
self._pump_thread = new_pump
new_pump.start()
self._spawn_in_progress = False
return True
adopt_pid(proc.pid) # bind to parent lifetime (Windows job / sweep)
except Exception:
logger.error("Failed to start training subprocess", exc_info = True)
return False
logger.info("Training subprocess started (pid=%s)", proc.pid)
# Reset state (old pump thread dead, proc.start() succeeded).
self.current_job_id = job_id
self._should_stop = False
self._cancel_requested = False
self._progress = TrainingProgress(
is_training = True, status_message = "Initializing training..."
)
self.loss_history.clear()
self.lr_history.clear()
self.step_history.clear()
self.grad_norm_history.clear()
self.grad_norm_step_history.clear()
self.eval_loss_history.clear()
self.eval_step_history.clear()
self.eval_enabled = False
self._output_dir = None
self._metric_buffer.clear()
self._run_finalized = False
self._db_run_created = False
self._db_total_steps_set = False
self._db_config = _sanitize_db_config(config)
self._db_started_at = datetime.now(timezone.utc).isoformat()
# Start each job Xet-first; keep config so a stall can respawn over HTTP.
self._last_full_config = config
self._in_model_load = False
self._xet_fallback_used = False
self._needs_xet_respawn = False
# Create the DB run row before the pump can consume events, so it appears
# in history during model loading and a fast terminal worker can't race the
# pump into a duplicate create/finalize. From here the pump only finalizes.
self._ensure_db_run_created()
# Assign handles and start the pump together under the lock so a concurrent
# poll can't see a live _proc with no pump and spawn a duplicate.
new_pump = threading.Thread(target = self._pump_loop, daemon = True)
with self._lock:
self._pump_running = False
self._event_queue = event_queue
self._stop_queue = stop_queue
self._proc = proc
self._pump_thread = new_pump
new_pump.start()
return True
self._spawn_in_progress = False
raise
def stop_training(self, save: bool = True) -> bool:
"""Send stop signal to the training subprocess."""
@ -953,15 +1024,212 @@ class TrainingBackend:
self._progress.status_message = (
"Stopping training and saving checkpoint..." if save else "Cancelling training..."
)
# Guarantee the run finalizes even if the worker wedges after saving.
self._start_stop_watchdog(cancel = not save)
return True
def force_terminate(self) -> None:
"""Force-kill the training subprocess so state can be reset immediately."""
def _start_stop_watchdog(self, cancel: bool) -> None:
"""Start a daemon that force-terminates the worker if a requested stop does not
exit on its own. No-op if no worker is alive or a live watchdog already watches
this proc (a stale watchdog on an old proc never blocks a new run's watcher)."""
with self._lock:
if self._proc is not None and self._proc.is_alive():
logger.info("Force-terminating training subprocess (pid=%s)", self._proc.pid)
self._proc.terminate()
proc = self._proc
if proc is None or not proc.is_alive():
return
if (
self._stop_watchdog is not None
and self._stop_watchdog.is_alive()
and self._stop_watchdog_proc is proc
):
return
watchdog = threading.Thread(
target = self._stop_watchdog_loop,
args = (proc, cancel, self.current_job_id),
name = f"stop-watchdog-{self.current_job_id or 'unknown'}",
daemon = True,
)
self._stop_watchdog = watchdog
self._stop_watchdog_proc = proc
watchdog.start()
def _stop_watchdog_loop(
self,
target_proc: "mp.Process",
cancel: bool,
watched_job_id: Optional[str] = None,
) -> None:
"""Escalate a stuck stop to force_terminate(): grace after "complete", else the
absolute backstop (see the module timeouts). No-ops on a clean exit; exits
silently if a new run replaces the worker."""
started = time.monotonic()
complete_at: Optional[float] = None
reason = ""
while True:
with self._lock:
superseded = self._proc is not target_proc
# A later cancel has nothing to save, so tighten an in-flight save
# watchdog to the shorter cancel cap.
cancelling = cancel or self._cancel_requested
if superseded or not target_proc.is_alive():
return
now = time.monotonic()
abs_timeout = _CANCEL_TIMEOUT_S if cancelling else _STOP_TIMEOUT_S
if complete_at is None and self._complete_seen.is_set():
complete_at = now
if complete_at is not None and now - complete_at >= _STOP_GRACE_S:
reason = "worker still alive after save"
break
if now - started >= abs_timeout:
reason = "worker did not exit within the absolute timeout"
break
time.sleep(0.5)
with self._lock:
superseded = self._proc is not target_proc
if superseded or not target_proc.is_alive():
return
if complete_at is None:
# Backstop fired pre-completion: a save may still be in progress.
logger.warning(
"Stop watchdog: absolute timeout with no completion signal; "
"force-terminating a possibly-mid-save worker: %s",
reason,
)
else:
logger.warning("Stop watchdog force-terminating stuck training worker: %s", reason)
# force_terminate can raise on a wedged child; finalize regardless.
try:
self.force_terminate(target_proc = target_proc)
except Exception:
logger.exception("Stop watchdog: force_terminate failed; finalizing anyway")
finally:
self._finalize_stopped_after_escalation(
target_proc = target_proc, watched_job_id = watched_job_id
)
def _finalize_stopped_after_escalation(
self,
target_proc: "Optional[mp.Process]" = None,
watched_job_id: Optional[str] = None,
) -> None:
"""Finalize parent state after a force-terminate so the UI leaves "Stopping..."
even if the worker is wedged in driver teardown; preserves output_dir so a saved
checkpoint is kept. No-ops if a new run already replaced the watched worker, so a
stale watchdog never marks a fresh run stopped or drops its handle.
Supersession is checked on both the watched proc and job id: start_training sets
current_job_id before it installs the new _proc, so a stale watchdog entering that
startup window still sees the old (dead) handle and is caught by the job-id guard.
The run's terminal DB state is recorded (create-if-needed + finish by captured id)
BEFORE _proc is dropped: a wedged worker still reports alive, so the pump never
reaches its own finalize and would bail on its _proc-is-None guard once the handle
is gone. While the handle is held is_training_active() stays true, so no new run can
start and current_job_id stays the watched run for the write. _proc is dropped last,
re-guarded on target_proc so a run that did replace the worker keeps its handle."""
with self._lock:
if target_proc is not None and self._proc is not target_proc:
return # a new run replaced the worker; never touch its state
if watched_job_id is not None and self.current_job_id != watched_job_id:
return # a new run is already starting up; leave its state alone
run_id = self.current_job_id # == watched_job_id
self._progress.is_training = False
self._progress.status_message = "Training stopped."
# Create the row if a start-time create failed (no-op otherwise; skips when the pump
# is mid-create, in which case its create-then-finalize records the run instead).
self._ensure_db_run_created()
with self._lock:
claim = (
bool(run_id)
and self.current_job_id == run_id
and self._db_run_created
and not self._run_finalized
)
batch: list = []
final_step = final_loss = duration = None
loss_history: list = []
output_dir = self._output_dir
if claim:
self._run_finalized = True # claim this run's finalize
batch = list(self._metric_buffer)
del self._metric_buffer[: len(batch)]
final_step = self._progress.step
final_loss = self._progress.loss
if final_loss is not None and not math.isfinite(final_loss):
final_loss = None
duration = self._progress.elapsed_seconds
loss_history = list(self.loss_history)
if claim:
self._finish_stopped_run(
run_id, output_dir, batch, final_step, final_loss, duration, loss_history
)
with self._lock:
if target_proc is None or self._proc is target_proc:
self._proc = None # drop only our handle, never a run that replaced it
def _finish_stopped_run(
self,
run_id: str,
output_dir: Optional[str],
batch: list,
final_step: Optional[int],
final_loss: Optional[float],
duration: Optional[float],
loss_history: list,
) -> None:
"""Record a force-stopped run finished by its captured id, from state snapshotted
under the lock. insert_metrics_batch upserts and finish_run is an idempotent UPDATE,
so a concurrent pump finalize of the same run is harmless and a different current run
is never touched. The watchdog is the sole finalizer once _proc is dropped, so a
transient DB error (e.g. a SQLite lock) is retried a few times; on final failure the
finalize is unclaimed (only if the run is still current) so the row is not left
claimed-but-unfinalized."""
for attempt in range(_DB_FINALIZE_RETRIES):
try:
from storage.studio_db import finish_run, insert_metrics_batch
from utils.downsample import downsample
if batch:
insert_metrics_batch(run_id, batch)
sparkline = downsample(loss_history, 50)
finish_run(
id = run_id,
status = "stopped",
ended_at = datetime.now(timezone.utc).isoformat(),
final_step = final_step,
final_loss = final_loss,
duration_seconds = duration,
loss_sparkline = _json.dumps(sparkline),
output_dir = output_dir,
error_message = None,
)
return
except Exception:
if attempt + 1 < _DB_FINALIZE_RETRIES:
time.sleep(_DB_FINALIZE_RETRY_S)
continue
logger.warning(
"Failed to finalize stopped run %s in DB after %d attempts",
run_id,
_DB_FINALIZE_RETRIES,
exc_info = True,
)
with self._lock:
# Only if still current; a new run's finalize state is never touched.
if self.current_job_id == run_id:
self._run_finalized = False
def force_terminate(self, target_proc: "Optional[mp.Process]" = None) -> None:
"""Force-kill the training subprocess so state can be reset immediately. With
``target_proc``, terminate only that handle and no-op if a new run has replaced
it, so the watchdog can never kill a fresh worker."""
with self._lock:
proc = self._proc
if target_proc is not None and proc is not target_proc:
return # superseded by a new run; do not touch the new worker
if proc is not None and proc.is_alive():
logger.info("Force-terminating training subprocess (pid=%s)", proc.pid)
proc.terminate()
cancelled = self._cancel_requested
output_dir = self._output_dir
@ -1038,50 +1306,84 @@ class TrainingBackend:
from .worker import run_training_process
try:
with native_path_secret_removed_for_child_start():
event_queue = _CTX.Queue()
stop_queue = _CTX.Queue()
new_proc = _CTX.Process(
target = run_without_native_path_secret,
args = (run_training_process,),
kwargs = {
"event_queue": event_queue,
"stop_queue": stop_queue,
"config": config,
},
daemon = True,
)
new_proc.start()
from utils.process_lifetime import adopt_pid
# This run is active, so an install request 409s rather than proceeds: a reservation seen here
# is transient (an aborting install or short lazy repair). Wait it out instead of stranding the
# stalled run; only a wedged reservation fails the respawn.
from utils.transformers_version import sidecar_swap_in_progress
adopt_pid(new_proc.pid) # bind to parent lifetime (Windows job / sweep)
except Exception:
logger.error("Failed to respawn training subprocess", exc_info = True)
with self._lock:
# No replacement pump will run; clear the flag so a later run can't
# inherit a stale _pump_running=True and spawn a duplicate.
self._pump_running = False
self._progress.is_training = False
self._progress.error = "Failed to recover stalled model download"
self._ensure_db_run_created()
self._finalize_run_in_db(
status = "error",
error_message = "Failed to recover stalled model download",
self._spawn_in_progress = True
_swap_wait_deadline = time.time() + 120
while sidecar_swap_in_progress() and time.time() < _swap_wait_deadline:
time.sleep(1)
if sidecar_swap_in_progress():
# Raising here would land in the pump's broad finalization catch and
# strand the run in a training state with no worker: finalize it as a
# failure explicitly instead.
self._spawn_in_progress = False
msg = (
"A transformers installation is replacing the latest sidecar; "
"cannot respawn the training worker."
)
logger.error(msg)
with self._lock:
self._progress.is_training = False
self._progress.error = msg
self._ensure_db_run_created()
self._finalize_run_in_db(status = "error", error_message = msg)
return
logger.info("Training subprocess respawned with Xet disabled (pid=%s)", new_proc.pid)
new_pump = threading.Thread(target = self._pump_loop, daemon = True)
with self._lock:
self._in_model_load = False
self._event_queue = event_queue
self._stop_queue = stop_queue
self._proc = new_proc
self._pump_thread = new_pump
# Start under the lock so _ensure_pump_alive can never observe the
# new pump as a not-yet-started (dead) thread and spawn a duplicate.
new_pump.start()
# Reset the handshake flag on any unexpected failure past this point, so a
# crashed respawn cannot wedge is_training_active until restart.
try:
try:
with native_path_secret_removed_for_child_start():
event_queue = _CTX.Queue()
stop_queue = _CTX.Queue()
new_proc = _CTX.Process(
target = run_without_native_path_secret,
args = (run_training_process,),
kwargs = {
"event_queue": event_queue,
"stop_queue": stop_queue,
"config": config,
},
daemon = True,
)
new_proc.start()
from utils.process_lifetime import adopt_pid
adopt_pid(new_proc.pid) # bind to parent lifetime (Windows job / sweep)
except Exception:
logger.error("Failed to respawn training subprocess", exc_info = True)
self._spawn_in_progress = False
with self._lock:
# No replacement pump will run; clear the flag so a later run can't
# inherit a stale _pump_running=True and spawn a duplicate.
self._pump_running = False
self._progress.is_training = False
self._progress.error = "Failed to recover stalled model download"
self._ensure_db_run_created()
self._finalize_run_in_db(
status = "error",
error_message = "Failed to recover stalled model download",
)
return
logger.info("Training subprocess respawned with Xet disabled (pid=%s)", new_proc.pid)
new_pump = threading.Thread(target = self._pump_loop, daemon = True)
with self._lock:
self._in_model_load = False
self._event_queue = event_queue
self._stop_queue = stop_queue
self._proc = new_proc
self._spawn_in_progress = False
self._pump_thread = new_pump
# Start under the lock so _ensure_pump_alive can never observe the
# new pump as a not-yet-started (dead) thread and spawn a duplicate.
new_pump.start()
except Exception:
self._spawn_in_progress = False
raise
def _ensure_pump_alive(self) -> bool:
"""Restart the event pump if it crashed, even after the worker exited.
@ -1114,6 +1416,10 @@ class TrainingBackend:
def is_training_active(self) -> bool:
"""Check if training is currently active."""
# A spawn past its sidecar-swap recheck counts as active even before _proc is recorded,
# so an install cannot slip in mid-spawn.
if getattr(self, "_spawn_in_progress", False):
return True
# Self-heal a crashed pump first: a dead pump must never leave the worker
# training invisibly behind a frozen UI. Cheap enough for per-second polls.
self._ensure_pump_alive()
@ -1468,6 +1774,8 @@ class TrainingBackend:
"training cancelled",
"training stopped",
}
# Save is done by now; let the stop watchdog start its grace timer.
self._complete_seen.set()
self._progress.is_training = False
self._progress.is_completed = not stopped
self._output_dir = event.get("output_dir")
@ -1531,91 +1839,167 @@ class TrainingBackend:
elif db_action == "finalize":
self._finalize_run_in_db(**db_action_kwargs)
def _ensure_db_run_created(self) -> None:
"""Create the DB row if it doesn't exist yet. Called outside the lock."""
if self._db_run_created or not self.current_job_id or not self._db_config:
if etype == "progress":
self._log_training_progress()
def _log_training_progress(self) -> None:
"""One throttled training-status line to the server log (the per-step stream
still goes to the UI via SSE): first step, then at most every 30s, plus the
final step; resyncs on a new run. Runs on the pump thread."""
p = self._progress
step = int(p.step or 0)
if step <= 0:
return
total = int(p.total_steps or 0)
is_final = total > 0 and step >= total
prev = self._last_progress_log_step
if step == prev:
return
now = time.monotonic()
if prev >= 0 and step > prev and not is_final and (now - self._last_progress_log_ts) < 30.0:
return
self._last_progress_log_ts = now
self._last_progress_log_step = step
logger.info(
"training_progress",
step = step,
total_steps = total or None,
percent = int(step * 100 / total) if total > 0 else None,
loss = round(p.loss, 4) if p.loss is not None else None,
epoch = round(p.epoch, 2) if p.epoch is not None else None,
eta_s = int(p.eta_seconds) if p.eta_seconds else None,
)
def _ensure_db_run_created(self) -> None:
"""Create the DB row if it doesn't exist yet. An in-progress flag lets only one
caller create at a time, and ``_db_run_created`` is published only after
``create_run`` commits, so a concurrent finalize never runs ``finish_run`` against a
not-yet-inserted row (a zero-row UPDATE that would leave the run stuck as running)."""
with self._lock:
if (
self._db_run_created
or self._db_create_in_progress
or not self.current_job_id
or not self._db_config
):
return
self._db_create_in_progress = True # only one caller creates
job_id = self.current_job_id
db_config = self._db_config
started_at = self._db_started_at or datetime.now(timezone.utc).isoformat()
total_steps = self._progress.total_steps or None
created = False
try:
from storage.studio_db import create_run
dataset_name = (
self._db_config.get("hf_dataset")
or next(iter(self._db_config.get("local_datasets") or []), None)
or _s3_dataset_name(self._db_config.get("s3_dataset"))
db_config.get("hf_dataset")
or next(iter(db_config.get("local_datasets") or []), None)
or _s3_dataset_name(db_config.get("s3_dataset"))
or "unknown"
)
create_run(
id = self.current_job_id,
model_name = self._db_config["model_name"],
id = job_id,
model_name = db_config["model_name"],
dataset_name = dataset_name,
config_json = _json.dumps(self._db_config),
started_at = self._db_started_at or datetime.now(timezone.utc).isoformat(),
total_steps = self._progress.total_steps or None,
config_json = _json.dumps(db_config),
started_at = started_at,
total_steps = total_steps,
)
self._db_run_created = True
created = True
except Exception:
logger.warning("Failed to create DB run record for early failure", exc_info = True)
finally:
with self._lock:
# Publish the flags only if this is still the current run. A killed worker
# lets a new /start proceed mid-create, and these flags are backend-wide, so
# a stale create for the captured job must not satisfy the new run's DB state
# (the row was still created by id; the new run owns/creates its own row).
if self.current_job_id == job_id:
if created:
self._db_run_created = True # publish only after the insert commits
self._db_create_in_progress = False
def _finalize_run_in_db(
self,
status: str,
error_message: Optional[str] = None,
output_dir: Optional[str] = None,
expected_job_id: Optional[str] = None,
) -> None:
"""Flush remaining metrics and mark a run as finished in the DB."""
if not self.current_job_id or not self._db_run_created or self._run_finalized:
return
self._flush_metrics_to_db()
"""Flush remaining metrics and mark a run finished in the DB. Claims the finalize
under the lock so the watchdog and pump can't double-finalize, and no-ops when
``expected_job_id`` no longer matches (a new run took over). The run id and final
progress are snapshotted under the lock and threaded through the flush/finish calls,
so a new run racing between this claim and the DB writes can't be flushed or marked
stopped under the old run's finalize."""
with self._lock:
if expected_job_id is not None and self.current_job_id != expected_job_id:
return
if not self.current_job_id or not self._db_run_created or self._run_finalized:
return
self._run_finalized = True
run_id = self.current_job_id
final_step = self._progress.step
final_loss = self._progress.loss
if final_loss is not None and not math.isfinite(final_loss):
final_loss = None
duration = self._progress.elapsed_seconds
loss_history = list(self.loss_history)
self._flush_metrics_to_db(run_id = run_id)
try:
from storage.studio_db import finish_run
from utils.downsample import downsample
sparkline = downsample(self.loss_history, 50)
sparkline = downsample(loss_history, 50)
finish_run(
id = self.current_job_id,
id = run_id,
status = status,
ended_at = datetime.now(timezone.utc).isoformat(),
final_step = self._progress.step,
final_loss = self._progress.loss
if (self._progress.loss is not None and math.isfinite(self._progress.loss))
else None,
duration_seconds = self._progress.elapsed_seconds,
final_step = final_step,
final_loss = final_loss,
duration_seconds = duration,
loss_sparkline = _json.dumps(sparkline),
output_dir = output_dir,
error_message = error_message,
)
self._run_finalized = True
except Exception:
with self._lock:
self._run_finalized = False # unclaim so a later flush can retry
logger.warning("Failed to finalize run in DB (status=%s)", status, exc_info = True)
def _flush_metrics_to_db(self) -> None:
"""Flush buffered metrics to the database and update live progress."""
if not self._metric_buffer or not self.current_job_id or not self._db_run_created:
return
# Cap buffer to bound memory growth.
if len(self._metric_buffer) > 500:
logger.warning(
"Metric buffer exceeded 500 entries (%d) — trimming oldest",
len(self._metric_buffer),
)
self._metric_buffer = self._metric_buffer[-500:]
# Snapshot before insert so metrics arriving during the write survive.
batch = list(self._metric_buffer)
def _flush_metrics_to_db(self, run_id: Optional[str] = None) -> None:
"""Flush buffered metrics to the DB and update live progress. The target run id,
metric batch, and progress snapshot are all taken under the lock, so a concurrent
flush can't double-remove metrics and a racing new run can't redirect the write to
a different job. A finalizer passes ``run_id`` to pin the target to its captured run."""
with self._lock:
target = run_id if run_id is not None else self.current_job_id
if not self._metric_buffer or not target or not self._db_run_created:
return
# Cap buffer to bound memory growth.
if len(self._metric_buffer) > 500:
logger.warning(
"Metric buffer exceeded 500 entries (%d) — trimming oldest",
len(self._metric_buffer),
)
del self._metric_buffer[:-500]
# Claim the batch under the lock so a concurrent flush can't re-remove it.
batch = list(self._metric_buffer)
del self._metric_buffer[: len(batch)]
step = self._progress.step
loss = self._progress.loss
if loss is not None and not math.isfinite(loss):
loss = None
duration = self._progress.elapsed_seconds
try:
from storage.studio_db import insert_metrics_batch, update_run_progress
insert_metrics_batch(self.current_job_id, batch)
del self._metric_buffer[: len(batch)]
update_run_progress(
id = self.current_job_id,
step = self._progress.step,
loss = self._progress.loss
if (self._progress.loss is not None and math.isfinite(self._progress.loss))
else None,
duration_seconds = self._progress.elapsed_seconds,
)
insert_metrics_batch(target, batch)
update_run_progress(id = target, step = step, loss = loss, duration_seconds = duration)
except Exception:
# Leave buffer intact for retry on next flush
# Re-queue the claimed batch at the front so it retries on the next flush.
with self._lock:
self._metric_buffer[:0] = batch
logger.warning("Failed to flush metrics to DB", exc_info = True)
@staticmethod

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

@ -1731,6 +1731,7 @@ def _run_mlx_training(event_queue, stop_queue, config):
# sharegpt+images) and text (alpaca/sharegpt/chatml → "text" column).
format_type = config.get("format_type", "")
custom_format_mapping = config.get("custom_format_mapping")
dataset_final_format = ""
try:
from utils.datasets import format_and_template_dataset
def _fmt_progress(status_message = "", **_kw):
@ -1796,6 +1797,7 @@ def _run_mlx_training(event_queue, stop_queue, config):
)
if info.get("success", True):
dataset = info.get("dataset", dataset)
dataset_final_format = str(info.get("final_format", "") or "").lower()
if eval_dataset is not None:
ev = format_and_template_dataset(
eval_dataset,
@ -1894,6 +1896,9 @@ def _run_mlx_training(event_queue, stop_queue, config):
eval_steps = eval_steps_val,
)
# Also gates the masking skip below, so defined outside the feature-detect block.
raw_text_mode = training_type == "Continued Pretraining" or format_type == "raw"
# Feature-detect optional fields so this PR works without the paired zoo bump.
_supported_fields = getattr(MLXTrainingConfig, "__dataclass_fields__", {})
if "cast_norm_output_to_input_dtype" in _supported_fields:
@ -1907,7 +1912,6 @@ def _run_mlx_training(event_queue, stop_queue, config):
if "max_grad_leaf_norm" in _supported_fields:
mlx_config_kwargs["max_grad_leaf_norm"] = max_grad_leaf_norm
if "append_eos" in _supported_fields:
raw_text_mode = training_type == "Continued Pretraining" or format_type == "raw"
# Studio SFT formatting owns rendered examples; raw/CPT text still
# needs MLX to append EOS like the CUDA raw-text path.
mlx_config_kwargs["append_eos"] = bool(raw_text_mode)
@ -1928,29 +1932,27 @@ def _run_mlx_training(event_queue, stop_queue, config):
_send("eval_configured")
# ── 7. Apply train_on_responses_only if requested ──
if config.get("train_on_completions", False):
# Auto-detect markers from the chat template first, manual table as
# fallback. Mirror the CUDA skips: raw/CPT text has no chat turns and
# Alpaca-rendered text lacks the chat markers. Also check the resolved
# format, since format_type="auto" can land on alpaca or raw text.
if (
config.get("train_on_completions", False)
and not raw_text_mode
and format_type != "alpaca"
and dataset_final_format not in ("alpaca", "raw_text")
):
_send("status", status_message = "Configuring response-only training...")
try:
from utils.datasets import (
MODEL_TO_TEMPLATE_MAPPER,
TEMPLATE_TO_RESPONSES_MAPPER,
)
template_name = MODEL_TO_TEMPLATE_MAPPER.get(model_name.lower())
markers = TEMPLATE_TO_RESPONSES_MAPPER.get(template_name) if template_name else None
if markers:
trainer = train_on_responses_only(
trainer,
instruction_part = markers["instruction"],
response_part = markers["response"],
)
else:
_send(
"status",
status_message = f"train_on_completions skipped (no template for {model_name})",
)
except Exception as e:
_send("status", status_message = f"train_on_completions failed: {e}")
# No catch: the helper handles detection failures and double misses, so
# an exception here is a real masking failure that must fail the run,
# not silently train on full sequences.
from utils.datasets.completion_masking import apply_completion_masking
trainer, _masking_applied = apply_completion_masking(
trainer,
model_name,
train_on_responses_only,
notify = lambda level, message: _send("status", status_message = message),
)
# ── 8. Setup wandb / tensorboard ──
wandb_run = None
@ -2188,7 +2190,11 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) ->
stop_queue: mp.Queue for stop commands from the parent.
config: Training config dict with all parameters.
"""
os.environ["TOKENIZERS_PARALLELISM"] = "false"
# Off on Linux (forked datasets map() workers deadlock otherwise); on spawn
# platforms map() is in-process, so keep tokenizer threads on for faster prep.
os.environ["TOKENIZERS_PARALLELISM"] = (
"true" if sys.platform in ("win32", "darwin") else "false"
)
os.environ["PYTHONWARNINGS"] = "ignore" # before imports
# HTTP-fallback respawn: disable Xet before any huggingface_hub import (the
@ -2577,6 +2583,15 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) ->
_bnb_rocm_ver,
)
# Setting BNB_ROCM_VERSION makes bitsandbytes log a benign override
# notice on import; drop only that record so real errors and mismatch
# warnings still show.
if os.environ.get("BNB_ROCM_VERSION"):
import logging as _logging
_logging.getLogger("bitsandbytes.cextension").addFilter(
lambda _r: "environment variable detected" not in _r.getMessage()
)
# Parse HIP version for the kernel-fix gate below, falling back to
# the rocm version embedded in torch.__version__ when version.hip is
# unset (AMD SDK / Radeon wheels).
@ -3008,11 +3023,24 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) ->
),
xet_disabled = os.environ.get("HF_HUB_DISABLE_XET") == "1",
)
# Latest-sidecar models load 16-bit here too: bnb 4-bit feeds quantized
# expert weights into unvalidated paths (same flip as the chat worker).
_train_load_in_4bit = config["load_in_4bit"]
if _train_load_in_4bit:
from utils.transformers_version import latest_tier_active_for
if latest_tier_active_for(model_name, hf_token):
_train_load_in_4bit = False
logger.info(
"Latest-transformers sidecar active for %s - forcing a 16-bit "
"training load (4-bit is disabled for brand-new architectures)",
model_name,
)
try:
success = trainer.load_model(
model_name = model_name,
max_seq_length = config["max_seq_length"],
load_in_4bit = config["load_in_4bit"],
load_in_4bit = _train_load_in_4bit,
full_finetuning = not use_lora,
hf_token = hf_token,
is_dataset_image = config.get("is_dataset_image", False),

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

@ -8,6 +8,7 @@ import os
import signal
import subprocess
import sys
import time
import threading
from pathlib import Path
from typing import Callable, Optional
@ -75,6 +76,10 @@ def spawn_worker(
env["HF_HUB_DISABLE_PROGRESS_BARS"] = "1"
env["HF_HUB_DISABLE_TELEMETRY"] = "1"
env["HF_HUB_DISABLE_XET"] = "0" if use_xet else "1"
# No token in Studio settings: fall back to the backend's own HF_TOKEN so
# private repos stay downloadable (needed while inkling repos are private).
if not hf_token:
hf_token = os.environ.get("HF_TOKEN") or None
env["HF_HUB_DISABLE_IMPLICIT_TOKEN"] = "0" if hf_token else "1"
# hf_transfer's parallel Range chunks can leave sparse partials even in
# "http" mode; disable so the worker's writer is always sequential.
@ -206,7 +211,9 @@ def finalize_worker_exit(
repo_type: Optional[RepoType] = None,
repo_id: Optional[str] = None,
transport: Optional[str] = None,
) -> None:
cancel_marker_transport: Optional[str] = None,
defer_error: bool = False,
) -> str:
"""Block until *proc* exits, then record the job's terminal state in
*registry*. Drains and scrubs stderr first, then classifies the exit code.
A no-op when the process was already dropped (e.g. superseded).
@ -218,7 +225,7 @@ def finalize_worker_exit(
rc = proc.wait()
cancel_requested = registry.cancel_requested(key)
if not registry.drop_process(key, proc):
return
return "idle"
stderr_text = download_registry.scrub_secrets(
(stderr_data or b"").decode("utf-8", "replace").strip(),
hf_token = hf_token,
@ -226,6 +233,8 @@ def finalize_worker_exit(
state = classify_exit(rc, cancel_requested = cancel_requested)
if state == "complete":
registry.set_job(key, "complete")
if transport == download_registry.TRANSPORT_HTTP:
registry.update_job_transport(key, download_registry.TRANSPORT_HTTP)
if stderr_text:
if download_manifest.MANIFEST_DEGRADED_MARKER in stderr_text:
logger.warning(
@ -258,18 +267,226 @@ def finalize_worker_exit(
metadata.variant
if metadata is not None and metadata.variant
else download_registry.variant_from_key(key),
transport,
cancel_marker_transport or transport,
logger = logger,
)
else:
registry.set_job(
key,
"error",
stderr_text or f"worker exited with code {rc}",
)
if not defer_error:
registry.set_job(
key,
"error",
stderr_text or f"worker exited with code {rc}",
)
logger.error(
f"{log_prefix} failed for {label} (rc={rc}): {stderr_text}",
)
return state
def _set_retry_failure_state(
registry: download_registry.DownloadRegistry,
key: str,
error: str,
*,
repo_type: RepoType,
repo_id: str,
fallback_variant: Optional[str],
fallback_transport: Optional[str],
logger,
) -> str:
state, metadata = registry.set_error_unless_cancelled(key, error)
if state == "cancelled":
download_registry.persist_cancel_marker(
repo_type,
repo_id,
metadata.variant if metadata is not None and metadata.variant else fallback_variant,
metadata.transport
if metadata is not None and metadata.transport
else fallback_transport,
logger = logger,
)
return state
def _try_http_retry(
registry: download_registry.DownloadRegistry,
key: str,
*,
hf_token: Optional[str],
label: str,
log_prefix: str,
logger,
repo_type: RepoType,
repo_id: str,
watch_name: str,
) -> bool:
"""Reclaim *key* with HTTP transport and spawn a recovery worker.
Returns ``True`` when the HTTP worker was successfully registered.
Caller is responsible for ensuring this is only called when: the job is
in ``"error"`` state, the original transport was XET, and HTTP is available.
Derives variant and blob-hash metadata from the registry entry written by
the original XET claim so callers do not re-construct worker arguments.
Re-queries peer protection hashes at spawn time to reflect any concurrent
sibling changes between the XET failure and this call.
"""
original_metadata = registry.get_job_metadata(key)
if original_metadata is None:
logger.debug("%s XET retry skipped for %s; metadata unavailable", log_prefix, label)
_set_retry_failure_state(
registry,
key,
"XET retry skipped: metadata unavailable",
repo_type = repo_type,
repo_id = repo_id,
fallback_variant = download_registry.variant_from_key(key),
fallback_transport = download_registry.TRANSPORT_XET,
logger = logger,
)
return False
if original_metadata.transport != download_registry.TRANSPORT_XET:
logger.debug(
"%s XET retry skipped for %s; original transport was %s",
log_prefix,
label,
original_metadata.transport,
)
_set_retry_failure_state(
registry,
key,
f"XET retry skipped: original transport was {original_metadata.transport}",
repo_type = repo_type,
repo_id = repo_id,
fallback_variant = original_metadata.variant,
fallback_transport = original_metadata.transport,
logger = logger,
)
return False
variant = original_metadata.variant
blob_hashes = original_metadata.blob_hashes
progress_blob_hashes = original_metadata.progress_blob_hashes
completed_baseline_bytes = (
download_registry.completed_blob_bytes(
repo_type,
repo_id,
progress_blob_hashes,
)
if progress_blob_hashes
else 0
)
generation = registry.current_generation(key)
registry.release_active_slot(key)
while True:
if registry.cancel_requested(key):
_set_retry_failure_state(
registry,
key,
"HTTP retry cancelled before reclaiming the download slot",
repo_type = repo_type,
repo_id = repo_id,
fallback_variant = variant,
fallback_transport = original_metadata.transport,
logger = logger,
)
return False
claimed, conflict_state = registry.claim(
key,
download_registry.TRANSPORT_HTTP,
repo_type = repo_type,
repo_id = repo_id,
variant = variant,
blob_hashes = blob_hashes,
progress_blob_hashes = progress_blob_hashes,
completed_baseline_bytes = completed_baseline_bytes,
generation = generation,
replace_active = True,
cancel_marker_transport = original_metadata.transport,
)
if claimed:
break
if conflict_state == "deleting":
logger.debug(
"%s XET retry claim rejected for %s; repo is being deleted",
log_prefix,
label,
)
_set_retry_failure_state(
registry,
key,
"HTTP retry could not reclaim the download slot",
repo_type = repo_type,
repo_id = repo_id,
fallback_variant = variant,
fallback_transport = original_metadata.transport,
logger = logger,
)
return False
logger.debug(
"%s XET retry claim blocked for %s by active sibling state %s; waiting",
log_prefix,
label,
conflict_state,
)
time.sleep(0.05)
args: list[str] = ["--repo-id", repo_id]
if repo_type == "dataset":
args.append("--dataset")
elif variant:
args.extend(["--variant", variant])
# Re-query at spawn time: sibling state may have changed since XET failed.
peer_hashes = registry.peer_blob_hashes(key) if variant else frozenset()
logger.warning(
"%s XET worker failed for %s; retrying over HTTP",
log_prefix,
label,
)
try:
proc = spawn_worker(
args,
hf_token,
use_xet = False,
protected_blob_hashes = peer_hashes or None,
)
except Exception as exc:
scrubbed = download_registry.scrub_secrets(str(exc), hf_token = hf_token)
logger.error(
"%s HTTP retry spawn failed for %s: %s",
log_prefix,
label,
scrubbed,
)
registry.update_job_transport(key, original_metadata.transport)
_set_retry_failure_state(
registry,
key,
scrubbed,
repo_type = repo_type,
repo_id = repo_id,
fallback_variant = variant,
fallback_transport = original_metadata.transport,
logger = logger,
)
return False
return register_worker(
registry,
key,
proc,
hf_token = hf_token,
label = label,
log_prefix = log_prefix,
logger = logger,
repo_type = repo_type,
repo_id = repo_id,
transport = download_registry.TRANSPORT_HTTP,
cancel_marker_transport = original_metadata.transport,
watch_name = watch_name,
)
def kill_and_reap_process(
@ -305,6 +522,7 @@ def register_worker(
repo_type: RepoType,
repo_id: str,
transport: str,
cancel_marker_transport: Optional[str] = None,
watch_name: str,
) -> bool:
if not registry.register_process(key, proc):
@ -315,7 +533,14 @@ def register_worker(
def _watch() -> None:
try:
finalize_worker_exit(
can_retry_http = (
transport == download_registry.TRANSPORT_XET
and download_registry.download_transport_unavailable_reason(
download_registry.TRANSPORT_HTTP
)
is None
)
state = finalize_worker_exit(
registry,
key,
proc,
@ -326,7 +551,25 @@ def register_worker(
repo_type = repo_type,
repo_id = repo_id,
transport = transport,
cancel_marker_transport = cancel_marker_transport,
defer_error = can_retry_http,
)
# XET-to-HTTP recovery: when a non-cancelled XET worker fails and
# HTTP is available, attempt one automatic retry over HTTP. The
# transport check is the recursion guard: an HTTP worker that errors
# never satisfies `transport == TRANSPORT_XET`, so it stays terminal.
if can_retry_http and state == "error":
_try_http_retry(
registry,
key,
hf_token = worker_token,
label = label,
log_prefix = log_prefix,
logger = logger,
repo_type = repo_type,
repo_id = repo_id,
watch_name = watch_name,
)
except Exception:
# finalize_worker_exit is the only thing that clears running/cancelling;
# if it raises, force a terminal state so claim() isn't blocked until restart.
@ -422,8 +665,19 @@ def cancel_worker(
return "cancelling"
return registry.get_job(key).state
# Worker already exited; let its watcher classify the real return code.
# Arming a pending cancel here could mislabel a genuine failure as a cancel.
if proc.poll() is not None:
get_metadata = getattr(registry, "get_job_metadata", None)
metadata = get_metadata(key) if get_metadata is not None else None
can_retry_http = (
metadata is not None
and metadata.transport == download_registry.TRANSPORT_XET
and download_registry.download_transport_unavailable_reason(
download_registry.TRANSPORT_HTTP
)
is None
)
if can_retry_http and registry.mark_pending_cancel(key, generation):
return "cancelling"
return registry.get_job(key).state
if not registry.request_cancel(key, proc, generation):

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

@ -49,6 +49,10 @@ _REPO_SIZE_NEG_TTL = 60.0
_MODEL_METADATA_TIMEOUT_SECONDS = 5.0
_repo_size_cache_lock = threading.Lock()
# Identity for a cached file with no HF blob (Windows without Developer Mode: hf
# moves the blob into snapshots/ and leaves blobs/ empty).
_LOCAL_SIZE_IDENTITY_PREFIX = "size:"
def get_repo_snapshot_metadata_cached(
repo_id: str, hf_token: Optional[str] = None
@ -135,23 +139,52 @@ def _cached_repo_file_name(file_obj) -> str:
return str(getattr(file_obj, "file_name", "")).replace("\\", "/")
def _is_real_cache_blob(blob: Optional[Path], repo_dir: Optional[Path]) -> bool:
"""True only for a real cache blob at ``<repo_dir>/blobs/<etag>``.
A no-symlink ``snapshots/`` file (name is the filename, not an etag) or a
repo's own ``blobs/`` subdir is not the cache blob store.
"""
if blob is None or repo_dir is None:
return False
try:
return blob.parent.resolve(strict = False) == (repo_dir / "blobs").resolve(strict = False)
except OSError:
return False
def _cached_blob_hash(blob_path, repo_path = None) -> Optional[str]:
"""The cache blob hash (etag) for a cached file, or None when there is no blob.
Only a real blob under the repo's ``blobs/`` dir has name == hash; a moved
no-symlink ``snapshots/`` file is "no blob", so the caller uses a size identity.
"""
path = Path(blob_path)
repo_dir = Path(repo_path) if repo_path is not None else None
return path.name if _is_real_cache_blob(path, repo_dir) else None
def local_size_identity(size: int) -> str:
"""Identity for a cached file whose blob hash is unknowable: its size.
Re-hashing multi-GB GGUFs on the inventory hot path is not viable, and a
``size:`` token never collides with a hex hash.
"""
return f"{_LOCAL_SIZE_IDENTITY_PREFIX}{int(size)}"
def _repo_gguf_blob_map(repo_info, *, include_companions: bool = False) -> dict[str, set[str]]:
"""Map each cached GGUF file's repo-relative name to the SET of its local
blob hashes across all cached revisions.
identities across all revisions.
HF names each local cache blob FILE by the file's etag (lfs.sha256 else
blob_id), so a local file's blob hash == ``Path(blob_path).name``. An updated
repo keeps BOTH the old and new revision snapshots until HF garbage-collects
them, so the same file resolves to several blobs; collecting them ALL (not
just the first one seen, since ``repo_info.revisions`` is a frozenset and
yields them in arbitrary order) lets the remote-vs-local diff treat the file
as current when the remote (``main``) blob is present in any cached revision.
Mirrors the ``cached_blob_ids`` membership test in routes/models.py.
By default this keeps the historical MAIN-GGUF-only behavior. GGUF update
checks opt into companions so a shared mmproj/MTP blob can be compared too.
An identity is the file's blob hash, or a size identity when the cache holds no
blob (Windows without Developer Mode). BOTH old and new revision blobs are kept
(a set), so the diff treats the file as current when the remote ``main`` blob is
in any cached revision. Main GGUF only by default; update checks opt into
companions to compare a shared mmproj/MTP blob too.
"""
blob_map: dict[str, set[str]] = {}
repo_path = getattr(repo_info, "repo_path", None)
for revision in repo_info.revisions:
for f in revision.files:
if include_companions:
@ -163,7 +196,13 @@ def _repo_gguf_blob_map(repo_info, *, include_companions: bool = False) -> dict[
if not blob_path:
continue
name = _cached_repo_file_name(f)
blob_map.setdefault(name, set()).add(Path(blob_path).name)
identity = _cached_blob_hash(blob_path, repo_path)
if identity is None:
size = int(getattr(f, "size_on_disk", 0) or 0)
if size <= 0:
continue
identity = local_size_identity(size)
blob_map.setdefault(name, set()).add(identity)
return blob_map

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

@ -408,7 +408,13 @@ def reclaim_replaced_gguf_variant(
and extract_quant_label(name).lower() == variant_key,
)
for snap, blob, name in matches:
blob_hash = _blob_hash_from_path(blob) if blob is not None else None
# Prune only a file we can identify as a real, stale cache blob. A
# no-symlink snapshot file has no identifiable blob hash, so keep it.
blob_hash = (
_blob_hash_from_path(blob)
if cache_inventory._is_real_cache_blob(blob, repo_dir)
else None
)
if blob_hash is None or blob_hash in keep_main_hashes:
continue
stale_matches.append((snap, blob, name))

30
studio/backend/hub/services/models/downloads.py
View file

@ -60,6 +60,30 @@ def _job_status(
return DownloadJobStatus(state = state, error = error, generation = generation)
def _load_in_flight(repo_id: str) -> bool:
try:
from core.inference.llama_cpp import hf_gguf_load_in_flight
return hf_gguf_load_in_flight(repo_id)
except Exception:
return False
def _load_in_flight_error(repo_id: str) -> HTTPException:
return HTTPException(
status_code = 409,
detail = (
f"A model load for '{repo_id}' is in progress and may be "
"downloading it. Wait for the load to finish (or cancel it), "
"then start the download."
),
)
def _reject_if_load_in_flight(repo_id: str) -> None:
if _load_in_flight(repo_id):
raise _load_in_flight_error(repo_id)
def _spawn_download_worker(
repo_id: str,
variant: Optional[str],
@ -89,6 +113,9 @@ async def download_model_response(body: DownloadModelRequest, hf_token: Optional
# Canonicalize so two different-cased paste-ins share one job + cache dir.
repo_id = await asyncio.to_thread(resolve_cached_repo_id_case, repo_id, repo_type = "model")
# Avoid concurrent writers to the same HF cache files.
_reject_if_load_in_flight(repo_id)
variant = (body.gguf_variant or "").strip() or None
if variant is not None and not _is_valid_gguf_variant(variant):
raise HTTPException(
@ -147,9 +174,12 @@ async def download_model_response(body: DownloadModelRequest, hf_token: Optional
blob_hashes = variant_blob_hashes,
progress_blob_hashes = variant_progress_blob_hashes,
completed_baseline_bytes = completed_baseline_bytes,
admission_check = lambda: not _load_in_flight(repo_id),
)
generation = _registry.current_generation(key)
if not claimed:
if claim_state == "admission_blocked":
raise _load_in_flight_error(repo_id)
# claim_state is the blocking job's state. The client can attach only
# when the blocker is this key's own in-flight job (adoptable); a
# cross-variant conflict or in-progress delete is not accepted.

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

@ -15,6 +15,7 @@ from loggers import get_logger
from hub.schemas.inventory import BrowseEntry, BrowseFoldersResponse
from hub.storage.scan_folders import (
contains_sensitive_path_component,
is_denied_system_path,
list_scan_folders,
)
from hub.utils.paths import (
@ -27,7 +28,10 @@ from hub.utils.paths import (
studio_root,
well_known_model_dirs,
)
from utils.paths.external_media import linux_run_media_mount_roots
from utils.paths.external_media import (
linux_run_media_mount_roots,
windows_drive_roots,
)
from hub.services.models.common import _safe_is_dir
from hub.services.models.local_inventory import _resolve_hf_cache_dir
@ -158,8 +162,14 @@ def _looks_like_model_dir(directory: Path) -> bool:
return False
def _build_browse_allowlist() -> list[Path]:
"""Root directories the browser may walk (also seeds the suggestion chips): HOME, resolved HF cache dirs, Studio outputs/exports/root, registered scan folders, and well-known local-LLM dirs. Each is added only if it resolves to a real directory so the sandbox has no dead boundary."""
def _build_browse_allowlist(
media_roots: Optional[list[Path]] = None, drive_roots: Optional[list[Path]] = None
) -> list[Path]:
"""Root directories the browser may walk (also seeds the suggestion chips): HOME, resolved HF cache dirs, Studio outputs/exports/root, registered scan folders, and well-known local-LLM dirs. Each is added only if it resolves to a real directory so the sandbox has no dead boundary.
*media_roots* / *drive_roots* let the caller pass already-probed
removable-media and Windows drive roots so they aren't scanned again (a
disconnected mapped drive can make each probe slow); probed here when ``None``."""
from hub.storage.scan_folders import list_scan_folders
candidates: list[Path] = []
@ -176,7 +186,13 @@ def _build_browse_allowlist() -> list[Path]:
candidates.append(resolved)
_add(Path.home())
for p in linux_run_media_mount_roots():
if media_roots is None:
media_roots = linux_run_media_mount_roots()
if drive_roots is None:
drive_roots = windows_drive_roots()
for p in media_roots:
_add(p)
for p in drive_roots:
_add(p)
_add(_resolve_hf_cache_dir())
try:
@ -218,7 +234,14 @@ def _build_browse_allowlist() -> list[Path]:
def _is_path_inside_allowlist(target: Path, allowed_roots: list[Path]) -> bool:
"""True if *target* equals or descends from any allowed root; uses ``os.path.realpath`` so symlinks cannot escape the sandbox."""
"""True if *target* equals or descends from any allowed root; uses ``os.path.realpath`` so symlinks cannot escape the sandbox.
A Windows drive root (``D:\\``) authorizes its descendants, but a bare POSIX
root (``/``) must NOT: a single ``/`` allowlist entry (e.g. a legacy scan
folder) would otherwise authorize every absolute path, reaching ``/var``,
``/root``, etc. the denylist does not cover. Mirrors the legacy browser so
both treat ``/`` identically.
"""
try:
target_real = os.path.normcase(os.path.realpath(str(target)))
except OSError:
@ -228,13 +251,25 @@ def _is_path_inside_allowlist(target: Path, allowed_roots: list[Path]) -> bool:
root_real = os.path.normcase(os.path.realpath(str(root)))
except OSError:
continue
if target_real == root_real:
return True
drive, tail = os.path.splitdrive(root_real)
if os.path.dirname(root_real) == root_real and not drive:
# Bare POSIX filesystem root ("/"): equality above is the only
# match; do not let it authorize arbitrary descendants.
continue
if drive.startswith(("\\\\", "//")) and not tail:
# Bare UNC share root (\\server\share): os.path.commonpath raises
# "can't mix absolute and relative" on it, so authorize its
# descendants with a boundary-safe prefix test (normcase applied).
if target_real.startswith(root_real.rstrip("\\/") + os.sep):
return True
continue
try:
if os.path.commonpath([target_real, root_real]) == root_real:
return True
except ValueError:
continue
if target_real == root_real:
return True
return False
@ -347,6 +382,11 @@ def _resolve_browse_target(path: Optional[str], allowed_roots: list[Path]) -> Pa
status_code = 403,
detail = "Credential or configuration directories are not browseable.",
)
if is_denied_system_path(str(resolved_child)):
raise HTTPException(
status_code = 403,
detail = "System directories are not browseable.",
)
current = resolved_child
if contains_sensitive_path_component(str(current)):
@ -354,6 +394,13 @@ def _resolve_browse_target(path: Optional[str], allowed_roots: list[Path]) -> Pa
status_code = 403,
detail = "Credential or configuration directories are not browseable.",
)
# Zero-component case: the requested path IS an allowlist root
# (e.g. a legacy-registered "/" or a Windows drive root).
if is_denied_system_path(str(current)):
raise HTTPException(
status_code = 403,
detail = "System directories are not browseable.",
)
if not current.is_dir():
raise HTTPException(
status_code = 400,
@ -382,9 +429,13 @@ def browse_folders_response(
"""
from hub.storage.scan_folders import list_scan_folders
# Probe removable-media and Windows drive roots once; the allowlist and
# chips reuse the result so a disconnected mapped drive isn't scanned twice.
media_roots = linux_run_media_mount_roots()
drive_roots = windows_drive_roots()
# Build the allowlist once -- the sandbox check and suggestion chips share
# it so chips are always navigable.
allowed_roots = _build_browse_allowlist()
allowed_roots = _build_browse_allowlist(media_roots, drive_roots)
try:
target = _resolve_browse_target(path, allowed_roots)
@ -440,6 +491,15 @@ def browse_folders_response(
# descending into them is refused and registration rejects them.
if contains_sensitive_path_component(name):
continue
# Same for denied system dirs (C:\Windows, /etc, ...): descent 403s,
# so don't render them as clickable rows. Resolve first so a
# symlink/junction into a denied dir is hidden too, not just a literal name.
try:
resolved_child = os.path.realpath(str(child))
except (OSError, ValueError):
resolved_child = str(child)
if is_denied_system_path(resolved_child):
continue
entries.append(
BrowseEntry(
name = name,
@ -487,13 +547,22 @@ def browse_folders_response(
return
if resolved in seen_sug:
return
# Drop a denied system dir (e.g. a stale scan-folder row) so it never
# becomes a chip that 403s on click. Drive roots stay: only their
# system subdirectories are denied, not the root itself.
if is_denied_system_path(resolved):
return
if _safe_is_dir(resolved):
seen_sug.add(resolved)
suggestions.append(resolved)
# Home first as the safe fallback.
_add_sug(Path.home())
for p in linux_run_media_mount_roots():
# Reuse the roots probed for the allowlist above (no second drive scan).
for p in media_roots:
_add_sug(p)
# Windows drive roots so the user can hop between C:, D:, E: ...
for p in drive_roots:
_add_sug(p)
# The HF cache root in use (honors HF_HOME / HF_HUB_CACHE), then the default.
try:

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

@ -337,6 +337,22 @@ def _local_main_gguf_blobs_by_quant(repo_id: str) -> dict[str, dict[str, set[str
return result
def _size_identity_matches(local_set: set[str], remote_size: int) -> bool:
"""Whether a cached file with NO blob hash is current, judged by size.
A size token only lands in ``local_set`` for a file the cache has no blob for,
so it never loosens the hash comparison for a normal file. Tradeoff: an
equal-size requant is missed, versus the status quo where every no-blob GGUF
shows a phantom update that no re-download clears.
"""
size = int(remote_size or 0)
if size <= 0:
return False
from hub.services.models import cache_inventory
return cache_inventory.local_size_identity(size) in local_set
def _variant_update_available_from_requirement(
local_blobs: dict[str, set[str]], requirement: Optional[_GgufVariantRequirement], variant: str
) -> bool:
@ -355,8 +371,13 @@ def _variant_update_available_from_requirement(
if not remote_blob:
continue
local_set = local_by_posix.get(path)
if not local_set or remote_blob not in local_set:
if not local_set:
return True
if remote_blob in local_set:
continue
if _size_identity_matches(local_set, expected.size):
continue
return True
return False

25
studio/backend/hub/services/snapshot_progress.py
View file

@ -12,6 +12,7 @@ summing stale blobs against the wrong total)."""
from __future__ import annotations
import asyncio
import threading
from pathlib import Path
from typing import Callable, Optional
@ -34,6 +35,28 @@ logger = get_logger(__name__)
# (repo_id, hf_token) -> (expected_total_bytes, expected_blob_hashes)
SnapshotMetadataResolver = Callable[[str, Optional[str]], "tuple[int, frozenset[str]]"]
# One progress log per 10% step per job, so an active download reports progress
# without emitting a line on every poll.
_progress_step_lock = threading.Lock()
_last_progress_step: dict[str, int] = {}
def _log_progress_step(job_key: str, repo_id: str, variant: Optional[str], progress: float) -> None:
step = int(progress * 10)
with _progress_step_lock:
last = _last_progress_step.get(job_key, -1)
if step == last:
return
_last_progress_step[job_key] = step
if step < last:
return # download restarted; resync without logging
logger.info(
"hub_download_progress",
repo_id = repo_id,
variant = variant or "",
percent = step * 10,
)
def _empty_progress(expected_bytes: int) -> dict:
return {
@ -215,6 +238,8 @@ def compute_snapshot_progress(
else 0
)
)
if force_active:
_log_progress_step(job_key, repo_id, variant, progress)
return {
"downloaded_bytes": display_downloaded_bytes,
"completed_bytes": display_completed_bytes,

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

@ -16,7 +16,7 @@ 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.external_media import is_linux_run_media_path, is_local_filesystem_root
from utils.paths.sensitive import (
contains_sensitive_path_component as _shared_contains_sensitive_path_component,
)
@ -52,6 +52,25 @@ def _denied_path_prefixes() -> list[str]:
return []
def is_denied_system_path(path: str) -> bool:
"""True if *path* is, or descends from, a denied system directory.
Mirrors the denylist add_scan_folder() enforces at registration so the
browser refuses /etc, /proc, C:\\Windows, etc. even when the allowlist holds
a broad root (a Windows drive root C:\\ or a legacy-registered / root). The
/run carve-out keeps Linux removable-media mounts browseable. Expects an
already-resolved (realpath) path so symlinks cannot escape into a denied subtree.
"""
is_win = platform.system() == "Windows"
check = os.path.normcase(path) if is_win else path
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
return True
return False
def _contains_sensitive_path_component(path: str) -> bool:
return _shared_contains_sensitive_path_component(path)
@ -108,8 +127,9 @@ 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 os.path.dirname(normalized) == normalized:
# Registering a filesystem root would expose denied system dirs via browse.
if is_local_filesystem_root(normalized):
# A local fs root ("/", "C:\\") would expose denied system dirs via browse;
# a UNC share root (\\server\share) has none under it and stays registerable.
raise ValueError("The filesystem root cannot be registered")
if _contains_sensitive_path_component(normalized):
raise ValueError("Credential or configuration directories are not allowed")

158
studio/backend/hub/tests/test_download_lifecycle.py
View file

@ -1,27 +1,147 @@
# 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 io
import logging
from hub.services import download_lifecycle
from hub.utils import download_registry, state_dir
def _set_xet_reason(monkeypatch, reason):
class _Proc:
pid = 4242
def __init__(
self,
rc,
stderr = b"",
):
self.rc = rc
self.stderr = io.BytesIO(stderr)
self.waited = False
def poll(self):
return self.rc if self.waited else None
def wait(self, timeout = None):
self.waited = True
return self.rc
def kill(self):
pass
class _ImmediateThread:
def __init__(self, *, target, **_kwargs):
self.target = target
def start(self):
self.target()
def test_resolve_effective_use_xet(monkeypatch):
for requested, unavailable_reason, expected in (
(False, "unused", False),
(True, None, True),
(True, "hf_xet is not installed", False),
):
monkeypatch.setattr(
download_lifecycle.download_registry,
"download_transport_unavailable_reason",
lambda _transport, reason = unavailable_reason: reason,
)
assert download_lifecycle.resolve_effective_use_xet(requested) is expected
def test_xet_failure_retries_over_http_for_model_and_dataset(monkeypatch, tmp_path):
monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path / "state")
monkeypatch.setattr(download_lifecycle.threading, "Thread", _ImmediateThread)
register_worker = download_lifecycle.register_worker
for repo_type, repo_id, variant, expected_args in (
("model", "Org/Model", "Q4_K_M", ["--repo-id", "Org/Model", "--variant", "Q4_K_M"]),
("dataset", "Org/Data", None, ["--repo-id", "Org/Data", "--dataset"]),
):
registry = download_registry.DownloadRegistry()
key = download_registry.normalize_job_key(f"{repo_id}::{variant}" if variant else repo_id)
assert registry.claim(
key,
download_registry.TRANSPORT_XET,
repo_type = repo_type,
repo_id = repo_id,
variant = variant,
blob_hashes = frozenset({"blob"}),
)[0]
generation = registry.current_generation(key)
spawned = []
def fake_spawn(
args,
_token,
*,
use_xet,
protected_blob_hashes = None,
):
spawned.append((args, use_xet, protected_blob_hashes))
return _Proc(0)
def fake_retry_register(*_args, **kwargs):
assert kwargs["transport"] == download_registry.TRANSPORT_HTTP
return True
monkeypatch.setattr(download_lifecycle, "spawn_worker", fake_spawn)
monkeypatch.setattr(download_lifecycle, "register_worker", fake_retry_register)
assert register_worker(
registry,
key,
_Proc(1, b"xet failed"),
hf_token = None,
label = repo_id,
log_prefix = "Download",
logger = logging.getLogger("test"),
repo_type = repo_type,
repo_id = repo_id,
transport = download_registry.TRANSPORT_XET,
watch_name = f"{repo_type}-watch",
)
metadata = registry.get_job_metadata(key)
assert spawned == [(expected_args, False, None)]
assert metadata.transport == download_registry.TRANSPORT_HTTP
assert metadata.blob_hashes == frozenset({"blob"})
assert registry.current_generation(key) == generation
def test_http_failure_remains_terminal(monkeypatch, tmp_path):
monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path / "state")
monkeypatch.setattr(download_lifecycle.threading, "Thread", _ImmediateThread)
register_worker = download_lifecycle.register_worker
registry = download_registry.DownloadRegistry()
key = download_registry.normalize_repo_key("Org/Data")
assert registry.claim(
key,
download_registry.TRANSPORT_HTTP,
repo_type = "dataset",
repo_id = "Org/Data",
)[0]
monkeypatch.setattr(
download_lifecycle.download_registry,
"download_transport_unavailable_reason",
lambda _transport: reason,
download_lifecycle,
"register_worker",
lambda *_args, **_kwargs: (_ for _ in ()).throw(
AssertionError("HTTP failures must not retry")
),
)
def test_resolve_effective_use_xet_keeps_http_when_not_requested(monkeypatch):
_set_xet_reason(monkeypatch, "should not be consulted")
assert download_lifecycle.resolve_effective_use_xet(False) is False
def test_resolve_effective_use_xet_keeps_xet_when_available(monkeypatch):
_set_xet_reason(monkeypatch, None)
assert download_lifecycle.resolve_effective_use_xet(True) is True
def test_resolve_effective_use_xet_downgrades_when_xet_unavailable(monkeypatch):
_set_xet_reason(monkeypatch, "Xet transport is unavailable because hf_xet is not installed.")
assert download_lifecycle.resolve_effective_use_xet(True) is False
assert register_worker(
registry,
key,
_Proc(1, b"http failed"),
hf_token = None,
label = "Org/Data",
log_prefix = "Download",
logger = logging.getLogger("test"),
repo_type = "dataset",
repo_id = "Org/Data",
transport = download_registry.TRANSPORT_HTTP,
watch_name = "dataset-watch",
)
assert registry.get_job(key).state == "error"

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

@ -36,6 +36,20 @@ from hub.utils import (
from hub.workers import hf_download
@pytest.fixture(autouse = True)
def _denylist_inert(monkeypatch):
# The browse tests here exercise allowlist containment, symlink safety and
# the sensitive-name filter, not the system-directory denylist (which has
# its own suite in tests/test_browse_denylist.py). On macOS tmp_path
# resolves under /private/var, a denied prefix, so _resolve_browse_target
# would 403 the fixture dirs before that logic runs. Keep the denylist inert
# so these assertions hold on every platform. folder_browser binds
# is_denied_system_path at import, so patch it on that module, not on
# scan_folders. The "rejects" cases still 403 via the allowlist/sensitive
# checks, and the non-browse tests never call it.
monkeypatch.setattr(folder_browser, "is_denied_system_path", lambda _p: False)
def _repo(repo_id: str, files: list[SimpleNamespace], repo_path: Path):
return SimpleNamespace(
repo_id = repo_id,
@ -228,7 +242,8 @@ def test_browse_folders_hides_sensitive_dirs(monkeypatch, tmp_path):
home = tmp_path / "home"
(home / ".ssh").mkdir(parents = True)
(home / "models").mkdir()
monkeypatch.setattr(folder_browser, "_build_browse_allowlist", lambda: [home])
# Accept and ignore the optional (media_roots, drive_roots) args the caller now passes.
monkeypatch.setattr(folder_browser, "_build_browse_allowlist", lambda *_a, **_k: [home])
response = folder_browser.browse_folders_response(str(home), show_hidden = True)

200
studio/backend/hub/utils/download_registry.py
View file

@ -45,9 +45,9 @@ import sys
import threading
import time
import weakref
from dataclasses import dataclass, field
from dataclasses import dataclass, field, replace
from pathlib import Path
from typing import Iterator, Literal, Optional
from typing import Callable, Iterator, Literal, Optional
from loggers import get_logger
@ -126,6 +126,9 @@ def write_worker_breadcrumb(key: str, pid: int, metadata: Optional["DownloadMeta
"repo_id": metadata.repo_id if metadata is not None else None,
"variant": metadata.variant if metadata is not None else None,
"transport": metadata.transport if metadata is not None else None,
"cancel_marker_transport": metadata.cancel_marker_transport
if metadata is not None
else None,
}
tmp = path.with_name(f".{path.name}.tmp-{pid}")
try:
@ -305,7 +308,7 @@ def reap_orphan_workers() -> None:
data.get("repo_type"),
repo_id,
data.get("variant"),
data.get("transport"),
data.get("cancel_marker_transport") or data.get("transport"),
)
except Exception as exc:
logger.debug("Reaper failed for breadcrumb %s: %s", entry, exc)
@ -699,6 +702,7 @@ class DownloadMetadata:
repo_id: str
variant: Optional[str]
transport: Optional[str]
cancel_marker_transport: Optional[str] = None
# GGUF variant main/writable hashes, identifying the variant-specific shards
# for concurrency decisions.
blob_hashes: frozenset[str] = field(default_factory = frozenset)
@ -801,6 +805,7 @@ class DownloadRegistry:
self._processes: dict[str, subprocess.Popen] = {}
self._repo_active: dict[str, set[str]] = {}
self._metadata: dict[str, DownloadMetadata] = {}
self._cancel_marker_transports: dict[str, str] = {}
self._pending_cancel: dict[str, Optional[int]] = {}
self._generations: dict[str, int] = {}
# Monotonic across keys so an evicted then re-claimed key never reuses a
@ -839,6 +844,7 @@ class DownloadRegistry:
if state in TERMINAL_STATES:
self._put_terminal_job_locked(key, state, error)
self._pending_cancel.pop(key, None)
self._cancel_marker_transports.pop(key, None)
repo = _repo_of_key(key)
active = self._repo_active.get(repo)
if active is not None:
@ -848,6 +854,57 @@ class DownloadRegistry:
else:
self._jobs[key] = DownloadState(state, error)
def set_error_unless_cancelled(
self, key: str, error: str
) -> tuple[JobState, Optional[DownloadMetadata]]:
key = normalize_job_key(key)
with self._lock:
current = self._jobs.get(key, DownloadState("idle")).state
has_pending_cancel = key in self._pending_cancel
pending_generation = self._pending_cancel.get(key)
metadata = self._metadata.get(key)
should_cancel = current == "cancelling" or (
has_pending_cancel and self._generation_matches_locked(key, pending_generation)
)
terminal_state: JobState = "cancelled" if should_cancel else "error"
marker_transport = self._cancel_marker_transports.pop(key, None)
if marker_transport is None and metadata is not None:
marker_transport = metadata.cancel_marker_transport
self._put_terminal_job_locked(
key,
terminal_state,
None if should_cancel else error,
)
self._pending_cancel.pop(key, None)
repo = _repo_of_key(key)
active = self._repo_active.get(repo)
if active is not None:
active.discard(key)
if not active:
self._repo_active.pop(repo, None)
if should_cancel and metadata is not None and marker_transport is not None:
metadata = replace(metadata, transport = marker_transport)
return terminal_state, metadata
def update_job_transport(self, key: str, transport: str) -> None:
key = normalize_job_key(key)
with self._lock:
metadata = self._metadata.get(key)
if metadata is None or metadata.transport == transport:
return
self._metadata[key] = replace(metadata, transport = transport)
def release_active_slot(self, key: str) -> None:
key = normalize_job_key(key)
repo = _repo_of_key(key)
with self._lock:
active = self._repo_active.get(repo)
if active is None:
return
active.discard(key)
if not active:
self._repo_active.pop(repo, None)
def get_job(self, key: str) -> DownloadState:
key = normalize_job_key(key)
with self._lock:
@ -884,6 +941,14 @@ class DownloadRegistry:
):
self._put_terminal_job_locked(key, "cancelled")
metadata_to_persist = self._metadata.pop(key, None)
marker_transport = self._cancel_marker_transports.pop(key, None)
if marker_transport is None and metadata_to_persist is not None:
marker_transport = metadata_to_persist.cancel_marker_transport
if metadata_to_persist is not None and marker_transport is not None:
metadata_to_persist = replace(
metadata_to_persist,
transport = marker_transport,
)
repo = _repo_of_key(key)
active = self._repo_active.get(repo)
if active is not None:
@ -963,12 +1028,24 @@ class DownloadRegistry:
blob_hashes: Optional[frozenset[str]] = None,
progress_blob_hashes: Optional[frozenset[str]] = None,
completed_baseline_bytes: int = 0,
admission_check: Optional[Callable[[], bool]] = None,
generation: Optional[int] = None,
replace_active: bool = False,
metadata_transport: Optional[str] = None,
cancel_marker_transport: Optional[str] = None,
) -> tuple[bool, str]:
key = normalize_job_key(key)
repo = _repo_of_key(key)
requested_hashes = blob_hashes or frozenset()
requested_progress_hashes = progress_blob_hashes or frozenset()
with self._lock:
# Run the final external admission check while the registry lock is
# held, immediately before inspecting and publishing active state.
# The GGUF load path establishes its marker before calling
# its active-job probe, so either this claim observes that marker
# or the load's later probe observes this claim.
if admission_check is not None and not admission_check():
return False, "admission_blocked"
deleting_scopes = self._deleting.get(repo)
if deleting_scopes is not None and (
None in deleting_scopes or variant_from_key(key) in deleting_scopes
@ -1007,10 +1084,13 @@ class DownloadRegistry:
if conflict_state is not None:
return False, conflict_state
current = self._jobs.get(key, DownloadState("idle")).state
if current in _ACTIVE_STATES:
if current in _ACTIVE_STATES and not replace_active:
return False, current
self._generation_seq += 1
self._generations[key] = self._generation_seq
if generation is None:
self._generation_seq += 1
self._generations[key] = self._generation_seq
else:
self._generations[key] = generation
self._jobs[key] = DownloadState("running")
self._repo_active.setdefault(repo, active).add(key)
if repo_type and repo_id:
@ -1018,7 +1098,8 @@ class DownloadRegistry:
repo_type = repo_type,
repo_id = repo_id,
variant = variant,
transport = transport,
transport = metadata_transport if metadata_transport is not None else transport,
cancel_marker_transport = cancel_marker_transport,
blob_hashes = requested_hashes,
progress_blob_hashes = requested_progress_hashes,
completed_baseline_bytes = max(
@ -1026,8 +1107,13 @@ class DownloadRegistry:
int(completed_baseline_bytes or 0),
),
)
if cancel_marker_transport is not None:
self._cancel_marker_transports[key] = cancel_marker_transport
else:
self._cancel_marker_transports.pop(key, None)
else:
self._metadata.pop(key, None)
self._cancel_marker_transports.pop(key, None)
return True, "running"
def adoptable(self, key: str) -> bool:
@ -1053,7 +1139,8 @@ class DownloadRegistry:
download. A variant delete conflicts only with that same variant or a
whole-repo download writing the shared snapshot; other quantizations
download concurrently and never block it."""
for key in self._repo_active.get(repo_id, set()):
active_keys = self._repo_active.get(repo_id, set())
for key in active_keys:
job = self._jobs.get(key)
if job is None or job.state not in _ACTIVE_STATES:
continue
@ -1062,6 +1149,16 @@ class DownloadRegistry:
other_variant = self._active_job_variant_locked(key)
if other_variant is None or other_variant == variant:
return True
for key, job in self._jobs.items():
if key in active_keys or _repo_of_key(key) != repo_id:
continue
if job.state not in _ACTIVE_STATES:
continue
if variant is None:
return True
other_variant = self._active_job_variant_locked(key)
if other_variant is None or other_variant == variant:
return True
return False
def peer_blob_hashes(self, key: str) -> frozenset[str]:
@ -1108,6 +1205,16 @@ class DownloadRegistry:
candidate_keys = list(self._repo_active.get(repo_key, set()))
else:
candidate_keys = [key for active in self._repo_active.values() for key in active]
# An XET->HTTP retry handoff briefly drops its key from _repo_active
# while its job stays active; include those released-but-active jobs
# so the waiting retry still lists and can be adopted or cancelled.
seen = set(candidate_keys)
for key, job in self._jobs.items():
if key in seen or job.state not in _ACTIVE_STATES:
continue
if repo_key is not None and _repo_of_key(key) != repo_key:
continue
candidate_keys.append(key)
refs: list[ActiveDownloadRef] = []
for key in candidate_keys:
job = self._jobs.get(key)
@ -1123,6 +1230,23 @@ class DownloadRegistry:
)
return refs
def has_active_variant(self, repo_id: str, variant: Optional[str]) -> bool:
"""Whether an active model job targets this exact GGUF variant.
Scans the job table rather than only ``_repo_active`` so an XET-to-HTTP
retry handoff remains visible while it has temporarily released its
active slot.
"""
repo_key = normalize_repo_key(repo_id)
target = (variant or "").strip().lower() or None
with self._lock:
for key, job in self._jobs.items():
if _repo_of_key(key) != repo_key or job.state not in _ACTIVE_STATES:
continue
if self._active_job_variant_locked(key) == target:
return True
return False
def begin_delete(
self,
repo_id: str,
@ -1169,12 +1293,25 @@ class DownloadRegistry:
repo_id = normalize_repo_key(repo_id)
target = (variant or "").strip().lower() or None
with self._lock:
for key in self._repo_active.get(repo_id, set()):
active_keys = self._repo_active.get(repo_id, set())
for key in active_keys:
job = self._jobs.get(key)
if job is None or job.state not in _ACTIVE_STATES:
continue
if self._active_job_variant_locked(key) != target:
return True
# An XET->HTTP retry peer between release_active_slot() and its reclaim
# is briefly absent from _repo_active while its job stays active and
# still owns the shared companion; mirror the released-but-active scan
# used by _delete_blocked_by_active_locked so it still blocks companion
# deletion of a different variant.
for key, job in self._jobs.items():
if key in active_keys or _repo_of_key(key) != repo_id:
continue
if job.state not in _ACTIVE_STATES:
continue
if self._active_job_variant_locked(key) != target:
return True
return False
def request_cancel(
@ -1198,17 +1335,58 @@ class DownloadRegistry:
return True
def terminate_all(self, kind: str = "download") -> None:
settled_no_proc: list[Optional[DownloadMetadata]] = []
with self._lock:
live = [
(key, proc, self._metadata.get(key))
for key, proc in self._processes.items()
if proc.poll() is None
]
live_keys = {key for key, _proc, _metadata in live}
# Flag as an intentional stop so the watcher's exit classification
# reports them cancelled rather than an OOM/crash once SIGKILL lands.
for key, _proc, _metadata in live:
if self._jobs.get(key, DownloadState("idle")).state == "running":
self._jobs[key] = DownloadState("cancelling")
# Settle active jobs without a live worker too. Two cases: an
# XET->HTTP retry parked in the reclaim wait loop has dropped its
# worker and slot guard, so it is absent from `live`; and a
# registered worker that already exited with an error but whose
# watcher has not yet run would otherwise stay `running` and spawn an
# HTTP retry after this shutdown snapshot. Skip a registered worker
# that exited cleanly (rc == 0): it completed and the watcher will
# mark it done, so marking it cancelling would strand a stale marker.
for key, job in list(self._jobs.items()):
if job.state not in _ACTIVE_STATES or key in live_keys:
continue
proc = self._processes.get(key)
if proc is not None:
if proc.poll() == 0:
continue
# A registered worker that exited nonzero on its own over HTTP
# is a genuine terminal download failure, not a shutdown cancel
# and not retry-capable: leave its error status intact rather
# than persisting a cancel marker that would read as
# cancelled/resumable after restart. Only an exited XET worker
# could still spawn a post-shutdown HTTP retry, so only that
# needs settling here.
metadata = self._metadata.get(key)
if metadata is not None and metadata.transport == TRANSPORT_HTTP:
continue
self._pending_cancel[key] = self._generations.get(key)
self._jobs[key] = DownloadState("cancelling")
settled_no_proc.append(self._metadata.get(key))
# Persist a cancel marker for each settled no-live-worker job outside the
# lock (mirroring the reaped path) so shutdown records resumable/cancelled
# state even if it returns before the daemon watcher wakes to do so.
for metadata in settled_no_proc:
if metadata is not None:
persist_cancel_marker(
metadata.repo_type,
metadata.repo_id,
metadata.variant,
metadata.cancel_marker_transport or metadata.transport,
)
reaped: list[tuple[str, subprocess.Popen, Optional[DownloadMetadata]]] = []
for key, proc, metadata in live:
try:
@ -1222,7 +1400,7 @@ class DownloadRegistry:
metadata.repo_type,
metadata.repo_id,
metadata.variant,
metadata.transport,
metadata.cancel_marker_transport or metadata.transport,
)
continue
reaped.append((key, proc, metadata))
@ -1242,7 +1420,7 @@ class DownloadRegistry:
metadata.repo_type,
metadata.repo_id,
metadata.variant,
metadata.transport,
metadata.cancel_marker_transport or metadata.transport,
)

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

@ -102,16 +102,17 @@ def preferred_mmproj_sibling(siblings: Sequence) -> Optional[object]:
def preferred_mtp_sibling(siblings: Sequence) -> Optional[object]:
"""The separate MTP drafter to fetch with every variant: the repo-root
``mtp-*.gguf`` copy unsloth ships for llama.cpp ``-hf`` auto-discovery
(Gemma 4). Same pick as the loader's drafter resolution (``mtp-`` basename
prefix, first in sort order) so download and load resolve the same file;
the higher-precision ``MTP/`` subdir copies are for explicit selection and
are not auto-fetched. None for repos with the head baked into the main
GGUF (Qwen)."""
(Gemma 4). Same pick as the loader's drafter resolution (root-level
``mtp-`` prefix, first in sort order) so download and load resolve the same
file; the higher-precision ``MTP/`` subdir copies are for explicit
selection and are not auto-fetched. None for repos with the head baked into
the main GGUF (Qwen)."""
# Root-level only: the MTP/ subdir copies now share the mtp- prefix too.
candidates = sorted(
(
s
for s in siblings
if (name := _gguf_rfilename(s)) and name.lower().rsplit("/", 1)[-1].startswith("mtp-")
if (name := _gguf_rfilename(s)) and "/" not in name and name.lower().startswith("mtp-")
),
key = lambda s: getattr(s, "rfilename"),
)

18
studio/backend/loggers/config.py
View file

@ -17,6 +17,15 @@ import structlog
from loggers.handlers import filter_sensitive_data
class _DropTorchDtypeDeprecation(logging.Filter):
"""Drop transformers' once-per-run "`torch_dtype` is deprecated" warning_once.
It is emitted via logging (not warnings), so a warnings filter cannot catch it."""
def filter(self, record: logging.LogRecord) -> bool:
msg = record.getMessage()
return not ("torch_dtype" in msg and "deprecated" in msg)
class LogConfig:
"""Structured logging configuration for the application."""
@ -72,4 +81,13 @@ class LogConfig:
cache_logger_on_first_use = True,
)
# Drop transformers' cosmetic "`torch_dtype` is deprecated" warning_once (see filter).
_dtype_filter = _DropTorchDtypeDeprecation()
for _name in (
"transformers.configuration_utils",
"transformers.modeling_utils",
"transformers.pipelines.base",
):
logging.getLogger(_name).addFilter(_dtype_filter)
return structlog.get_logger(service_name)

81
studio/backend/loggers/handlers.py
View file

@ -28,19 +28,26 @@ def _env_int(name: str, default: int) -> int:
return default
# Drop duplicate successful-GET access logs repeated within the window: the SPA
# fans one cache invalidation into many identical list fetches; only the first
# informs. Loading polls, mutations, and errors are unaffected. 0 = log all.
# Collapse identical GET/2xx logs within the window (the SPA fans one invalidation
# into many list fetches). Mutations and errors always log. 0 = off.
_ACCESS_LOG_DEDUP_MS = _env_int("UNSLOTH_STUDIO_ACCESS_LOG_DEDUP_MS", 300)
# Pure-liveness/UI polls whose access line carries no signal beyond "client still
# polling" (state changes are logged by their own modules). Collapsed to a longer
# heartbeat instead of one line per poll; first hit and any error still log. 0 = off.
# Liveness/UI polls whose line means only "still polling"; collapse to a longer
# heartbeat. First hit and errors still log. 0 = off.
_QUIET_POLL_DEDUP_MS = _env_int("UNSLOTH_STUDIO_ACCESS_LOG_POLL_DEDUP_MS", 10000)
_QUIET_POLL_PATHS = {
"/api/health",
"/api/auth/status",
"/api/inference/status",
"/api/inference/monitor",
# List polls the tabs refetch on a timer and on every tab switch.
"/api/train/runs",
"/api/models/checkpoints",
"/api/models/local",
"/api/rag/knowledge-bases",
# Legacy download polls emit no progress events (unlike /api/hub/*), so heartbeat them.
"/api/models/download-progress",
"/api/models/gguf-download-progress",
"/api/datasets/download-progress",
}
_DEDUP_MAP_MAX = 4096
_NATIVE_PATH_LEASE_RE = re.compile(
@ -62,6 +69,46 @@ _EXCLUDED_SUFFIXES = (
".woff2",
".ttf",
)
# GET polls whose 2xx line carries no signal (their progress/phase events and the UI
# do), so drop it entirely; non-2xx still logs. Only /api/hub download polls emit
# events; the legacy /api/models and /api/datasets ones heartbeat via _QUIET_POLL_PATHS.
_QUIET_SUCCESS_PATHS = {
"/api/inference/load-progress",
"/api/llama/update-status",
"/api/export/logs",
"/api/export/status",
"/api/hub/download-status",
"/api/hub/download-progress",
"/api/hub/gguf-download-progress",
"/api/hub/active-downloads",
"/api/hub/transport-status",
"/api/hub/datasets/download-status",
"/api/hub/datasets/download-progress",
"/api/hub/datasets/active-downloads",
"/api/hub/datasets/transport-status",
}
# The token-refresh route. Its first 2xx means the client has obtained a valid
# session, so from then on chat 401s are real failures and must stay visible.
_AUTH_REFRESH_PATH = "/api/auth/refresh"
# High-frequency chat list polls; their 2xx is covered by generation/tool-call/stats
# events. Exact paths only, so detail/message reads (/threads/{id}, .../messages,
# /projects/{id}) keep their logs. The pre-auth 401 race also fires on these polls.
_CHAT_LIST_PATHS = {
"/api/chat/threads",
"/api/chat/projects",
}
def _is_quiet_success(method: str, path: str, status_code: int, pre_auth: bool) -> bool:
"""GET-only. Suppress a 2xx poll line that carries no signal, plus a chat list
poll's transient pre-auth 401 (only in the bootstrap window before the first
successful token refresh). Mutations, real (post-refresh) auth failures, and
all other errors always log."""
if method != "GET":
return False
if 200 <= status_code < 300:
return path in _QUIET_SUCCESS_PATHS or path in _CHAT_LIST_PATHS
return pre_auth and status_code == 401 and path in _CHAT_LIST_PATHS
class LoggingMiddleware:
@ -71,14 +118,16 @@ class LoggingMiddleware:
self.app = app
# (method, path, query, status_code) -> monotonic ts of the last EMITTED log.
self._last_log: dict[tuple[str, str, bytes, int], float] = {}
# Flips True after the first successful /api/auth/refresh; before that, chat
# list-poll 401s are the transient bootstrap race and are suppressed.
self._auth_refreshed = False
def _is_redundant_repeat(
self, method: str, path: str, query: bytes, status_code: int, now: float
) -> bool:
"""True if an identical GET/2xx log fired < window ago. The query string
is part of the identity, so distinct query-driven GETs are not collapsed.
Mutations and non-2xx are never deduped. Quiet-poll paths use a longer
heartbeat window. Stamps only on emit, so steady polls still log."""
"""True if an identical GET/2xx log fired < window ago (query string is part
of the identity). Non-GET/non-2xx never dedup; quiet-poll paths use the longer
heartbeat. Stamps only on emit, so steady polls still log."""
if method != "GET" or not (200 <= status_code < 300):
return False
window_ms = _QUIET_POLL_DEDUP_MS if path in _QUIET_POLL_PATHS else _ACCESS_LOG_DEDUP_MS
@ -129,8 +178,16 @@ class LoggingMiddleware:
raise
else:
end_time = time.perf_counter()
if not excluded and not self._is_redundant_repeat(
scope["method"], path, scope.get("query_string", b""), status_code, end_time
if 200 <= status_code < 300 and path == _AUTH_REFRESH_PATH:
self._auth_refreshed = True
if (
not excluded
and not _is_quiet_success(
scope["method"], path, status_code, not self._auth_refreshed
)
and not self._is_redundant_repeat(
scope["method"], path, scope.get("query_string", b""), status_code, end_time
)
):
logger.info(
"request_completed",

108
studio/backend/main.py
View file

@ -159,6 +159,14 @@ if sys.platform == "win32":
_bnb_rocm_ver_final,
)
# Setting BNB_ROCM_VERSION makes bitsandbytes log a benign override notice on
# import; drop only that record so real errors and mismatch warnings show.
if os.environ.get("BNB_ROCM_VERSION"):
import logging as _logging
_logging.getLogger("bitsandbytes.cextension").addFilter(
lambda _r: "environment variable detected" not in _r.getMessage()
)
# ── WSL AMD Strix Halo (gfx1151): enable ROCDXG before any torch import ──────
# In WSL the AMD GPU is reached via the ROCDXG bridge (librocdxg.so over
# /dev/dxg), which HSA loads only when HSA_ENABLE_DXG_DETECTION=1 is set BEFORE
@ -226,6 +234,7 @@ if _STUDIO_ROOT_RESOLVED != _LEGACY_STUDIO_ROOT:
os.environ.setdefault("UNSLOTH_IS_PRESENT", "1")
import hashlib
import ipaddress
import mimetypes
import re as _re
import shutil
@ -551,8 +560,12 @@ async def lifespan(app: FastAPI):
(_time.perf_counter() - _lifespan_started) * 1000,
)
# run_server's pre-bind gate sets suppress_bootstrap_injection when a public
# URL is about to serve with the default credential active: never (re)capture
# the bootstrap password into app.state, or the HTML would hand it out.
_suppress_bootstrap = getattr(app.state, "suppress_bootstrap_injection", False)
if storage.ensure_default_admin():
bootstrap_pw = storage.get_bootstrap_password()
bootstrap_pw = None if _suppress_bootstrap else storage.get_bootstrap_password()
app.state.bootstrap_password = bootstrap_pw
bootstrap_path = storage.DB_PATH.parent / ".bootstrap_password"
@ -563,7 +576,9 @@ async def lifespan(app: FastAPI):
print(" Open the Studio UI to sign in and change it.")
print("=" * 60 + "\n")
else:
app.state.bootstrap_password = storage.get_bootstrap_password()
app.state.bootstrap_password = (
None if _suppress_bootstrap else storage.get_bootstrap_password()
)
_lifespan_log.info(
"lifespan startup completed in %.1fms",
@ -597,6 +612,22 @@ app = FastAPI(
lifespan = lifespan,
)
# The MCP surface is opt-in because it can start GPU jobs and write model
# artifacts. Mount it only when explicitly enabled by the Studio process.
if os.environ.get("UNSLOTH_STUDIO_ENABLE_MCP") == "1":
from fastmcp.utilities.lifespan import combine_lifespans
from mcp_server import BearerTokenMiddleware, create_studio_mcp
_studio_mcp_app = create_studio_mcp().http_app(path = "/")
_studio_mcp_lifespan = _studio_mcp_app.lifespan
_mcp_token = os.environ.get("UNSLOTH_STUDIO_MCP_TOKEN")
if not _mcp_token:
raise RuntimeError("UNSLOTH_STUDIO_MCP_TOKEN is required when MCP is enabled")
_studio_mcp_app = BearerTokenMiddleware(_studio_mcp_app, _mcp_token)
app.router.lifespan_context = combine_lifespans(lifespan, _studio_mcp_lifespan)
app.mount("/mcp", _studio_mcp_app)
from loggers.config import LogConfig
from loggers.handlers import LoggingMiddleware
@ -737,6 +768,7 @@ _BODY_PROTECTED_PREFIXES = (
"/api/settings",
"/api/train",
"/api/export",
"/mcp",
)
_DATASET_UPLOAD_PASSTHROUGH_PREFIX = "/api/datasets/upload"
_DATA_RECIPE_UNSTRUCTURED_UPLOAD_PASSTHROUGH_PREFIX = (
@ -1355,6 +1387,61 @@ def _canonical_origin(scheme: str, netloc: str) -> Optional[tuple[str, str, int]
return (scheme, host, port)
def _is_loopback_ip(host: Optional[str]) -> bool:
"""Return whether ``host`` is a loopback IP, including IPv4-mapped IPv6."""
if not host or "%" in host: # a scope id (::1%eth0) is never a plain loopback
return False
try:
ip = ipaddress.ip_address(host)
except (TypeError, ValueError):
return False
mapped = getattr(ip, "ipv4_mapped", None)
return ip.is_loopback or (mapped is not None and mapped.is_loopback)
# A loopback peer carrying any of these is a proxy/tunnel relaying a remote
# client, so the peer is the proxy, not the caller: cloudflared sets
# cf-connecting-ip, reverse proxies set the rest (uvicorn only consumes
# x-forwarded-for, so the others survive to here).
_PROXIED_CLIENT_HEADERS = (
"cf-connecting-ip",
"forwarded",
"x-forwarded-for",
"x-forwarded-host",
"x-real-ip",
)
def _host_header_is_loopback(host_header: Optional[str]) -> bool:
"""Loopback/localhost check on the raw Host header.
Reads the header directly so a malformed or absent Host cannot fall back to
``request.url.hostname``'s (loopback) ASGI server address.
"""
if not host_header:
return False
host = host_header.strip()
if host.startswith("["): # [IPv6] or [IPv6]:port
end = host.find("]")
if end == -1 or (host[end + 1 :] and not host[end + 1 :].startswith(":")):
return False # unclosed bracket or junk after ] (e.g. [::1]evil)
host = host[1:end]
elif host.count(":") == 1: # host:port
host = host.split(":", 1)[0]
host = host.lower().rstrip(".")
return host == "localhost" or _is_loopback_ip(host)
def _is_local_bootstrap_request(request: Request) -> bool:
"""Allow bootstrap injection only through a direct loopback authority."""
client = request.client
if client is None or not _is_loopback_ip(client.host):
return False
if any(request.headers.get(h) is not None for h in _PROXIED_CLIENT_HEADERS):
return False
return _host_header_is_loopback(request.headers.get("host"))
def _is_same_origin_request(request: Request) -> bool:
"""True when Origin is missing or matches request's scheme://host:port.
@ -1390,6 +1477,17 @@ def _is_same_origin_request(request: Request) -> bool:
return origin_canon == self_canon
def _should_inject_bootstrap(request: Request) -> bool:
"""Whether to embed the seeded bootstrap password in index.html."""
if not _is_same_origin_request(request):
return False
if _IS_COLAB:
# Single-user notebook proxy: allow autofill, but never a public
# shareable tunnel (a Colab Cloudflare link sets cf-connecting-ip).
return request.headers.get("cf-connecting-ip") is None
return _is_local_bootstrap_request(request)
def setup_frontend(app: FastAPI, build_path: Path):
"""Mount frontend static files (optional)"""
if not build_path.exists():
@ -1402,8 +1500,10 @@ def setup_frontend(app: FastAPI, build_path: Path):
def _build_index_response(request: Request) -> Response:
content = (build_path / "index.html").read_bytes()
content = _strip_crossorigin(content)
# Bootstrap pw is same-origin only; Vary: Origin keeps caches honest.
if _is_same_origin_request(request):
# Bootstrap pw goes only to a same-origin, direct-loopback client (or
# Colab's single-user notebook proxy): a wildcard bind must not serve it
# in-page to a LAN or proxied peer. Vary: Origin keeps caches honest.
if _should_inject_bootstrap(request):
content, nonce = _inject_bootstrap(content, app)
else:
nonce = None

259
studio/backend/mcp_server.py Normal file
View file

@ -0,0 +1,259 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Curated MCP tools for driving an Unsloth Studio instance.
The MCP surface deliberately wraps the existing Studio services instead of
duplicating training or export logic. It is opt-in because several tools can
start GPU work or write model artifacts.
"""
from __future__ import annotations
import hmac
from typing import Any
from fastmcp import FastMCP
class BearerTokenMiddleware:
"""Require an exact bearer token when Studio MCP is exposed remotely."""
def __init__(self, app: Any, token: str) -> None:
if not token or not token.strip():
raise ValueError("Studio MCP bearer token must be a non-empty value")
if not token.isascii():
# A non-ASCII token cannot be sent in an HTTP header; reject it here.
raise ValueError("Studio MCP bearer token must contain ASCII characters only")
self.app = app
# Compare on raw header bytes: str hmac.compare_digest raises on non-ASCII
# input, which would surface as a 500 instead of a clean 401.
self.expected = token.encode("utf-8")
async def __call__(self, scope: dict[str, Any], receive: Any, send: Any) -> None:
scope_type = scope.get("type")
if scope_type not in ("http", "websocket"):
await self.app(scope, receive, send)
return
headers = dict(scope.get("headers", []))
raw_auth = headers.get(b"authorization", b"")
scheme, _, supplied = raw_auth.partition(b" ")
if scheme.lower() != b"bearer" or not hmac.compare_digest(supplied, self.expected):
await _send_unauthorized(send, scope_type)
return
await self.app(scope, receive, send)
async def _send_unauthorized(send: Any, scope_type: str) -> None:
if scope_type == "websocket":
await send({"type": "websocket.close", "code": 4401})
return
await send(
{
"type": "http.response.start",
"status": 401,
"headers": [(b"content-type", b"application/json"), (b"www-authenticate", b"Bearer")],
}
)
await send(
{
"type": "http.response.body",
"body": b'{"detail":"MCP bearer token required"}',
}
)
def _dump(value: Any) -> Any:
"""Convert Pydantic responses to plain JSON values for MCP clients."""
if hasattr(value, "model_dump"):
return value.model_dump(mode = "json")
return value
def _clamp(value: int, low: int, high: int) -> int:
"""Clamp an MCP-supplied integer into an inclusive range.
MCP tools call the Studio route functions directly, which skips FastAPI's
Query(ge=, le=) validation, so we re-apply the same bounds here.
"""
return max(low, min(value, high))
def create_studio_mcp() -> FastMCP:
"""Create the Studio MCP server and register the high-value tools."""
mcp = FastMCP(
"Unsloth Studio",
instructions = (
"Use read tools to inspect the local Studio state before starting GPU work. "
"Training and export tools can consume substantial VRAM and write files. "
"Never expose tokens or local paths from tool results unless the user asks."
),
)
@mcp.tool
async def studio_status() -> dict[str, Any]:
"""Return the current training, export, inference, and GPU state."""
from routes.export import get_export_status
from routes.inference import get_status as get_inference_status
from routes.training import get_training_status
from utils.hardware import get_gpu_utilization
training, export, inference = await _gather_status(
get_training_status(current_subject = "mcp"),
get_export_status(current_subject = "mcp"),
get_inference_status(current_subject = "mcp"),
)
return {
"training": _dump(training),
"export": _dump(export),
"inference": _dump(inference),
"hardware": get_gpu_utilization(),
}
@mcp.tool
async def list_local_models(models_dir: str = "./models") -> dict[str, Any]:
"""List local and cached models available to Studio."""
from routes.models import list_local_models as list_models
return _dump(await list_models(models_dir = models_dir, current_subject = "mcp"))
@mcp.tool
async def get_training_status() -> dict[str, Any]:
"""Read the active training job, phase, progress, and recent metrics."""
from routes.training import get_training_status as get_status
return _dump(await get_status(current_subject = "mcp"))
@mcp.tool
async def start_training(config: dict[str, Any]) -> dict[str, Any]:
"""Start a validated Studio training job from a TrainingStartRequest-shaped object.
The config is validated by the same Pydantic model used by the Studio UI.
Call get_training_status first and do not start work while another job runs.
"""
from models import TrainingStartRequest
from routes.training import start_training as start
request = TrainingStartRequest.model_validate(config)
# Pass via_api_key explicitly (a direct call leaves it a Depends object).
# MCP drives Studio like the UI session, so it coexists and frees VRAM.
return _dump(await start(request, current_subject = "mcp", via_api_key = False))
@mcp.tool
async def stop_training(save: bool = True) -> dict[str, Any]:
"""Ask the active training process to stop at its next safe checkpoint."""
from routes.training import TrainingStopRequest, stop_training as stop
return _dump(await stop(TrainingStopRequest(save = save), current_subject = "mcp"))
@mcp.tool
async def list_training_runs(limit: int = 50, offset: int = 0) -> dict[str, Any]:
"""List completed and stopped training runs, newest first."""
from routes.training_history import list_training_runs as list_runs
# Clamp here (direct call skips Query bounds); a negative LIMIT = no limit.
limit = _clamp(limit, 1, 200)
offset = max(0, offset)
return _dump(await list_runs(limit = limit, offset = offset, current_subject = "mcp"))
@mcp.tool
def validate_recipe(recipe: dict[str, Any]) -> dict[str, Any]:
"""Validate a Data Recipe with the same validator used by Studio."""
from models.data_recipe import RecipePayload
from routes.data_recipe.validate import validate
return _dump(validate(RecipePayload(recipe = recipe)))
@mcp.tool
def get_recipe_job_status(job_id: str) -> dict[str, Any]:
"""Read the status of a Data Recipe job."""
from routes.data_recipe.jobs import job_status
return _dump(job_status(job_id))
@mcp.tool
def get_recipe_job_dataset(
job_id: str,
limit: int = 20,
offset: int = 0,
) -> dict[str, Any]:
"""Read a bounded page of generated Data Recipe rows."""
from routes.data_recipe.jobs import job_dataset
# Clamp here (direct call skips FastAPI's Query bounds).
limit = _clamp(limit, 1, 500)
offset = max(0, offset)
return _dump(job_dataset(job_id, limit = limit, offset = offset))
@mcp.tool
async def load_checkpoint(
checkpoint_path: str,
max_seq_length: int = 2048,
load_in_4bit: bool = True,
trust_remote_code: bool = False,
approved_remote_code_fingerprint: str | None = None,
hf_token: str | None = None,
) -> dict[str, Any]:
"""Load a checkpoint into the export backend.
Export runs in its own subprocess and coexists with training and
inference; it does not unload them, so a load can fail with a clear
out-of-memory error if the GPU is already full. Pass hf_token to load a
gated checkpoint, and approved_remote_code_fingerprint to retry a
trust_remote_code load that was blocked pending review.
"""
from models import LoadCheckpointRequest
from routes.export import load_checkpoint as load
request = LoadCheckpointRequest(
checkpoint_path = checkpoint_path,
max_seq_length = max_seq_length,
load_in_4bit = load_in_4bit,
trust_remote_code = trust_remote_code,
approved_remote_code_fingerprint = approved_remote_code_fingerprint,
hf_token = hf_token,
)
return _dump(await load(request, current_subject = "mcp"))
@mcp.tool
async def export_gguf(
save_directory: str,
quantization_method: str | list[str] = "Q4_K_M",
push_to_hub: bool = False,
repo_id: str | None = None,
hf_token: str | None = None,
imatrix: bool = False,
imatrix_path: str | None = None,
) -> dict[str, Any]:
"""Export the loaded model to GGUF using Studio's existing path validation.
quantization_method may be a single method or a list to produce several
GGUFs from one load. Pass hf_token when push_to_hub is set (the backend
rejects a Hub upload without it). Set imatrix (or imatrix_path) for the
IQ low-bit quants that require an importance matrix.
"""
from models import ExportGGUFRequest
from routes.export import export_gguf as export
request = ExportGGUFRequest(
save_directory = save_directory,
quantization_method = quantization_method,
push_to_hub = push_to_hub,
repo_id = repo_id,
hf_token = hf_token,
imatrix = imatrix,
imatrix_path = imatrix_path,
)
return _dump(await export(request, current_subject = "mcp"))
return mcp
async def _gather_status(*coroutines: Any) -> tuple[Any, ...]:
"""Gather independent status calls without letting one optional backend fail all state."""
import asyncio
results = await asyncio.gather(*coroutines, return_exceptions = True)
return tuple(
{"error": str(result)} if isinstance(result, Exception) else result for result in results
)

10
studio/backend/models/auth.py
View file

@ -7,6 +7,8 @@ from typing import Optional
from pydantic import BaseModel, Field
from auth.storage import MIN_PASSWORD_LENGTH
class AuthLoginRequest(BaseModel):
"""Login payload: username/password to obtain a JWT."""
@ -45,10 +47,14 @@ class ChangePasswordRequest(BaseModel):
"""Change the current user's password, typically on first login."""
current_password: str = Field(
..., min_length = 8, description = "Existing password for the authenticated user"
...,
min_length = MIN_PASSWORD_LENGTH,
description = "Existing password for the authenticated user",
)
new_password: str = Field(
..., min_length = 8, description = "Replacement password (minimum 8 characters)"
...,
min_length = MIN_PASSWORD_LENGTH,
description = f"Replacement password (minimum {MIN_PASSWORD_LENGTH} characters)",
)

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

@ -140,6 +140,27 @@ class ValidateModelRequest(BaseModel):
)
class TransformersUpgradeInfo(BaseModel):
"""A model architecture no installed transformers ships, but a newer release does."""
model_type: str = Field(
..., description = "config.json model_type unknown to every installed transformers"
)
pypi_version: Optional[str] = Field(
None, description = "Latest transformers release on PyPI at check time"
)
supported_in_pypi: bool = Field(
False,
description = "True if the latest PyPI release ships this model_type; Studio can "
"install it into a persistent sidecar after user consent.",
)
supported_in_main: bool = Field(
False,
description = "True if transformers GitHub main ships this model_type (dev-only; "
"not installable through Studio yet).",
)
class ValidateModelResponse(BaseModel):
"""Result of model validation.
@ -167,6 +188,48 @@ class ValidateModelResponse(BaseModel):
description = "Native training context length, read from the GGUF header when the file "
"is already downloaded locally; None for non-GGUF, gated, or not-yet-downloaded models.",
)
# Additive fields; the consuming consent dialog ships in a follow-up frontend PR.
requires_transformers_upgrade: bool = Field(
False,
description = "True when the model's architecture is unknown to every installed "
"transformers but a newer transformers ships it; the UI should offer the "
"install-latest-transformers consent dialog (or the dev-only notice).",
)
transformers_upgrade: Optional[TransformersUpgradeInfo] = Field(
None,
description = "Details for the transformers-upgrade dialog; set only when "
"requires_transformers_upgrade is true.",
)
class InstallLatestTransformersRequest(BaseModel):
"""Consented request to install the latest transformers release into a sidecar."""
version: str = Field(
...,
min_length = 1,
max_length = 64,
description = "Exact transformers version to install; must match the current "
"latest PyPI release reported by /validate.",
)
class InstallLatestTransformersResponse(BaseModel):
"""Result of the consented latest-transformers sidecar install."""
success: bool = Field(..., description = "Whether the sidecar was provisioned")
version: str = Field(..., description = "The requested transformers version")
message: str = Field(..., description = "Human-readable result")
model_unloaded: bool = Field(
False,
description = "Whether the active chat model was unloaded before the swap "
"(reported even on failure, so the client can restore its state)",
)
latest_version: Optional[str] = Field(
None,
description = "On a version-mismatch failure: the release that superseded "
"the requested one, so the client can retry with it",
)
class GenerateRequest(BaseModel):
@ -632,6 +695,23 @@ class ThinkingConfig(BaseModel):
type: Literal["disabled", "enabled"] = "disabled"
# Recognized permission_mode values. The field accepts a plain string rather than
# a Literal so an unrecognized value from a newer UI/client degrades to the
# safest gate ("ask") instead of a 422; the tool loops apply the same unknown ->
# ask fallback, so normalizing here keeps that forward-compat path reachable at
# the API boundary. None stays unset ("behaves as 'ask'" without self-enabling
# the confirm gate).
_KNOWN_PERMISSION_MODES = ("ask", "auto", "off", "full")
def _normalize_permission_mode(value: Any) -> Any:
if value is None:
return None
if value not in _KNOWN_PERMISSION_MODES:
return "ask"
return value
class ChatCompletionRequest(BaseModel):
"""OpenAI-compatible chat completion request.
@ -777,6 +857,19 @@ class ChatCompletionRequest(BaseModel):
False,
description = "[x-unsloth] Bypass Permissions: when true, skip the tool-call confirmation gate AND disable the python/terminal execution sandbox (safety checks, command blocklist, resource limits). Secret env vars are still stripped. Takes precedence over confirm_tool_calls.",
)
permission_mode: Optional[str] = Field(
None,
description = (
"[x-unsloth] Permission level for local tool calls. 'ask' pauses every "
"call for approval; 'ask'/'auto' enable the confirmation gate on their "
"own (needs a streaming request to deliver prompts). 'auto' ('Approve for "
"me') only pauses calls detected as potentially unsafe (state-mutating "
"terminal/python/MCP calls); read-only calls run immediately, and the "
"sandbox stays on. 'full' is equivalent to bypass_permissions=true (no "
"confirmation, no sandbox). Unset behaves as 'ask'. An unrecognized value "
"(e.g. from a newer client) is treated as 'ask'."
),
)
auto_heal_tool_calls: Optional[bool] = Field(
True,
description = "[x-unsloth] Auto-detect and fix malformed tool calls from model output.",
@ -815,6 +908,10 @@ class ChatCompletionRequest(BaseModel):
None,
description = "[x-unsloth] Session/thread ID for scoping tool execution sandbox.",
)
thread_id: Optional[str] = Field(
None,
description = "[x-unsloth] Conversation ID for scoping stateful tool sessions (e.g. stdio MCP); stays per-thread where session_id may be shared project-wide.",
)
rag_scope: Optional[dict] = Field(
None,
description = (
@ -1036,6 +1133,52 @@ class ChatCompletionRequest(BaseModel):
self.enable_thinking = self.thinking.type == "enabled"
return self
@field_validator("permission_mode", mode = "before")
@classmethod
def _coerce_permission_mode(cls, value: Any) -> Any:
# Accept any string so an unknown mode degrades to 'ask' instead of a
# 422; mirrors the tool loops' unknown -> ask fallback.
return _normalize_permission_mode(value)
@model_validator(mode = "after")
def _fold_full_permission_into_bypass(self) -> "ChatCompletionRequest":
"""permission_mode='full' is the documented equivalent of
bypass_permissions=true, so fold it in before any route guard reads
the flag (else a full request would trip the confirm-gate rejections)."""
if self.permission_mode == "full":
self.bypass_permissions = True
elif self.bypass_permissions:
# Legacy bypass callers map onto Full access (mirrors the tool loop).
self.permission_mode = "full"
elif self.permission_mode == "off":
# "Off" never prompts, so route guards must see confirm disabled.
self.confirm_tool_calls = False
elif (
self.permission_mode == "ask"
and self.confirm_tool_calls is None
and not (self.provider_id or self.provider_type)
and (self.enable_tools is True or bool(self.mcp_enabled))
):
# "Ask" gates every call, so a direct API caller that omits the legacy
# confirm flag must still hit the confirmation gate for Studio's own
# tool loop. An explicit confirm_tool_calls=False wins over the mode
# (mirrors _permission_mode_confirm and the Anthropic pre-switch guard),
# so only self-enable when the flag is unset. Only self-enable when that
# loop is actually requested
# (enable_tools / mcp_enabled) -- the router enters the loop on those
# signals, not on enabled_tools alone (which merely filters which tools
# run). A plain client-tool passthrough (client-supplied `tools` that
# Studio does not execute) must route verbatim, and external-provider
# routing rejects confirm_tool_calls with tools, so skip the fold there.
#
# "auto" is deliberately NOT folded: it only prompts for a call the
# classifier flags, so leaving confirm_tool_calls unset lets the route's
# _confirm_gate_needs_stream apply the safe-only exception (a safe-only
# auto selection needs no stream) instead of an explicit-confirm forcing
# stream=true. The mode still drives the loop's per-call gate.
self.confirm_tool_calls = True
return self
class ToolConfirmRequest(BaseModel):
session_id: Optional[str] = None
@ -1533,12 +1676,41 @@ class AnthropicToolResultBlock(BaseModel):
tool_use_id: str
content: Union[str, list] = ""
@field_validator("content", mode = "before")
@classmethod
def _coerce_null_content(cls, v):
# Some clients send null content for an empty tool result; the str|list
# union would 400 on it, so treat null as "".
return "" if v is None else v
# Block types the converter translates explicitly. Anything else (thinking /
# redacted_thinking, a provider block a resumed session replays, or a future type)
# is accepted as an unknown block and dropped by the converter, rather than 400-ing
# the whole request on strict validation.
_KNOWN_ANTHROPIC_BLOCK_TYPES = frozenset({"text", "image", "tool_use", "tool_result"})
class AnthropicUnknownBlock(BaseModel):
type: str
model_config = {"extra": "allow"}
@field_validator("type")
@classmethod
def _only_unknown_types(cls, v):
# Known types parse as their typed models above (so a malformed known block
# still fails cleanly); this fallback only catches the rest.
if v in _KNOWN_ANTHROPIC_BLOCK_TYPES:
raise ValueError("known block type handled by its typed model")
return v
AnthropicContentBlock = Union[
AnthropicTextBlock,
AnthropicImageBlock,
AnthropicToolUseBlock,
AnthropicToolResultBlock,
AnthropicUnknownBlock,
]
@ -1583,6 +1755,40 @@ class AnthropicMessage(BaseModel):
role: Literal["user", "assistant"]
content: Union[str, list[AnthropicContentBlock]]
@model_validator(mode = "before")
@classmethod
def _normalize_content(cls, data):
# Role-aware leniency that never silently drops real user input:
# - assistant: a resumed tool-only turn's null content -> "" (str|list would
# 400 on null; "" keeps the converter's `for block in content` safe).
# Unknown blocks (thinking / future types) validate via
# AnthropicUnknownBlock and are dropped by the converter.
# - user: keep strict. Null user content stays None so str|list rejects it
# (400) rather than forwarding an empty prompt; and reject block types the
# converter cannot translate, since it silently skips unknown user blocks
# -- a user turn made only of them would validate yet send no content
# (silent data loss).
if not isinstance(data, dict):
return data
content = data.get("content")
if data.get("role") == "assistant":
# Coerce only an explicit null (resumed tool-only turn). A missing
# content key stays malformed so the required-field check still 400s.
if "content" in data and content is None:
return {**data, "content": ""}
return data
if isinstance(content, list):
for block in content:
btype = (
block.get("type") if isinstance(block, dict) else getattr(block, "type", None)
)
# Guard the value: a non-string type is unsupported too, and a
# membership test on an unhashable value would raise TypeError
# (escaping as a 500 instead of a clean 400).
if not isinstance(btype, str) or btype not in _KNOWN_ANTHROPIC_BLOCK_TYPES:
raise ValueError(f"unsupported content block type {btype!r} in a user message")
return data
class AnthropicTool(BaseModel):
# Client tools have input_schema; server tools may only have type/name.
@ -1619,11 +1825,19 @@ class AnthropicMessagesRequest(BaseModel):
enable_tools: Optional[bool] = None
enabled_tools: Optional[list[str]] = None
session_id: Optional[str] = None
thread_id: Optional[str] = Field(
None,
description = "[x-unsloth] Conversation ID for scoping stateful tool sessions (e.g. stdio MCP); stays per-thread where session_id may be shared project-wide.",
)
cancel_id: Optional[str] = None
bypass_permissions: Optional[bool] = Field(
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.",
)
permission_mode: Optional[str] = Field(
None,
description = "[x-unsloth] Permission level for local tool calls: 'ask' pauses every call, 'auto' only pauses calls detected as potentially unsafe, 'off' never pauses (sandbox stays on), 'full' equals bypass_permissions=true. Unset behaves as 'ask'; an unrecognized value (e.g. from a newer client) is treated as 'ask'. Declared explicitly so omitted requests default to None 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).",
@ -1665,6 +1879,27 @@ class AnthropicMessagesRequest(BaseModel):
normalized["system"] = _merge_anthropic_system(normalized.get("system"), system_additions)
return normalized
@field_validator("permission_mode", mode = "before")
@classmethod
def _coerce_permission_mode(cls, value: Any) -> Any:
# Accept any string so an unknown mode degrades to 'ask' instead of a
# 422; mirrors the tool loops' unknown -> ask fallback.
return _normalize_permission_mode(value)
@model_validator(mode = "after")
def _fold_full_permission_into_bypass(self) -> "AnthropicMessagesRequest":
"""permission_mode='full' equals bypass_permissions=true (mirrors the
Chat Completions request)."""
if self.permission_mode == "full":
self.bypass_permissions = True
elif self.bypass_permissions:
# Legacy bypass callers map onto Full access (mirrors the tool loop).
self.permission_mode = "full"
elif self.permission_mode == "off":
# "Off" never prompts, so route guards must see confirm disabled.
self.confirm_tool_calls = False
return self
# ── Response models ────────────────────────────────────────────

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

@ -500,8 +500,9 @@ async def change_password(
detail = "New password must be different from the current password",
)
storage.update_password(current_subject, payload.new_password)
storage.revoke_user_refresh_tokens(current_subject)
# Single transaction: a separate refresh-token purge could fail after the
# password commit, leaving pre-change tokens able to mint access tokens.
storage.update_password(current_subject, payload.new_password, revoke_refresh_tokens = True)
try:
request.app.state.bootstrap_password = None
except AttributeError:

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

@ -51,7 +51,17 @@ def _ensure_export_supported() -> None:
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.
Also refuses (409) while a latest-transformers install is swapping .venv_t5_latest: an
export worker spawned mid-swap could activate a half-replaced sidecar.
"""
from utils.transformers_latest import is_install_in_progress
if is_install_in_progress():
raise HTTPException(
status_code = 409,
detail = "A transformers installation is in progress. Retry when it completes.",
)
from utils.hardware import export_capability
cap = export_capability()
@ -97,6 +107,11 @@ async def load_checkpoint(
except HTTPException:
raise
except Exception as e:
from utils.transformers_version import SidecarSwapInProgress
if isinstance(e, SidecarSwapInProgress):
# Expected loss of the race against a sidecar install: retryable 409.
raise HTTPException(status_code = 409, detail = str(e))
logger.error(f"Error loading checkpoint: {e}", exc_info = True)
raise HTTPException(
status_code = 500,
@ -308,6 +323,11 @@ async def export_merged_model(
except HTTPException:
raise
except Exception as e:
from utils.transformers_version import SidecarSwapInProgress
if isinstance(e, SidecarSwapInProgress):
# Expected loss of the race against a sidecar install: retryable 409.
raise HTTPException(status_code = 409, detail = str(e))
logger.error(f"Error exporting merged model: {e}", exc_info = True)
raise HTTPException(
status_code = 500,
@ -347,6 +367,11 @@ async def export_base_model(
except HTTPException:
raise
except Exception as e:
from utils.transformers_version import SidecarSwapInProgress
if isinstance(e, SidecarSwapInProgress):
# Expected loss of the race against a sidecar install: retryable 409.
raise HTTPException(status_code = 409, detail = str(e))
logger.error(f"Error exporting base model: {e}", exc_info = True)
raise HTTPException(
status_code = 500,
@ -388,6 +413,11 @@ async def export_gguf(
except HTTPException:
raise
except Exception as e:
from utils.transformers_version import SidecarSwapInProgress
if isinstance(e, SidecarSwapInProgress):
# Expected loss of the race against a sidecar install: retryable 409.
raise HTTPException(status_code = 409, detail = str(e))
logger.error(f"Error exporting GGUF model: {e}", exc_info = True)
raise HTTPException(
status_code = 500,
@ -428,6 +458,11 @@ async def export_lora_adapter(
except HTTPException:
raise
except Exception as e:
from utils.transformers_version import SidecarSwapInProgress
if isinstance(e, SidecarSwapInProgress):
# Expected loss of the race against a sidecar install: retryable 409.
raise HTTPException(status_code = 409, detail = str(e))
logger.error(f"Error exporting LoRA adapter: {e}", exc_info = True)
raise HTTPException(
status_code = 500,

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

File diff suppressed because it is too large Load diff

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

@ -14,14 +14,17 @@ never blocks on a missing marker / offline GitHub.
from __future__ import annotations
import asyncio
import threading
from typing import Optional
from fastapi import APIRouter, Depends, Query
from pydantic import BaseModel, Field
from auth.authentication import get_current_subject
from loggers import get_logger
from utils.llama_cpp_update import get_update_status, start_update
logger = get_logger(__name__)
router = APIRouter()
@ -69,6 +72,27 @@ class LlamaUpdateActionResponse(BaseModel):
job: LlamaUpdateJob = Field(default_factory = LlamaUpdateJob)
_llama_update_lock = threading.Lock()
_last_llama_update_step = -1
def _log_llama_update_progress(job: LlamaUpdateJob) -> None:
"""One llama_update_progress line per 10% step so a prebuilt update reports
progress without a line per poll. Resyncs when a new update starts."""
global _last_llama_update_step
if job.state != "running" or job.progress is None:
return
step = int(max(0.0, min(float(job.progress), 1.0)) * 10)
with _llama_update_lock:
prev = _last_llama_update_step
if step == prev:
return
_last_llama_update_step = step
if step < prev:
return # new update; resync without logging
logger.info("llama_update_progress", to_tag = job.to_tag or "", percent = step * 10)
@router.get("/update-status", response_model = LlamaUpdateStatusResponse)
async def llama_update_status(
force_refresh: bool = Query(
@ -78,7 +102,9 @@ async def llama_update_status(
) -> LlamaUpdateStatusResponse:
# Off the event loop: detection may probe the host and read GitHub.
status = await asyncio.to_thread(get_update_status, force_refresh = force_refresh)
return LlamaUpdateStatusResponse(**status)
resp = LlamaUpdateStatusResponse(**status)
_log_llama_update_progress(resp.job)
return resp
@router.post("/update", response_model = LlamaUpdateActionResponse)

15
studio/backend/routes/mcp_servers.py
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 asyncio
import json
import uuid
from urllib.parse import urlparse
@ -13,6 +14,7 @@ from core.inference.mcp_client import (
TOOL_CACHE_INVALIDATING_FIELDS,
cache_tools,
clear_oauth_tokens_async,
close_stdio_sessions,
invalidate_tool_cache,
is_stdio,
list_tools_async,
@ -206,11 +208,15 @@ async def update_mcp_server(
):
await clear_oauth_tokens_async(old["url"])
mcp_servers_db.update_server(server_id, changes)
# A new endpoint/auth makes cached tools wrong and disabling makes them
# unreachable, so drop them and let the next send re-probe; a rename
# leaves them valid.
if changes.keys() & TOOL_CACHE_INVALIDATING_FIELDS:
# A new endpoint/auth makes cached tools wrong and disabling makes them unreachable, so drop
# them and let the next send re-probe; a rename leaves them valid. Live stdio sessions for the
# old endpoint close too. Gate on a real value change, not mere presence: the edit dialog
# resends url/headers/oauth unchanged on a rename, which must not drop the session.
if any(changes[k] != old.get(k) for k in changes.keys() & TOOL_CACHE_INVALIDATING_FIELDS):
invalidate_tool_cache(server_id)
# Narrow to this row's env: another server row sharing the command but
# with a different env keeps its live sessions.
await asyncio.to_thread(close_stdio_sessions, old["url"], parse_server_headers(old))
return _row_to_response(mcp_servers_db.get_server(server_id))
@ -223,6 +229,7 @@ async def delete_mcp_server(server_id: str, current_subject: str = Depends(get_c
await clear_oauth_tokens_async(old["url"])
mcp_servers_db.delete_server(server_id)
invalidate_tool_cache(server_id)
await asyncio.to_thread(close_stdio_sessions, old["url"], parse_server_headers(old))
@router.post("/{server_id}/refresh", response_model = McpServerProbeResult)

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

@ -1188,7 +1188,9 @@ def _looks_like_model_dir(directory: Path) -> bool:
return False
def _build_browse_allowlist() -> list[Path]:
def _build_browse_allowlist(
media_roots: Optional[list[Path]] = None, drive_roots: Optional[list[Path]] = None
) -> list[Path]:
"""Return the root directories the folder browser may walk.
The same list seeds the sidebar suggestion chips, so chip targets are
@ -1196,13 +1198,20 @@ def _build_browse_allowlist() -> list[Path]:
outputs/exports/studio root, registered scan folders, and well-known
local-LLM dirs (LM Studio, Ollama, ``~/models``); each added only if
it resolves to a real directory.
*media_roots* / *drive_roots* let the caller pass already-probed
removable-media and Windows drive roots so they aren't scanned again (a
disconnected mapped drive can make each probe slow); probed here when ``None``.
"""
from utils.paths import (
hf_default_cache_dir,
legacy_hf_cache_dir,
well_known_model_dirs,
)
from utils.paths.external_media import linux_run_media_mount_roots
from utils.paths.external_media import (
linux_run_media_mount_roots,
windows_drive_roots,
)
from storage.studio_db import list_scan_folders
candidates: list[Path] = []
@ -1218,7 +1227,13 @@ def _build_browse_allowlist() -> list[Path]:
candidates.append(resolved)
_add(Path.home())
for p in linux_run_media_mount_roots():
if media_roots is None:
media_roots = linux_run_media_mount_roots()
if drive_roots is None:
drive_roots = windows_drive_roots()
for p in media_roots:
_add(p)
for p in drive_roots:
_add(p)
_add(_resolve_hf_cache_dir())
try:
@ -1269,19 +1284,43 @@ def _build_browse_allowlist() -> list[Path]:
def _is_path_inside_allowlist(target: Path, allowed_roots: list[Path]) -> bool:
"""True if *target* equals or descends from any allowed root.
Uses ``os.path.realpath`` so symlinks can't escape the sandbox.
Uses ``os.path.realpath`` (symlinks can't escape the sandbox) and
``os.path.commonpath`` for a component-wise containment test, so a string
prefix like ``/home/u`` never matches a sibling ``/home/user2`` while a
drive root ``D:\\`` still contains ``D:\\models``. A Windows drive root
authorizes its descendants, but a bare POSIX root ``/`` must NOT, else one
``/`` allowlist entry would authorize every absolute path. ``normcase`` keeps
the drive-letter comparison case-insensitive, matching the hub browser.
"""
try:
target_real = os.path.realpath(str(target))
target_real = os.path.normcase(os.path.realpath(str(target)))
except OSError:
return False
for root in allowed_roots:
try:
root_real = os.path.realpath(str(root))
root_real = os.path.normcase(os.path.realpath(str(root)))
except OSError:
continue
if target_real == root_real or target_real.startswith(root_real + os.sep):
if target_real == root_real:
return True
drive, tail = os.path.splitdrive(root_real)
if os.path.dirname(root_real) == root_real and not drive:
# Bare POSIX filesystem root ("/"): equality above is the only
# match; do not let it authorize arbitrary descendants.
continue
if drive.startswith(("\\\\", "//")) and not tail:
# Bare UNC share root (\\server\share): os.path.commonpath raises
# "can't mix absolute and relative" on it, so authorize its
# descendants with a boundary-safe prefix test (normcase applied).
if target_real.startswith(root_real.rstrip("\\/") + os.sep):
return True
continue
try:
if os.path.commonpath([target_real, root_real]) == root_real:
return True
except ValueError:
# Different drives / mixed absolute-relative: not contained.
continue
return False
@ -1339,7 +1378,10 @@ 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
from storage.studio_db import (
contains_sensitive_path_component,
is_denied_system_path,
)
requested_path = _normalize_browse_request_path(path)
resolved_roots: list[Path] = []
@ -1396,6 +1438,11 @@ def _resolve_browse_target(path: Optional[str], allowed_roots: list[Path]) -> Pa
status_code = 403,
detail = "Credential or configuration directories are not browseable.",
)
if is_denied_system_path(str(resolved_child)):
raise HTTPException(
status_code = 403,
detail = "System directories are not browseable.",
)
current = resolved_child
if contains_sensitive_path_component(str(current)):
@ -1403,6 +1450,13 @@ def _resolve_browse_target(path: Optional[str], allowed_roots: list[Path]) -> Pa
status_code = 403,
detail = "Credential or configuration directories are not browseable.",
)
# Zero-component case: the requested path IS an allowlist root
# (e.g. a legacy-registered "/" or a Windows drive root).
if is_denied_system_path(str(current)):
raise HTTPException(
status_code = 403,
detail = "System directories are not browseable.",
)
if not current.is_dir():
raise HTTPException(
status_code = 400,
@ -1420,8 +1474,12 @@ def _resolve_browse_target(path: Optional[str], allowed_roots: list[Path]) -> Pa
)
# Sync (def, not async) so FastAPI runs the blocking filesystem I/O (drive
# probes, iterdir, realpath) in the threadpool: a disconnected mapped drive can
# make the probe wait out its timeout, which on the event loop would stall every
# other request. Matches the hub browse endpoint.
@router.get("/browse-folders", response_model = BrowseFoldersResponse)
async def browse_folders(
def browse_folders(
path: Optional[str] = Query(
None,
description = (
@ -1450,11 +1508,22 @@ async def browse_folders(
then hidden (if ``show_hidden=true``).
"""
from utils.paths import hf_default_cache_dir, well_known_model_dirs
from utils.paths.external_media import linux_run_media_mount_roots
from storage.studio_db import contains_sensitive_path_component, list_scan_folders
from utils.paths.external_media import (
linux_run_media_mount_roots,
windows_drive_roots,
)
from storage.studio_db import (
contains_sensitive_path_component,
is_denied_system_path,
list_scan_folders,
)
# Probe removable-media and Windows drive roots once; the allowlist and
# chips reuse the result so a disconnected mapped drive isn't scanned twice.
media_roots = linux_run_media_mount_roots()
drive_roots = windows_drive_roots()
# Build once; the sandbox check and suggestion chips share it.
allowed_roots = _build_browse_allowlist()
allowed_roots = _build_browse_allowlist(media_roots, drive_roots)
try:
target = _resolve_browse_target(path, allowed_roots)
@ -1506,6 +1575,15 @@ async def browse_folders(
continue
if contains_sensitive_path_component(name):
continue
# Hide denied system dirs (C:\Windows, /etc, ...) so they don't
# render as clickable rows that then 403 on descent. Resolve first
# so a symlink/junction into a denied dir is hidden too, not just a literal name.
try:
resolved_child = os.path.realpath(str(child))
except (OSError, ValueError):
resolved_child = str(child)
if is_denied_system_path(resolved_child):
continue
entries.append(
BrowseEntry(
name = name,
@ -1553,13 +1631,22 @@ async def browse_folders(
return
if resolved in seen_sug:
return
# Drop a denied system dir (e.g. a stale scan-folder row) so it never
# becomes a chip that 403s on click. Drive roots stay: only their
# system subdirectories are denied, not the root itself.
if is_denied_system_path(resolved):
return
if _safe_is_dir(resolved):
seen_sug.add(resolved)
suggestions.append(resolved)
# Home first -- the safe fallback when everything else is cold.
_add_sug(Path.home())
for p in linux_run_media_mount_roots():
# Reuse the roots probed for the allowlist above (no second drive scan).
for p in media_roots:
_add_sug(p)
# Windows drive roots so the user can hop between C:, D:, E: ...
for p in drive_roots:
_add_sug(p)
# The HF cache root the process is actually using.
try:

204
studio/backend/routes/settings.py
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 re
from typing import Literal, Optional
from urllib.parse import unquote, urlsplit
@ -551,6 +552,7 @@ class PersonalizationProfile(BaseModel):
nickname: str = Field("", max_length = 200)
avatarDataUrl: Optional[str] = Field(None, max_length = MAX_AVATAR_DATA_URL_BYTES)
avatarShape: Literal["circle", "rounded"] = "circle"
showGreetingSloth: bool = True
@field_validator("avatarDataUrl")
@classmethod
@ -562,11 +564,179 @@ class PersonalizationProfile(BaseModel):
return value
class PersonalizationCustomColors(BaseModel):
model_config = ConfigDict(extra = "ignore")
accent: Optional[str] = Field(None, pattern = r"^#[0-9a-fA-F]{6}$")
background: Optional[str] = Field(None, pattern = r"^#[0-9a-fA-F]{6}$")
foreground: Optional[str] = Field(None, pattern = r"^#[0-9a-fA-F]{6}$")
class PersonalizationCustomColorModes(BaseModel):
model_config = ConfigDict(extra = "ignore")
light: PersonalizationCustomColors = Field(default_factory = PersonalizationCustomColors)
dark: PersonalizationCustomColors = Field(default_factory = PersonalizationCustomColors)
MAX_IMPORTED_FONTS = 3
# ~1.5 MB font file as base64; matches MAX_IMPORTED_FONT_DATA_URL_LENGTH in
# the frontend appearance-custom-store.
MAX_FONT_DATA_URL_LENGTH = 2_200_000
# Aggregate cap across all imported fonts; matches
# MAX_TOTAL_IMPORTED_FONT_DATA_URL_LENGTH in the frontend so a synced payload
# always fits the browser's localStorage quota.
MAX_TOTAL_FONT_DATA_URL_LENGTH = 4_400_000
# Characters that could terminate a CSS declaration, escape the quoted
# font-family value (backslash), or smuggle extra fallbacks/comments (comma,
# slash) if a stored name ever reached a stylesheet. The server is the
# authoritative gate; the frontend strips the same set before use.
_FONT_NAME_FORBIDDEN = set(";{}()<>\"'\\/,`")
def _check_font_name(value: str) -> str:
if any(c in _FONT_NAME_FORBIDDEN or ord(c) < 0x20 for c in value):
raise ValueError("Font name contains invalid characters.")
return value
# Matches FONT_DATA_URL_PATTERN in the frontend appearance-custom-store.
_FONT_DATA_URL_PATTERN = re.compile(
r"^data:(?:font/(?:woff2?|ttf|otf|sfnt)"
r"|application/(?:octet-stream|x-font-\w+|font-\w+));base64,[A-Za-z0-9+/=]+$"
)
class PersonalizationImportedFont(BaseModel):
model_config = ConfigDict(extra = "ignore")
name: str = Field(..., min_length = 1, max_length = 100)
dataUrl: str = Field(..., max_length = MAX_FONT_DATA_URL_LENGTH)
@field_validator("name")
@classmethod
def _validate_font_name(cls, value: str) -> str:
return _check_font_name(value)
@field_validator("dataUrl")
@classmethod
def _validate_font_data_url(cls, value: str) -> str:
# fullmatch, not match: re's ``$`` also matches just before a trailing
# newline, so ``match`` would accept "data:font/woff2;base64,AAAA\n",
# which the frontend's JS pattern (``$`` = end of string) rejects.
if not _FONT_DATA_URL_PATTERN.fullmatch(value):
raise ValueError("dataUrl must be a base64 font data URL.")
return value
# Optional user-menu items; the boolean is each id's default visibility.
# Settings-tab shortcuts ship hidden.
SIDEBAR_MENU_ITEM_DEFAULTS = {
"api": True,
"darkMode": True,
"guidedTour": True,
"profile": False,
"appearance": False,
"resources": False,
"chat": False,
"connections": False,
}
# The sidebarMenu validator below dedupes ids and re-fills any missing ones, so
# the stored list is always exactly one entry per id. Cap the *incoming* list at
# a generous multiple rather than len(defaults): a stale or duplicated payload
# (more items than distinct ids) must reach the validator so it can normalize,
# instead of being rejected by the length constraint before dedupe runs. A
# pathologically long list is still refused.
MAX_SIDEBAR_MENU_INPUT_ITEMS = 4 * len(SIDEBAR_MENU_ITEM_DEFAULTS)
class PersonalizationSidebarMenuItem(BaseModel):
model_config = ConfigDict(extra = "ignore")
id: Literal[
"api",
"darkMode",
"guidedTour",
"profile",
"appearance",
"resources",
"chat",
"connections",
]
visible: bool = True
def _default_sidebar_menu() -> "list[PersonalizationSidebarMenuItem]":
return [
PersonalizationSidebarMenuItem(id = item_id, visible = visible)
for item_id, visible in SIDEBAR_MENU_ITEM_DEFAULTS.items()
]
class PersonalizationCustomization(BaseModel):
model_config = ConfigDict(extra = "ignore")
colors: PersonalizationCustomColorModes = Field(default_factory = PersonalizationCustomColorModes)
uiFont: Optional[str] = Field(None, max_length = 200)
headingFont: Optional[str] = Field(None, max_length = 200)
chatFont: Optional[str] = Field(None, max_length = 200)
codeFont: Optional[str] = Field(None, max_length = 200)
importedFonts: list[PersonalizationImportedFont] = Field(
default_factory = list, max_length = MAX_IMPORTED_FONTS
)
@field_validator("importedFonts")
@classmethod
def _validate_total_font_size(
cls, value: list[PersonalizationImportedFont]
) -> list[PersonalizationImportedFont]:
if sum(len(f.dataUrl) for f in value) > MAX_TOTAL_FONT_DATA_URL_LENGTH:
raise ValueError("Imported fonts exceed the total size limit.")
return value
@field_validator("uiFont", "headingFont", "chatFont", "codeFont")
@classmethod
def _validate_selected_fonts(cls, value: Optional[str]) -> Optional[str]:
# Selected font names reach CSS the same way imported names do.
return value if value is None else _check_font_name(value)
uiFontSize: Optional[int] = Field(None, ge = 12, le = 20)
codeFontSize: Optional[int] = Field(None, ge = 10, le = 20)
contrast: int = Field(50, ge = 0, le = 100)
pointerCursors: bool = False
reduceMotion: Literal["system", "on", "off"] = "system"
fontSmoothing: bool = True
sidebarMenu: list[PersonalizationSidebarMenuItem] = Field(
default_factory = _default_sidebar_menu,
max_length = MAX_SIDEBAR_MENU_INPUT_ITEMS,
)
@field_validator("sidebarMenu")
@classmethod
def _validate_sidebar_menu(
cls, value: list[PersonalizationSidebarMenuItem]
) -> list[PersonalizationSidebarMenuItem]:
# Drop duplicate ids (keep the first) and re-append any missing ids so
# the stored list always covers every optional menu item exactly once.
seen: set[str] = set()
items = [item for item in value if not (item.id in seen or seen.add(item.id))]
for item_id, visible in SIDEBAR_MENU_ITEM_DEFAULTS.items():
if item_id not in seen:
items.append(PersonalizationSidebarMenuItem(id = item_id, visible = visible))
return items
class PersonalizationAppearance(BaseModel):
model_config = ConfigDict(extra = "ignore")
theme: Literal["light", "dark", "system"] = "system"
palette: Literal["standard", "classic", "minimal"] = "standard"
language: Optional[str] = Field(None, max_length = 20)
customization: PersonalizationCustomization = Field(
default_factory = PersonalizationCustomization
)
class PersonalizationPayload(BaseModel):
@ -579,6 +749,11 @@ class PersonalizationPayload(BaseModel):
class PersonalizationResponse(PersonalizationPayload):
saved: bool = False
# False when the stored record predates a field, so the client keeps local
# overrides instead of treating a server-filled default as an explicit value.
customizationSaved: bool = False
paletteSaved: bool = False
greetingSlothSaved: bool = False
@router.get("/personalization", response_model = PersonalizationResponse)
@ -588,15 +763,38 @@ def get_personalization_settings(
stored = get_personalization()
response = PersonalizationResponse.model_validate(stored or {})
response.saved = bool(stored)
appearance = stored.get("appearance") if isinstance(stored, dict) else None
profile = stored.get("profile") if isinstance(stored, dict) else None
response.customizationSaved = isinstance(appearance, dict) and "customization" in appearance
response.paletteSaved = isinstance(appearance, dict) and "palette" in appearance
response.greetingSlothSaved = isinstance(profile, dict) and "showGreetingSloth" in profile
return response
def _merge_personalization(base: dict, overlay: dict) -> dict:
# Recursively overlay only the request's set fields onto the stored record,
# so a stale client that omits newer keys (palette, customization) does not
# materialize their defaults and defeat the *Saved legacy detection.
merged = dict(base)
for key, value in overlay.items():
existing = merged.get(key)
if isinstance(value, dict) and isinstance(existing, dict):
merged[key] = _merge_personalization(existing, value)
else:
merged[key] = value
return merged
@router.put("/personalization", response_model = PersonalizationPayload)
def update_personalization_settings(
payload: PersonalizationPayload, current_subject: str = Depends(get_current_subject)
) -> PersonalizationPayload:
try:
set_personalization(payload.model_dump())
# exclude_unset so absent fields are not persisted as defaults; merge so
# fields the request omits keep whatever the record already stored.
incoming = payload.model_dump(exclude_unset = True)
merged = _merge_personalization(get_personalization(), incoming)
set_personalization(merged)
except ValueError as exc:
raise log_and_http_error(
exc,
@ -605,4 +803,6 @@ def update_personalization_settings(
event = "settings.update_personalization_failed",
log = logger,
) from exc
return payload
# Return the stored record, not the defaults-filled request, so the response
# matches storage (and the next GET) for fields the client omitted.
return PersonalizationPayload.model_validate(merged)

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

@ -146,6 +146,16 @@ async def start_training(
# No in-process ensure_transformers_version(): the subprocess
# (worker.py) activates the correct version before importing ML libs.
# A consented latest-transformers install stage-and-swaps .venv_t5_latest;
# a worker spawned mid-swap could activate a half-replaced sidecar.
from utils.transformers_latest import is_install_in_progress
if is_install_in_progress():
raise HTTPException(
status_code = 409,
detail = ("A transformers installation is in progress. Retry when it completes."),
)
backend = get_training_backend()
# S3 dataset loading needs the optional boto3 dependency. Reject early
@ -341,6 +351,24 @@ async def start_training(
"s3_config": request.s3_config.model_dump() if request.s3_config else None,
}
# Latest-sidecar models size and train 16-bit (same flip as chat load):
# 4-bit is disabled for brand-new architectures, so VRAM coexistence
# checks must not underestimate against a load the worker will refuse.
if training_kwargs["load_in_4bit"]:
from utils.transformers_version import latest_tier_active_for
if await asyncio.to_thread(
latest_tier_active_for,
training_kwargs["model_name"],
training_kwargs["hf_token"] or None,
):
training_kwargs["load_in_4bit"] = False
logger.info(
"Latest-transformers sidecar active for %s - sizing and "
"training in 16-bit (4-bit is disabled for brand-new "
"architectures)",
training_kwargs["model_name"],
)
# Training page has no trust_remote_code toggle, so honor the YAML default
# -- but only for genuine first-party (unsloth/nvidia) Hub repos, never a
# local path or a name merely starting with "unsloth/".
@ -426,9 +454,16 @@ async def start_training(
logger.warning("Chat/training VRAM coordination failed; proceeding: %s", e)
# The hook runs only once start guards pass -> VRAM freed iff training starts.
success = backend.start_training(
job_id = job_id, before_spawn = _free_vram_for_training, **training_kwargs
)
from utils.transformers_version import SidecarSwapInProgress
try:
success = backend.start_training(
job_id = job_id, before_spawn = _free_vram_for_training, **training_kwargs
)
except SidecarSwapInProgress as exc:
# Expected loss of the race against a sidecar install: a retryable
# 409 matching the route-entry guard, not an internal error.
raise HTTPException(status_code = 409, detail = str(exc))
if not success:
progress_error = backend.trainer.training_progress.error
@ -698,7 +733,9 @@ async def stream_training_progress(
if last_event_id is not None:
try:
resume_from_step = int(last_event_id)
logger.info(f"SSE reconnect: resuming from step {resume_from_step}")
# Fires on every reconnect (each tab switch); the meaningful signal is
# the "replayed N missed steps" line below, logged only when N > 0.
logger.debug(f"SSE reconnect: resuming from step {resume_from_step}")
except ValueError:
logger.warning(f"Invalid Last-Event-ID: {last_event_id}")

337
studio/backend/run.py
View file

@ -10,7 +10,7 @@ import os
import sys
import time
from pathlib import Path
from typing import Optional
from typing import Optional, Tuple
def _fix_torch_cuda_ld_path():
@ -616,24 +616,28 @@ def _print_cloudflare_line(secure: bool = False, loopback_host: str = "127.0.0.1
f"bind {loopback_host} or close firewall access to keep Studio private.",
warn,
)
elif not _cloudflare_flag:
elif _cloudflare_flag is False or _cloudflare_flag is None:
# None = off by default (no flag); False = explicit --no-cloudflare.
_reason = "default" if _cloudflare_flag is None else "--no-cloudflare"
if _public_reachable is True:
_emit(
" Cloudflare tunnel: OFF (--no-cloudflare). The raw port is still "
f" Cloudflare tunnel: OFF ({_reason}). The raw port is still "
"reachable from the public internet (see the reachability check above): "
"--no-cloudflare disables only the Cloudflare link, not the public bind.",
"pass --cloudflare to also expose a public Cloudflare HTTPS link, or "
f"bind {loopback_host} to keep Studio private.",
warn,
)
elif _public_reachable is False:
_emit(
" Cloudflare tunnel: OFF (--no-cloudflare). Studio is reachable on your "
"local network only. Omit --no-cloudflare to expose a public "
f" Cloudflare tunnel: OFF ({_reason}). Studio is reachable on your "
"local network only. Pass --cloudflare to expose a public "
"Cloudflare HTTPS link."
)
else:
_emit(
" Cloudflare tunnel: OFF (--no-cloudflare). There is no Cloudflare "
"public link. Raw port reachability was not verified; "
f" Cloudflare tunnel: OFF ({_reason}). There is no Cloudflare "
"public link. Raw port reachability was not verified; pass --cloudflare "
"to expose a public Cloudflare HTTPS link, or "
f"bind {loopback_host} or close firewall access to keep Studio private.",
warn,
)
@ -874,7 +878,9 @@ _cloudflare_url = None
_public_reachable = None
_cloudflare_requested = False
_cloudflare_flag = True
# Opt-in tri-state (mirrors the CLI): None = off by default, True = on,
# False = explicit --no-cloudflare. run_server overwrites it before the banner.
_cloudflare_flag = None
_DEFAULT_FRONTEND_PATH = Path(__file__).resolve().parent.parent / "frontend" / "dist"
@ -1057,6 +1063,199 @@ def _cloudflare_tunnel_should_start(
return host in ("0.0.0.0", "::") and not api_only
def _stream_isatty(stream) -> bool:
"""isatty() that treats broken streams as non-interactive.
isatty() can raise under service wrappers (closed stdin -> ValueError;
sys.stdin None in Windows GUI -> AttributeError); such a stream can't host a
prompt, which is a fallback, not an error.
"""
try:
return stream.isatty()
except (AttributeError, ValueError):
return False
def _terminal_password_gate(
*,
tunnel_will_start: bool,
host: str,
secure: bool,
api_only: bool,
frontend_served: bool,
is_colab: bool = False,
) -> Tuple[bool, bool]:
"""Force a terminal password change before the public tunnel goes up.
When the tunnel is about to publish Studio and the seeded admin password was
never changed, ask for a new one (masked, confirmed) before any public URL
exists. The CLI normally does this before re-exec'ing the backend; this is
the backstop for direct `python run.py` launches and older-CLI installs.
Must run BEFORE the uvicorn socket binds: on a wildcard bind the served HTML
injects the bootstrap credential, so a pre-gate listener would hand the
default password to anyone reaching the raw port while the operator types.
Returns (proceed, drop_bootstrap_injection):
proceed False -> abort the launch (interactive refusal, or a headless
public launch nothing would protect); fail closed.
drop_bootstrap_injection True -> caller must null
app.state.bootstrap_password: the password just changed (stale), or a
public URL is about to serve the default credential and must not leak it.
Without a usable terminal the prompt is skipped: proceed if the bootstrap
deadline (armed later) will protect the launch; if even that is disabled
(api-only, timeout 0) nothing protects it, so refuse. NOT wrapped in a broad
try/except: an auth storage failure must abort rather than expose the default.
"""
if not tunnel_will_start:
return True, False
from auth import hashing as _auth_hashing
from auth import storage as _auth_storage
from auth.bootstrap_timeout import (
bootstrap_timeout_seconds,
should_arm_bootstrap_timeout,
)
from auth.terminal_prompt import (
prompt_for_password_change,
should_prompt_password_change,
)
_admin = _auth_storage.DEFAULT_ADMIN_USERNAME
# Gate can run before lifespan: seed the admin row here (idempotent).
_auth_storage.ensure_default_admin()
requires_change = _auth_storage.requires_password_change(_admin)
if not requires_change:
return True, False
if not should_prompt_password_change(
tunnel_will_start = tunnel_will_start,
requires_change = requires_change,
stdin_isatty = _stream_isatty(sys.stdin),
stderr_isatty = _stream_isatty(sys.stderr),
):
# No terminal: only proceed if the bootstrap deadline will arm; api-only
# and TIMEOUT=0 never arm it, leaving the default credential public.
deadline_arms = should_arm_bootstrap_timeout(
host = host,
secure = secure,
api_only = api_only,
frontend_served = frontend_served,
is_colab = is_colab,
requires_change = True,
timeout_seconds = bootstrap_timeout_seconds(),
)
if not deadline_arms:
print(
"Refusing to publish Studio on a public Cloudflare URL: the "
"default admin password was never changed, no terminal is "
"attached to change it here, and the bootstrap shutdown "
"deadline does not apply to this launch (api-only, or "
"UNSLOTH_STUDIO_BOOTSTRAP_TIMEOUT=0). Change the password "
"first (run `unsloth studio` locally and log in, or re-run "
"with a terminal attached), then retry.",
file = sys.stderr,
flush = True,
)
return False, False
# The public page won't auto-fill the bootstrap credential (suppressed
# below) and the seeded file may already be gone, so point recovery at a
# terminal-attached run / reset-password instead of reading it from disk.
print(
" WARNING: the default admin password is still active while "
"Studio is about to be published on a public Cloudflare URL, and "
"no terminal is attached to change it here. The public page will "
"NOT auto-fill the bootstrap credential. Set a new password by "
"running `unsloth studio` locally with a terminal attached, or "
"`unsloth studio reset-password`. Studio shuts down after the "
"bootstrap deadline (UNSLOTH_STUDIO_BOOTSTRAP_TIMEOUT, default 1h) "
"unless the password is changed.",
file = sys.stderr,
flush = True,
)
# Never serve the default credential in HTML over a public URL.
return True, True
def _is_current_password(candidate: str) -> bool:
record = _auth_storage.get_user_and_secret(_admin)
if record is None:
return False
salt, pwd_hash, _jwt_secret, _must_change = record
return _auth_hashing.verify_password(candidate, salt, pwd_hash)
def _apply_change(new_password: str) -> None:
# Same effects as routes/auth.py change_password: rehash, rotate the JWT
# secret, revoke refresh tokens in the SAME transaction.
_auth_storage.update_password(_admin, new_password, revoke_refresh_tokens = True)
changed = prompt_for_password_change(
min_length = _auth_storage.MIN_PASSWORD_LENGTH,
is_current_password = _is_current_password,
apply_change = _apply_change,
out = sys.stderr,
)
return (True, True) if changed else (False, False)
def _apply_supplied_password(password_value: "Optional[str]") -> None:
"""Non-interactively set the INITIAL admin password before the socket binds,
for a direct ``python run.py`` launch (the CLI does this in its own parent).
Value comes from --password / UNSLOTH_STUDIO_PASSWORD / stdin.
Only ever sets the FIRST password: an already-set one is a hard error, an
invalid value fails closed. NOT wrapped in a broad try/except: an auth
storage failure must abort rather than expose the default credential.
"""
from auth import hashing as _auth_hashing
from auth import storage as _auth_storage
from auth.terminal_prompt import SUPPLIED_PASSWORD_ENV, resolve_supplied_password
supplied = resolve_supplied_password(password_value)
# Strip the env var once read so child subprocesses (cloudflared, llama-server,
# code-exec tools) can't inherit the plaintext via /proc/PID/environ. Mirrors
# the CLI. Unconditional: strips a leftover value even when a literal --password won.
os.environ.pop(SUPPLIED_PASSWORD_ENV, None)
if not supplied:
return
_admin = _auth_storage.DEFAULT_ADMIN_USERNAME
_auth_storage.ensure_default_admin()
if not _auth_storage.requires_password_change(_admin):
print(
"Error: a Studio admin password is already set; --password only sets "
"the initial password. Run `unsloth studio reset-password` first.",
file = sys.stderr,
flush = True,
)
sys.exit(1)
def _is_current_password(candidate: str) -> bool:
record = _auth_storage.get_user_and_secret(_admin)
if record is None:
return False
salt, pwd_hash, _jwt_secret, _must_change = record
return _auth_hashing.verify_password(candidate, salt, pwd_hash)
if len(supplied) < _auth_storage.MIN_PASSWORD_LENGTH:
print(
f"Error: password must be at least {_auth_storage.MIN_PASSWORD_LENGTH} "
"characters; not starting.",
file = sys.stderr,
flush = True,
)
sys.exit(1)
if _is_current_password(supplied):
print(
"Error: the new password must differ from the current bootstrap "
"password; not starting.",
file = sys.stderr,
flush = True,
)
sys.exit(1)
_auth_storage.update_password(_admin, supplied, revoke_refresh_tokens = True)
print(f"Password updated for '{_admin}'.", file = sys.stderr, flush = True)
def _apply_cli_tool_policy(enable_tools: "Optional[bool]") -> None:
"""Honor an explicit --enable-tools/--disable-tools; None leaves the policy
unset (tools default on, per-request enable_tools honored). Host is never
@ -1075,9 +1274,10 @@ def run_server(
silent: bool = False,
api_only: bool = False,
llama_parallel_slots: int = 1,
cloudflare: bool = True,
cloudflare: "Optional[bool]" = None,
secure: bool = False,
enable_tools: "Optional[bool]" = None,
password: "Optional[str]" = None,
emit_tauri_port: bool = True,
):
"""
@ -1090,6 +1290,9 @@ def run_server(
silent: Suppress startup messages
api_only: API server only, no frontend (for Tauri desktop app)
llama_parallel_slots: parallel slots for llama-server
cloudflare: opt in to the public Cloudflare HTTPS tunnel for a wildcard
bind. Tri-state: None (unset) and False both mean off; True enables it.
--secure implies it (True) and rejects an explicit False.
enable_tools: explicit --enable-tools/--disable-tools policy; None leaves
the default (tools on, per-request enable_tools honored)
emit_tauri_port: print the machine-readable TAURI_PORT line the desktop
@ -1111,13 +1314,16 @@ def run_server(
initialize_parent_lifetime()
# --secure exposes only the Cloudflare link: force a loopback bind so the raw
# port is never public (even with -H 0.0.0.0), and reject the contradictory combo.
if secure and not cloudflare:
raise SystemExit(
"A secure Cloudflare link is not allowed, use --no-secure which provides a 0.0.0.0 link"
)
# --secure exposes ONLY the Cloudflare link: reject --secure --no-cloudflare,
# then force a loopback bind so the raw port is never public (even -H 0.0.0.0).
# Otherwise keep the tri-state so the banner distinguishes "off by default"
# from an explicit --no-cloudflare.
if secure:
if cloudflare is False:
raise SystemExit(
"--secure requires the Cloudflare tunnel; do not combine it with --no-cloudflare."
)
cloudflare = True
host = "127.0.0.1"
# `unsloth studio run` installs its own resolved policy and passes None here.
@ -1148,14 +1354,35 @@ def run_server(
if secure:
os.environ["UNSLOTH_SECURE"] = "1"
import nest_asyncio
nest_asyncio.apply()
import asyncio
# nest_asyncio is for Colab/IPython, where the main thread already runs a loop
# the blocking waits below would collide with. Apply it only with a loop running
# (a plain CLI start has nothing to nest) and only on Python <= 3.13: on 3.14+
# its global Task patch leaves asyncio.current_task() None (tracking moved into
# C), which also breaks the background uvicorn loop and 500s every request. It
# is archived upstream, so no 3.14 fix is coming; skip it there.
if sys.version_info < (3, 14):
try:
asyncio.get_running_loop()
except RuntimeError:
pass
else:
import nest_asyncio
nest_asyncio.apply()
from threading import Thread, Event
import uvicorn
# `from main import app` below loads torch/unsloth/transformers (~2 min cold,
# silent), so print a flushed heads-up (piped stdout is block-buffered).
if not silent:
print(
"Loading Unsloth Studio, please wait... (this can take a few minutes)",
flush = True,
)
print(" - loading PyTorch, Unsloth and Transformers...", flush = True)
import_started = time.perf_counter()
from main import app, setup_frontend, _IS_COLAB
@ -1164,6 +1391,8 @@ def run_server(
"Imported FastAPI app in %.1fms",
(time.perf_counter() - import_started) * 1000,
)
if not silent:
print(" - Starting server...", flush = True)
from utils.paths import ensure_studio_directories
# Allow local stdio MCP servers on a loopback bind (the user's own machine),
@ -1191,7 +1420,7 @@ def run_server(
print("=" * 50)
if blocker:
pid, name = blocker
print(f"Port {original_port} is already in use by " f"{name} (PID {pid}).")
print(f"Port {original_port} is already in use by {name} (PID {pid}).")
else:
print(f"Port {original_port} is already in use.")
print(f"Unsloth Studio will use port {port} instead.")
@ -1304,6 +1533,44 @@ def run_server(
app.state.trigger_shutdown = _trigger_shutdown
# A supplied --password / UNSLOTH_STUDIO_PASSWORD / stdin sets the initial
# admin password before the gate and socket bind (direct `python run.py`;
# the CLI applies it in its own parent).
_apply_supplied_password(password)
# Never publish with the seeded default password active: prompt first (or
# warn / fail closed headless; see _terminal_password_gate). Runs BEFORE the
# socket binds so a pre-gate listener can't hand out the injected credential.
_pw_proceed, _pw_drop_bootstrap = _terminal_password_gate(
tunnel_will_start = _cloudflare_tunnel_should_start(
cloudflare = cloudflare,
host = host,
secure = secure,
api_only = api_only,
is_colab = _IS_COLAB,
),
host = host,
secure = secure,
api_only = api_only,
frontend_served = bool(frontend_path) and not api_only,
is_colab = _IS_COLAB,
)
if not _pw_proceed:
print(
"Not starting Studio; set a new admin password first, or launch "
"without --secure/--cloudflare.",
file = sys.stderr,
flush = True,
)
sys.exit(1)
if _pw_drop_bootstrap:
# Password just changed (stale) or a public URL is about to serve the
# default credential: don't leak it in the HTML. Lifespan runs AFTER this
# and re-reads the bootstrap password, so the flag (not a plain None)
# makes it skip that re-read.
app.state.suppress_bootstrap_injection = True
app.state.bootstrap_password = None
# Run server in a daemon thread with explicit new_event_loop() +
# run_until_complete() (not asyncio.run) so nest_asyncio's patches don't
# interfere when Colab/IPython already runs a loop on the main thread.
@ -1373,6 +1640,7 @@ def run_server(
is_colab = _IS_COLAB,
)
_cloudflare_requested = _cloudflare_enabled
if _cloudflare_enabled:
try: # best-effort: any failure must not block startup
from cloudflare_tunnel import start_studio_tunnel, stop_studio_tunnel
@ -1460,6 +1728,14 @@ def _build_arg_parser():
default = "127.0.0.1",
help = "Host to bind to (default: 127.0.0.1; use 0.0.0.0 for network/cloud access)",
)
parser.add_argument(
"--password",
default = None,
help = "Set the INITIAL admin password non-interactively (headless), only when "
"none is set yet. Also reads UNSLOTH_STUDIO_PASSWORD, or --password - for stdin. "
"A literal value is visible in the process list. Rotate later via "
"`unsloth studio reset-password`.",
)
parser.add_argument("--port", type = int, default = 8888, help = "Port to bind to")
parser.add_argument(
"--frontend",
@ -1476,11 +1752,13 @@ def _build_arg_parser():
parser.add_argument(
"--cloudflare",
action = argparse.BooleanOptionalAction,
default = True,
help = "Auto-create a free Cloudflare HTTPS tunnel for non-api-only wildcard "
"binds (0.0.0.0 or ::), exposing Studio on a PUBLIC internet URL (default on). "
"Pass --no-cloudflare to disable that Cloudflare URL; it does not change a "
"public wildcard bind. --api-only keeps it off unless paired with --secure.",
default = None,
help = "Expose Studio on a PUBLIC internet URL via a free Cloudflare HTTPS "
"tunnel, for non-api-only wildcard binds (0.0.0.0 or ::). Off by default; "
"pass --cloudflare to enable it (--secure implies it), --no-cloudflare to "
"force it off. It does not change a raw wildcard bind. If the admin "
"password was never changed, Studio asks for a new one in the terminal "
"before publishing the URL.",
)
parser.add_argument(
"--secure",
@ -1488,7 +1766,9 @@ def _build_arg_parser():
default = False,
help = "Expose ONLY a Cloudflare HTTPS link: bind localhost and fail closed "
"if the tunnel can't start. Without it, --no-secure also serves the raw "
"0.0.0.0 port, which is reachable from anywhere on the network",
"0.0.0.0 port, which is reachable from anywhere on the network. If the "
"admin password was never changed, Studio asks for a new one in the "
"terminal before publishing the URL.",
)
# Back-compat: accept --not-secure as a hidden alias for --no-secure.
parser.add_argument(
@ -1550,7 +1830,7 @@ if __name__ == "__main__":
args = parser.parse_args()
if not _PARALLEL_MIN <= args.parallel <= _PARALLEL_MAX:
parser.error(f"--parallel must be between {_PARALLEL_MIN} and {_PARALLEL_MAX}")
if args.secure and not args.cloudflare:
if args.secure and args.cloudflare is False:
parser.error(
"--secure requires the Cloudflare tunnel; do not combine it with --no-cloudflare"
)
@ -1564,6 +1844,7 @@ if __name__ == "__main__":
cloudflare = args.cloudflare,
secure = args.secure,
enable_tools = args.enable_tools,
password = args.password,
)
if args.frontend is not None:
kwargs["frontend_path"] = Path(args.frontend)

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

@ -27,7 +27,7 @@ from utils.paths import (
project_workspaces_root,
studio_db_path,
)
from utils.paths.external_media import is_linux_run_media_path
from utils.paths.external_media import is_linux_run_media_path, is_local_filesystem_root
from utils.paths.sensitive import (
contains_sensitive_path_component as _shared_contains_sensitive_path_component,
)
@ -69,6 +69,25 @@ def _denied_path_prefixes() -> list[str]:
return []
def is_denied_system_path(path: str) -> bool:
"""True if *path* is, or descends from, a denied system directory.
Mirrors the denylist add_scan_folder() enforces at registration so the
browser refuses /etc, /proc, C:\\Windows, etc. even when the allowlist holds
a broad root (a Windows drive root C:\\ or a legacy-registered / root). The
/run carve-out keeps Linux removable-media mounts browseable. Expects an
already-resolved (realpath) path so symlinks cannot escape into a denied subtree.
"""
is_win = platform.system() == "Windows"
check = os.path.normcase(path) if is_win else path
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
return True
return False
def _contains_sensitive_path_component(path: str) -> bool:
return _shared_contains_sensitive_path_component(path)
@ -931,6 +950,12 @@ 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")
# Reject a local filesystem root ("/", or a bare Windows drive root "C:\\"):
# registering one seeds the browse allowlist with a root above denied system
# dirs. A UNC share root (\\server\share) has none under it and was
# registerable before this guard, so it stays allowed. Mirrors scan_folders.py.
if is_local_filesystem_root(normalized):
raise ValueError("The filesystem root cannot be registered")
if _contains_sensitive_path_component(normalized):
raise ValueError("Credential or configuration directories are not allowed")

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

@ -1631,6 +1631,48 @@ class TestAnthropicMessagesToolRouting:
assert entry["status"] == "cancelled"
assert monitor.active_count() == 0
@staticmethod
def _sse_blob(chunks):
# StreamingResponse may hand back str or already-encoded bytes.
return "".join(c.decode() if isinstance(c, (bytes, bytearray)) else c for c in chunks)
def test_plain_streaming_unclassified_error_emits_error_event(self, monkeypatch):
# An unclassified mid-stream failure must surface as an SSE `error` event
# and stop, not a message_stop that masks a truncated turn as clean.
def _gen_boom(**_kwargs):
yield "partial"
raise RuntimeError("llama-server crashed mid-decode")
_mock_backend(monkeypatch, generate_chat_completion = _gen_boom)
payload = _basic_payload(stream = True)
response = _drive(anthropic_messages(payload, request = self._Request(), current_subject = "t"))
blob = self._sse_blob(self._consume_response(response))
assert "event: error" in blob
assert '"type": "error"' in blob
assert "event: message_stop" not in blob
def test_tool_streaming_unclassified_error_emits_error_event(self, monkeypatch):
# Same guarantee on the tool-calling stream path.
def _gen_tools_boom(**_kwargs):
yield {"type": "content", "text": "partial"}
raise RuntimeError("llama-server crashed mid-decode")
_mock_backend(monkeypatch, generate_chat_completion_with_tools = _gen_tools_boom)
payload = _basic_payload(
stream = True,
enable_tools = True,
tools = [{"type": "web_search_20250305", "name": "web_search"}],
)
response = _drive(anthropic_messages(payload, request = self._Request(), current_subject = "t"))
blob = self._sse_blob(self._consume_response(response))
assert "event: error" in blob
assert '"type": "error"' in blob
assert "event: message_stop" not in blob
def test_mixed_server_and_client_tools_rejected_with_400(self, monkeypatch):
_mock_backend(monkeypatch)
payload = _basic_payload(
@ -1739,7 +1781,9 @@ class TestAnthropicMessagesToolRouting:
assert backend.calls[0][0] == "plain"
def test_server_tool_alias_enters_tool_path_when_policy_unset(self, monkeypatch):
# Mirror of the previous test for the default (None) policy.
# Mirror of the previous test for the default (None) policy. An omitted
# permission_mode still runs here because web_search is a safe server tool
# (only a selected terminal/python would require the missing gate).
backend = _mock_backend(monkeypatch)
payload = _basic_payload(
tools = [{"type": "web_search_20250305", "name": "web_search"}],
@ -1761,6 +1805,126 @@ class TestAnthropicMessagesToolRouting:
assert "confirm_tool_calls is not supported" in exc.value.detail["error"]["message"]
assert backend.calls == []
def test_permission_mode_gating_for_server_tools(self, monkeypatch):
# ask is a request for a per-call pause this channel cannot honor, so it is
# always rejected, even for a safe-only server tool (web_search).
safe_tools = [{"type": "web_search_20250305", "name": "web_search"}]
backend = _mock_backend(monkeypatch)
payload = _basic_payload(tools = safe_tools, permission_mode = "ask")
with pytest.raises(HTTPException) as exc:
_drive(anthropic_messages(payload, request = None, current_subject = "t"))
assert exc.value.status_code == 400
assert "no confirmation channel" in exc.value.detail["error"]["message"]
assert backend.calls == []
# auto only gates unsafe calls, so a safe-only selection runs (nothing to
# gate), like the omitted default. Both keep existing callers working.
for extra in ({"permission_mode": "auto"}, {}):
backend = _mock_backend(monkeypatch)
payload = _basic_payload(tools = safe_tools, **extra)
_drive(anthropic_messages(payload, request = None, current_subject = "t"))
assert backend.calls[0][0] == "tools"
# But auto or an omitted mode that would run a local tool (terminal/python,
# via a bare Anthropic tool type or enabled_tools) is rejected, since that
# tool could need the gate this channel lacks.
for local_payload in (
_basic_payload(tools = [{"type": "terminal", "name": "terminal"}]),
_basic_payload(
tools = [{"type": "terminal", "name": "terminal"}], permission_mode = "auto"
),
_basic_payload(tools = safe_tools, enable_tools = True, enabled_tools = ["python"]),
):
backend = _mock_backend(monkeypatch)
with pytest.raises(HTTPException) as exc:
_drive(anthropic_messages(local_payload, request = None, current_subject = "t"))
assert exc.value.status_code == 400
assert "terminal" in exc.value.detail["error"]["message"]
assert backend.calls == []
# off, full, and a legacy confirm_tool_calls=False opt-out all run, even
# with a local tool selected. The explicit opt-out wins over the mode
# (mirrors _permission_mode_confirm and the GGUF path), so it runs even
# under ask, which otherwise always rejects.
for extra in (
{"tools": safe_tools, "permission_mode": "off"},
{"tools": safe_tools, "permission_mode": "full"},
{"tools": safe_tools, "enabled_tools": ["python"], "confirm_tool_calls": False},
{"tools": safe_tools, "permission_mode": "ask", "confirm_tool_calls": False},
{
"tools": [{"type": "terminal", "name": "terminal"}],
"permission_mode": "ask",
"confirm_tool_calls": False,
},
):
backend = _mock_backend(monkeypatch)
payload = _basic_payload(**extra)
_drive(anthropic_messages(payload, request = None, current_subject = "t"))
assert backend.calls[0][0] == "tools"
def test_render_html_gated_for_server_tools(self, monkeypatch):
# render_html is no longer unconditionally safe: a networked canvas prompts
# in auto and this channel cannot present that gate, so selecting it under
# ask/auto/omitted rejects like terminal/python; off/full (and an explicit
# confirm opt-out) run it.
rh = {"enable_tools": True, "enabled_tools": ["render_html"]}
for mode in ("ask", "auto", None):
backend = _mock_backend(monkeypatch)
fields = dict(rh)
if mode is not None:
fields["permission_mode"] = mode
payload = _basic_payload(**fields)
with pytest.raises(HTTPException) as exc:
_drive(anthropic_messages(payload, request = None, current_subject = "t"))
assert exc.value.status_code == 400
assert "no confirmation channel" in exc.value.detail["error"]["message"]
assert backend.calls == []
for extra in (
{"permission_mode": "off"},
{"permission_mode": "full"},
{"confirm_tool_calls": False},
):
backend = _mock_backend(monkeypatch)
payload = _basic_payload(**{**rh, **extra})
_drive(anthropic_messages(payload, request = None, current_subject = "t"))
assert backend.calls[0][0] == "tools"
def test_permission_mode_rejected_before_auto_switch(self, monkeypatch):
# The unsupported-mode rejection must run before _maybe_auto_switch_model,
# so an invalid confirm-gated request never evicts the resident model
# (mirrors the pre-switch malformed- and mixed-tool guards).
import routes.inference as inf_mod
switch_calls = []
async def _rec_switch(*_args, **_kwargs):
switch_calls.append(1)
monkeypatch.setattr(inf_mod, "_maybe_auto_switch_model", _rec_switch)
safe_tools = [{"type": "web_search_20250305", "name": "web_search"}]
local_tools = [{"type": "terminal", "name": "terminal"}]
# ask (any server tool), auto with a local tool, and an omitted mode
# selecting a local tool are all rejected up front, before the switch runs.
for payload in (
_basic_payload(tools = safe_tools, permission_mode = "ask"),
_basic_payload(tools = local_tools, permission_mode = "auto"),
_basic_payload(tools = local_tools),
):
switch_calls.clear()
_mock_backend(monkeypatch)
with pytest.raises(HTTPException) as exc:
_drive(anthropic_messages(payload, request = None, current_subject = "t"))
assert exc.value.status_code == 400
assert switch_calls == [], "rejection must precede the auto-switch"
# A supported request (off) still reaches the switch and runs the loop.
switch_calls.clear()
_mock_backend(monkeypatch)
payload = _basic_payload(tools = safe_tools, permission_mode = "off")
_drive(anthropic_messages(payload, request = None, current_subject = "t"))
assert switch_calls == [1]
def test_per_request_enable_tools_false_blocks_server_tool_alias(self, monkeypatch):
backend = _mock_backend(monkeypatch)
payload = _basic_payload(
@ -1770,3 +1934,523 @@ class TestAnthropicMessagesToolRouting:
_drive(anthropic_messages(payload, request = None, current_subject = "t"))
assert backend.calls[0][0] == "plain"
def test_resumed_session_thinking_and_null_content_do_not_400():
# A resumed session replays assistant turns with `thinking` (and sometimes null)
# content. Those must be accepted (thinking dropped by the converter), not 400ed.
from pydantic import ValidationError
req = AnthropicMessagesRequest(
model = "x",
max_tokens = 16,
messages = [
{"role": "user", "content": "hi"},
{
"role": "assistant",
"content": [
{"type": "thinking", "thinking": "secret reasoning", "signature": "s"},
{"type": "text", "text": "the answer"},
{"type": "tool_use", "id": "t1", "name": "f", "input": {}},
],
},
{"role": "assistant", "content": None}, # tool-only turn serialized as null
],
)
# Known blocks still parse as their typed models; only the unknown one is loose.
assert type(req.messages[1].content[0]).__name__ == "AnthropicUnknownBlock"
assert type(req.messages[1].content[1]).__name__ == "AnthropicTextBlock"
assert req.messages[2].content == "" # null coerced
openai = anthropic_messages_to_openai([m.model_dump() for m in req.messages])
assistant = next(m for m in openai if m["role"] == "assistant" and m.get("content"))
assert assistant["content"] == "the answer"
assert "secret reasoning" not in json.dumps(openai) # thinking never forwarded
# A malformed KNOWN block still fails cleanly instead of being swallowed.
with pytest.raises(ValidationError):
AnthropicMessagesRequest(
model = "x",
max_tokens = 16,
messages = [{"role": "assistant", "content": [{"type": "tool_use", "name": "f"}]}],
)
def test_user_null_content_rejected():
# The null->"" leniency is assistant-only; a null user content must be rejected
# at the boundary, not coerced into an empty prompt and forwarded to the model.
from pydantic import ValidationError
with pytest.raises(ValidationError):
AnthropicMessagesRequest(
model = "x",
max_tokens = 16,
messages = [{"role": "user", "content": None}],
)
def test_user_unknown_block_rejected_not_silently_dropped():
# The converter skips user blocks it cannot translate, so a user turn whose only
# block is unknown would validate yet forward no content. Reject at the boundary
# to avoid that silent data loss (the assistant fallback is unaffected).
from pydantic import ValidationError
with pytest.raises(ValidationError):
AnthropicMessagesRequest(
model = "x",
max_tokens = 16,
messages = [
{"role": "user", "content": [{"type": "document", "source": {}}]},
],
)
def test_user_translatable_blocks_still_accepted():
# text / image / tool_result are translatable, so a real user message built from
# them must still pass; the unknown-block guard only trips on other types.
req = AnthropicMessagesRequest(
model = "x",
max_tokens = 16,
messages = [
{
"role": "user",
"content": [
{"type": "text", "text": "What is this?"},
{
"type": "image",
"source": {"type": "base64", "media_type": "image/png", "data": "AA"},
},
{"type": "tool_result", "tool_use_id": "t1", "content": "ok"},
],
}
],
)
assert [type(b).__name__ for b in req.messages[0].content] == [
"AnthropicTextBlock",
"AnthropicImageBlock",
"AnthropicToolResultBlock",
]
openai = anthropic_messages_to_openai([m.model_dump() for m in req.messages])
assert any(m["role"] == "tool" and m["tool_call_id"] == "t1" for m in openai)
def test_user_malformed_known_block_still_rejected():
# The guard only allow-lists a user block's *type*; the union still validates its
# shape, so a known-but-malformed block (tool_result without tool_use_id) fails.
from pydantic import ValidationError
with pytest.raises(ValidationError):
AnthropicMessagesRequest(
model = "x",
max_tokens = 16,
messages = [
{"role": "user", "content": [{"type": "tool_result", "content": "x"}]},
],
)
def test_user_content_block_non_string_type_rejected_cleanly():
# A user block whose `type` is a non-string (unhashable list / dict, or a stray
# int) must fail as a clean validation error, not raise TypeError from the
# frozenset membership test and escape as a 500.
from pydantic import ValidationError
for bad_type in ([], {}, 5):
with pytest.raises(ValidationError):
AnthropicMessagesRequest(
model = "x",
max_tokens = 16,
messages = [{"role": "user", "content": [{"type": bad_type}]}],
)
def test_assistant_missing_content_key_still_rejected():
# The null -> "" leniency is only for an EXPLICIT null. An assistant message that
# omits content entirely stays malformed and must fail required-field validation.
from pydantic import ValidationError
with pytest.raises(ValidationError):
AnthropicMessagesRequest(
model = "x",
max_tokens = 16,
messages = [{"role": "assistant"}],
)
# An explicit null is still accepted and coerced (regression guard).
req = AnthropicMessagesRequest(
model = "x",
max_tokens = 16,
messages = [
{"role": "user", "content": "hi"},
{"role": "assistant", "content": None},
],
)
assert req.messages[1].content == ""
def test_resumed_null_assistant_between_users_coalesced_on_messages_route(monkeypatch):
# user -> assistant(null) -> user is now accepted: the null assistant turn coerces
# to "" and is dropped. The route must then coalesce the two remaining user turns
# so a strict GGUF chat template does not 400 on non-alternating roles.
backend = _mock_backend(monkeypatch, context_length = 2048)
class _Req:
state = SimpleNamespace()
url = SimpleNamespace(path = "/v1/messages")
method = "POST"
async def is_disconnected(self):
return False
payload = AnthropicMessagesRequest(
model = "x",
max_tokens = 16,
messages = [
{"role": "user", "content": "first question"},
{"role": "assistant", "content": None},
{"role": "user", "content": "please continue"},
],
)
response = _drive(anthropic_messages(payload, request = _Req(), current_subject = "t"))
assert response.status_code == 200
[(_path, kwargs)] = backend.calls
user_turns = [m for m in kwargs["messages"] if m.get("role") == "user"]
assert len(user_turns) == 1 # the two user turns were merged, not left adjacent
merged = user_turns[0]["content"]
if isinstance(merged, list):
merged = " ".join(p.get("text", "") for p in merged if isinstance(p, dict))
assert "first question" in merged and "please continue" in merged
def test_disable_parallel_tool_use_forwards_heartbeats_while_dropping():
"""Heartbeats from a parallel-disabled, dropped tool call must still reach
the client as SSE keepalives: the dropped call runs server-side and the
stall keepalive never fires while the generator keeps producing events, so
swallowing them recreates the silent window keepalives exist to prevent."""
import threading as _threading
from routes.inference import (
_OPENAI_PASSTHROUGH_SSE_KEEPALIVE,
_anthropic_tool_stream,
)
def run_gen():
def gen():
yield {
"type": "tool_start",
"tool_name": "python",
"tool_call_id": "call_0",
"arguments": {},
}
yield {"type": "heartbeat"}
yield {
"type": "tool_end",
"tool_name": "python",
"tool_call_id": "call_0",
"result": "r1",
}
# Second call: dropped by disable_parallel_tool_use, still executed
# server-side (heartbeats + live output).
yield {
"type": "tool_start",
"tool_name": "python",
"tool_call_id": "call_1",
"arguments": {},
}
yield {"type": "heartbeat"}
yield {
"type": "tool_output",
"tool_name": "python",
"tool_call_id": "call_1",
"text": "x",
}
yield {"type": "heartbeat"}
yield {
"type": "tool_end",
"tool_name": "python",
"tool_call_id": "call_1",
"result": "r2",
}
yield {"type": "content", "text": "final answer"}
return gen()
async def _drive():
async def _is_disconnected():
return False
request = SimpleNamespace(is_disconnected = _is_disconnected)
resp = await _anthropic_tool_stream(
request,
_threading.Event(),
run_gen,
"msg_hb",
"m",
disable_parallel_tool_use = True,
)
return [chunk async for chunk in resp.body_iterator]
chunks = asyncio.run(_drive())
keepalives = [c for c in chunks if c == _OPENAI_PASSTHROUGH_SSE_KEEPALIVE]
# One heartbeat inside the kept call, two inside the dropped window.
assert len(keepalives) >= 3
# The dropped call must not surface as a second tool_use block.
tool_use_starts = [c for c in chunks if "content_block_start" in c and '"tool_use"' in c]
assert len(tool_use_starts) == 1
def test_dropped_tool_output_events_emit_rate_limited_keepalives(monkeypatch):
"""A chatty tool streaming tool_output/tool_args with no heartbeats keeps the
generator busy (stall keepalive never fires); the Anthropic path can't
translate those events and drops them. Dropping silently would let an idle
proxy kill the stream, so the drop branch emits a rate-limited keepalive."""
import threading as _threading
import routes.inference as inf_mod
from routes.inference import (
_OPENAI_PASSTHROUGH_SSE_KEEPALIVE,
_anthropic_tool_stream,
)
# Deterministic clock: only the drop-branch keepalive uses time.monotonic
# here, so jumping past the stall window per call makes each dropped event
# cross the rate-limit threshold. asyncio.wait uses the loop clock and
# next(gen) returns promptly, so the outer stall keepalive never fires --
# every keepalive here is from the drop branch.
_real_time = inf_mod.time
_tick = {"v": 0.0}
def _fast_monotonic():
_tick["v"] += 100.0
return _tick["v"]
fake_time = SimpleNamespace(
monotonic = _fast_monotonic,
sleep = _real_time.sleep,
time = _real_time.time,
perf_counter = _real_time.perf_counter,
)
monkeypatch.setattr(inf_mod, "time", fake_time)
n_output = 4
def run_gen():
def gen():
yield {
"type": "tool_start",
"tool_name": "python",
"tool_call_id": "call_0",
"arguments": {},
}
# Chatty streamed stdout, no heartbeats.
for i in range(n_output):
yield {
"type": "tool_output",
"tool_name": "python",
"tool_call_id": "call_0",
"text": f"line {i}\n",
}
yield {
"type": "tool_end",
"tool_name": "python",
"tool_call_id": "call_0",
"result": "done",
}
yield {"type": "content", "text": "final answer"}
return gen()
async def _drive():
async def _is_disconnected():
return False
request = SimpleNamespace(is_disconnected = _is_disconnected)
resp = await _anthropic_tool_stream(
request,
_threading.Event(),
run_gen,
"msg_drop_ka",
"m",
)
return [chunk async for chunk in resp.body_iterator]
chunks = asyncio.run(_drive())
keepalives = [c for c in chunks if c == _OPENAI_PASSTHROUGH_SSE_KEEPALIVE]
assert len(keepalives) == n_output
# Final answer still reaches the client (drop is transport-only).
assert any("final answer" in c for c in chunks)
def test_parallel_disabled_dropped_call_output_emits_rate_limited_keepalives(monkeypatch):
"""Under disable_parallel_tool_use a chatty second call is dropped whole
(drop_until_tool_end). Its tool_output/tool_args events must still emit
rate-limited keepalives: the drop window can last minutes with no heartbeats
and no stall keepalive, so swallowing them silently would let an idle proxy
kill the stream. The keepalive branch runs before the drop skip."""
import threading as _threading
import routes.inference as inf_mod
from routes.inference import (
_OPENAI_PASSTHROUGH_SSE_KEEPALIVE,
_anthropic_tool_stream,
)
# Deterministic clock: jumps past the stall window per call (see sibling test).
_real_time = inf_mod.time
_tick = {"v": 0.0}
def _fast_monotonic():
_tick["v"] += 100.0
return _tick["v"]
fake_time = SimpleNamespace(
monotonic = _fast_monotonic,
sleep = _real_time.sleep,
time = _real_time.time,
perf_counter = _real_time.perf_counter,
)
monkeypatch.setattr(inf_mod, "time", fake_time)
n_output = 4
def run_gen():
def gen():
# First (kept) call.
yield {
"type": "tool_start",
"tool_name": "python",
"tool_call_id": "call_0",
"arguments": {},
}
yield {
"type": "tool_end",
"tool_name": "python",
"tool_call_id": "call_0",
"result": "r1",
}
# Second call: dropped whole by disable_parallel_tool_use but still
# executed server-side, streaming chatty stdout with no heartbeats.
yield {
"type": "tool_start",
"tool_name": "python",
"tool_call_id": "call_1",
"arguments": {},
}
for i in range(n_output):
yield {
"type": "tool_output",
"tool_name": "python",
"tool_call_id": "call_1",
"text": f"line {i}\n",
}
yield {
"type": "tool_end",
"tool_name": "python",
"tool_call_id": "call_1",
"result": "r2",
}
yield {"type": "content", "text": "final answer"}
return gen()
async def _drive():
async def _is_disconnected():
return False
request = SimpleNamespace(is_disconnected = _is_disconnected)
resp = await _anthropic_tool_stream(
request,
_threading.Event(),
run_gen,
"msg_drop_ka2",
"m",
disable_parallel_tool_use = True,
)
return [chunk async for chunk in resp.body_iterator]
chunks = asyncio.run(_drive())
keepalives = [c for c in chunks if c == _OPENAI_PASSTHROUGH_SSE_KEEPALIVE]
assert len(keepalives) == n_output
# The dropped call must not surface as a second tool_use block.
tool_use_starts = [c for c in chunks if "content_block_start" in c and '"tool_use"' in c]
assert len(tool_use_starts) == 1
assert any("final answer" in c for c in chunks)
def test_plain_stream_emits_keepalive_during_prompt_stall(monkeypatch):
"""No-tool Anthropic stream must emit SSE keepalives while a long prompt
prefill blocks next(gen), matching the tool stream (finding 5). The old
single unbounded to_thread(next, ...) could sit silent past a proxy idle cap."""
import threading as _threading
import time as _time
from routes import inference as inf_mod
from routes.inference import _OPENAI_PASSTHROUGH_SSE_KEEPALIVE, _anthropic_plain_stream
monkeypatch.setattr(inf_mod, "_LOCAL_TOOL_STREAM_STALL_KEEPALIVE_S", 0.05)
def run_gen():
def gen():
_time.sleep(0.24) # stall past several shortened keepalive windows
yield "hello world"
return gen()
async def _drive():
async def _is_disconnected():
return False
request = SimpleNamespace(is_disconnected = _is_disconnected)
resp = await _anthropic_plain_stream(
request, _threading.Event(), run_gen, "msg_plain_ka", "m"
)
return [chunk async for chunk in resp.body_iterator]
chunks = asyncio.run(_drive())
keepalives = [c for c in chunks if c == _OPENAI_PASSTHROUGH_SSE_KEEPALIVE]
assert len(keepalives) >= 2
assert any("hello world" in c for c in chunks)
def test_plain_stream_closes_generator_on_disconnect():
"""On disconnect the no-tool teardown must drain any pending worker and close
the generator (finding 6). The old finally only stopped the disconnect
watcher, leaking the generator. A fake generator records close() so the
teardown is asserted deterministically, not via GC."""
import threading as _threading
from routes.inference import _anthropic_plain_stream
closed = _threading.Event()
class _FakeGen:
def __init__(self):
self._items = iter(["tok0", "tok1", "tok2", "tok3"])
def __next__(self):
return next(self._items)
def close(self):
closed.set()
def run_gen():
return _FakeGen()
state = {"disconnected": False}
async def _drive():
async def _is_disconnected():
return state["disconnected"]
request = SimpleNamespace(is_disconnected = _is_disconnected)
resp = await _anthropic_plain_stream(
request, _threading.Event(), run_gen, "msg_plain_close", "m"
)
out = []
async for chunk in resp.body_iterator:
out.append(chunk)
if "tok0" in chunk:
# Client drops after the first token; the next loop turn tears down.
state["disconnected"] = True
return out
asyncio.run(_drive())
assert closed.is_set()

371
studio/backend/tests/test_browse_denylist.py Normal file
View file

@ -0,0 +1,371 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""System-directory denylist enforcement for the folder browser.
Once the allowlist can hold a whole Windows drive root (C:\\) or a legacy /
root, the browse endpoints must re-apply the ``_denied_path_prefixes()`` policy
``add_scan_folder`` enforces, so /etc, /proc, C:\\Windows, C:\\Program Files stay
unbrowseable even under an allowlisted root. Windows/macOS branches run on this
POSIX host by AST-extracting the pure helper with ``ntpath`` / a mocked ``platform``.
"""
from __future__ import annotations
import ast
import ntpath
import os
import posixpath
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.external_media import is_local_filesystem_root
_BACKEND_ROOT = Path(__file__).resolve().parent.parent
class _HTTPException(Exception):
def __init__(self, status_code: int, detail: str):
super().__init__(detail)
self.status_code = status_code
self.detail = detail
def _extract_is_denied_windows():
"""is_denied_system_path (+ _denied_path_prefixes) from studio_db.py under Windows semantics (ntpath) on a POSIX host."""
src = (_BACKEND_ROOT / "storage" / "studio_db.py").read_text(encoding = "utf-8")
tree = ast.parse(src)
funcs = [
n
for n in tree.body
if isinstance(n, ast.FunctionDef)
and n.name in {"_denied_path_prefixes", "is_denied_system_path"}
]
module = ast.Module(body = funcs, type_ignores = [])
ast.fix_missing_locations(module)
win_os = SimpleNamespace(
sep = "\\",
environ = {
"SystemRoot": r"C:\Windows",
"ProgramFiles": r"C:\Program Files",
"ProgramFiles(x86)": r"C:\Program Files (x86)",
},
path = SimpleNamespace(normcase = ntpath.normcase),
)
ns = {
"os": win_os,
"platform": SimpleNamespace(system = lambda: "Windows"),
# /run has no Windows analog, so the carve-out is never reached.
"is_linux_run_media_path": lambda _p: False,
}
exec(compile(module, "<extracted studio_db.py>", "exec"), ns)
return ns["is_denied_system_path"]
# is_denied_system_path -- Linux (real helper, this host)
@pytest.mark.parametrize(
"path",
[
"/etc",
"/etc/ssl/private",
"/proc",
"/proc/1",
"/sys",
"/dev",
"/boot",
"/run",
"/run/systemd/private",
"/run/media",
"/run/media/dspofu",
],
)
def test_is_denied_system_path_linux_denies_system_dirs(monkeypatch, path):
monkeypatch.setattr(studio_db.platform, "system", lambda: "Linux")
assert studio_db.is_denied_system_path(path) is True
@pytest.mark.parametrize(
"path",
["/run/media/dspofu/nvmeB", "/run/media/dspofu/nvmeB/models"],
)
def test_is_denied_system_path_linux_allows_run_media_mounts(monkeypatch, path):
# The /run/media/<user>/<volume> carve-out keeps removable media browseable.
monkeypatch.setattr(studio_db.platform, "system", lambda: "Linux")
assert studio_db.is_denied_system_path(path) is False
@pytest.mark.parametrize(
"path",
["/etc-backup", "/etcetera", "/home/u/models", "/mnt/data", "/devices", "/", "/opt/models"],
)
def test_is_denied_system_path_linux_allows_non_system(monkeypatch, path):
monkeypatch.setattr(studio_db.platform, "system", lambda: "Linux")
assert studio_db.is_denied_system_path(path) is False
def test_legacy_and_hub_denylist_agree(monkeypatch):
monkeypatch.setattr(studio_db.platform, "system", lambda: "Linux")
monkeypatch.setattr(scan_folders.platform, "system", lambda: "Linux")
for p in ["/etc", "/proc/1", "/home/u", "/boot", "/opt/x"]:
assert studio_db.is_denied_system_path(p) == scan_folders.is_denied_system_path(p)
# is_denied_system_path -- Windows (ntpath-backed), case-insensitive + collisions
@pytest.mark.parametrize(
"path",
[
r"C:\Windows",
r"C:\Windows\System32",
r"c:\windows",
r"C:\WINDOWS\Temp",
r"C:\Program Files",
r"C:\Program Files\x",
r"C:\Program Files (x86)\y",
r"c:\program files",
],
)
def test_is_denied_system_path_windows_denies_system_dirs(path):
is_denied = _extract_is_denied_windows()
assert is_denied(path) is True
@pytest.mark.parametrize(
"path",
[
r"C:\Models",
r"D:\models",
r"C:\WindowsApps",
r"C:\ProgramData",
r"C:\Program Files Extra",
r"E:\gguf",
r"C:\Users\me\models",
],
)
def test_is_denied_system_path_windows_allows_non_system(path):
is_denied = _extract_is_denied_windows()
assert is_denied(path) is False
# _resolve_browse_target -- real-FS integration (legacy browser)
def _extract_resolver():
"""Extract the legacy browse resolver; its inline imports use the real storage.studio_db policy."""
src = (_BACKEND_ROOT / "routes" / "models.py").read_text(encoding = "utf-8")
tree = ast.parse(src)
names = {
"_is_path_inside_allowlist",
"_normalize_browse_request_path",
"_browse_relative_parts",
"_match_browse_child",
"_resolve_browse_target",
}
funcs = [n for n in tree.body if isinstance(n, ast.FunctionDef) and n.name in names]
module = ast.Module(body = funcs, type_ignores = [])
ast.fix_missing_locations(module)
ns = {
"os": os,
"Path": Path,
"Optional": Optional,
"HTTPException": _HTTPException,
"logger": SimpleNamespace(warning = lambda *a, **k: None, debug = lambda *a, **k: None),
}
exec(compile(module, "<extracted routes/models.py>", "exec"), ns)
return ns["_resolve_browse_target"]
def test_resolve_browse_target_blocks_etc_via_root():
# Registering "/" must not make /etc browsable (Codex #3 regression guard).
resolve = _extract_resolver()
with pytest.raises(_HTTPException) as exc:
resolve("/etc", [Path("/")])
assert exc.value.status_code == 403
def test_resolve_browse_target_blocks_stale_denied_root(tmp_path, monkeypatch):
# A stale scan-folder row pointing at a denied dir is refused by the
# browse-time denylist even though it is its own allowlist root. A tmp-based
# denied prefix (+ Linux compare) keeps the assertion OS-agnostic: on macOS
# tmp lives under the already-denied /private/var, masking the message.
denied = (tmp_path / "sysfake").resolve()
denied.mkdir()
monkeypatch.setattr(studio_db.platform, "system", lambda: "Linux")
monkeypatch.setattr(studio_db, "_denied_path_prefixes", lambda: [str(denied)])
resolve = _extract_resolver()
with pytest.raises(_HTTPException) as exc:
resolve(str(denied), [denied])
assert exc.value.status_code == 403
assert "System directories" in exc.value.detail
def test_resolve_browse_target_allows_root_itself():
resolve = _extract_resolver()
assert resolve("/", [Path("/")]) == Path("/")
def test_resolve_browse_target_allows_legit_nested_dir(tmp_path, monkeypatch):
# Force the Linux denylist so the macOS temp location (under the denied
# /private/var) doesn't reject the tmp fixture; a normal nested dir must not be over-blocked.
monkeypatch.setattr(studio_db.platform, "system", lambda: "Linux")
resolve = _extract_resolver()
base = tmp_path / "allowed"
sub = base / "models" / "gguf"
sub.mkdir(parents = True)
assert resolve(str(sub), [base]) == sub.resolve()
def test_resolve_browse_target_symlink_escape_blocked(tmp_path):
resolve = _extract_resolver()
base = tmp_path / "allowed"
base.mkdir()
link = base / "escape"
try:
link.symlink_to("/etc", target_is_directory = True)
except OSError:
pytest.skip("symlinks unsupported on this host")
with pytest.raises(_HTTPException) as exc:
resolve(str(link), [base])
assert exc.value.status_code == 403
# _is_path_inside_allowlist -- bare POSIX root parity (legacy == hub)
def _extract_is_inside(rel_parts, *, os_module = os):
"""Extract a standalone _is_path_inside_allowlist (os/Path only) so both browsers' copies compare without importing their heavy modules."""
src = _BACKEND_ROOT.joinpath(*rel_parts).read_text(encoding = "utf-8")
tree = ast.parse(src)
funcs = [
n
for n in tree.body
if isinstance(n, ast.FunctionDef) and n.name == "_is_path_inside_allowlist"
]
module = ast.Module(body = funcs, type_ignores = [])
ast.fix_missing_locations(module)
ns = {"os": os_module, "Path": Path}
exec(compile(module, f"<extracted {'/'.join(rel_parts)}>", "exec"), ns)
return ns["_is_path_inside_allowlist"]
# ntpath semantics with a no-FS realpath, so UNC containment can be driven on a
# POSIX CI (the real realpath cannot resolve \\server\share off Windows).
_WIN_OS = SimpleNamespace(
sep = ntpath.sep,
path = SimpleNamespace(
realpath = lambda p: ntpath.normpath(str(p)),
normcase = ntpath.normcase,
splitdrive = ntpath.splitdrive,
dirname = ntpath.dirname,
commonpath = ntpath.commonpath,
),
)
def test_legacy_and_hub_allowlist_agree_on_posix_root():
# A bare "/" allowlist entry must authorize only "/" itself in BOTH
# browsers, never descend into /var, /root, /home (which the denylist does
# not cover). Guards the hub browser against authorizing every absolute path.
legacy = _extract_is_inside(["routes", "models.py"])
hub = _extract_is_inside(["hub", "services", "models", "folder_browser.py"])
roots = [Path("/")]
for tgt in ["/var", "/root", "/home", "/usr", "/opt", "/etc"]:
assert legacy(Path(tgt), roots) is False
assert hub(Path(tgt), roots) is False
# "/" itself stays browseable; only its descendants are withheld.
assert legacy(Path("/"), roots) is True
assert hub(Path("/"), roots) is True
def test_hub_allowlist_authorizes_normal_nested_dir(tmp_path):
# The bare-root special case must not over-block a normal allowlist root's descendants.
hub = _extract_is_inside(["hub", "services", "models", "folder_browser.py"])
base = tmp_path / "allowed"
sub = base / "models" / "gguf"
sub.mkdir(parents = True)
assert hub(sub, [base]) is True
assert hub(base, [base]) is True
# add_scan_folder -- filesystem-root rejection parity (legacy == hub)
def test_legacy_add_scan_folder_rejects_filesystem_root(monkeypatch):
monkeypatch.setattr(studio_db.platform, "system", lambda: "Linux")
with pytest.raises(ValueError, match = "filesystem root"):
studio_db.add_scan_folder("/")
def test_hub_add_scan_folder_rejects_filesystem_root(monkeypatch):
monkeypatch.setattr(scan_folders.platform, "system", lambda: "Linux")
with pytest.raises(ValueError, match = "filesystem root"):
scan_folders.add_scan_folder("/")
# is_local_filesystem_root: reject "/" and "C:\\" (roots above denied system dirs),
# but NOT a UNC share root -- registering \\server\share was allowed before this
# guard and has no system dirs under it. _pathmod drives Windows semantics on POSIX CI.
@pytest.mark.parametrize(
"path, pathmod, expected",
[
# Local filesystem roots -> rejected (True).
("/", posixpath, True),
("C:\\", ntpath, True),
("c:\\", ntpath, True),
("D:\\", ntpath, True),
# UNC share roots -> NOT a local root, stay registerable (False).
(r"\\server\share", ntpath, False),
(r"\\nas\models", ntpath, False),
("//server/share", ntpath, False),
# Device / extended-length volume roots -> still local roots (rejected),
# so neither \\?\C:\ nor a drive-letter-less \\?\Volume{GUID}\ can slip
# past the guard as if it were a share root.
(r"\\?\C:" + "\\", ntpath, True),
(r"\\.\C:" + "\\", ntpath, True),
(r"\\?\C:", ntpath, True),
(r"\\.\C:", ntpath, True),
(r"\\?\Volume{2f8e6d31-0000-0000-0000-100000000000}" + "\\", ntpath, True),
(r"\\.\Volume{2f8e6d31-0000-0000-0000-100000000000}", ntpath, True),
# Device-namespace UNC share root -> stays registerable (False).
(r"\\?\UNC\server\share", ntpath, False),
# Non-root paths (incl. deep device / extended-length) -> not a root (False).
("C:\\Models", ntpath, False),
(r"\\server\share\models", ntpath, False),
(r"\\?\C:\Users\me\models", ntpath, False),
(r"\\?\Volume{2f8e6d31-0000-0000-0000-100000000000}\models", ntpath, False),
("/home/user", posixpath, False),
],
)
def test_is_local_filesystem_root(path, pathmod, expected):
assert is_local_filesystem_root(path, _pathmod = pathmod) is expected
def test_both_guards_use_the_shared_local_root_helper():
# Register-root parity: both browsers reject the same roots via one helper, so a
# UNC-share exemption can never drift between the legacy and hub code paths.
legacy_src = (_BACKEND_ROOT / "storage" / "studio_db.py").read_text(encoding = "utf-8")
hub_src = (_BACKEND_ROOT / "hub" / "storage" / "scan_folders.py").read_text(encoding = "utf-8")
assert "is_local_filesystem_root(normalized)" in legacy_src
assert "is_local_filesystem_root(normalized)" in hub_src
# A registered UNC share root must authorize its own descendants in both browsers.
# os.path.commonpath raises "can't mix absolute and relative" on a bare
# \\server\share, so containment falls back to a boundary-safe prefix test; without
# it, registering a UNC share (now allowed) would 403 every folder under it.
@pytest.mark.parametrize(
"rel_parts",
[
["routes", "models.py"],
["hub", "services", "models", "folder_browser.py"],
],
)
def test_unc_share_root_authorizes_its_descendants(rel_parts):
is_inside = _extract_is_inside(rel_parts, os_module = _WIN_OS)
root = [Path(r"\\server\share")]
assert is_inside(Path(r"\\server\share"), root) is True # the root itself
assert is_inside(Path(r"\\server\share\models"), root) is True # direct child
assert is_inside(Path(r"\\server\share\a\b\c"), root) is True # deep descendant
assert is_inside(Path(r"\\SERVER\SHARE\Models"), root) is True # case-insensitive
assert is_inside(Path(r"\\server\share2\models"), root) is False # sibling share
assert is_inside(Path(r"C:\models"), root) is False # different volume

12
studio/backend/tests/test_browse_folders_route.py
View file

@ -22,6 +22,18 @@ if "structlog" not in sys.modules:
)
import routes.models as models_route
import storage.studio_db as studio_db
@pytest.fixture(autouse = True)
def _denylist_inert(monkeypatch):
# These tests exercise allowlist containment and the file-vs-directory guard,
# not the system-directory denylist (which has its own suite in
# test_browse_denylist.py). On macOS tmp_path resolves under /private/var, a
# denied prefix, so _resolve_browse_target would 403 the fixture dirs before
# the containment logic runs. Keep the denylist inert here so these
# assertions hold on every platform.
monkeypatch.setattr(studio_db, "is_denied_system_path", lambda _p: False)
def test_resolve_browse_target_returns_allowed_directory(tmp_path):

22
studio/backend/tests/test_bypass_permissions.py
View file

@ -13,6 +13,7 @@ never gating under bypass.
Run with: ``PYTHONPATH=studio/backend python -m pytest studio/backend/tests/test_bypass_permissions.py -q``
"""
import io
import os
import sys
@ -94,10 +95,23 @@ def test_safe_env_excludes_host_and_secret(monkeypatch, tmp_path):
class _FakeProc:
returncode = 0
"""A subprocess.Popen double for the drain path (``tools._drain_process_output``):
a readable ``stdout`` pipe yielding the fake output then EOF, plus
``wait()`` / ``poll()`` / ``pid``. The pid is non-existent so
``_capture_process_group``'s ``os.getpgid`` returns None; ``wait`` returns
immediately so the drain never kills.
"""
def communicate(self, timeout = None):
return ("FAKEOUT", None)
returncode = 0
# Unlikely-to-exist pid: os.getpgid raises ProcessLookupError (caught) -> None.
pid = 2**22
def __init__(self):
# Readable stdout: iter(readline, "") yields "FAKEOUT" then hits EOF.
self.stdout = io.StringIO("FAKEOUT")
def wait(self, timeout = None):
return 0
def poll(self):
return 0
@ -257,6 +271,7 @@ def test_loop_forwards_disable_sandbox_and_does_not_gate():
cancel_event = None,
timeout = None,
session_id = None,
thread_id = None,
rag_scope = None,
disable_sandbox = False,
):
@ -290,6 +305,7 @@ def test_loop_bypass_overrides_confirm_for_direct_callers():
cancel_event = None,
timeout = None,
session_id = None,
thread_id = None,
rag_scope = None,
disable_sandbox = False,
):

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

@ -651,11 +651,19 @@ class TestLoadModelGuardIntegration(unittest.TestCase):
inf._shutdown_subprocess = MagicMock()
llama = SimpleNamespace(is_loaded = False, model_identifier = None, hf_variant = None)
llama.unload_model = MagicMock()
cfg = SimpleNamespace(is_gguf = False, is_lora = False, path = None, base_model = None)
cfg = SimpleNamespace(
is_gguf = False,
is_lora = False,
path = None,
base_model = None,
identifier = "unsloth/Qwen3-1.7B",
)
request = LoadRequest(model_path = "unsloth/Qwen3-1.7B")
info = {"required_gb": 40.0, "usable_gb": 5.0, "needed_gb": 50.0, "mode": "auto"}
with (
# Pin the latest-sidecar tier check so the guard path stays offline.
patch("utils.transformers_version.latest_tier_active_for", return_value = False),
patch.object(self.route, "validate_extra_args", return_value = None),
patch.object(
self.route,

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

@ -691,13 +691,14 @@ def _argparse_default(source, option):
return None
def test_run_server_cloudflare_default_true():
def test_run_server_cloudflare_default_off():
defaults = _func_param_defaults(_RUN_PY.read_text(), "run_server")
assert defaults.get("cloudflare") is True
assert "cloudflare" in defaults
assert defaults["cloudflare"] is None
def test_argparse_cloudflare_default_true():
assert _argparse_default(_RUN_PY.read_text(), "--cloudflare") is True
def test_argparse_cloudflare_default_off():
assert _argparse_default(_RUN_PY.read_text(), "--cloudflare") is None
def test_verify_global_reachability_marks_private_address_unreachable():
@ -832,6 +833,31 @@ def test_cloudflare_line_states_disabled_when_off(monkeypatch):
assert "local network only" in out
def test_cloudflare_line_labels_unset_as_default(monkeypatch):
# None = off by default (no flag) -> banner says "(default)", not "(--no-cloudflare)".
out = _run_print_cloudflare_line(
monkeypatch,
cloudflare_url = None,
public_reachable = False,
cloudflare_requested = False,
cloudflare_flag = None,
)
assert "Cloudflare tunnel: OFF (default)" in out
assert "--no-cloudflare" not in out
def test_cloudflare_line_labels_explicit_no_cloudflare(monkeypatch):
# False = explicit --no-cloudflare -> banner says "(--no-cloudflare)".
out = _run_print_cloudflare_line(
monkeypatch,
cloudflare_url = None,
public_reachable = False,
cloudflare_requested = False,
cloudflare_flag = False,
)
assert "Cloudflare tunnel: OFF (--no-cloudflare)" in out
def test_cloudflare_line_states_failed_when_requested_but_no_url(monkeypatch):
out = _run_print_cloudflare_line(
monkeypatch,

314
studio/backend/tests/test_completion_masking.py Normal file
View file

@ -0,0 +1,314 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Completion-only masking policy: auto-detect first, manual table fallback.
Covers utils.datasets.completion_masking.apply_completion_masking, shared by
the CUDA trainer (core/training/trainer.py) and the MLX worker
(core/training/worker.py):
- unmapped models use chat template auto-detection (previously masking was
silently disabled),
- gpt-oss goes auto-first too (its quantized checkpoints ship a template
the manual markers cannot match),
- an auto-detection failure falls back to the template table markers,
- a table miss after an auto failure warns and leaves the trainer unchanged.
"""
from __future__ import annotations
import pytest
from utils.datasets.completion_masking import apply_completion_masking, lookup_manual_markers
from utils.datasets.model_mappings import TEMPLATE_TO_RESPONSES_MAPPER
class _Trainer:
"""Sentinel trainer; train_fn wraps it in a new object when applied."""
class _Recorder:
"""Fake train_on_responses_only that records calls."""
def __init__(self):
self.calls = []
def __call__(self, trainer, **kwargs):
self.calls.append(kwargs)
wrapped = _Trainer()
wrapped.wrapped_from = trainer
return wrapped
def _detect_ok(processor):
return "<INS>", "<RES>"
def _detect_fail(processor):
raise ValueError(
"Unsloth: Could not reliably auto-detect response_part - "
"pass instruction_part and response_part."
)
_AUTO = {"instruction_part": "<INS>", "response_part": "<RES>"}
class _Notes:
def __init__(self):
self.messages = []
def __call__(self, level, message):
self.messages.append((level, message))
def warnings(self):
return [m for level, m in self.messages if level == "warning"]
def test_unmapped_model_uses_auto_detection():
# Unmapped model: the auto path applies masking (was silently disabled).
trainer = _Trainer()
train_fn = _Recorder()
notes = _Notes()
result, applied = apply_completion_masking(
trainer, "LiquidAI/LFM2-8B-A1B", train_fn, notify = notes, detect_fn = _detect_ok
)
assert applied is True
assert result.wrapped_from is trainer
assert train_fn.calls == [dict(_AUTO)] # applied with the detected markers
assert notes.warnings() == []
def test_mapped_model_prefers_auto_detection():
trainer = _Trainer()
train_fn = _Recorder()
_, applied = apply_completion_masking(
trainer, "unsloth/Qwen3-0.6B", train_fn, detect_fn = _detect_ok
)
assert applied is True
assert train_fn.calls == [dict(_AUTO)]
def test_gpt_oss_uses_auto_detection_first():
# The quantized gpt-oss checkpoints ship a template without the
# <|channel|>final header, where the manual markers match nothing; auto
# derives markers from the template the checkpoint actually ships.
trainer = _Trainer()
train_fn = _Recorder()
_, applied = apply_completion_masking(
trainer, "unsloth/gpt-oss-20b", train_fn, detect_fn = _detect_ok
)
assert applied is True
assert train_fn.calls == [dict(_AUTO)]
def test_gpt_oss_detection_failure_falls_back_to_manual_markers():
trainer = _Trainer()
train_fn = _Recorder()
_, applied = apply_completion_masking(
trainer, "unsloth/gpt-oss-20b", train_fn, detect_fn = _detect_fail
)
assert applied is True
expected = TEMPLATE_TO_RESPONSES_MAPPER["gpt-oss"]
assert train_fn.calls == [
{
"instruction_part": expected["instruction"],
"response_part": expected["response"],
}
]
def test_auto_failure_falls_back_to_template_table():
trainer = _Trainer()
train_fn = _Recorder()
notes = _Notes()
result, applied = apply_completion_masking(
trainer, "unsloth/Qwen3-0.6B", train_fn, notify = notes, detect_fn = _detect_fail
)
assert applied is True
assert result.wrapped_from is trainer
expected = TEMPLATE_TO_RESPONSES_MAPPER["qwen3"]
assert train_fn.calls == [
{
"instruction_part": expected["instruction"],
"response_part": expected["response"],
},
]
assert any("falling back to the template table" in m for m in notes.warnings())
def test_application_failure_propagates_not_fallback():
# Detection succeeds; a failure while APPLYING the masking must propagate,
# never silently fall back to full-sequence training.
def train_fn(trainer, **kwargs):
raise RuntimeError("dataset map worker crashed")
with pytest.raises(RuntimeError, match = "dataset map worker crashed"):
apply_completion_masking(_Trainer(), "LiquidAI/LFM2-8B-A1B", train_fn, detect_fn = _detect_ok)
def test_preset_tokenizer_markers_used_directly():
# Preset unsloth marker attrs skip detection; zoo reuses them on a bare call.
class _Tok:
_unsloth_input_part = "<I>"
_unsloth_output_part = "<O>"
trainer = _Trainer()
trainer.processing_class = _Tok()
train_fn = _Recorder()
_, applied = apply_completion_masking(
trainer, "LiquidAI/LFM2-8B-A1B", train_fn, detect_fn = _detect_fail
)
assert applied is True
assert train_fn.calls == [{}] # bare call, stored parts
def test_table_miss_warns_and_disables_without_crashing():
trainer = _Trainer()
train_fn = _Recorder()
notes = _Notes()
result, applied = apply_completion_masking(
trainer, "some-org/not-in-any-mapper", train_fn, notify = notes, detect_fn = _detect_fail
)
assert applied is False
assert result is trainer # unchanged: full sequence training
assert train_fn.calls == [] # detection failed; nothing applied
assert any("could not be applied" in m for m in notes.warnings())
assert any("full sequences" in m for m in notes.warnings())
def test_num_proc_forwarded_only_when_given():
# CUDA path passes num_proc; the MLX path omits it.
train_fn = _Recorder()
apply_completion_masking(
_Trainer(), "unsloth/Qwen3-0.6B", train_fn, num_proc = 4, detect_fn = _detect_ok
)
assert train_fn.calls == [dict(_AUTO, num_proc = 4)]
train_fn = _Recorder()
apply_completion_masking(
_Trainer(), "unsloth/Qwen3-0.6B", train_fn, num_proc = 4, detect_fn = _detect_fail
)
assert train_fn.calls[0]["num_proc"] == 4
train_fn = _Recorder()
apply_completion_masking(_Trainer(), "unsloth/Qwen3-0.6B", train_fn, detect_fn = _detect_ok)
assert train_fn.calls == [dict(_AUTO)]
def test_manual_fallback_failure_propagates_to_caller():
# Errors while applying the manual fallback must propagate to the caller.
def train_fn(trainer, **kwargs):
raise RuntimeError("boom")
with pytest.raises(RuntimeError, match = "boom"):
apply_completion_masking(_Trainer(), "unsloth/gpt-oss-20b", train_fn)
def test_notify_is_optional():
train_fn = _Recorder()
_, applied = apply_completion_masking(
_Trainer(), "some-org/not-in-any-mapper", train_fn, detect_fn = _detect_fail
)
assert applied is False
def test_lookup_manual_markers():
template, instruction, response = lookup_manual_markers("unsloth/Qwen3-0.6B")
assert template == "qwen3"
assert instruction == TEMPLATE_TO_RESPONSES_MAPPER["qwen3"]["instruction"]
assert response == TEMPLATE_TO_RESPONSES_MAPPER["qwen3"]["response"]
template, instruction, response = lookup_manual_markers("some-org/unknown")
assert (template, instruction, response) == (None, None, None)
template, instruction, response = lookup_manual_markers(None)
assert (template, instruction, response) == (None, None, None)
def test_renamed_gpt_oss_gets_template_markers():
# Name-detected as gpt-oss but not in the exact-name table: must use the
# gpt-oss markers, not fall through to full-sequence training.
trainer = _Trainer()
train_fn = _Recorder()
_, applied = apply_completion_masking(
trainer, "some-org/gpt-oss-20b-sft", train_fn, detect_fn = _detect_fail
)
assert applied is True
expected = TEMPLATE_TO_RESPONSES_MAPPER["gpt-oss"]
assert train_fn.calls == [
{
"instruction_part": expected["instruction"],
"response_part": expected["response"],
}
]
class _FakeTokenizerWrapper:
"""mlx-lm TokenizerWrapper semantics: plain reads delegate to the wrapped
tokenizer, underscore attrs do not (so preset markers are hidden)."""
def __init__(self, tokenizer):
object.__setattr__(self, "_tokenizer", tokenizer)
def __getattr__(self, attr):
if attr.startswith("_"):
return object.__getattribute__(self, attr)
return getattr(object.__getattribute__(self, "_tokenizer"), attr)
_FakeTokenizerWrapper.__name__ = "TokenizerWrapper"
def test_mlx_tokenizer_wrapper_unwrapped_for_preset_markers():
# Markers live on the inner HF tokenizer that the wrapper hides; the helper
# must unwrap so the preset bare-call path still fires on MLX.
class _Tok:
_unsloth_input_part = "<I>"
_unsloth_output_part = "<O>"
trainer = _Trainer()
trainer.tokenizer = _FakeTokenizerWrapper(_Tok())
train_fn = _Recorder()
_, applied = apply_completion_masking(
trainer, "LiquidAI/LFM2-8B-A1B", train_fn, detect_fn = _detect_fail
)
assert applied is True
assert train_fn.calls == [{}] # bare call, stored parts
def test_mlx_tokenizer_wrapper_unwrapped_for_detection():
# Detection must see the real tokenizer, not the wrapper, so it does not
# depend on the loader's __call__ patch.
class _Tok:
pass
inner = _Tok()
trainer = _Trainer()
trainer.tokenizer = _FakeTokenizerWrapper(inner)
train_fn = _Recorder()
seen = []
def detect(processor):
seen.append(processor)
return "<INS>", "<RES>"
_, applied = apply_completion_masking(
trainer, "LiquidAI/LFM2-8B-A1B", train_fn, detect_fn = detect
)
assert applied is True
assert seen == [inner]

17
studio/backend/tests/test_deepseek_v4_thinking_effort.py
View file

@ -16,7 +16,6 @@ from __future__ import annotations
import sys
from pathlib import Path
from types import SimpleNamespace
import pytest
@ -128,15 +127,13 @@ def _kwargs_for(flags: dict, enable_thinking, reasoning_effort):
"""Drive the real backend method with a shim carrying the detected flags."""
from core.inference.llama_cpp import LlamaCppBackend
shim = SimpleNamespace(
_supports_reasoning = flags["supports_reasoning"],
_reasoning_always_on = flags["reasoning_always_on"],
_reasoning_style = flags["reasoning_style"],
_reasoning_effort_levels = flags["reasoning_effort_levels"],
_supports_preserve_thinking = flags["supports_preserve_thinking"],
)
build = LlamaCppBackend._request_reasoning_kwargs.__get__(shim)
return build(enable_thinking, reasoning_effort, None) or {}
shim = object.__new__(LlamaCppBackend)
shim._supports_reasoning = flags["supports_reasoning"]
shim._reasoning_always_on = flags["reasoning_always_on"]
shim._reasoning_style = flags["reasoning_style"]
shim._reasoning_effort_levels = flags["reasoning_effort_levels"]
shim._supports_preserve_thinking = flags["supports_preserve_thinking"]
return shim._request_reasoning_kwargs(enable_thinking, reasoning_effort, None) or {}
def _flags():

740
studio/backend/tests/test_gguf_load_cache_reuse.py Normal file
View file

@ -0,0 +1,740 @@
# 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 cached GGUF reuse and load/download exclusion.
No GPU, network, or subprocesses are required.
"""
from __future__ import annotations
import asyncio
import sys
import threading
import types as _types
from pathlib import Path
from unittest.mock import patch
import pytest
_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
if _BACKEND_DIR not in sys.path:
sys.path.insert(0, _BACKEND_DIR)
# Stub optional dependencies before importing the modules under test.
_loggers_stub = _types.ModuleType("loggers")
_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name)
sys.modules.setdefault("loggers", _loggers_stub)
_structlog_stub = _types.ModuleType("structlog")
sys.modules.setdefault("structlog", _structlog_stub)
try:
import httpx # noqa: F401
except ImportError:
_httpx_stub = _types.ModuleType("httpx")
for _exc_name in (
"ConnectError",
"TimeoutException",
"ReadTimeout",
"ReadError",
"RemoteProtocolError",
"CloseError",
"HTTPError",
"RequestError",
"HTTPStatusError",
):
setattr(_httpx_stub, _exc_name, type(_exc_name, (Exception,), {}))
_httpx_stub.Response = type("Response", (), {})
_httpx_stub.Request = type("Request", (), {})
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 huggingface_hub import constants as hf_constants
from core.inference.llama_cpp import (
LlamaCppBackend,
cached_gguf_for_load,
gguf_load_in_flight,
hf_gguf_load_in_flight,
)
REPO = "unsloth/gemma-test-GGUF"
VARIANT = "UD-Q4_K_XL"
MAIN = f"gemma-test-{VARIANT}.gguf"
def _build_cache(
root: Path,
repo_id: str,
files: dict[str, int],
*,
snapshot_sha: str = "a" * 40,
) -> Path:
"""Create ``$root/models--<repo>/snapshots/<sha>/<rel>`` for each entry."""
repo_dir = root / f"models--{repo_id.replace('/', '--')}"
(repo_dir / "blobs").mkdir(parents = True, exist_ok = True)
snap = repo_dir / "snapshots" / snapshot_sha
snap.mkdir(parents = True, exist_ok = True)
for rel, size in files.items():
full = snap / rel
full.parent.mkdir(parents = True, exist_ok = True)
full.write_bytes(b"\0" * size)
return snap
@pytest.fixture
def hf_cache(tmp_path, monkeypatch):
monkeypatch.setattr(hf_constants, "HF_HUB_CACHE", str(tmp_path))
monkeypatch.delenv("HF_HUB_OFFLINE", raising = False)
monkeypatch.delenv("TRANSFORMERS_OFFLINE", raising = False)
return tmp_path
def _fail_download(*_args, **_kwargs):
raise AssertionError("must reuse the cached GGUF instead of downloading")
def _fail_get_paths_info(*_args, **_kwargs):
raise AssertionError("cached reuse must return before the sizing preflight")
class TestLoadReusesCachedCopy:
def test_online_reuse_after_revision_bump(self, hf_cache):
"""A new repo revision does not replace a complete cached model."""
backend = LlamaCppBackend()
snap = _build_cache(hf_cache, REPO, {MAIN: 4})
with (
patch("huggingface_hub.list_repo_files", lambda *_a, **_k: [MAIN]),
patch("huggingface_hub.get_paths_info", _fail_get_paths_info),
patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", _fail_download),
):
out = backend._download_gguf(hf_repo = REPO, hf_variant = VARIANT)
assert out == str(snap / MAIN)
def test_reuse_size_check_uses_cached_snapshot_revision(self, hf_cache):
"""Current-revision size changes do not invalidate an older complete copy."""
backend = LlamaCppBackend()
snap = _build_cache(hf_cache, REPO, {MAIN: 4})
revisions: list[str | None] = []
def fake_get_paths_info(
_repo,
paths,
*,
revision = None,
token = None,
):
revisions.append(revision)
size = 4 if revision == snap.name else 8
return [_types.SimpleNamespace(path = path, size = size) for path in paths]
with (
patch("huggingface_hub.list_repo_files", lambda *_a, **_k: [MAIN]),
patch("huggingface_hub.get_paths_info", fake_get_paths_info),
patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", _fail_download),
):
out = backend._download_gguf(hf_repo = REPO, hf_variant = VARIANT)
assert out == str(snap / MAIN)
assert revisions == [snap.name]
def test_reuse_when_cached_revision_vanished_from_hub(self, hf_cache):
"""The Hub answers an unknown revision with an empty result, not an error."""
backend = LlamaCppBackend()
snap = _build_cache(hf_cache, REPO, {MAIN: 4})
with (
patch("huggingface_hub.list_repo_files", lambda *_a, **_k: [MAIN]),
patch("huggingface_hub.get_paths_info", lambda *_a, **_k: []),
patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", _fail_download),
):
out = backend._download_gguf(hf_repo = REPO, hf_variant = VARIANT)
assert out == str(snap / MAIN)
def test_truncated_cached_file_is_not_reused(self, hf_cache):
backend = LlamaCppBackend()
_build_cache(hf_cache, REPO, {MAIN: 4})
downloaded: list[str] = []
def fake_get_paths_info(
_repo,
paths,
*,
revision = None,
token = None,
):
return [_types.SimpleNamespace(path = path, size = 8) for path in paths]
def fake_download(
repo_id,
filename,
token = None,
**_kwargs,
):
downloaded.append(filename)
return f"/fake/{repo_id}/{filename}"
with (
patch("huggingface_hub.list_repo_files", lambda *_a, **_k: [MAIN]),
patch("huggingface_hub.get_paths_info", fake_get_paths_info),
patch("huggingface_hub.try_to_load_from_cache", lambda *_a, **_k: None),
patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", fake_download),
):
out = backend._download_gguf(hf_repo = REPO, hf_variant = VARIANT)
assert downloaded == [MAIN]
assert out == f"/fake/{REPO}/{MAIN}"
def test_truncated_cached_split_shard_is_not_reused(self, hf_cache):
backend = LlamaCppBackend()
shard1 = f"gemma-test-{VARIANT}-00001-of-00002.gguf"
shard2 = f"gemma-test-{VARIANT}-00002-of-00002.gguf"
_build_cache(hf_cache, REPO, {shard1: 8, shard2: 4})
downloaded: list[str] = []
def fake_get_paths_info(
_repo,
paths,
*,
revision = None,
token = None,
):
return [_types.SimpleNamespace(path = path, size = 8) for path in paths]
def fake_download(
repo_id,
filename,
token = None,
**_kwargs,
):
downloaded.append(filename)
return f"/fake/{repo_id}/{filename}"
with (
patch("huggingface_hub.list_repo_files", lambda *_a, **_k: [shard1, shard2]),
patch("huggingface_hub.get_paths_info", fake_get_paths_info),
patch("huggingface_hub.try_to_load_from_cache", lambda *_a, **_k: None),
patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", fake_download),
):
out = backend._download_gguf(hf_repo = REPO, hf_variant = VARIANT)
assert downloaded == [shard1, shard2]
assert out == f"/fake/{REPO}/{shard1}"
def test_online_reuse_when_reupload_renamed_the_file(self, hf_cache):
"""A renamed variant still reuses its cached file."""
backend = LlamaCppBackend()
old_name = f"gemma-test-old-{VARIANT}.gguf"
snap = _build_cache(hf_cache, REPO, {old_name: 4})
with (
patch("huggingface_hub.list_repo_files", lambda *_a, **_k: [MAIN]),
patch("huggingface_hub.get_paths_info", _fail_get_paths_info),
patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", _fail_download),
):
out = backend._download_gguf(hf_repo = REPO, hf_variant = VARIANT)
assert out == str(snap / old_name)
def test_downloads_when_nothing_cached(self, hf_cache):
backend = LlamaCppBackend()
downloaded: list[str] = []
def fake_download(
repo_id,
filename,
token = None,
**_kwargs,
):
downloaded.append(filename)
return f"/fake/{repo_id}/{filename}"
def fake_get_paths_info(
_repo_id,
paths,
token = None,
):
return [_types.SimpleNamespace(path = p, size = 1) for p in paths if p is not None]
with (
patch("huggingface_hub.list_repo_files", lambda *_a, **_k: [MAIN]),
patch("huggingface_hub.get_paths_info", fake_get_paths_info),
patch("huggingface_hub.try_to_load_from_cache", lambda *_a, **_k: None),
patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", fake_download),
):
out = backend._download_gguf(hf_repo = REPO, hf_variant = VARIANT)
assert downloaded == [MAIN]
assert out == f"/fake/{REPO}/{MAIN}"
def test_force_redownloads_despite_cache(self, hf_cache):
"""A forced download ignores a complete cached copy."""
backend = LlamaCppBackend()
_build_cache(hf_cache, REPO, {MAIN: 4})
downloaded: list[str] = []
def fake_download(
repo_id,
filename,
token = None,
**kwargs,
):
assert kwargs.get("force_download") is True
downloaded.append(filename)
return f"/fake/{repo_id}/{filename}"
def fake_get_paths_info(
_repo_id,
paths,
token = None,
):
return [_types.SimpleNamespace(path = p, size = 1) for p in paths if p is not None]
with (
patch("huggingface_hub.list_repo_files", lambda *_a, **_k: [MAIN]),
patch("huggingface_hub.get_paths_info", fake_get_paths_info),
patch("huggingface_hub.try_to_load_from_cache", lambda *_a, **_k: None),
patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", fake_download),
):
out = backend._download_gguf(hf_repo = REPO, hf_variant = VARIANT, force = True)
assert downloaded == [MAIN]
assert out == f"/fake/{REPO}/{MAIN}"
def test_split_reused_only_when_colocated(self, hf_cache):
backend = LlamaCppBackend()
shard1 = f"gemma-test-{VARIANT}-00001-of-00002.gguf"
shard2 = f"gemma-test-{VARIANT}-00002-of-00002.gguf"
snap = _build_cache(hf_cache, REPO, {shard1: 4, shard2: 4})
with (
patch("huggingface_hub.list_repo_files", lambda *_a, **_k: [shard1, shard2]),
patch("huggingface_hub.get_paths_info", _fail_get_paths_info),
patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", _fail_download),
):
out = backend._download_gguf(hf_repo = REPO, hf_variant = VARIANT)
assert out == str(snap / shard1)
def test_partial_split_set_downloads(self, hf_cache):
"""A partial split set is not reused."""
backend = LlamaCppBackend()
shard1 = f"gemma-test-{VARIANT}-00001-of-00002.gguf"
shard2 = f"gemma-test-{VARIANT}-00002-of-00002.gguf"
_build_cache(hf_cache, REPO, {shard1: 4})
downloaded: list[str] = []
def fake_download(
repo_id,
filename,
token = None,
**_kwargs,
):
downloaded.append(filename)
return f"/fake/{repo_id}/{filename}"
def fake_get_paths_info(
_repo_id,
paths,
token = None,
):
return [_types.SimpleNamespace(path = p, size = 4) for p in paths if p is not None]
with (
patch("huggingface_hub.list_repo_files", lambda *_a, **_k: [shard1, shard2]),
patch("huggingface_hub.get_paths_info", fake_get_paths_info),
patch("huggingface_hub.try_to_load_from_cache", lambda *_a, **_k: None),
patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", fake_download),
):
out = backend._download_gguf(hf_repo = REPO, hf_variant = VARIANT)
assert downloaded == [shard1, shard2]
assert out == f"/fake/{REPO}/{shard1}"
def test_reuse_prefers_newest_snapshot_after_update(self, hf_cache):
"""Loads prefer the newest complete snapshot."""
import os
backend = LlamaCppBackend()
old_snap = _build_cache(hf_cache, REPO, {MAIN: 4}, snapshot_sha = "a" * 40)
new_snap = _build_cache(hf_cache, REPO, {MAIN: 6}, snapshot_sha = "b" * 40)
os.utime(old_snap, (1_000_000, 1_000_000))
os.utime(new_snap, (2_000_000, 2_000_000))
with (
patch("huggingface_hub.list_repo_files", lambda *_a, **_k: [MAIN]),
patch("huggingface_hub.get_paths_info", _fail_get_paths_info),
patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", _fail_download),
):
out = backend._download_gguf(hf_repo = REPO, hf_variant = VARIANT)
assert out == str(new_snap / MAIN)
def test_low_disk_fallback_reuses_cached_copy(self, hf_cache):
backend = LlamaCppBackend()
fallback = "gemma-test-Q2_K.gguf"
snap = _build_cache(hf_cache, REPO, {fallback: 4})
def fake_get_paths_info(
_repo,
paths,
*,
revision = None,
token = None,
):
size = 4 if revision == snap.name else 100
return [_types.SimpleNamespace(path = path, size = size) for path in paths]
with (
patch("huggingface_hub.list_repo_files", lambda *_a, **_k: [MAIN]),
patch("huggingface_hub.get_paths_info", fake_get_paths_info),
patch("huggingface_hub.try_to_load_from_cache", lambda *_a, **_k: None),
patch("shutil.disk_usage", lambda *_a, **_k: _types.SimpleNamespace(free = 10)),
patch.object(
backend,
"_find_smallest_fitting_variant",
lambda *_a, **_k: (fallback, 4, []),
),
patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", _fail_download),
):
out = backend._download_gguf(hf_repo = REPO, hf_variant = VARIANT)
assert out == str(snap / fallback)
def test_companion_prefers_main_snapshot_sibling(self, hf_cache):
"""A cached mmproj is reused from the main model's snapshot."""
backend = LlamaCppBackend()
snap = _build_cache(hf_cache, REPO, {MAIN: 4, "mmproj-F16.gguf": 2})
def _fail_list(*_args, **_kwargs):
raise AssertionError("snapshot sibling must resolve without a repo listing")
with patch("huggingface_hub.list_repo_files", _fail_list):
out = backend._download_mmproj(hf_repo = REPO, near_path = str(snap / MAIN))
assert out == str(snap / "mmproj-F16.gguf")
def test_companion_finds_snapshot_through_hf_symlink(self, hf_cache):
backend = LlamaCppBackend()
snap = _build_cache(hf_cache, REPO, {})
blobs = snap.parent.parent / "blobs"
main_blob = blobs / "main"
mmproj_blob = blobs / "mmproj"
main_blob.write_bytes(b"main")
mmproj_blob.write_bytes(b"mmproj")
try:
(snap / MAIN).symlink_to(main_blob)
(snap / "mmproj-F16.gguf").symlink_to(mmproj_blob)
except OSError as exc:
pytest.skip(f"symlinks unavailable: {exc}")
with patch("huggingface_hub.list_repo_files", _fail_download):
out = backend._download_mmproj(hf_repo = REPO, near_path = str(snap / MAIN))
assert out == str(snap / "mmproj-F16.gguf")
def test_companion_does_not_download_during_hub_job(self, hf_cache):
backend = LlamaCppBackend()
snap = _build_cache(hf_cache, REPO, {MAIN: 4})
registry = _types.SimpleNamespace(active_job_refs = lambda _repo: [object()])
with (
patch("huggingface_hub.list_repo_files", _fail_download),
patch("hub.utils.download_registry.get_models_registry", lambda: registry),
patch("core.inference.llama_cpp.hf_hub_download_with_xet_fallback", _fail_download),
):
out = backend._download_mmproj(hf_repo = REPO, near_path = str(snap / MAIN))
assert out is None
class TestCachedGgufForLoadProbe:
def test_complete_copy_found(self, hf_cache):
snap = _build_cache(hf_cache, REPO, {MAIN: 4})
assert cached_gguf_for_load(REPO, VARIANT) == str(snap / MAIN)
def test_absent_copy_is_none(self, hf_cache):
assert cached_gguf_for_load(REPO, VARIANT) is None
def test_partial_split_is_none(self, hf_cache):
shard1 = f"gemma-test-{VARIANT}-00001-of-00002.gguf"
_build_cache(hf_cache, REPO, {shard1: 4})
assert cached_gguf_for_load(REPO, VARIANT) is None
def test_partial_new_snapshot_does_not_hide_complete_split(self, hf_cache):
import os
shard1 = f"gemma-test-{VARIANT}-00001-of-00002.gguf"
shard2 = f"gemma-test-{VARIANT}-00002-of-00002.gguf"
old = _build_cache(
hf_cache,
REPO,
{shard1: 4, shard2: 4},
snapshot_sha = "a" * 40,
)
new = _build_cache(hf_cache, REPO, {shard1: 4}, snapshot_sha = "b" * 40)
os.utime(old, (1_000_000, 1_000_000))
os.utime(new, (2_000_000, 2_000_000))
assert cached_gguf_for_load(REPO, VARIANT) == str(old / shard1)
def test_split_requires_every_declared_shard(self, hf_cache):
shard1 = f"gemma-test-{VARIANT}-00001-of-00003.gguf"
shard2 = f"gemma-test-{VARIANT}-00002-of-00003.gguf"
_build_cache(hf_cache, REPO, {shard1: 4, shard2: 4})
assert cached_gguf_for_load(REPO, VARIANT) is None
def test_required_mmproj_must_share_main_snapshot(self, hf_cache):
snap = _build_cache(hf_cache, REPO, {MAIN: 4})
assert cached_gguf_for_load(REPO, VARIANT) == str(snap / MAIN)
assert cached_gguf_for_load(REPO, VARIANT, require_mmproj = True) is None
(snap / "mmproj-F16.gguf").write_bytes(b"mmproj")
assert cached_gguf_for_load(REPO, VARIANT, require_mmproj = True) == str(snap / MAIN)
def test_required_mmproj_scans_past_newer_main_only_snapshot(self, hf_cache):
import os
old = _build_cache(
hf_cache,
REPO,
{MAIN: 4, "mmproj-F16.gguf": 2},
snapshot_sha = "a" * 40,
)
new = _build_cache(hf_cache, REPO, {MAIN: 4}, snapshot_sha = "b" * 40)
os.utime(old, (1_000_000, 1_000_000))
os.utime(new, (2_000_000, 2_000_000))
assert cached_gguf_for_load(REPO, VARIANT, require_mmproj = True) == str(old / MAIN)
class TestLoadHubDownloadExclusion:
def test_in_flight_marker_counts_and_normalizes_case(self):
assert not hf_gguf_load_in_flight(REPO)
with gguf_load_in_flight(REPO):
assert hf_gguf_load_in_flight(REPO.upper())
with gguf_load_in_flight(REPO.lower()):
assert hf_gguf_load_in_flight(REPO)
assert hf_gguf_load_in_flight(REPO)
assert not hf_gguf_load_in_flight(REPO)
def test_marker_noops_for_local_loads(self):
with gguf_load_in_flight(None):
assert not hf_gguf_load_in_flight("")
def test_marker_cleared_on_exception(self):
with pytest.raises(RuntimeError):
with gguf_load_in_flight(REPO):
raise RuntimeError("boom")
assert not hf_gguf_load_in_flight(REPO)
def test_hub_download_refused_while_load_in_flight(self):
from fastapi import HTTPException
from hub.schemas.downloads import DownloadModelRequest
from hub.services.models import downloads as dl
body = DownloadModelRequest(repo_id = REPO, gguf_variant = VARIANT)
with (
patch.object(dl, "resolve_cached_repo_id_case", lambda repo_id, repo_type: repo_id),
gguf_load_in_flight(REPO),
):
with pytest.raises(HTTPException) as exc_info:
asyncio.run(dl.download_model_response(body))
assert exc_info.value.status_code == 409
assert "load" in exc_info.value.detail.lower()
def test_hub_download_rechecks_marker_before_claim(self):
from fastapi import HTTPException
from hub.schemas.downloads import DownloadModelRequest
from hub.services.models import downloads as dl
scope = None
def mark_load(*_args, **_kwargs):
nonlocal scope
if scope is None:
scope = gguf_load_in_flight(REPO)
scope.__enter__()
return frozenset()
class _Registry:
def claim(self, *_args, admission_check, **_kwargs):
assert admission_check() is False
return False, "admission_blocked"
def current_generation(self, _key):
return 0
registry = _Registry()
body = DownloadModelRequest(repo_id = REPO, gguf_variant = VARIANT)
try:
with (
patch.object(dl, "resolve_cached_repo_id_case", lambda repo_id, repo_type: repo_id),
patch.object(dl.gguf_variants, "gguf_variant_blob_hashes", mark_load),
patch.object(dl, "_registry", registry),
):
with pytest.raises(HTTPException) as exc_info:
asyncio.run(dl.download_model_response(body))
finally:
if scope is not None:
scope.__exit__(None, None, None)
assert exc_info.value.status_code == 409
def test_registry_admission_check_prevents_claim(self):
from hub.utils.download_registry import DownloadRegistry, TRANSPORT_HTTP
registry = DownloadRegistry()
claimed, state = registry.claim(
f"{REPO}::{VARIANT}",
TRANSPORT_HTTP,
repo_type = "model",
repo_id = REPO,
variant = VARIANT,
admission_check = lambda: False,
)
assert claimed is False
assert state == "admission_blocked"
assert registry.active_jobs(REPO) == {}
def test_same_variant_job_stays_visible_during_retry_handoff(self):
from hub.utils.download_registry import DownloadRegistry, TRANSPORT_XET
from core.inference.llama_cpp import _hub_download_blocks_gguf_load
registry = DownloadRegistry()
key = f"{REPO}::{VARIANT}"
claimed, _ = registry.claim(
key,
TRANSPORT_XET,
repo_type = "model",
repo_id = REPO,
variant = VARIANT,
)
assert claimed is True
assert registry.has_active_variant(REPO, VARIANT.lower()) is True
registry.release_active_slot(key)
assert registry.active_jobs(REPO) == {}
assert registry.active_job_refs(REPO)
assert registry.has_active_variant(REPO, VARIANT) is True
with (
patch("hub.utils.download_registry.get_models_registry", lambda: registry),
patch(
"core.inference.llama_cpp.cached_gguf_for_load",
side_effect = AssertionError("same-variant jobs must block before cache reuse"),
),
):
assert _hub_download_blocks_gguf_load(REPO, VARIANT) is True
registry.set_job(key, "complete")
assert registry.has_active_variant(REPO, VARIANT) is False
def test_other_variant_job_still_allows_complete_cached_load(self):
from core.inference.llama_cpp import _hub_download_blocks_gguf_load
from hub.utils.download_registry import DownloadRegistry, TRANSPORT_HTTP
registry = DownloadRegistry()
registry.claim(
f"{REPO}::Q8_0",
TRANSPORT_HTTP,
repo_type = "model",
repo_id = REPO,
variant = "Q8_0",
)
with (
patch("hub.utils.download_registry.get_models_registry", lambda: registry),
patch(
"core.inference.llama_cpp.cached_gguf_for_load",
return_value = "/cached/model.gguf",
) as cached_probe,
):
assert _hub_download_blocks_gguf_load(REPO, VARIANT) is False
cached_probe.assert_called_once_with(
REPO,
VARIANT,
require_mmproj = False,
verify_sizes = True,
hf_token = None,
)
def test_cancelled_request_keeps_marker_until_load_thread_finishes(self):
from core.inference.llama_cpp import _with_gguf_load_marker
started = threading.Event()
release = threading.Event()
finished = threading.Event()
class FakeBackend:
@_with_gguf_load_marker
def load_model(self, *, hf_repo):
started.set()
release.wait(timeout = 2)
finished.set()
return True
async def scenario():
with patch(
"core.inference.llama_cpp._hub_download_blocks_gguf_load",
return_value = False,
):
task = asyncio.create_task(
asyncio.to_thread(FakeBackend().load_model, hf_repo = REPO)
)
assert await asyncio.to_thread(started.wait, 1)
task.cancel()
with pytest.raises(asyncio.CancelledError):
await task
assert hf_gguf_load_in_flight(REPO)
release.set()
assert await asyncio.to_thread(finished.wait, 1)
for _ in range(100):
if not hf_gguf_load_in_flight(REPO):
break
await asyncio.sleep(0.001)
assert not hf_gguf_load_in_flight(REPO)
asyncio.run(scenario())
def test_load_marker_precedes_hub_guard_and_unload(self):
source = (Path(__file__).resolve().parent.parent / "routes" / "inference.py").read_text()
gguf_branch = source[source.index("if config.is_gguf:") :]
assert (
gguf_branch.index("enter_context(gguf_load_in_flight")
< gguf_branch.index("if request.llama_extra_args is None")
< gguf_branch.index("_hub_download_blocks_gguf_load")
< gguf_branch.index("unsloth_backend.unload_model")
)
llama_source = (
Path(__file__).resolve().parent.parent / "core" / "inference" / "llama_cpp.py"
).read_text()
assert "@_with_gguf_load_marker\n def load_model(" in llama_source

123
studio/backend/tests/test_index_bootstrap_loopback.py Normal file
View file

@ -0,0 +1,123 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Regression coverage for bootstrap password exposure to remote clients."""
from types import SimpleNamespace
def _request(
client_host,
request_host = "127.0.0.1",
headers = None,
):
"""Build a minimal request; ``None`` models an unresolved peer / absent Host."""
client = None if client_host is None else SimpleNamespace(host = client_host, port = 0)
hdrs = {}
if request_host is not None:
hdrs["host"] = request_host
hdrs.update(headers or {})
return SimpleNamespace(client = client, headers = hdrs, url = SimpleNamespace(hostname = request_host))
def test_loopback_peers_are_local():
from main import _is_local_bootstrap_request
cases = (
("127.0.0.1", "127.0.0.1"),
("::1", "::1"),
("::ffff:127.0.0.1", "::ffff:127.0.0.1"),
("127.0.0.1", "localhost"),
)
for peer, host in cases:
assert _is_local_bootstrap_request(_request(peer, host)) is True, (peer, host)
def test_non_loopback_peers_are_remote():
from main import _is_local_bootstrap_request
# ::1%eth0 is a scope-id'd address, which ipaddress treats as loopback on
# 3.9+; it must not count as a direct local peer.
for host in ("192.168.1.10", "::ffff:192.168.1.10", "::1%eth0"):
assert _is_local_bootstrap_request(_request(host)) is False, host
def test_absent_or_unparseable_peer_fails_safe():
from main import _is_local_bootstrap_request
for host in (None, "localhost"):
assert _is_local_bootstrap_request(_request(host)) is False, host
def test_cloudflare_tunnel_clients_are_remote_despite_loopback_peer():
from main import _is_local_bootstrap_request
for client_ip in ("203.0.113.7", ""):
request = _request("127.0.0.1", headers = {"cf-connecting-ip": client_ip})
assert _is_local_bootstrap_request(request) is False, client_ip
def test_dns_rebinding_host_is_remote_despite_loopback_peer():
from main import _is_local_bootstrap_request
for host in ("attacker.example", "192.168.1.10", None):
assert _is_local_bootstrap_request(_request("127.0.0.1", host)) is False, host
def test_unparseable_request_host_fails_safe():
"""A Host that makes ``request.url.hostname`` raise must fall to remote."""
from main import _is_local_bootstrap_request
class _RaisingURL:
@property
def hostname(self):
raise ValueError("malformed host")
request = SimpleNamespace(
client = SimpleNamespace(host = "127.0.0.1", port = 0), headers = {}, url = _RaisingURL()
)
assert _is_local_bootstrap_request(request) is False
def test_reverse_proxy_forwarded_headers_are_remote():
"""A loopback proxy relaying a remote client (non-Cloudflare headers) is remote."""
from main import _is_local_bootstrap_request
for header in ("forwarded", "x-forwarded-for", "x-forwarded-host", "x-real-ip"):
request = _request("127.0.0.1", "localhost", headers = {header: "203.0.113.7"})
assert _is_local_bootstrap_request(request) is False, header
def test_malformed_or_absent_host_is_remote():
"""A malformed/absent/scope-id Host must not fall back to the loopback server address."""
from main import _is_local_bootstrap_request
# incl. bracket smuggling: [::1]evil / unclosed [::1 must not reduce to ::1
for host in (
"e_vil",
"[malformed",
"",
None,
"[::1%25eth0]:8888",
"[::1]attacker",
"[::1]evil.com",
"[::1",
"[::1]x",
):
assert _is_local_bootstrap_request(_request("127.0.0.1", host)) is False, host
def test_colab_allows_notebook_proxy_but_not_shareable_tunnel(monkeypatch):
"""Colab autofills its single-user proxy, but not a public Cloudflare link."""
import main
monkeypatch.setattr(main, "_IS_COLAB", True)
# In-notebook proxy: same-origin, no tunnel header, injects off-loopback too.
assert main._should_inject_bootstrap(_request("10.0.0.2", "colab.proxy")) is True
# Shareable Cloudflare link marks visitors with cf-connecting-ip; withhold.
tunnel = _request("127.0.0.1", "localhost", headers = {"cf-connecting-ip": "203.0.113.7"})
assert main._should_inject_bootstrap(tunnel) is False
def test_non_colab_gate_requires_local_client(monkeypatch):
"""Outside Colab the gate injects only for a direct loopback client."""
import main
monkeypatch.setattr(main, "_IS_COLAB", False)
assert main._should_inject_bootstrap(_request("127.0.0.1", "localhost")) is True
assert main._should_inject_bootstrap(_request("192.168.1.10", "localhost")) is False

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

@ -112,7 +112,7 @@ def test_route_llama_streaming_async_clients_disable_proxy_env():
continue
calls.append(node)
assert len(calls) == 4
assert len(calls) == 5
for call in calls:
assert any(
kw.arg == "trust_env" and isinstance(kw.value, ast.Constant) and kw.value.value is False

404
studio/backend/tests/test_install_resolve_prebuilt.py
View file

@ -55,6 +55,27 @@ def _host(**kw):
return ilp.HostInfo(**base)
def test_force_cpu_clears_all_gpu_attributes_including_intel():
# --cpu-fallback is the "select the CPU prebuilt even when a GPU is present"
# escape hatch. It must drop EVERY GPU attribute, including has_intel_gpu, or
# the planner still prepends the Vulkan asset on an Intel-GPU host.
host = _host(
is_linux = True,
is_x86_64 = True,
has_usable_nvidia = True,
has_physical_nvidia = True,
has_rocm = True,
rocm_gfx_target = "gfx1100",
has_intel_gpu = True,
)
forced = ilp._apply_host_overrides(host, force_cpu = True)
assert forced.has_usable_nvidia is False
assert forced.has_physical_nvidia is False
assert forced.has_rocm is False
assert forced.rocm_gfx_target is None
assert forced.has_intel_gpu is False
def test_macos_upstream_pin_only_for_explicit_pre26_upstream():
pre26 = _host(
system = "Darwin",
@ -313,3 +334,386 @@ def test_sm103_host_drops_cuda128_windows_build():
)
kept_b200 = ilp._drop_blackwell_incapable_windows_cuda(b200, [cuda128, cuda129])
assert [a.name for a in kept_b200] == [cuda128.name, cuda129.name]
def _upstream_release(tag, asset_names):
return {
"tag_name": tag,
"assets": [
{"name": n, "browser_download_url": f"https://example/{n}"} for n in asset_names
],
}
def test_direct_upstream_arm64_intel_prefers_vulkan():
# Auto-detected Intel GPU on Linux arm64 -> Vulkan prebuilt first, CPU
# second (mirrors the x86_64 branch; ggml-org ships the arm64 Vulkan asset).
host = _host(is_linux = True, is_arm64 = True, machine = "aarch64", has_intel_gpu = True)
rel = _upstream_release(
"b9925",
["llama-b9925-bin-ubuntu-vulkan-arm64.tar.gz", "llama-b9925-bin-ubuntu-arm64.tar.gz"],
)
plan = ilp.direct_upstream_release_plan(rel, host, UPSTREAM, "latest")
kinds = [a.install_kind for a in plan.attempts]
assert kinds[0] == "linux-vulkan", kinds
assert "linux-arm64" in kinds
assert plan.attempts[0].name == "llama-b9925-bin-ubuntu-vulkan-arm64.tar.gz"
def test_direct_upstream_intel_with_hidden_nvidia_is_cpu_only():
# A host with a physical NVIDIA hidden via CUDA_VISIBLE_DEVICES (physical
# True, usable False) + an Intel iGPU must NOT get the Vulkan archive even
# when planning directly against upstream: Vulkan ignores CUDA_VISIBLE_DEVICES
# and could grab the reserved card. It falls through to the CPU asset.
host = _host(
is_linux = True,
is_x86_64 = True,
has_intel_gpu = True,
has_physical_nvidia = True,
has_usable_nvidia = False,
)
rel = _upstream_release(
"b9925",
["llama-b9925-bin-ubuntu-vulkan-x64.tar.gz", "llama-b9925-bin-ubuntu-x64.tar.gz"],
)
plan = ilp.direct_upstream_release_plan(rel, host, UPSTREAM, "latest")
assert [a.install_kind for a in plan.attempts] == ["linux-cpu"]
def test_direct_upstream_arm64_without_intel_is_cpu_only():
host = _host(is_linux = True, is_arm64 = True, machine = "aarch64")
rel = _upstream_release(
"b9925",
["llama-b9925-bin-ubuntu-vulkan-arm64.tar.gz", "llama-b9925-bin-ubuntu-arm64.tar.gz"],
)
plan = ilp.direct_upstream_release_plan(rel, host, UPSTREAM, "latest")
assert [a.install_kind for a in plan.attempts] == ["linux-arm64"]
def test_direct_upstream_x86_intel_prefers_vulkan():
host = _host(is_linux = True, is_x86_64 = True, has_intel_gpu = True)
rel = _upstream_release(
"b9925",
["llama-b9925-bin-ubuntu-vulkan-x64.tar.gz", "llama-b9925-bin-ubuntu-x64.tar.gz"],
)
plan = ilp.direct_upstream_release_plan(rel, host, UPSTREAM, "latest")
kinds = [a.install_kind for a in plan.attempts]
assert kinds[0] == "linux-vulkan", kinds
assert "linux-cpu" in kinds
def test_linux_vulkan_health_glob_matches_bare_cpu_lib():
# The widened glob must cover both arch-suffixed (x64) and bare (arm64) CPU
# libs so a valid Vulkan install is not re-flagged unhealthy every check.
choice = ilp.AssetChoice(
repo = UPSTREAM,
tag = "b9925",
name = "llama-b9925-bin-ubuntu-vulkan-arm64.tar.gz",
url = "https://example/x",
source_label = "upstream",
install_kind = "linux-vulkan",
)
groups = ilp.runtime_payload_health_groups(choice)
assert ["libggml-cpu*.so*"] in groups
assert ["libggml-cpu-*.so*"] not in groups
def test_route_to_vulkan_prebuilt_auto_intel_goes_upstream_and_drops_fork_pin():
# Routing fork -> upstream also drops the fork release pin, which is in a
# different tag namespace and would make the upstream resolver miss.
host = _host(is_linux = True, is_x86_64 = True, has_intel_gpu = True)
routed, repo, tag = ilp._route_to_vulkan_prebuilt(host, FORK, "b9596-mix-abc", force_cpu = False)
assert repo == UPSTREAM
assert tag == ""
assert routed.has_intel_gpu is True
def test_route_to_vulkan_prebuilt_preserves_explicit_upstream_pin():
# A pin set WITH an explicit upstream repo is already on upstream -> kept.
host = _host(is_linux = True, is_x86_64 = True, has_intel_gpu = True)
_routed, repo, tag = ilp._route_to_vulkan_prebuilt(host, UPSTREAM, "b9596", force_cpu = False)
assert repo == UPSTREAM
assert tag == "b9596"
def test_route_to_vulkan_prebuilt_cpu_fallback_wins():
# --cpu-fallback suppresses Vulkan routing even for an Intel host.
host = _host(is_linux = True, is_x86_64 = True, has_intel_gpu = True)
routed, repo, tag = ilp._route_to_vulkan_prebuilt(host, FORK, "b9596-mix-abc", force_cpu = True)
assert repo == FORK
assert tag == "b9596-mix-abc"
assert routed is host
def test_route_to_vulkan_prebuilt_hidden_nvidia_not_rerouted():
# A mixed NVIDIA+Intel host that hid NVIDIA (CUDA_VISIBLE_DEVICES=""/-1):
# physical NVIDIA present but not usable. Must NOT auto-route to Vulkan, or
# Vulkan (which ignores CUDA_VISIBLE_DEVICES) could grab the reserved GPU.
host = _host(
is_linux = True,
is_x86_64 = True,
has_intel_gpu = True,
has_physical_nvidia = True,
has_usable_nvidia = False,
)
_routed, repo, _tag = ilp._route_to_vulkan_prebuilt(host, FORK, "", force_cpu = False)
assert repo == FORK
def test_route_to_vulkan_prebuilt_rocm_host_not_rerouted():
# An Intel iGPU alongside a usable ROCm GPU stays on its ROCm/fork path.
host = _host(is_linux = True, is_x86_64 = True, has_intel_gpu = True, has_rocm = True)
_routed, repo, _tag = ilp._route_to_vulkan_prebuilt(host, FORK, "", force_cpu = False)
assert repo == FORK
def test_route_to_vulkan_prebuilt_non_intel_unchanged():
host = _host(is_linux = True, is_x86_64 = True)
routed, repo, _tag = ilp._route_to_vulkan_prebuilt(host, FORK, "", force_cpu = False)
assert repo == FORK
assert routed is host
def test_resolve_prebuilt_intel_host_routes_to_upstream(monkeypatch, capsys):
# The --resolve-prebuilt probe must agree with the install path: an
# auto-detected Intel host resolves against upstream (Vulkan), not the fork.
monkeypatch.setattr(
ilp, "detect_host", lambda: _host(is_linux = True, is_x86_64 = True, has_intel_gpu = True)
)
seen, out = _run_resolve_capture_host(monkeypatch, capsys)
assert seen["repo"] == UPSTREAM
assert out["repo"] == UPSTREAM
# ---------------------------------------------------------------------------
# windows_intel_gpu_in_registry: the in-process Windows Intel probe. A fake
# winreg module stands in for the real registry so the walk runs anywhere.
# ---------------------------------------------------------------------------
class _FakeRegKey:
def __init__(
self,
subkeys = None,
values = None,
denied = False,
):
self.subkeys = subkeys or {}
self.values = values or {}
self.denied = denied
def __enter__(self):
return self
def __exit__(self, *exc):
return False
class _FakeWinreg:
HKEY_LOCAL_MACHINE = object()
def __init__(self, root_key):
self._root_key = root_key
def OpenKey(self, parent, name):
if parent is self.HKEY_LOCAL_MACHINE:
# Pin the production constant: a typo'd class GUID must fail here,
# not silently return the fake tree.
if name != ilp._WINDOWS_DISPLAY_CLASS_KEY:
raise FileNotFoundError(name)
if self._root_key is None:
raise FileNotFoundError(name)
return self._root_key
key = parent.subkeys.get(name)
if key is None:
# Real winreg raises OSError, never KeyError, for a missing key.
raise FileNotFoundError(name)
if key.denied:
raise PermissionError(name)
return key
def QueryInfoKey(self, key):
return (len(key.subkeys), len(key.values), 0)
def EnumKey(self, key, index):
return list(key.subkeys)[index]
def QueryValueEx(self, key, value_name):
if value_name not in key.values:
raise FileNotFoundError(value_name)
return (key.values[value_name], 1)
def _probe_with_display_class(monkeypatch, adapters):
# The helper lazily does `import winreg`; plant the fake in sys.modules the
# same way unsloth_cli/tests/test_start.py fakes it for _refresh_windows_path.
monkeypatch.setitem(sys.modules, "winreg", _FakeWinreg(_FakeRegKey(subkeys = adapters)))
return ilp.windows_intel_gpu_in_registry()
def test_windows_intel_registry_matches_vendor_id(monkeypatch):
assert (
_probe_with_display_class(
monkeypatch,
{
"0000": _FakeRegKey(
values = {
"MatchingDeviceId": r"PCI\VEN_8086&DEV_56A0&SUBSYS_12345678",
"DriverDesc": "Intel(R) Arc(TM) A770 Graphics",
}
),
},
)
is True
)
def test_windows_intel_registry_matches_driver_desc_without_device_id(monkeypatch):
assert (
_probe_with_display_class(
monkeypatch,
{
"0000": _FakeRegKey(values = {"DriverDesc": "Intel(R) UHD Graphics 630"}),
},
)
is True
)
def test_windows_intel_registry_ignores_non_intel_adapters(monkeypatch):
assert (
_probe_with_display_class(
monkeypatch,
{
"0000": _FakeRegKey(
values = {
"MatchingDeviceId": r"PCI\VEN_10DE&DEV_2684",
"DriverDesc": "NVIDIA GeForce RTX 4090",
}
),
"0001": _FakeRegKey(
values = {
"MatchingDeviceId": r"PCI\VEN_1002&DEV_744C",
"DriverDesc": "AMD Radeon RX 7900 XTX",
}
),
},
)
is False
)
def test_windows_intel_registry_skips_restricted_properties_subkey(monkeypatch):
# The real class key carries an ACL-restricted "Properties" subkey and can
# deny access to individual adapter keys; neither may abort the walk.
assert (
_probe_with_display_class(
monkeypatch,
{
"Properties": _FakeRegKey(denied = True),
"0000": _FakeRegKey(denied = True),
"0001": _FakeRegKey(
values = {
"MatchingDeviceId": r"PCI\VEN_8086&DEV_56A0",
}
),
},
)
is True
)
def test_windows_intel_registry_missing_class_key_is_false(monkeypatch):
monkeypatch.setitem(sys.modules, "winreg", _FakeWinreg(None))
assert ilp.windows_intel_gpu_in_registry() is False
def _detect_windows_host(
monkeypatch,
winreg_fake,
powershell_stdout = "",
):
"""Drive the real detect_host() as a GPU-less Windows host with a fake
registry, recording every run_capture invocation. Pins the wiring the
unit tests above cannot see: registry-first, CIM only on a registry miss."""
monkeypatch.setitem(sys.modules, "winreg", winreg_fake)
monkeypatch.setattr(ilp.platform, "system", lambda: "Windows")
monkeypatch.setattr(ilp.platform, "machine", lambda: "AMD64")
for _env in (
"CUDA_VISIBLE_DEVICES",
"HIP_VISIBLE_DEVICES",
"ROCR_VISIBLE_DEVICES",
"HIP_PATH",
"ROCM_PATH",
):
monkeypatch.delenv(_env, raising = False)
monkeypatch.setattr(
ilp.shutil,
"which",
lambda name: "powershell" if name in ("powershell", "pwsh") else None,
)
captured = []
def _fake_run_capture(command, **kwargs):
captured.append(command[0])
if command[0] == "powershell":
return SimpleNamespace(returncode = 0, stdout = powershell_stdout, stderr = "")
return SimpleNamespace(returncode = 1, stdout = "", stderr = "")
monkeypatch.setattr(ilp, "run_capture", _fake_run_capture)
return ilp.detect_host(), captured
def test_detect_host_registry_intel_skips_cim_probe(monkeypatch):
winreg = _FakeWinreg(
_FakeRegKey(
subkeys = {
"0000": _FakeRegKey(values = {"MatchingDeviceId": r"PCI\VEN_8086&DEV_56A0"}),
}
)
)
host, captured = _detect_windows_host(monkeypatch, winreg)
assert host.has_intel_gpu is True
assert "powershell" not in captured
def test_detect_host_cim_fallback_fires_on_registry_miss(monkeypatch):
winreg = _FakeWinreg(
_FakeRegKey(
subkeys = {
"0000": _FakeRegKey(values = {"MatchingDeviceId": r"PCI\VEN_10DE&DEV_2684"}),
}
)
)
host, captured = _detect_windows_host(
monkeypatch, winreg, powershell_stdout = "Intel(R) Arc(TM) A770 Graphics"
)
assert host.has_intel_gpu is True
assert "powershell" in captured
def test_windows_intel_registry_unexpected_error_is_false(monkeypatch):
# The probe is advisory: even a non-OSError bug in the walk must return
# False (deferring to the CIM fallback), never crash detect_host.
class _ExplodingWinreg:
HKEY_LOCAL_MACHINE = object()
def OpenKey(self, parent, name):
raise TypeError(name)
monkeypatch.setitem(sys.modules, "winreg", _ExplodingWinreg())
assert ilp.windows_intel_gpu_in_registry() is False
def test_detect_host_cim_rescues_exploding_registry(monkeypatch):
class _ExplodingWinreg:
HKEY_LOCAL_MACHINE = object()
def OpenKey(self, parent, name):
raise TypeError(name)
host, captured = _detect_windows_host(
monkeypatch, _ExplodingWinreg(), powershell_stdout = "Intel(R) Arc(TM) A770 Graphics"
)
assert host.has_intel_gpu is True
assert "powershell" in captured

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

@ -252,10 +252,17 @@ def test_legacy_browse_allowlist_includes_linux_run_media_mounts(monkeypatch, tm
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_external_media = SimpleNamespace(
linux_run_media_mount_roots = lambda: [media_root],
windows_drive_roots = lambda: [],
)
fake_studio_db = SimpleNamespace(
list_scan_folders = lambda: [],
contains_sensitive_path_component = studio_db.contains_sensitive_path_component,
# The media root is a legitimate mount, not denied; the .ssh 403 below
# comes from the credential check. A False stub keeps this OS-independent
# (on macOS tmp_path lives under the denied /private/var).
is_denied_system_path = lambda _p: False,
)
monkeypatch.setitem(sys.modules, "utils.paths", fake_paths)
monkeypatch.setitem(sys.modules, "utils.paths.external_media", fake_external_media)

320
studio/backend/tests/test_llama_admission.py Normal file
View file

@ -0,0 +1,320 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
import asyncio
import os
import sys
import threading
import pytest
_backend = os.path.join(os.path.dirname(__file__), "..")
sys.path.insert(0, _backend)
from core.inference import llama_admission
from core.inference.llama_admission import (
ADMISSION_CONTROL_ENV,
ADMISSION_KEEPALIVE_INTERVAL_ENV,
ADMISSION_MAX_QUEUE_ENV,
ADMISSION_QUEUE_TIMEOUT_ENV,
DEFAULT_ADMISSION_KEEPALIVE_INTERVAL_S,
DEFAULT_ADMISSION_MAX_QUEUE,
DEFAULT_ADMISSION_QUEUE_TIMEOUT_S,
LlamaAdmissionConfig,
LlamaAdmissionQueueFull,
get_llama_admission_queue,
llama_admission_config_from_env,
reset_llama_admission_queues,
)
@pytest.fixture(autouse = True)
def _reset_queues():
reset_llama_admission_queues()
yield
reset_llama_admission_queues()
def test_admission_config_defaults(monkeypatch):
for name in (
ADMISSION_CONTROL_ENV,
ADMISSION_QUEUE_TIMEOUT_ENV,
ADMISSION_KEEPALIVE_INTERVAL_ENV,
ADMISSION_MAX_QUEUE_ENV,
):
monkeypatch.delenv(name, raising = False)
config = llama_admission_config_from_env()
assert config.enabled is True
assert config.queue_timeout_s == DEFAULT_ADMISSION_QUEUE_TIMEOUT_S
assert config.keepalive_interval_s == DEFAULT_ADMISSION_KEEPALIVE_INTERVAL_S
assert config.max_queue == DEFAULT_ADMISSION_MAX_QUEUE
def test_admission_config_env_overrides(monkeypatch):
monkeypatch.setenv(ADMISSION_CONTROL_ENV, "off")
monkeypatch.setenv(ADMISSION_QUEUE_TIMEOUT_ENV, "0")
monkeypatch.setenv(ADMISSION_KEEPALIVE_INTERVAL_ENV, "0.25")
monkeypatch.setenv(ADMISSION_MAX_QUEUE_ENV, "0")
config = llama_admission_config_from_env()
assert config.enabled is False
assert config.queue_timeout_s is None
assert config.keepalive_interval_s == 0.25
assert config.max_queue is None
def test_admission_config_positive_queue_timeout_env(monkeypatch):
monkeypatch.setenv(ADMISSION_QUEUE_TIMEOUT_ENV, "600")
config = llama_admission_config_from_env()
assert config.queue_timeout_s == 600.0
def test_fifo_capacity_one_grants_next_waiter_on_release():
async def _run():
queue = get_llama_admission_queue("http://llama.test")
config = LlamaAdmissionConfig()
first = queue.reserve(capacity = 1, config = config)
second = queue.reserve(capacity = 1, config = config)
third = queue.reserve(capacity = 1, config = config)
first_lease = first.lease_nowait()
assert first_lease is not None
assert second.lease_nowait() is None
assert third.lease_nowait() is None
assert queue.snapshot().queued == 2
first_lease.release()
second_lease = await second.wait(0.1)
assert second_lease is not None
assert third.lease_nowait() is None
second_lease.release()
third_lease = await third.wait(0.1)
assert third_lease is not None
third_lease.release()
snapshot = queue.snapshot()
assert snapshot.active == 0
assert snapshot.queued == 0
asyncio.run(_run())
def test_queue_full_rejects_excess_waiter():
async def _run():
queue = get_llama_admission_queue("http://llama.test")
config = LlamaAdmissionConfig(max_queue = 1)
first = queue.reserve(capacity = 1, config = config)
queued = queue.reserve(capacity = 1, config = config)
assert first.lease_nowait() is not None
assert queued.lease_nowait() is None
with pytest.raises(LlamaAdmissionQueueFull):
queue.reserve(capacity = 1, config = config)
asyncio.run(_run())
def test_disabled_admission_bypasses_active_slot_limit():
async def _run():
queue = get_llama_admission_queue("http://llama.test")
config = LlamaAdmissionConfig(enabled = False)
first = queue.reserve(capacity = 1, config = config)
second = queue.reserve(capacity = 1, config = config)
assert first.lease_nowait() is not None
assert second.lease_nowait() is not None
assert queue.snapshot().active == 0
assert queue.snapshot().queued == 0
asyncio.run(_run())
def test_cancelling_promoted_waiter_releases_slot():
async def _run():
queue = get_llama_admission_queue("http://llama.test")
config = LlamaAdmissionConfig()
first = queue.reserve(capacity = 1, config = config)
second = queue.reserve(capacity = 1, config = config)
first_lease = first.lease_nowait()
first_lease.release()
await asyncio.sleep(0)
second.cancel()
snapshot = queue.snapshot()
assert snapshot.active == 0
assert snapshot.queued == 0
asyncio.run(_run())
def test_cancelling_promoted_waiter_before_delivery_releases_slot():
async def _run():
queue = get_llama_admission_queue("http://llama.test")
config = LlamaAdmissionConfig()
first = queue.reserve(capacity = 1, config = config)
second = queue.reserve(capacity = 1, config = config)
first_lease = first.lease_nowait()
first_lease.release()
second.cancel()
snapshot = queue.snapshot()
assert snapshot.active == 0
assert snapshot.queued == 0
asyncio.run(_run())
def test_external_waiter_future_cancel_invalidates_reservation():
async def _run():
queue = get_llama_admission_queue("http://llama.test")
config = LlamaAdmissionConfig()
first = queue.reserve(capacity = 1, config = config)
second = queue.reserve(capacity = 1, config = config)
first_lease = first.lease_nowait()
assert first_lease is not None
assert second._waiter is not None
second._waiter.future.cancel()
assert second.lease_nowait() is None
assert second.is_cancelled is True
assert await second.wait(0.01) is None
first_lease.release()
snapshot = queue.snapshot()
assert snapshot.active == 0
assert snapshot.queued == 0
asyncio.run(_run())
def test_wait_returns_none_when_waiter_future_cancelled_during_wait():
async def _run():
queue = get_llama_admission_queue("http://llama.test")
config = LlamaAdmissionConfig()
first = queue.reserve(capacity = 1, config = config)
second = queue.reserve(capacity = 1, config = config)
first_lease = first.lease_nowait()
assert first_lease is not None
assert second._waiter is not None
wait_task = asyncio.create_task(second.wait(1.0))
await asyncio.sleep(0)
second._waiter.future.cancel()
assert await asyncio.wait_for(wait_task, timeout = 0.1) is None
assert second.is_cancelled is True
first_lease.release()
snapshot = queue.snapshot()
assert snapshot.active == 0
assert snapshot.queued == 0
asyncio.run(_run())
def test_capacity_increase_promotes_existing_waiter_fifo():
async def _run():
queue = get_llama_admission_queue("http://llama.test")
config = LlamaAdmissionConfig()
first = queue.reserve(capacity = 1, config = config)
second = queue.reserve(capacity = 1, config = config)
first_lease = first.lease_nowait()
assert first_lease is not None
assert second.lease_nowait() is None
assert queue.snapshot().active == 1
assert queue.snapshot().queued == 1
third = queue.reserve(capacity = 2, config = config)
second_lease = await second.wait(0.1)
assert second_lease is not None
assert third.lease_nowait() is None
snapshot = queue.snapshot()
assert snapshot.capacity == 2
assert snapshot.active == 2
assert snapshot.queued == 1
first_lease.release()
third_lease = await third.wait(0.1)
assert third_lease is not None
second_lease.release()
third_lease.release()
snapshot = queue.snapshot()
assert snapshot.active == 0
assert snapshot.queued == 0
asyncio.run(_run())
def test_lease_release_is_idempotent_under_concurrent_calls():
async def _run():
queue = get_llama_admission_queue("http://llama.test")
config = LlamaAdmissionConfig()
reservation = queue.reserve(capacity = 1, config = config)
lease = reservation.lease_nowait()
assert lease is not None
threads = [threading.Thread(target = lease.release) for _ in range(16)]
for thread in threads:
thread.start()
for thread in threads:
thread.join()
snapshot = queue.snapshot()
assert snapshot.active == 0
assert snapshot.queued == 0
asyncio.run(_run())
def test_new_key_evicts_idle_prior_load_queues():
# Each model load carries a fresh ephemeral port, so a new base_url key must
# not leave the drained queues from earlier loads accumulating forever.
get_llama_admission_queue("http://127.0.0.1:1001")
get_llama_admission_queue("http://127.0.0.1:1002")
assert set(llama_admission._QUEUES) == {"http://127.0.0.1:1002"}
get_llama_admission_queue("http://127.0.0.1:1003")
assert set(llama_admission._QUEUES) == {"http://127.0.0.1:1003"}
def test_new_key_retains_in_flight_prior_load_queue():
config = LlamaAdmissionConfig()
busy = get_llama_admission_queue("http://127.0.0.1:2001")
async def _run():
reservation = busy.reserve(capacity = 1, config = config)
lease = reservation.lease_nowait()
assert lease is not None
# A new load must not drop a queue that still has an in-flight request.
get_llama_admission_queue("http://127.0.0.1:2002")
assert set(llama_admission._QUEUES) == {"http://127.0.0.1:2001", "http://127.0.0.1:2002"}
# Once it drains, the next load reclaims it.
lease.release()
get_llama_admission_queue("http://127.0.0.1:2003")
assert set(llama_admission._QUEUES) == {"http://127.0.0.1:2003"}
asyncio.run(_run())

53
studio/backend/tests/test_llama_cpp_effective_parallel_slots.py Normal file
View file

@ -0,0 +1,53 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
import os
import sys
import pytest
_backend = os.path.join(os.path.dirname(__file__), "..")
sys.path.insert(0, _backend)
from core.inference import llama_cpp as llama_cpp_module
from core.inference.llama_cpp import LlamaCppBackend
@pytest.fixture
def backend(monkeypatch):
monkeypatch.setattr(LlamaCppBackend, "_kill_orphaned_servers", lambda self: 0)
monkeypatch.setattr(llama_cpp_module.atexit, "register", lambda *_args, **_kwargs: None)
return LlamaCppBackend()
def test_effective_parallel_slots_initial_value_is_one(backend):
assert backend.effective_parallel_slots == 1
def test_effective_parallel_slots_commit_uses_final_positive_parallel(backend):
backend._commit_effective_parallel_slots(3)
assert backend.effective_parallel_slots == 3
@pytest.mark.parametrize("value", [None, 0, -2, "not-an-int"])
def test_effective_parallel_slots_commit_invalid_value_falls_back_to_one(backend, value):
backend._commit_effective_parallel_slots(value)
assert backend.effective_parallel_slots == 1
def test_effective_parallel_slots_reset_returns_to_one(backend):
backend._commit_effective_parallel_slots(4)
backend._reset_effective_parallel_slots()
assert backend.effective_parallel_slots == 1
def test_effective_parallel_slots_unload_resets_to_one(backend):
backend._commit_effective_parallel_slots(4)
backend.unload_model()
assert backend.effective_parallel_slots == 1

192
studio/backend/tests/test_llama_cpp_stream_cancel.py Normal file
View file

@ -0,0 +1,192 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
import contextlib
import os
import socket
import sys
import threading
import time
import httpx
import pytest
_backend = os.path.join(os.path.dirname(__file__), "..")
sys.path.insert(0, _backend)
from core.inference.llama_cpp import LlamaCppBackend, _LlamaStreamCancelled
def _backend_stub() -> LlamaCppBackend:
backend = LlamaCppBackend.__new__(LlamaCppBackend)
backend._process = object()
backend._healthy = True
backend._port = 48848
backend._effective_context_length = 4096
backend._supports_reasoning = False
backend._reasoning_always_on = False
backend._reasoning_style = "enable_thinking"
backend._supports_preserve_thinking = False
return backend
def test_stream_cancel_uses_internal_exception_not_generator_exit():
class FakeResponse:
status_code = 200
def close(self):
pass
class FakeStream:
def __enter__(self):
return FakeResponse()
def __exit__(self, *_args):
return False
class FakeClient:
def stream(self, *_args, **_kwargs):
return FakeStream()
cancel_event = threading.Event()
with pytest.raises(Exception) as exc_info:
with LlamaCppBackend._stream_with_retry(
FakeClient(),
"http://llama.test/v1/chat/completions",
{},
cancel_event,
):
cancel_event.set()
raise httpx.ReadError("client closed")
assert exc_info.type is _LlamaStreamCancelled
assert not issubclass(exc_info.type, GeneratorExit)
def test_generate_chat_completion_swallows_internal_stream_cancel(monkeypatch):
backend = _backend_stub()
@contextlib.contextmanager
def fake_open_stream(*_args, **_kwargs):
raise _LlamaStreamCancelled
monkeypatch.setattr(backend, "_open_stream", fake_open_stream)
chunks = list(
backend.generate_chat_completion(
[{"role": "user", "content": "hi"}],
cancel_event = threading.Event(),
)
)
assert chunks == []
class _StallUpstream:
"""Raw HTTP/1.1 server that streams one chunked SSE chunk, then holds the
socket open and silent so the client's next read blocks in recv() until its
side is torn down. Reproduces a mid-stream stall (llama-server goes quiet)."""
def __init__(self):
self._sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
self._sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
self._sock.bind(("127.0.0.1", 0))
self._sock.listen(1)
self.port = self._sock.getsockname()[1]
self._stop = threading.Event()
self._thread = threading.Thread(target = self._serve, daemon = True)
@property
def url(self) -> str:
return f"http://127.0.0.1:{self.port}/v1/chat/completions"
def __enter__(self):
self._thread.start()
return self
def __exit__(self, *_exc):
self._stop.set()
try:
self._sock.close()
except OSError:
pass
self._thread.join(timeout = 5)
def _serve(self) -> None:
try:
conn, _ = self._sock.accept()
except OSError:
return
with conn:
conn.settimeout(5)
try:
buf = b""
while b"\r\n\r\n" not in buf:
data = conn.recv(4096)
if not data:
return
buf += data
head, _, body = buf.partition(b"\r\n\r\n")
content_length = 0
for line in head.split(b"\r\n"):
if line.lower().startswith(b"content-length:"):
content_length = int(line.split(b":", 1)[1].strip())
break
while len(body) < content_length:
data = conn.recv(4096)
if not data:
break
body += data
except OSError:
return
conn.sendall(
b"HTTP/1.1 200 OK\r\n"
b"Content-Type: text/event-stream\r\n"
b"Transfer-Encoding: chunked\r\n"
b"\r\n"
)
chunk = b"data: hello\n\n"
conn.sendall(b"%x\r\n%s\r\n" % (len(chunk), chunk))
# Stall: stay open and silent until the client shuts its side down.
while not self._stop.wait(timeout = 0.05):
try:
conn.settimeout(0.05)
if conn.recv(1) == b"":
return
except socket.timeout:
continue
except OSError:
return
def test_cancel_interrupts_a_read_blocked_on_a_mid_stream_stall():
# Mid-stream stall: the reader is parked in recv() on a long bound read timeout,
# so response.close() alone can't wake it; the watcher must shut the socket down.
# Assert cancel lands in seconds, not at the far-off deadline (pre-fix: hung ~30s).
with _StallUpstream() as server:
cancel_event = threading.Event()
def _cancel_soon():
time.sleep(0.3)
cancel_event.set()
threading.Thread(target = _cancel_soon, daemon = True).start()
started = time.monotonic()
with httpx.Client(
limits = httpx.Limits(max_keepalive_connections = 0), trust_env = False
) as client:
with pytest.raises(_LlamaStreamCancelled):
with LlamaCppBackend._stream_with_retry(
client,
server.url,
{},
cancel_event,
first_token_deadline = started + 30,
) as response:
for _chunk in response.iter_text():
pass # first chunk arrives, then the read blocks silently
elapsed = time.monotonic() - started
assert elapsed < 10, f"cancel took {elapsed:.1f}s; the blocked read was not interrupted"

361
studio/backend/tests/test_llama_cpp_tool_loop.py
View file

@ -1061,6 +1061,80 @@ def test_same_turn_duplicate_web_search_is_internal_noop(monkeypatch):
]
def test_same_turn_duplicate_does_not_drop_later_parallel_call(monkeypatch):
# One batch: search(a), search(a) [duplicate], search(b). The duplicate is an
# internal no-op, but the distinct search(b) after it must still run, and the
# no-op nudge must land after the tool results rather than splitting them.
batch = [
_sse(
{
"tool_calls": [
{
"index": 0,
"id": "call_a1",
"type": "function",
"function": {"name": "web_search", "arguments": json.dumps({"query": "a"})},
},
{
"index": 1,
"id": "call_a2",
"type": "function",
"function": {"name": "web_search", "arguments": json.dumps({"query": "a"})},
},
{
"index": 2,
"id": "call_b",
"type": "function",
"function": {"name": "web_search", "arguments": json.dumps({"query": "b"})},
},
]
}
),
_done(),
]
final_stream = [_sse({"content": "Final answer."}), _done()]
payloads: list[dict] = []
backend = _make_backend(monkeypatch, [batch, final_stream], payloads)
calls: list[dict] = []
def fake_execute_tool(name, arguments, **_kwargs):
calls.append(arguments)
return "search-result"
monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool)
events = list(
backend.generate_chat_completion_with_tools(
messages = [{"role": "user", "content": "search"}],
tools = [{"type": "function", "function": {"name": "web_search"}}],
max_tool_iterations = 3,
)
)
# Both distinct calls ran; the duplicate did not (old `break` dropped search(b)).
assert calls == [{"query": "a"}, {"query": "b"}]
assert [e.get("tool_call_id") for e in events if e.get("type") == "tool_end"] == [
"call_a1",
"call_b",
]
# The next generation's conversation must be well-formed: the assistant lists
# only the executed calls (no orphan for the duplicate), the two tool results
# follow contiguously, and the no-op nudge lands after them, never between.
conv = payloads[1]["messages"]
asst = next(m for m in conv if m["role"] == "assistant" and m.get("tool_calls"))
assert [tc.get("id") for tc in asst["tool_calls"]] == ["call_a1", "call_b"]
after = conv[conv.index(asst) + 1 :]
assert [m["role"] for m in after[:2]] == ["tool", "tool"]
assert [m.get("tool_call_id") for m in after[:2]] == ["call_a1", "call_b"]
assert after[2]["role"] == "user" # deferred duplicate nudge, after the results
assert after[2]["content"].startswith(
"One earlier request to call tool 'web_search' in this batch was not executed"
)
assert "previous tool request" not in after[2]["content"].lower()
def test_same_turn_repeated_render_html_does_not_emit_second_provisional_start(monkeypatch):
same_turn_render_calls = [
_sse(
@ -1498,6 +1572,59 @@ def test_textual_mistral_marker_not_leaked_when_inline_with_preface(monkeypatch)
assert any("Let me search." in t for t in content_texts)
def test_textual_explicit_id_reuses_provisional_card(monkeypatch):
# A textual Mistral-style call with an explicit ``id`` must reconcile onto the
# open provisional TEXT card (keyed "call_0"), not spawn a duplicate under the
# explicit id (which the parser keeps for execution).
big_query = "cats " * 80 # push the drained call past the provisional floor
call = "[TOOL_CALLS]" + json.dumps(
[{"name": "web_search", "arguments": {"query": big_query}, "id": "explicit-42"}]
)
assert len(call) > 256
# Small chunks so the provisional card opens mid-generation (a single-shot
# delta parses instantly and never shows a provisional to exercise).
chunks = [call[i : i + 24] for i in range(0, len(call), 24)]
streams = [
[_sse({"content": c}) for c in chunks] + [_done()],
[_sse({"content": "done"}), _done()],
]
payloads: list[dict] = []
backend = _make_backend(monkeypatch, streams, payloads)
calls: list[tuple[str, dict]] = []
def fake_execute_tool(name, arguments, **_kwargs):
calls.append((name, arguments))
return "result"
monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool)
events = list(
backend.generate_chat_completion_with_tools(
messages = [{"role": "user", "content": "search"}],
tools = [{"type": "function", "function": {"name": "web_search"}}],
max_tool_iterations = 1,
)
)
assert calls == [("web_search", {"query": big_query})]
tool_starts = [e for e in events if e.get("type") == "tool_start"]
# Empty-args card = provisional open; full-args card = reconciled real start.
provisional = [e for e in tool_starts if not e.get("arguments")]
real = [e for e in tool_starts if e.get("arguments", {}).get("query")]
assert len(provisional) == 1, tool_starts # provisional actually opened
prov_id = provisional[0]["tool_call_id"]
# Exactly one real card, sharing the provisional id, not a duplicate under
# the explicit "explicit-42" id.
assert len(real) == 1, tool_starts
assert real[0]["tool_call_id"] == prov_id
assert real[0]["tool_name"] == "web_search"
assert {e["tool_call_id"] for e in tool_starts} == {prov_id}
# A single tool_end reconciles the card; no stale empty-result close.
ends = [e for e in events if e.get("type") == "tool_end"]
assert [e["tool_call_id"] for e in ends] == [prov_id]
assert ends[0]["result"] == "result"
def test_textual_llama_python_tag_marker_not_leaked(monkeypatch):
# Same leak class for the Llama-3 built-in ``<|python_tag|>NAME.call(...)`` form.
streams = [
@ -1826,6 +1953,39 @@ def test_large_python_tool_call_emits_early_provisional_start(monkeypatch):
assert any(e.get("type") == "tool_end" and e.get("tool_name") == "python" for e in events)
def test_auto_mode_render_html_suppresses_provisional_card_under_confirm(monkeypatch):
"""render_html is no longer unconditionally safe (a networked canvas asks), so
with confirm_tool_calls set under permission_mode="auto" its early provisional
card is suppressed; the real full-argument tool_start still fires and a static
canvas runs without a prompt."""
args = {"code": "<html>" + "x" * 80 + "</html>"}
first_stream = _streamed_structured_tool_call("render_html", args, "call_rh")
final_stream = [_sse({"content": "Done."}), _done()]
payloads: list[dict] = []
backend = _make_backend(monkeypatch, [first_stream, final_stream], payloads)
monkeypatch.setattr("core.inference.tools.execute_tool", lambda name, arguments, **_k: "OK")
events = list(
backend.generate_chat_completion_with_tools(
messages = [{"role": "user", "content": "make a card"}],
tools = [{"type": "function", "function": {"name": "render_html"}}],
confirm_tool_calls = True,
permission_mode = "auto",
max_tool_iterations = 1,
)
)
tool_starts = [e for e in events if e.get("type") == "tool_start"]
provisional = [e for e in tool_starts if not e.get("arguments")]
# The confirm gate now suppresses the early provisional card for render_html.
assert provisional == [], tool_starts
real = [e for e in tool_starts if e.get("arguments")]
assert real and real[0]["tool_name"] == "render_html"
# A static canvas is classified safe, so it still runs without an approval gate.
assert real[0].get("awaiting_confirmation") in (False, None)
def test_small_python_tool_call_has_no_provisional_start(monkeypatch):
"""A small tool-call argument finishes streaming instantly, so it keeps the
existing behavior of a single (real) tool_start with no provisional card."""
@ -2883,3 +3043,204 @@ def test_gguf_valid_tool_calls_respect_max_tool_iterations(monkeypatch):
m.get("role") == "user" and "used all available tool calls" in m.get("content", "")
for m in payloads[2]["messages"]
), payloads[2]["messages"]
# ── Live tool-call argument streaming (tool_args events) ─────────────────────
def _python_tool_schema() -> list[dict]:
return [
{
"type": "function",
"function": {
"name": "python",
"description": "Run python code.",
"parameters": {
"type": "object",
"properties": {"code": {"type": "string"}},
"required": ["code"],
},
},
}
]
def test_structured_tool_args_stream_to_provisional_card(monkeypatch):
"""A large structured tool call must stream its arguments as tool_args events
to the provisional card (backlog that triggered the card, then each
fragment), while the executed call and the model's view stay exactly what the
accumulator built."""
code = "print('x')\n" + ("# pad\n" * 80)
args_json = json.dumps({"code": code})
call_id = "call_live_args"
split = _PROVISIONAL_ARGS_MIN_CHARS + 16
frag1, frag2, frag3 = (
args_json[:split],
args_json[split : split + 40],
args_json[split + 40 :],
)
def _tc_delta(fragment: str, with_header: bool) -> str:
entry: dict = {"index": 0, "function": {"arguments": fragment}}
if with_header:
entry.update({"id": call_id, "type": "function"})
entry["function"]["name"] = "python"
return _sse({"tool_calls": [entry]})
first_stream = [
_tc_delta(frag1, with_header = True),
_tc_delta(frag2, with_header = False),
_tc_delta(frag3, with_header = False),
_done(),
]
second_stream = [_sse({"content": "Done."}), _done()]
payloads: list[dict] = []
backend = _make_backend(monkeypatch, [first_stream, second_stream], payloads)
executed: list[tuple[str, dict]] = []
def fake_execute_tool(name, arguments, **_kwargs):
executed.append((name, arguments))
return "ok"
monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool)
events = list(
backend.generate_chat_completion_with_tools(
messages = [{"role": "user", "content": "run it"}],
tools = _python_tool_schema(),
max_tool_iterations = 1,
)
)
starts = [e for e in events if e.get("type") == "tool_start"]
assert starts and starts[0]["tool_call_id"] == call_id
args_events = [e for e in events if e.get("type") == "tool_args"]
assert args_events, "no tool_args events were streamed"
assert all(e["tool_call_id"] == call_id for e in args_events)
# First event is the backlog, the rest raw fragments; together the args JSON.
assert args_events[0]["text"] == frag1
assert "".join(e["text"] for e in args_events) == args_json
# The streamed display path must not perturb execution or the model view.
assert executed == [("python", {"code": code})]
assistant_messages = [m for m in payloads[1]["messages"] if m.get("role") == "assistant"]
tc = assistant_messages[-1]["tool_calls"][0]
assert tc["id"] == call_id
# Controller re-serializes args (normalized JSON); parsed payload unchanged.
assert json.loads(tc["function"]["arguments"]) == {"code": code}
def test_text_tool_call_streams_args_and_reconciles_card(monkeypatch):
"""A TEXT (XML) tool call must stream its raw call text as tool_args under the
id the stream-end parser assigns ("call_0"), so the provisional card and the
final tool_start reconcile."""
code = "print('hello')\n" + ("# filler\n" * 60)
call_json = json.dumps({"name": "python", "arguments": {"code": code}})
call_text = f"<tool_call>{call_json}</tool_call>"
chunks = [call_text[i : i + 48] for i in range(0, len(call_text), 48)]
first_stream = [_sse({"content": chunk}) for chunk in chunks] + [_done()]
second_stream = [_sse({"content": "Done."}), _done()]
payloads: list[dict] = []
backend = _make_backend(monkeypatch, [first_stream, second_stream], payloads)
executed: list[tuple[str, dict]] = []
def fake_execute_tool(name, arguments, **_kwargs):
executed.append((name, arguments))
return "ok"
monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool)
events = list(
backend.generate_chat_completion_with_tools(
messages = [{"role": "user", "content": "run it"}],
tools = _python_tool_schema(),
max_tool_iterations = 1,
)
)
starts = [e for e in events if e.get("type") == "tool_start"]
assert starts, "no tool_start emitted"
# Provisional card first (parser's first-call id), then the reconciling start.
assert starts[0]["tool_call_id"] == "call_0"
assert starts[0]["arguments"] == {}
assert starts[-1]["tool_call_id"] == "call_0"
args_events = [e for e in events if e.get("type") == "tool_args"]
assert args_events, "no tool_args events for the text call"
assert all(e["tool_call_id"] == "call_0" for e in args_events)
streamed = "".join(e["text"] for e in args_events)
# Streamed text is the drained call (display only); it must never leak into
# content events.
assert '"name": "python"' in streamed
assert executed == [("python", {"code": code})]
content_events = [e for e in events if e.get("type") == "content"]
assert not any("<tool_call>" in e["text"] for e in content_events)
def test_ordinary_json_answer_streams_no_tool_args(monkeypatch):
"""A large ordinary JSON answer (no enabled tool name) must not spawn a
provisional card or tool_args events; it stays a normal content answer."""
answer = json.dumps({"result": "fine", "data": ["x" * 40] * 12, "note": "not a tool call"})
chunks = [answer[i : i + 64] for i in range(0, len(answer), 64)]
stream = [_sse({"content": chunk}) for chunk in chunks] + [_done()]
payloads: list[dict] = []
backend = _make_backend(monkeypatch, [stream], payloads)
events = list(
backend.generate_chat_completion_with_tools(
messages = [{"role": "user", "content": "give me json"}],
tools = _python_tool_schema(),
max_tool_iterations = 1,
)
)
assert not [e for e in events if e.get("type") == "tool_args"]
assert not [e for e in events if e.get("type") == "tool_start"]
content_events = [e for e in events if e.get("type") == "content"]
assert content_events and answer in content_events[-1]["text"]
def test_provisional_text_card_closed_when_parse_fails(monkeypatch):
"""A >=256-char enabled-name text sniff opens a provisional card; if the
drained text then fails to parse (auto-heal off, truncated call), the
DRAINING false-positive path must close the card with a tool_end instead of
leaving it spinning forever."""
# Truncated mid-arguments and never closed: unparseable without healing.
call_text = '<tool_call>{"name": "python", "arguments": {"code": "' + "x" * (
_PROVISIONAL_ARGS_MIN_CHARS + 64
)
chunks = [call_text[i : i + 48] for i in range(0, len(call_text), 48)]
stream = [_sse({"content": chunk}) for chunk in chunks] + [_done()]
payloads: list[dict] = []
backend = _make_backend(monkeypatch, [stream], payloads)
executed: list[tuple[str, dict]] = []
def fake_execute_tool(name, arguments, **_kwargs):
executed.append((name, arguments))
return "ok"
monkeypatch.setattr("core.inference.tools.execute_tool", fake_execute_tool)
events = list(
backend.generate_chat_completion_with_tools(
messages = [{"role": "user", "content": "run it"}],
tools = _python_tool_schema(),
max_tool_iterations = 1,
auto_heal_tool_calls = False,
)
)
starts = [e for e in events if e.get("type") == "tool_start"]
ends = [e for e in events if e.get("type") == "tool_end"]
assert starts and starts[0]["tool_call_id"] == "call_0"
assert executed == [] # nothing parsed, nothing ran
assert ends, "provisional card left dangling (no tool_end)"
assert ends[-1]["tool_call_id"] == "call_0"

123
studio/backend/tests/test_llama_cpp_update.py
View file

@ -393,6 +393,9 @@ def test_start_update_source_build_installs_prebuilt(monkeypatch, tmp_path):
assert "--llama-tag" in cmd and "latest" in cmd
assert cmd[cmd.index("--rocm-gfx") + 1] == "gfx110x"
assert "--simple-policy" not in cmd and "--cpu-fallback" not in cmd
# No pin: source-build detection and the unpinned apply share the same
# "latest" resolver, so they already agree.
assert "--published-release-tag" not in cmd
def test_start_update_happy_path(monkeypatch, tmp_path):
@ -448,6 +451,48 @@ def test_start_update_happy_path(monkeypatch, tmp_path):
assert popen_kwargs["env"]["UNSLOTH_PROGRESS_PERCENT_STEP"] == "5"
def test_start_update_preserves_vulkan_via_env(monkeypatch, tmp_path):
# A Vulkan install (marker asset carries 'vulkan') must re-assert
# UNSLOTH_FORCE_VULKAN on update, or detect_host on a GPU box re-routes to
# CUDA/ROCm and silently replaces the Vulkan build.
install_dir = tmp_path / "llama.cpp"
binary = _write_install(
install_dir,
"b9493",
repo = "ggml-org/llama.cpp",
asset = "llama-b9493-bin-ubuntu-vulkan-x64.tar.gz",
)
monkeypatch.setattr(upd, "_find_binary", lambda: binary)
monkeypatch.setattr(upd, "_installer_script", lambda: tmp_path / "install_llama_prebuilt.py")
monkeypatch.setattr(freshness, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9518")
def _on_start(cmd):
_write_install(
install_dir,
"b9518",
repo = "ggml-org/llama.cpp",
asset = "llama-b9518-bin-ubuntu-vulkan-x64.tar.gz",
)
popen_kwargs: dict = {}
_patch_installer_popen(
monkeypatch,
lines = ["installed\n"],
on_start = _on_start,
captured_kwargs = popen_kwargs,
)
assert upd.start_update()["started"] is True
deadline = time.time() + 10
while time.time() < deadline:
job = upd.get_update_status()["job"]
if job["state"] in ("success", "error"):
break
time.sleep(0.05)
assert job["state"] == "success", job
assert popen_kwargs["env"]["UNSLOTH_FORCE_VULKAN"] == "1"
def test_start_update_reports_full_release_tag(monkeypatch, tmp_path):
install_dir = tmp_path / "llama.cpp"
binary = _write_install(install_dir, "b9595")
@ -477,6 +522,57 @@ def test_start_update_reports_full_release_tag(monkeypatch, tmp_path):
assert "Updated llama.cpp to b9596-mix-e6f2453." in job["message"]
def _run_start_update_to_completion():
res = upd.start_update()
assert res["started"] is True
deadline = time.time() + 10
while time.time() < deadline:
job = upd.get_update_status()["job"]
if job["state"] in ("success", "error"):
return job
time.sleep(0.05)
return upd.get_update_status()["job"]
def test_start_update_pinned_tag_mismatch_fails(monkeypatch, tmp_path):
# Installer stays on the pinned repo but produces a different tag -> it
# ignored the pin (the silent mismatch this pin exists to prevent). Fail loud.
monkeypatch.setattr(sys, "platform", "linux")
install_dir = tmp_path / "llama.cpp"
binary = _write_install(install_dir, "b9595")
monkeypatch.setattr(upd, "_find_binary", lambda: binary)
monkeypatch.setattr(upd, "_installer_script", lambda: tmp_path / "install_llama_prebuilt.py")
monkeypatch.setattr(
freshness, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9601-mix-a0e2906"
)
_patch_installer_popen(
monkeypatch,
on_start = lambda cmd: _write_install(install_dir, "b9500", release_tag = "b9500-mix-deadbee"),
)
job = _run_start_update_to_completion()
assert job["state"] == "error", job
assert "b9601-mix-a0e2906" in (job["error"] or "")
def test_start_update_pinned_reroute_to_other_repo_ok(monkeypatch, tmp_path):
# A Vulkan/Intel host reroutes fork->upstream and drops the pin, installing a
# different-repo tag. Legitimate: the pin check must not flag the repo switch.
monkeypatch.setattr(sys, "platform", "linux")
install_dir = tmp_path / "llama.cpp"
binary = _write_install(install_dir, "b9595", repo = "unslothai/llama.cpp")
monkeypatch.setattr(upd, "_find_binary", lambda: binary)
monkeypatch.setattr(upd, "_installer_script", lambda: tmp_path / "install_llama_prebuilt.py")
monkeypatch.setattr(
freshness, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9601-mix-a0e2906"
)
_patch_installer_popen(
monkeypatch,
on_start = lambda cmd: _write_install(install_dir, "b9601", repo = "ggml-org/llama.cpp"),
)
job = _run_start_update_to_completion()
assert job["state"] == "success", job
def test_start_update_installer_failure_reports_error(monkeypatch, tmp_path):
install_dir = tmp_path / "llama.cpp"
binary = _write_install(install_dir, "b9493")
@ -621,6 +717,33 @@ def test_install_cmd_cuda_marker_minimal_and_backward_compatible(monkeypatch, tm
assert "--cpu-fallback" not in cmd
def test_install_cmd_pins_offered_release_tag(monkeypatch, tmp_path):
# Apply must install exactly the release the banner offered. The installer's
# own "latest" comes from commit-date-ordered sources, which can lag the
# published_at-newest tag detection picked; unpinned, that lag makes Update
# reinstall the current build while the banner never clears.
monkeypatch.setattr(sys, "platform", "linux")
cmd = _capture_install_cmd(monkeypatch, tmp_path, latest = "b9601-mix-a0e2906")
# The full release identity is pinned, not the bare upstream base.
assert cmd[cmd.index("--published-release-tag") + 1] == "b9601-mix-a0e2906"
def test_install_cmd_pins_on_windows(monkeypatch, tmp_path):
# The darwin exemption must not leak to other platforms.
monkeypatch.setattr(sys, "platform", "win32")
cmd = _capture_install_cmd(monkeypatch, tmp_path)
assert cmd[cmd.index("--published-release-tag") + 1] == "b9518"
def test_install_cmd_does_not_pin_on_macos(monkeypatch, tmp_path):
# A pinned tag disables the installer's older-release walk-back, which macOS
# needs to skip prebuilts built for a newer macOS than the host.
monkeypatch.setattr(sys, "platform", "darwin")
cmd = _capture_install_cmd(monkeypatch, tmp_path)
assert "--published-release-tag" not in cmd
assert "--llama-tag" in cmd and "latest" in cmd
# --- refusal + maintenance-state coordination ---

193
studio/backend/tests/test_llama_cpp_vulkan_probe.py Normal file
View file

@ -0,0 +1,193 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Vulkan free-VRAM reader regression tests on a synthetic probe output.
Covers the post-probe handling in
``LlamaCppBackend._get_gpu_free_memory_vulkan``:
* integrated GPUs (probe reports is_igpu=1) leave a flat per-device host
margin matching llama.cpp's --fit-target, so context auto-sizing can't
over-commit shared RAM, and report total 0 (shared RAM is not a budget),
* discrete GPUs (is_igpu=0) keep their free untouched and pass their real
total through so the fit can reserve absolute headroom,
* an inherited ``GGML_VK_VISIBLE_DEVICES`` is passed through to ggml unchanged
(ggml applies it), not stripped or filtered in Python -- the probe reports
ggml's compact ordinal, which load_model pins with ``--device Vulkan<i>``.
The ggml Vulkan library is never loaded: subprocess.run is mocked to emit
the tab-separated lines the real ``_vulkan_probe.py`` would print.
"""
from __future__ import annotations
import subprocess
import sys
import types as _types
from pathlib import Path
from unittest import mock
import pytest
_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
if _BACKEND_DIR not in sys.path:
sys.path.insert(0, _BACKEND_DIR)
import importlib as _importlib # noqa: E402
def _maybe_stub(name: str, builder):
try:
_importlib.import_module(name)
except ImportError:
sys.modules[name] = builder()
def _build_loggers_stub():
m = _types.ModuleType("loggers")
m.get_logger = lambda name: __import__("logging").getLogger(name)
return m
_maybe_stub("loggers", _build_loggers_stub)
_maybe_stub("structlog", lambda: _types.ModuleType("structlog"))
from core.inference import llama_cpp as _llama_mod # noqa: E402
from core.inference.llama_cpp import ( # noqa: E402
LlamaCppBackend,
_llama_lib_dir,
_vulkan_lib_filename,
)
MIB = 1024 * 1024
GIB = 1024 * MIB
def _make_vulkan_install(tmp_path: Path) -> str:
"""A binary whose sibling dir holds the Vulkan ggml lib, so the
reader's ``is_vulkan_backend`` sibling-file check passes."""
bindir = tmp_path / "build" / "bin"
bindir.mkdir(parents = True)
binary = bindir / ("llama-server.exe" if sys.platform == "win32" else "llama-server")
binary.write_bytes(b"stub")
(bindir / _vulkan_lib_filename()).write_bytes(b"stub")
return str(binary)
def _mock_probe(rows: list[str], captured_env: dict | None = None):
"""Patch subprocess.run so the _vulkan_probe.py call returns ``rows``
(already tab-formatted), recording the env it was launched with."""
real_run = subprocess.run
def fake_run(cmd, *args, **kwargs):
if isinstance(cmd, list) and any("_vulkan_probe" in str(c) for c in cmd):
if captured_env is not None:
captured_env.clear()
captured_env.update(kwargs.get("env") or {})
return subprocess.CompletedProcess(
args = cmd, returncode = 0, stdout = "\n".join(rows), stderr = ""
)
return real_run(cmd, *args, **kwargs)
return mock.patch("subprocess.run", side_effect = fake_run)
def _row(
idx: int,
free_bytes: int,
is_igpu: int,
total_bytes: int = 0,
) -> str:
return f"{idx}\t{free_bytes}\t{is_igpu}\t{total_bytes}"
def test_integrated_gpu_leaves_host_margin(tmp_path):
binary = _make_vulkan_install(tmp_path)
# iGPU with 30 GiB free; reserve a flat 1024 MiB (llama.cpp --fit-target).
# total stays 0: shared system RAM is not a VRAM budget for the fit.
rows = [_row(0, 30 * GIB, is_igpu = 1, total_bytes = 32 * GIB)]
with _mock_probe(rows):
gpus = LlamaCppBackend._get_gpu_free_memory_vulkan(binary)
assert gpus == [(0, 30 * 1024 - 1024, 0)], gpus
def test_discrete_gpu_free_is_untouched_and_total_passed_through(tmp_path):
binary = _make_vulkan_install(tmp_path)
# 6 GiB free on a partially occupied 24 GiB card: free is untouched and the
# real total flows through so the fit reserves absolute headroom (CUDA/ROCm
# parity) instead of the looser free*frac budget.
rows = [_row(0, 6 * GIB, is_igpu = 0, total_bytes = 24 * GIB)]
with _mock_probe(rows):
gpus = LlamaCppBackend._get_gpu_free_memory_vulkan(binary)
assert gpus == [(0, 6 * 1024, 24 * 1024)], gpus
def test_large_discrete_gpu_is_untouched(tmp_path):
binary = _make_vulkan_install(tmp_path)
# A 48 GiB discrete card stays untouched regardless of size; only the
# iGPU flag triggers the host margin, never a VRAM/RAM ratio.
rows = [_row(0, 47 * GIB, is_igpu = 0, total_bytes = 48 * GIB)]
with _mock_probe(rows):
gpus = LlamaCppBackend._get_gpu_free_memory_vulkan(binary)
assert gpus == [(0, 47 * 1024, 48 * 1024)], gpus
def test_inherited_visible_devices_mask_is_passed_through_to_probe(tmp_path, monkeypatch):
# The mask is NOT stripped or filtered in Python: ggml parses it in raw
# physical-device space while this probe reports the compact post-filter
# ordinal, so mixing spaces would be wrong. It is passed through unchanged
# so ggml applies it to the same device list the launch will enumerate.
binary = _make_vulkan_install(tmp_path)
monkeypatch.setenv("GGML_VK_VISIBLE_DEVICES", "1")
captured: dict = {}
rows = [_row(0, 23 * GIB, is_igpu = 0, total_bytes = 24 * GIB)]
with _mock_probe(rows, captured_env = captured):
LlamaCppBackend._get_gpu_free_memory_vulkan(binary)
assert captured.get("GGML_VK_VISIBLE_DEVICES") == "1", captured
def test_vulkan_pin_args_uses_device_names_not_env_mask():
# Pin by compact device name via --device (the space the probe reports and
# the registry names), never by writing a compact ordinal into the raw
# GGML_VK_VISIBLE_DEVICES index space.
assert LlamaCppBackend._vulkan_pin_args([0]) == ["--device", "Vulkan0"]
assert LlamaCppBackend._vulkan_pin_args([1, 2]) == ["--device", "Vulkan1,Vulkan2"]
assert LlamaCppBackend._vulkan_pin_args(None) == []
assert LlamaCppBackend._vulkan_pin_args([]) == []
def test_vulkan_only_build_is_detected(tmp_path):
binary = _make_vulkan_install(tmp_path)
assert LlamaCppBackend._is_vulkan_backend(binary) is True
def test_multi_backend_build_is_not_vulkan_only(tmp_path):
# A custom build that ships CUDA (or HIP) alongside Vulkan must NOT be
# treated as Vulkan-only, or its CUDA GPU would be probed/pinned as a Vulkan
# device; defer to the CUDA/HIP path instead.
binary = _make_vulkan_install(tmp_path)
cuda = "ggml-cuda.dll" if sys.platform == "win32" else "libggml-cuda.so"
(_llama_lib_dir(binary) / cuda).write_bytes(b"stub")
assert LlamaCppBackend._is_vulkan_backend(binary) is False
@pytest.mark.skipif(sys.platform == "win32", reason = "shell wrapper fallback is POSIX")
def test_shell_wrapper_entrypoint_resolves_to_real_lib_dir(tmp_path):
# create_exec_entrypoint falls back to a #!/bin/sh wrapper at the install root
# when it cannot symlink; _find_llama_server_binary returns that root entrypoint,
# so _llama_lib_dir must follow the wrapper's exec target to build/bin -- else
# _is_vulkan_backend misses libggml-vulkan.so and the Vulkan probe/pin silently
# never engage on a valid Vulkan install.
import os
binary = _make_vulkan_install(tmp_path) # tmp_path/build/bin/llama-server + vulkan lib
bindir = Path(binary).parent
wrapper = tmp_path / "llama-server"
wrapper.write_text('#!/bin/sh\nexec "$(dirname "$0")/build/bin/llama-server" "$@"\n')
os.chmod(wrapper, 0o755)
assert _llama_lib_dir(str(wrapper)) == bindir
assert LlamaCppBackend._is_vulkan_backend(str(wrapper)) is True
if __name__ == "__main__":
raise SystemExit(pytest.main([__file__, "-v"]))

43
studio/backend/tests/test_llama_cpp_wait_for_vram_settle.py
View file

@ -373,6 +373,7 @@ def test_kill_orphaned_servers_returns_count():
with (
patch.dict(sys.modules, {"psutil": fake_psutil}),
patch.dict(os.environ, {"LLAMA_SERVER_PATH": fake_path}),
patch.object(LlamaCppBackend, "_pid_parent_is_alive", staticmethod(lambda pid: False)),
):
n = LlamaCppBackend._kill_orphaned_servers()
assert n == 1, "only the Studio-owned orphan should be counted"
@ -384,11 +385,53 @@ def test_kill_orphaned_servers_returns_count():
with (
patch.dict(sys.modules, {"psutil": fake_psutil}),
patch.dict(os.environ, {"LLAMA_SERVER_PATH": fake_path}),
patch.object(LlamaCppBackend, "_pid_parent_is_alive", staticmethod(lambda pid: False)),
):
assert LlamaCppBackend._kill_orphaned_servers() == 0
assert killed == []
def test_kill_orphaned_servers_spares_live_parent():
"""A Studio-owned llama-server whose parent is still running is not an
orphan (a live Studio or the user's shell owns it) and must never be
killed; only the true orphan (parent gone) is reaped."""
import os
mypid = os.getpid()
fake_path = "/tmp/unsloth-test-llama/llama-server"
killed: list[int] = []
class _FakeProc:
def __init__(self, pid, name, exe):
self.info = {"pid": pid, "name": name, "exe": exe}
def kill(self):
killed.append(self.info["pid"])
live_parent = _FakeProc(mypid + 1, "llama-server", fake_path)
true_orphan = _FakeProc(mypid + 2, "llama-server", fake_path)
fake_psutil = _types.ModuleType("psutil")
fake_psutil.NoSuchProcess = type("NoSuchProcess", (Exception,), {})
fake_psutil.AccessDenied = type("AccessDenied", (Exception,), {})
fake_psutil.ZombieProcess = type("ZombieProcess", (Exception,), {})
fake_psutil.process_iter = lambda attrs = None: [live_parent, true_orphan]
with (
patch.dict(sys.modules, {"psutil": fake_psutil}),
patch.dict(os.environ, {"LLAMA_SERVER_PATH": fake_path}),
patch.object(LlamaCppBackend, "_reap_recorded_pid", staticmethod(lambda: 0)),
patch.object(
LlamaCppBackend,
"_pid_parent_is_alive",
staticmethod(lambda pid: pid == mypid + 1),
),
):
n = LlamaCppBackend._kill_orphaned_servers()
assert n == 1, "only the true orphan should be reaped"
assert killed == [mypid + 2], "the live-parent server must be spared"
def test_startup_reaper_arms_settle_timestamp():
"""__init__ arms ``_last_kill_monotonic`` when the startup reaper kills an
orphan (so the first load_model waits for VRAM to settle), and leaves the

84
studio/backend/tests/test_llama_route_timeouts.py
View file

@ -203,3 +203,87 @@ def test_preheader_send_cleanup_on_disconnect_and_cancel():
asyncio.run(_run(False))
asyncio.run(_run(True))
def test_stream_stall_timeout_callable_re_resolved_each_read():
# The OpenAI passthrough passes a callable so the stall bound can switch to
# the short post-terminal grace mid-stream; it must be re-resolved per read,
# not captured once at generator start.
async def _run():
response = SimpleNamespace(request = SimpleNamespace(extensions = {"timeout": {}}))
values = iter([100.0, 2.0])
seen = []
class _Request:
async def is_disconnected(self):
return False
class _Items:
def __init__(self):
self.count = 0
async def __anext__(self):
self.count += 1
if self.count > 3:
raise StopAsyncIteration
return "data: {}"
async for _ in inf_mod._aiter_llama_stream_items(
_Items(),
cancel_event = threading.Event(),
request = _Request(),
response = response,
first_token_deadline = time.monotonic() + 1,
post_first_item_read_timeout_s = lambda: next(values, 5.0),
):
seen.append(response.request.extensions["timeout"].get("read"))
assert len(seen) == 3
# The callable is resolved right after the first item (arming the
# post-first window) and again before each later read, consuming
# successive values.
assert seen[0] == 100.0
assert 1.0 <= seen[1] <= 2.0
assert 4.0 <= seen[2] <= 5.0
asyncio.run(_run())
def test_stream_stall_timeout_disabled_clears_read_timeout():
# UNSLOTH_OPENAI_COMPAT_STREAM_STALL_TIMEOUT=0 disables the stall guard, so
# the callable returns None. Once a chunk has arrived the leftover
# first-token read timeout must be cleared, else a long post-first-chunk gap
# trips a stale deadline the operator asked to turn off.
async def _run():
response = SimpleNamespace(request = SimpleNamespace(extensions = {"timeout": {}}))
seen = []
class _Request:
async def is_disconnected(self):
return False
class _Items:
def __init__(self):
self.count = 0
async def __anext__(self):
self.count += 1
if self.count > 2:
raise StopAsyncIteration
return "data: {}"
async for _ in inf_mod._aiter_llama_stream_items(
_Items(),
cancel_event = threading.Event(),
request = _Request(),
response = response,
first_token_deadline = time.monotonic() + 5,
post_first_item_read_timeout_s = lambda: None,
):
seen.append(response.request.extensions["timeout"].get("read"))
# The first-token path armed a finite read timeout; after the first chunk
# with the guard disabled, it is cleared to None on every subsequent read.
assert seen == [None, None], seen
asyncio.run(_run())

48
studio/backend/tests/test_load_progress_throttle.py Normal file
View file

@ -0,0 +1,48 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""The /api/inference/load-progress throttle: one line per 10% step, reset per load."""
import pytest
import routes.inference as ri
class _Capture:
def __init__(self):
self.events = []
def info(self, event, **kw):
self.events.append((event, kw))
@pytest.fixture
def cap(monkeypatch):
capture = _Capture()
monkeypatch.setattr(ri, "logger", capture)
ri._reset_load_progress_step()
return capture
def _percents(cap):
return [kw["percent"] for _event, kw in cap.events]
def test_new_load_first_step_logs_after_reset(cap):
# Load A reaches 100%.
ri._log_load_progress_step(1.0, "ready")
assert _percents(cap) == [100]
# Same value keeps deduping (steady poll on a finished load stays quiet).
ri._log_load_progress_step(1.0, "ready")
assert _percents(cap) == [100]
# A new load arms the throttle, so a cached load B that reports 100% on its
# first poll still emits its progress line instead of hitting step == prev.
ri._reset_load_progress_step()
ri._log_load_progress_step(1.0, "ready")
assert _percents(cap) == [100, 100]
def test_steady_poll_dedups_within_a_load(cap):
for _ in range(3):
ri._log_load_progress_step(0.3, "mmap")
assert _percents(cap) == [30] # one line per 10% step, not one per poll

121
studio/backend/tests/test_logging_middleware.py
View file

@ -135,7 +135,7 @@ def test_duplicate_get_within_window_deduped(logs, monkeypatch):
mw = LoggingMiddleware(app)
for _ in range(3):
_run(mw(_http_scope("/api/chat/projects"), _noop_receive, send))
_run(mw(_http_scope("/api/models/browse-folders"), _noop_receive, send))
# Only the first of the identical GET/200 burst is logged.
assert len(logs.events) == 1
@ -183,11 +183,11 @@ def test_quiet_poll_paths_use_longer_heartbeat_window(logs, monkeypatch):
for _ in range(3):
_run(mw(_http_scope("/api/inference/monitor"), _noop_receive, send)) # quiet
for _ in range(3):
_run(mw(_http_scope("/api/chat/projects"), _noop_receive, send)) # normal
_run(mw(_http_scope("/api/models/browse-folders"), _noop_receive, send)) # normal
paths = [e[2]["path"] for e in logs.events]
assert paths.count("/api/inference/monitor") == 1 # collapsed to one heartbeat
assert paths.count("/api/chat/projects") == 3 # base dedup off -> all logged
assert paths.count("/api/models/browse-folders") == 3 # base dedup off -> all logged
def test_distinct_query_strings_are_not_deduped(logs, monkeypatch):
@ -242,3 +242,118 @@ def test_fastapi_static_asset_success_skips_log(tmp_path, logs):
assert response.status_code == 200
assert response.text == "body { color: black; }"
assert len(logs.events) == log_count
def _status_app(status):
async def app(scope, receive, send):
await send({"type": "http.response.start", "status": status, "headers": []})
await send({"type": "http.response.body", "body": b""})
return app
async def _drop(message):
pass
def _paths_logged(logs):
return [e[2]["path"] for e in logs.events]
def test_quiet_success_get_2xx_suppressed(logs):
# A GET/2xx poll on a quiet-success path logs nothing; the signal is in events.
for path in ("/api/chat/threads", "/api/export/status", "/api/hub/download-status"):
_run(LoggingMiddleware(_status_app(200))(_http_scope(path), _noop_receive, _drop))
assert logs.events == []
def test_chat_detail_and_message_reads_still_log(logs):
# Only the exact list polls are suppressed; detail/message reads carry latency
# signal and keep their access line.
for path in (
"/api/chat/threads/abc123",
"/api/chat/threads/abc123/messages",
"/api/chat/threads/abc123/messages/m1",
"/api/chat/projects/p1",
):
_run(LoggingMiddleware(_status_app(200))(_http_scope(path), _noop_receive, _drop))
assert _paths_logged(logs) == [
"/api/chat/threads/abc123",
"/api/chat/threads/abc123/messages",
"/api/chat/threads/abc123/messages/m1",
"/api/chat/projects/p1",
]
def test_quiet_success_is_get_only(logs):
# Mutations on the same paths still log (suppression is GET-only).
for method in ("POST", "PUT", "DELETE"):
_run(
LoggingMiddleware(_status_app(200))(
_http_scope("/api/chat/threads", method = method), _noop_receive, _drop
)
)
assert len(logs.events) == 3
def test_chat_pre_auth_401_suppressed_other_errors_logged(logs):
# The transient bootstrap 401 on a chat list GET is dropped, but a 500 (or any
# other status) still logs so real failures stay visible.
_run(
LoggingMiddleware(_status_app(401))(_http_scope("/api/chat/projects"), _noop_receive, _drop)
)
assert logs.events == []
_run(
LoggingMiddleware(_status_app(500))(_http_scope("/api/chat/projects"), _noop_receive, _drop)
)
assert _paths_logged(logs) == ["/api/chat/projects"]
def test_chat_401_logged_after_first_auth_refresh(logs):
# A chat 401 before any successful token refresh is the bootstrap race and is
# dropped, but once /api/auth/refresh has succeeded on this instance later chat
# 401s are real failures and stay visible.
responses: dict[tuple[str, str], int] = {}
async def app(scope, receive, send):
status = responses.get((scope["method"], scope["path"]), 200)
await send({"type": "http.response.start", "status": status, "headers": []})
await send({"type": "http.response.body", "body": b""})
mw = LoggingMiddleware(app)
responses[("GET", "/api/chat/threads")] = 401
_run(mw(_http_scope("/api/chat/threads"), _noop_receive, _drop))
assert logs.events == [] # bootstrap race: suppressed
# A successful refresh (POST, always logged) closes the bootstrap window.
responses[("POST", "/api/auth/refresh")] = 200
_run(mw(_http_scope("/api/auth/refresh", method = "POST"), _noop_receive, _drop))
assert _paths_logged(logs) == ["/api/auth/refresh"]
# Now the same chat 401 is a real failure and logs.
_run(mw(_http_scope("/api/chat/threads"), _noop_receive, _drop))
assert _paths_logged(logs) == ["/api/auth/refresh", "/api/chat/threads"]
def test_export_status_error_still_logs(logs):
# 2xx suppressed, but an HTTP-level error on export status remains visible.
_run(
LoggingMiddleware(_status_app(200))(_http_scope("/api/export/status"), _noop_receive, _drop)
)
assert logs.events == []
_run(
LoggingMiddleware(_status_app(500))(_http_scope("/api/export/status"), _noop_receive, _drop)
)
assert _paths_logged(logs) == ["/api/export/status"]
def test_legacy_download_progress_heartbeats_not_suppressed(logs, monkeypatch):
# Legacy /api/models download polls emit no progress events, so they heartbeat
# (first hit logs, the burst collapses) rather than vanish entirely.
monkeypatch.setattr(hmod, "_ACCESS_LOG_DEDUP_MS", 0)
monkeypatch.setattr(hmod, "_QUIET_POLL_DEDUP_MS", 1000)
mw = LoggingMiddleware(_status_app(200))
for _ in range(3):
_run(mw(_http_scope("/api/models/download-progress"), _noop_receive, _drop))
assert _paths_logged(logs) == ["/api/models/download-progress"]

177
studio/backend/tests/test_mcp_flatten_result.py Normal file
View file

@ -0,0 +1,177 @@
# 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 contextlib
import json
import sys
from pathlib import Path
from types import SimpleNamespace
_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
if _BACKEND_DIR not in sys.path:
sys.path.insert(0, _BACKEND_DIR)
from core.inference import mcp_client
from core.inference.mcp_client import (
MAX_IMAGE_PAYLOAD_CHARS,
MCP_IMAGES_SENTINEL,
_flatten_result,
call_tool_sync,
)
from core.inference.tool_loop_controller import is_tool_error, strip_result_for_model
PNG_B64 = "iVBORw0KGgoAAAANSUhEUg=="
def _text(value: str) -> SimpleNamespace:
return SimpleNamespace(type = "text", text = value)
def _image(data: str = PNG_B64, mime: str = "image/png") -> SimpleNamespace:
return SimpleNamespace(type = "image", data = data, mimeType = mime)
def _result(
*blocks,
is_error = False,
structured = None,
) -> SimpleNamespace:
return SimpleNamespace(
content = list(blocks),
is_error = is_error,
structured_content = structured,
)
def test_text_only_result_unchanged():
assert _flatten_result(_result(_text("hello"))) == "hello"
def test_image_only_result_keeps_image_and_notes_model():
flat = _flatten_result(_result(_image()))
body, payload = flat.split("\n" + MCP_IMAGES_SENTINEL, 1)
assert body == "[1 image attached; displayed to the user]"
assert json.loads(payload) == [{"data": PNG_B64, "mimeType": "image/png"}]
def test_text_plus_image_keeps_both():
flat = _flatten_result(_result(_text("Took a screenshot"), _image()))
body, payload = flat.split("\n" + MCP_IMAGES_SENTINEL, 1)
assert body == "Took a screenshot\n[1 image attached; displayed to the user]"
assert json.loads(payload)[0]["mimeType"] == "image/png"
def test_multiple_images_pluralized():
flat = _flatten_result(_result(_image(), _image(mime = "image/jpeg")))
body, payload = flat.split("\n" + MCP_IMAGES_SENTINEL, 1)
assert "[2 images attached; displayed to the user]" in body
assert [img["mimeType"] for img in json.loads(payload)] == ["image/png", "image/jpeg"]
def test_strip_result_for_model_drops_image_payload():
flat = _flatten_result(_result(_text("Took a screenshot"), _image()))
stripped = strip_result_for_model(flat)
assert stripped == "Took a screenshot\n[1 image attached; displayed to the user]"
assert PNG_B64 not in stripped
def test_strip_preserves_literal_mcp_sentinel_in_text():
# A tool that legitimately returns text containing the marker (e.g. reading
# source/docs that quote it) must not be truncated: the suffix is not a
# valid JSON image array.
text = "before\n__MCP_IMAGES__: literal from source\nafter"
assert strip_result_for_model(text) == text
def test_strip_preserves_non_image_json_after_marker():
text = 'log line\n__MCP_IMAGES__:["not", "image", "dicts"]'
assert strip_result_for_model(text) == text
def test_strip_removes_only_valid_terminal_envelope():
text = (
"Earlier mention: __MCP_IMAGES__: is documented here"
"\n[1 image attached; displayed to the user]"
'\n__MCP_IMAGES__:[{"data": "AAAA", "mimeType": "image/png"}]'
)
assert strip_result_for_model(text) == (
"Earlier mention: __MCP_IMAGES__: is documented here"
"\n[1 image attached; displayed to the user]"
)
def test_strip_still_handles_images_and_rag_sentinels():
assert strip_result_for_model("output\n__IMAGES__:['a.png']") == "output"
assert strip_result_for_model("answer\n__RAG_SOURCES__:[{}]") == "answer"
def test_error_result_keeps_error_prefix_and_images():
flat = _flatten_result(_result(_text("boom"), _image(), is_error = True))
assert flat.startswith("Error: boom")
assert is_tool_error(flat)
assert MCP_IMAGES_SENTINEL in flat
def test_image_only_error_no_longer_reports_no_content():
flat = _flatten_result(_result(_image(), is_error = True))
assert flat.startswith("Error: [1 image attached")
assert "tool returned no content" not in flat
def test_oversized_image_omitted_with_note():
huge = "A" * (MAX_IMAGE_PAYLOAD_CHARS + 1)
flat = _flatten_result(_result(_image(data = huge)))
assert flat == "[1 image omitted (too large)]"
assert MCP_IMAGES_SENTINEL not in flat
def test_oversized_budget_shared_across_images():
big = "A" * (MAX_IMAGE_PAYLOAD_CHARS - 10)
flat = _flatten_result(_result(_image(data = big), _image()))
body, payload = flat.split("\n" + MCP_IMAGES_SENTINEL, 1)
assert "1 image attached" in body
assert "1 image omitted (too large)" in body
images = json.loads(payload)
assert len(images) == 1 and images[0]["data"] == big
def test_non_image_binary_block_still_ignored():
flat = _flatten_result(
_result(SimpleNamespace(type = "audio", data = PNG_B64, mimeType = "audio/wav"))
)
assert flat == ""
def test_structured_content_fallback_still_used():
flat = _flatten_result(_result(structured = {"ok": True}))
assert flat == "{'ok': True}"
def test_call_tool_sync_passes_raise_on_error_false_and_keeps_error_images(monkeypatch):
# Guards that call_tool_sync passes raise_on_error=False, so an is_error result
# with image content reaches _flatten_result instead of FastMCP raising ToolError.
seen = {}
class _FakeClient:
async def call_tool(
self,
name,
args,
raise_on_error = True,
):
seen["raise_on_error"] = raise_on_error
return _result(_text("boom"), _image(), is_error = True)
@contextlib.asynccontextmanager
async def _fake_client(url, headers, use_oauth):
yield _FakeClient()
monkeypatch.setattr(mcp_client, "_client", _fake_client)
out = call_tool_sync("http://x", None, "take_screenshot", {})
assert seen["raise_on_error"] is False
assert out.startswith("Error: boom")
assert MCP_IMAGES_SENTINEL in out
assert is_tool_error(out)

290
studio/backend/tests/test_mcp_server.py Normal file
View file

@ -0,0 +1,290 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import asyncio
import sys
import types
import pytest
from mcp_server import BearerTokenMiddleware, _clamp, _dump, create_studio_mcp
def _get_tool(name):
tools = asyncio.run(create_studio_mcp().list_tools())
return {tool.name: tool for tool in tools}[name]
def test_studio_mcp_registers_control_plane_tools():
tools = asyncio.run(create_studio_mcp().list_tools())
assert {tool.name for tool in tools} == {
"studio_status",
"list_local_models",
"get_training_status",
"start_training",
"stop_training",
"list_training_runs",
"validate_recipe",
"get_recipe_job_status",
"get_recipe_job_dataset",
"load_checkpoint",
"export_gguf",
}
def test_dump_serializes_pydantic_values():
class Response:
def model_dump(self, *, mode):
assert mode == "json"
return {"ok": True}
assert _dump(Response()) == {"ok": True}
assert _dump({"already": "json"}) == {"already": "json"}
def test_bearer_token_middleware_rejects_wrong_token():
events = []
async def app(scope, receive, send):
events.append("app")
async def send(message):
events.append(message)
middleware = BearerTokenMiddleware(app, "secret")
asyncio.run(
middleware(
{"type": "http", "headers": [(b"authorization", b"Bearer wrong")]},
None,
send,
)
)
assert events[0]["status"] == 401
assert "app" not in events
def test_bearer_token_middleware_closes_unauthorized_websocket():
events = []
async def app(scope, receive, send):
events.append("app")
async def send(message):
events.append(message)
middleware = BearerTokenMiddleware(app, "secret")
asyncio.run(
middleware(
{"type": "websocket", "headers": []},
None,
send,
)
)
assert events == [{"type": "websocket.close", "code": 4401}]
def test_bearer_token_middleware_rejects_non_ascii_authorization():
# A non-ASCII bearer value must produce a clean 401, not a 500. Comparing on
# bytes avoids the str hmac.compare_digest TypeError on non-ASCII input.
events = []
async def app(scope, receive, send):
events.append("app")
async def send(message):
events.append(message)
middleware = BearerTokenMiddleware(app, "secret")
asyncio.run(
middleware(
{"type": "http", "headers": [(b"authorization", b"Bearer \xff\xff")]},
None,
send,
)
)
assert events[0]["status"] == 401
assert "app" not in events
def test_bearer_token_middleware_accepts_correct_token():
events = []
async def app(scope, receive, send):
events.append("app")
async def send(message):
events.append(message)
middleware = BearerTokenMiddleware(app, "secret")
asyncio.run(
middleware(
{"type": "http", "headers": [(b"authorization", b"Bearer secret")]},
None,
send,
)
)
assert events == ["app"]
def test_bearer_token_middleware_requires_non_empty_token():
async def app(scope, receive, send):
pass
for bad in ("", " "):
with pytest.raises(ValueError):
BearerTokenMiddleware(app, bad)
def test_bearer_token_middleware_rejects_non_ascii_token():
async def app(scope, receive, send):
pass
# non-ASCII tokens cannot be transmitted in an HTTP header by a standard
# client, so they are rejected at construction instead of locking out.
for bad in ("töken", "\U0001f600"):
with pytest.raises(ValueError):
BearerTokenMiddleware(app, bad)
def test_bearer_token_middleware_passes_through_non_http_scopes():
events = []
async def app(scope, receive, send):
events.append("app")
async def send(message):
events.append(message)
middleware = BearerTokenMiddleware(app, "secret")
asyncio.run(middleware({"type": "lifespan"}, None, send))
assert events == ["app"]
def test_clamp_restricts_to_inclusive_bounds():
assert _clamp(5, 1, 200) == 5
assert _clamp(-10, 1, 200) == 1
assert _clamp(10_000, 1, 200) == 200
assert _clamp(0, 1, 500) == 1
assert _clamp(1_000, 1, 500) == 500
def test_export_and_checkpoint_tools_expose_forwarded_fields():
export_props = set(_get_tool("export_gguf").parameters["properties"])
assert {"hf_token", "imatrix", "imatrix_path"} <= export_props
checkpoint_props = set(_get_tool("load_checkpoint").parameters["properties"])
assert {"hf_token", "approved_remote_code_fingerprint"} <= checkpoint_props
def _stub_module(monkeypatch, name, **attrs):
module = types.ModuleType(name)
for key, value in attrs.items():
setattr(module, key, value)
if "." in name:
module.__path__ = [] # mark package-like so submodule imports resolve
monkeypatch.setitem(sys.modules, name, module)
return module
def test_export_gguf_forwards_hf_token_and_imatrix(monkeypatch):
captured = {}
class FakeExportGGUFRequest:
def __init__(self, **kwargs):
captured.update(kwargs)
async def fake_export(request, current_subject):
return {"current_subject": current_subject}
_stub_module(monkeypatch, "models", ExportGGUFRequest = FakeExportGGUFRequest)
_stub_module(monkeypatch, "routes")
_stub_module(monkeypatch, "routes.export", export_gguf = fake_export)
tool = _get_tool("export_gguf")
result = asyncio.run(
tool.fn(
save_directory = "/tmp/out",
quantization_method = ["Q4_K_M", "Q8_0"],
push_to_hub = True,
repo_id = "me/model",
hf_token = "hf_secret",
imatrix = True,
imatrix_path = "/tmp/imatrix.dat",
)
)
assert captured["hf_token"] == "hf_secret"
assert captured["imatrix"] is True
assert captured["imatrix_path"] == "/tmp/imatrix.dat"
assert captured["quantization_method"] == ["Q4_K_M", "Q8_0"]
assert result["current_subject"] == "mcp"
def test_load_checkpoint_forwards_token_and_fingerprint(monkeypatch):
captured = {}
class FakeLoadCheckpointRequest:
def __init__(self, **kwargs):
captured.update(kwargs)
async def fake_load(request, current_subject):
return {"current_subject": current_subject}
_stub_module(monkeypatch, "models", LoadCheckpointRequest = FakeLoadCheckpointRequest)
_stub_module(monkeypatch, "routes")
_stub_module(monkeypatch, "routes.export", load_checkpoint = fake_load)
tool = _get_tool("load_checkpoint")
asyncio.run(
tool.fn(
checkpoint_path = "/tmp/ckpt",
approved_remote_code_fingerprint = "sha256:abc",
hf_token = "hf_secret",
)
)
assert captured["hf_token"] == "hf_secret"
assert captured["approved_remote_code_fingerprint"] == "sha256:abc"
def test_list_training_runs_clamps_pagination(monkeypatch):
captured = {}
async def fake_list_runs(limit, offset, current_subject):
captured["limit"] = limit
captured["offset"] = offset
return {"ok": True}
_stub_module(monkeypatch, "routes")
_stub_module(monkeypatch, "routes.training_history", list_training_runs = fake_list_runs)
tool = _get_tool("list_training_runs")
asyncio.run(tool.fn(limit = 10_000, offset = -5))
assert captured["limit"] == 200
assert captured["offset"] == 0
def test_get_recipe_job_dataset_clamps_pagination(monkeypatch):
captured = {}
def fake_job_dataset(job_id, limit, offset):
captured["limit"] = limit
captured["offset"] = offset
return {"ok": True}
_stub_module(monkeypatch, "routes")
_stub_module(monkeypatch, "routes.data_recipe")
_stub_module(monkeypatch, "routes.data_recipe.jobs", job_dataset = fake_job_dataset)
tool = _get_tool("get_recipe_job_dataset") # this tool is synchronous
tool.fn(job_id = "job-1", limit = -1, offset = -9)
assert captured["limit"] == 1
assert captured["offset"] == 0

76
studio/backend/tests/test_mcp_servers.py
View file

@ -198,7 +198,12 @@ def test_call_tool_sync_respects_pre_set_cancel_event(monkeypatch):
async def __aexit__(self, *args):
return False
async def call_tool(self, name, args):
async def call_tool(
self,
name,
args,
raise_on_error = True,
):
import asyncio as _asyncio
await _asyncio.sleep(30) # never finishes during the test
@ -520,7 +525,12 @@ def test_call_tool_sync_short_circuits_on_pre_set_cancel(monkeypatch):
async def __aexit__(self, *args):
return False
async def call_tool(self, name, args):
async def call_tool(
self,
name,
args,
raise_on_error = True,
):
return "ran"
monkeypatch.setattr(mcp_client, "_client", lambda *a, **kw: _StubClient())
@ -787,6 +797,68 @@ def test_update_display_name_keeps_tool_cache(tmp_path, monkeypatch):
assert mcp_client.get_cached_tools("s1") == cached
def test_update_rename_keeps_stdio_session(tmp_path, monkeypatch):
"""The edit dialog resends url/headers/oauth unchanged on a rename, so gating
the close on field presence would drop the live stdio session. Only a real
endpoint/auth change may close it."""
import asyncio
import json
_reset_db(tmp_path, monkeypatch)
from models.mcp_servers import McpServerUpdate
import routes.mcp_servers as routes_mcp
closed: list = []
monkeypatch.setattr(routes_mcp, "stdio_mcp_enabled", lambda: True)
monkeypatch.setattr(routes_mcp, "close_stdio_sessions", lambda *a, **k: closed.append(a))
mcp_servers_db.create_server(
id = "s1",
display_name = "A",
url = "npx demo-server",
headers_json = json.dumps({"API_KEY": "x"}),
is_enabled = True,
)
asyncio.run(
routes_mcp.update_mcp_server(
"s1",
McpServerUpdate(
display_name = "B",
url = "npx demo-server",
headers = {"API_KEY": "x"},
use_oauth = False,
),
current_subject = "u",
)
)
assert closed == []
assert mcp_servers_db.get_server("s1")["display_name"] == "B"
def test_update_stdio_command_change_closes_session(tmp_path, monkeypatch):
"""A real command change must still close the old stdio session."""
import asyncio
_reset_db(tmp_path, monkeypatch)
from models.mcp_servers import McpServerUpdate
import routes.mcp_servers as routes_mcp
closed: list = []
monkeypatch.setattr(routes_mcp, "stdio_mcp_enabled", lambda: True)
monkeypatch.setattr(routes_mcp, "close_stdio_sessions", lambda *a, **k: closed.append(a))
mcp_servers_db.create_server(
id = "s1",
display_name = "A",
url = "npx demo-server",
is_enabled = True,
)
asyncio.run(
routes_mcp.update_mcp_server(
"s1", McpServerUpdate(url = "npx other-server"), current_subject = "u"
)
)
assert len(closed) == 1
def test_update_disable_evicts_tool_cache(tmp_path, monkeypatch):
"""Disabling a server must drop its cached tools, not leave them unread."""
import asyncio

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