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

Merge branch 'main' into fix/6860-attn-mask-compat

Browse source
This commit is contained in:
Etherl 2026-07-23 20:02:12 +03:00 • committed by GitHub
commit e16b9ac5f1
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
982 changed files with 170368 additions and 19555 deletions
Download patch file Download diff file Expand all files Collapse all files

2
.gitattributes vendored
View file

@ -6,7 +6,7 @@
# them when run in WSL/Linux (e.g. `set -e` -> "set: Illegal option -").
*.sh text eol=lf
# Normalize Studio frontend sources to LF. Scoped to the frontend tree (rather
# Normalize Unsloth frontend sources to LF. Scoped to the frontend tree (rather
# than repo-wide *.ts/*.tsx/... rules) so the policy can't force LF on files
# elsewhere. text=auto lets Git detect and leave binary assets (logos, fonts)
# untouched while text files (.ts/.tsx/.json/.html/.svg/...) are stored as LF.

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

@ -36,6 +36,23 @@ AGENT="${2:?usage: agent-guides-drive.sh <mode> <agent>}"
# Determinism (seed/temp) is applied at the server level by
# serve-unsloth-run.sh --extra; agents inherit it through the API.
TIMEOUT="${AGENT_INVOKE_TIMEOUT:-180}"
# opencode is the slow outlier. Unlike the print-mode agents (claude -p, codex
# exec) it runs a full turn AND a separate small_model call to name the session,
# so one connection reply takes ~8 min on a CPU-served 4B -- right at the shared
# 600s cap, so the cell flaked when a run drifted past a ~480s success. Give it
# headroom (still well under the 40-min job budget); the fast agents keep the
# tight cap that still catches a real headless-TTY hang.
case "$AGENT" in
opencode)
# Double it, but only for a bare-integer seconds value. A GNU timeout(1)
# duration suffix (s/m/h/d, including floats like 0.5s) is left unchanged so
# the arithmetic never sees a non-number; timeout(1) parses it directly.
case "$TIMEOUT" in
*[!0-9]*) ;;
*) TIMEOUT=$(( TIMEOUT * 2 )) ;;
esac
;;
esac
# Claude refuses --dangerously-skip-permissions outside a sandbox; the CI runner
# IS the sandbox, so declare it (mirrors unslothai/scripts launcher.sh). Harmless
@ -154,15 +171,20 @@ raw_env() { # $1 = var name -> value (one shlex-quote layer stripped)
# writers as a side effect (it writes each agent's relocated session config).
parse_connect() {
local raw="$LOGS_DIR/connect-${AGENT}.txt"
if ! unsloth start "$AGENT" --no-launch --api-key "$UNSLOTH_API_KEY" > "$raw" 2>&1; then
# CONNECT_YOLO=1 adds --yolo. opencode/openclaw gate tool approval through their
# config (which now prompts by default), so the file-edit test opts into auto-approval
# here, the same intent as claude/codex's per-call bypass flags.
local yolo=()
[ -n "${CONNECT_YOLO:-}" ] && yolo=(--yolo)
if ! unsloth start "$AGENT" --no-launch "${yolo[@]}" --api-key "$UNSLOTH_API_KEY" > "$raw" 2>&1; then
cat_redacted "$raw"
guide_fail "'unsloth start ${AGENT} --no-launch' exited non-zero"
fi
echo "[$AGENT] connect --no-launch printed:"; cat_redacted "$raw"
CONNECT_ENV="$(grep -E '^(export |unset )' "$raw" || true)"
# The launch command is the last non-export, non-status line. start.py
# prints "Studio <url> · model <id>" and "Updated ..." status lines first.
CONNECT_CMD="$(grep -vE '^(export |unset |Studio |Updated |Disabled |Warning|Loading)' "$raw" \
# prints "Unsloth <url> · model <id>" and "Updated ..." status lines first.
CONNECT_CMD="$(grep -vE '^(export |unset |Unsloth |Updated |Disabled |Warning|Loading)' "$raw" \
| grep -E '[^[:space:]]' | tail -1)"
[ -n "$CONNECT_CMD" ] || guide_fail "could not parse a launch command from connect --no-launch output"
redact "$raw"
@ -371,7 +393,7 @@ case "$MODE" in
hermes) patch_hermes_tools none
invoke_via_connect "$OUT" -z "$PROMPT" ;;
openclaw) patch_openclaw_agent notools
invoke_via_connect "$OUT" agent --local --agent ci \
CONNECT_CMD_OVERRIDE=openclaw invoke_via_connect "$OUT" agent --local --agent ci \
--model "unsloth/${UNSLOTH_MODEL_ID}" --message "$PROMPT" ;;
*) invoke_via_connect "$OUT" "$PROMPT" ;;
esac
@ -394,7 +416,10 @@ case "$MODE" in
T2='Run hello.py with python and show me the exact output.'
# The start.py recipe writers + crosscheck must see the repo; run them
# from the repo root BEFORE cd-ing into the scratch work dir.
# from the repo root BEFORE cd-ing into the scratch work dir. opencode/openclaw
# gate tool approval through their config (prompting by default), so file-edit
# opts them into auto-approval to run edits/commands headlessly.
case "$AGENT" in opencode|openclaw) CONNECT_YOLO=1 ;; esac
parse_connect
crosscheck_contract
# File-edit needs real tools, so we cannot zero them as in connection.
@ -441,7 +466,7 @@ case "$MODE" in
fi ;;
opencode) invoke_via_connect "$out" run "$prompt" ;;
hermes) invoke_via_connect "$out" -z "$prompt" ;;
openclaw) invoke_via_connect "$out" agent --local --agent ci \
openclaw) CONNECT_CMD_OVERRIDE=openclaw invoke_via_connect "$out" agent --local --agent ci \
--model "unsloth/${UNSLOTH_MODEL_ID}" --message "$prompt" ;;
*) invoke_via_connect "$out" "$prompt" ;;
esac
@ -519,6 +544,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

2
.github/scripts/assert-llama-loads.sh vendored
View file

@ -2,7 +2,7 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
#
# Assert Studio installed a llama.cpp that loads and runs on THIS macOS. Tests
# Assert Unsloth installed a llama.cpp that loads and runs on THIS macOS. Tests
# the contract that matters (binaries load and their minimum-OS is <= this host)
# instead of the old "did install.sh fall back to a source build?" grep, since a
# source build with a correct deployment target is a valid outcome.

2
.github/scripts/assert-prompt-cache.sh vendored
View file

@ -31,7 +31,7 @@
# (llama_cpp.py:337-340). So default: ~/.unsloth/studio/logs/llama-server/.
#
# <P> is the INTERNAL llama-server port (self._find_free_port(),
# llama_cpp.py:3489 / :4641) -- a RANDOM port, NOT the Studio port. So we must
# llama_cpp.py:3489 / :4641) -- a RANDOM port, NOT the Unsloth port. So we must
# NOT filter the log glob by STUDIO_PORT (the brief's `port-<STUDIO_PORT>`
# glob would never match). We pick the newest llama-*.log instead.
#

4
.github/scripts/hf-download-with-retry.sh vendored
View file

@ -3,7 +3,7 @@
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
#
# Download a single file from a Hugging Face repo with a stall-retry
# watchdog. Used by the Studio CI workflows so a hung hf-xet transfer
# watchdog. Used by the Unsloth CI workflows so a hung hf-xet transfer
# kills + retries instead of silently consuming the job's timeout.
#
# Usage: hf-download-with-retry.sh REPO FILE LOCAL_DIR
@ -35,7 +35,7 @@ REPO="${1:?usage: hf-download-with-retry.sh REPO FILE [LOCAL_DIR]}"
FILE="${2:?usage: hf-download-with-retry.sh REPO FILE [LOCAL_DIR]}"
# LOCAL_DIR is optional. If empty, hf falls back to HF_HUB_CACHE
# (~/.cache/huggingface/hub) which is the desired path for callers
# that populate HF_HOME for a downstream Studio model load.
# that populate HF_HOME for a downstream Unsloth model load.
LOCAL_DIR="${3:-}"
# Stall threshold per attempt, in seconds. Override with

69
.github/scripts/run-studio-permission-browser.sh vendored Executable file
View file

@ -0,0 +1,69 @@
#!/usr/bin/env bash
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
set -euo pipefail
port="${1:?usage: $0 PORT BROWSER [CHANNEL]}"
browser="${2:?usage: $0 PORT BROWSER [CHANNEL]}"
channel="${3:-}"
slug="$browser${channel:+-$channel}"
artifact_dir="logs/playwright-permissions-$slug"
server_log="logs/studio-permissions-$slug.log"
studio_home="${UNSLOTH_STUDIO_HOME:-$HOME/.unsloth/studio}"
set --
if [ -n "${STUDIO_PERMISSION_FRONTEND:-}" ]; then
set -- -f "$STUDIO_PERMISSION_FRONTEND"
fi
mkdir -p "$artifact_dir"
unsloth studio reset-password
UNSLOTH_API_ONLY=1 unsloth studio -H 127.0.0.1 -p "$port" "$@" \
>"$server_log" 2>&1 &
studio_pid=$!
cleanup() {
kill "$studio_pid" 2>/dev/null || true
wait "$studio_pid" 2>/dev/null || true
}
trap cleanup EXIT
healthy=0
for _ in $(seq 1 180); do
if curl -fs "http://127.0.0.1:$port/api/health" >/dev/null; then
healthy=1
break
fi
if ! kill -0 "$studio_pid" 2>/dev/null; then
tail -100 "$server_log" || true
exit 1
fi
sleep 1
done
if [ "$healthy" -ne 1 ]; then
tail -100 "$server_log" || true
exit 1
fi
old_password=$(cat "$studio_home/auth/.bootstrap_password")
new_password="CIPerm-$(python -c 'import secrets; print(secrets.token_urlsafe(16))')"
if [ "${GITHUB_ACTIONS:-}" = "true" ]; then
echo "::add-mask::$old_password"
echo "::add-mask::$new_password"
fi
export BASE_URL="http://127.0.0.1:$port"
export STUDIO_OLD_PW="$old_password"
export STUDIO_NEW_PW="$new_password"
export STUDIO_UI_STRICT=1
export STUDIO_UI_PERMISSION_ONLY=1
export STUDIO_UI_WALL_TIMEOUT_S=240
export STUDIO_PLAYWRIGHT_BROWSER="$browser"
export PW_ART_DIR="$artifact_dir"
if [ -n "$channel" ]; then
export STUDIO_PLAYWRIGHT_CHANNEL="$channel"
else
unset STUDIO_PLAYWRIGHT_CHANNEL || true
fi
python tests/studio/playwright_chat_ui.py

7
.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,13 +359,18 @@ 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 \
tests/test_bad_mappings_redirect.py \
tests/test_prefetch_snapshot_scope.py \
tests/test_gemma_2b_mapper_key.py \
--deselect 'tests/utils/test_attention_masks.py::test_run_attention_flash_varlen_receives_window_and_softcap'
# The deselected test monkeypatches flash_attn_varlen_func, which is
# only bound on the module when `flash_attn` is importable. flash_attn

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

@ -1,18 +1,16 @@
# 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 across all three platforms.
#
# 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
# studio-backend-ci.yml's "Repo tests (CPU)" job). The test reads both
# installer scripts, and on Windows Path.read_text() defaults to the
# cp1252 locale encoding, so a non-cp1252 byte in install.sh (it already
# contains a U+274C) raises UnicodeDecodeError there even though Linux and
# macOS default to UTF-8. The reads were pinned to encoding="utf-8" in
# #6166; this job keeps that from silently regressing by exercising the
# test on the platforms it claims parity for. Pure pytest, no GPU,
# sub-second, so the matrix is cheap.
# Why: the parity test guards that install.sh and install.ps1 stay in sync.
# It originally ran only on ubuntu-latest through studio-backend-ci.yml.
# On Windows, Path.read_text() defaults to the cp1252 locale encoding, so a
# non-cp1252 byte in install.sh raises UnicodeDecodeError even though Linux
# and macOS default to UTF-8. The reads were pinned to encoding="utf-8" in
# #6166; this matrix keeps that from silently regressing. Pure pytest, no GPU,
# sub-second, so the matrix is cheap. Linux also runs the POSIX rollback test
# under dash, matching the supported curl-to-sh installer path.
name: Cross-platform parity
@ -21,14 +19,20 @@ on:
paths:
- 'install.sh'
- 'install.ps1'
- 'tests/test_installer_skip_autostart.py'
- 'tests/python/test_cross_platform_parity.py'
- 'tests/sh/test_install_rollback_lifecycle.sh'
- 'tests/studio/test_install_rollback_lifecycle.ps1'
- '.github/workflows/cross-platform-parity-ci.yml'
push:
branches: [main]
paths:
- 'install.sh'
- 'install.ps1'
- 'tests/test_installer_skip_autostart.py'
- 'tests/python/test_cross_platform_parity.py'
- 'tests/sh/test_install_rollback_lifecycle.sh'
- 'tests/studio/test_install_rollback_lifecycle.ps1'
- '.github/workflows/cross-platform-parity-ci.yml'
workflow_dispatch:
@ -45,7 +49,7 @@ jobs:
strategy:
fail-fast: false
matrix:
os: [windows-latest, macos-latest]
os: [ubuntu-latest, windows-latest, macos-latest]
runs-on: ${{ matrix.os }}
timeout-minutes: 10
steps:
@ -57,5 +61,18 @@ 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
- name: PowerShell rollback lifecycle tests
if: runner.os == 'Windows'
shell: pwsh
run: pwsh -NoProfile -File tests/studio/test_install_rollback_lifecycle.ps1
- name: POSIX rollback lifecycle tests
if: runner.os == 'Linux'
run: sh tests/sh/test_install_rollback_lifecycle.sh

4
.github/workflows/lint-ci.yml vendored
View file

@ -13,10 +13,10 @@
# committed YAML / JSON config.
#
# TypeScript and Rust are NOT duplicated here on purpose:
# - Studio Frontend CI runs `npm run typecheck` (= `tsc --noEmit`)
# - Unsloth Frontend CI runs `npm run typecheck` (= `tsc --noEmit`)
# and `npm run build` (vite/swc) on every studio/frontend/**
# change, which is a full TS AST + type check.
# - Studio Tauri CI runs `tauri build --debug --no-bundle` on
# - Unsloth Tauri CI runs `tauri build --debug --no-bundle` on
# every studio/src-tauri/** or studio/frontend/** change, which
# compiles the Rust crate (= cargo check + cargo build).
# Each is a stricter check than a parse-only step would be, so a

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

@ -154,7 +154,7 @@ jobs:
path: gguf-cache
key: ${{ runner.os }}-gguf-${{ env.GGUF_REPO }}-${{ env.GGUF_FILE }}-v1
- name: Install Studio (--local, --no-torch)
- name: Install Unsloth (--local, --no-torch)
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
# Gated off PR (see note above); public GGUF still downloads.
@ -256,7 +256,7 @@ jobs:
done
fi
- name: Stop Studio
- name: Stop Unsloth
if: always()
run: |
# Guard the PID: an unset/zero UNSLOTH_SERVER_PID would make
@ -359,7 +359,7 @@ jobs:
path: gguf-cache
key: ${{ runner.os }}-gguf-${{ env.GGUF_REPO }}-${{ env.GGUF_FILE }}-v1
- name: Install Studio (--local, --no-torch)
- name: Install Unsloth (--local, --no-torch)
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
# Gated off PR (see note above); public GGUF still downloads.
@ -448,7 +448,7 @@ jobs:
done
fi
- name: Stop Studio
- name: Stop Unsloth
if: always()
run: |
# Guard the PID: an unset/zero UNSLOTH_SERVER_PID would make
@ -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 Unsloth (--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 Unsloth
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
@ -536,7 +706,7 @@ jobs:
path: hf-cache
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2
- name: Install Studio (--local, --no-torch)
- name: Install Unsloth (--local, --no-torch)
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
# Gated off PR (see note above); public GGUF still downloads.
@ -594,7 +764,7 @@ jobs:
done
fi
- name: Stop Studio
- name: Stop Unsloth
if: always()
run: |
# Guard the PID: an unset/zero UNSLOTH_SERVER_PID would make

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

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

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

@ -130,7 +130,7 @@ jobs:
# MLX support landed after the most recent unsloth-zoo PyPI
# release; the wheel still raises NotImplementedError on
# Apple Silicon when device_type.get_device_type() runs
# unguarded. Studio's own install.sh overlays unsloth-zoo
# unguarded. Unsloth's own install.sh overlays unsloth-zoo
# from git main for the same reason. Pulling deps lets pip
# resolve the platform-conditional MLX-only wheels (mlx,
# mlx-lm, mlx-vlm gated on darwin+arm64 in unsloth-zoo's
@ -317,13 +317,13 @@ jobs:
echo
done
# Validates the macOS prebuilt path Studio's setup.sh uses (#5963): install the
# Validates the macOS prebuilt path Unsloth's setup.sh uses (#5963): install the
# unslothai/llama.cpp fork's latest release, download a small public GGUF, and
# check llama-server /completion end to end. Split and placed last so the
# untrusted binary runs only in the final smoke step, after every HF_TOKEN step,
# leaving no token-bearing step or shared workspace for a tampered prebuilt to
# corrupt. GH_TOKEN: releases API; HF_TOKEN (withheld on PR): probe + GGUF fetch.
- name: Studio prebuilt llama.cpp install + GGUF download (Mac M1)
- name: Unsloth prebuilt llama.cpp install + GGUF download (Mac M1)
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
@ -344,12 +344,12 @@ jobs:
# Final step: runs the downloaded binaries with no secrets present, and clears
# the GitHub Actions command files so a tampered prebuilt cannot influence the job.
- name: Studio prebuilt llama.cpp GGUF inference smoke (Mac M1)
- name: Unsloth prebuilt llama.cpp GGUF inference smoke (Mac M1)
run: |
set -euo pipefail
unset GITHUB_ENV GITHUB_PATH GITHUB_OUTPUT GITHUB_STEP_SUMMARY
INSTALL_DIR="$HOME/.unsloth-studio-prebuilt-test/llama.cpp"
# Studio bundles only llama-server + llama-quantize (not llama-cli);
# Unsloth bundles only llama-server + llama-quantize (not llama-cli);
# inference goes through llama-server's HTTP /completion endpoint.
LLAMA_SERVER="$INSTALL_DIR/build/bin/llama-server"
LLAMA_QUANT="$INSTALL_DIR/build/bin/llama-quantize"
@ -400,4 +400,4 @@ jobs:
tail -40 /tmp/llama-server.log
exit 1
fi
echo "OK: Studio prebuilt llama.cpp on Mac M1 + GGUF /completion works"
echo "OK: Unsloth prebuilt llama.cpp on Mac M1 + GGUF /completion works"

78
.github/workflows/ossf.yml vendored Normal file
View file

@ -0,0 +1,78 @@
# This workflow uses actions that are not certified by GitHub. They are provided
# by a third-party and are governed by separate terms of service, privacy
# policy, and support documentation.
name: Scorecard supply-chain security
on:
# For Branch-Protection check. Only the default branch is supported. See
# https://github.com/ossf/scorecard/blob/main/docs/checks.md#branch-protection
branch_protection_rule:
# To guarantee Maintained check is occasionally updated. See
# https://github.com/ossf/scorecard/blob/main/docs/checks.md#maintained
schedule:
- cron: '21 20 * * 0'
push:
branches: [ "main" ]
# Declare default permissions as read only.
permissions: read-all
jobs:
analysis:
name: Scorecard analysis
runs-on: ubuntu-latest
# `publish_results: true` only works when run from the default branch. conditional can be removed if disabled.
if: github.event.repository.default_branch == github.ref_name || github.event_name == 'pull_request'
permissions:
# Needed to upload the results to code-scanning dashboard.
security-events: write
# Needed to publish results and get a badge (see publish_results below).
id-token: write
# Uncomment the permissions below if installing in a private repository.
# contents: read
# actions: read
steps:
- name: "Checkout code"
uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2
with:
persist-credentials: false
- name: "Run analysis"
uses: ossf/scorecard-action@f49aabe0b5af0936a0987cfb85d86b75731b0186 # v2.4.1
with:
results_file: results.sarif
results_format: sarif
# (Optional) "write" PAT token. Uncomment the `repo_token` line below if:
# - you want to enable the Branch-Protection check on a *public* repository, or
# - you are installing Scorecard on a *private* repository
# To create the PAT, follow the steps in https://github.com/ossf/scorecard-action?tab=readme-ov-file#authentication-with-fine-grained-pat-optional.
# repo_token: ${{ secrets.SCORECARD_TOKEN }}
# Public repositories:
# - Publish results to OpenSSF REST API for easy access by consumers
# - Allows the repository to include the Scorecard badge.
# - See https://github.com/ossf/scorecard-action#publishing-results.
# For private repositories:
# - `publish_results` will always be set to `false`, regardless
# of the value entered here.
publish_results: true
# (Optional) Uncomment file_mode if you have a .gitattributes with files marked export-ignore
# file_mode: git
# Upload the results as artifacts (optional). Commenting out will disable uploads of run results in SARIF
# format to the repository Actions tab.
- name: "Upload artifact"
uses: actions/upload-artifact@4cec3d8aa04e39d1a68397de0c4cd6fb9dce8ec1 # v4.6.1
with:
name: SARIF file
path: results.sarif
retention-days: 5
# Upload the results to GitHub's code scanning dashboard (optional).
# Commenting out will disable upload of results to your repo's Code Scanning dashboard
- name: "Upload to code-scanning"
uses: github/codeql-action/upload-sarif@v3
with:
sarif_file: results.sarif

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

@ -4,7 +4,7 @@ on:
workflow_dispatch:
inputs:
studio_version:
description: 'Studio version tag to release (for example, v0.1.39-beta)'
description: 'Unsloth version tag to release (for example, v0.1.39-beta)'
type: string
required: true
pypi_version:
@ -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
@ -56,7 +69,7 @@ jobs:
if not studio_version:
sys.exit('studio_version is required, for example v0.1.39-beta')
if re.fullmatch(r'v?20\d{2}\.\d+\.\d+(?:[-+][0-9A-Za-z.-]+)?', studio_version):
sys.exit(f'studio_version must be a Studio SemVer tag, not a date-style backend version: {studio_version}')
sys.exit(f'studio_version must be an Unsloth SemVer tag, not a date-style backend version: {studio_version}')
semver_tag = re.compile(
r'^v(0|[1-9]\d*)\.(0|[1-9]\d*)\.(0|[1-9]\d*)'
@ -133,7 +146,7 @@ jobs:
print(f'pypi_version={pypi_version}', file=output)
PY
- name: Verify PyPI package and Studio stamp
- name: Verify PyPI package and Unsloth stamp
shell: bash
env:
STUDIO_VERSION: ${{ steps.prepare.outputs.studio_version }}
@ -198,7 +211,7 @@ jobs:
fi
python3 scripts/stamp_studio_release.py --verify-dist "$RUNNER_TEMP/pypi-unsloth-dist" --expected "$STUDIO_VERSION"
else
echo "scripts/stamp_studio_release.py not found; release-desktop requires #5308 to verify the PyPI Studio stamp." >&2
echo "scripts/stamp_studio_release.py not found; release-desktop requires #5308 to verify the PyPI Unsloth stamp." >&2
exit 1
fi
@ -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 }}

36
.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,
@ -36,8 +36,8 @@
# - unsloth `huggingfacenotorch` extras (the canonical install path
# for fine-tuning users; pulls transformers / peft / accelerate /
# trl / datasets / diffusers / sentence-transformers / etc.)
# - all six Studio backend requirements files
# - Studio frontend (npm) and Tauri shell (cargo)
# - all six Unsloth backend requirements files
# - Unsloth frontend (npm) and Tauri shell (cargo)
# Each Python step builds a filtered dep list from pyproject.toml +
# requirements/*.txt before auditing. We do NOT install any of these
# -- pip-audit resolves through PyPI metadata, scan_packages.py
@ -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]
@ -216,7 +218,7 @@ jobs:
# on the runner). A comment line is left in place so the
# skipped specs are obvious in the artifact.
# The `huggingface` extra is `huggingfacenotorch` plus torch /
# torchvision / triton, deliberately skipped: Studio backend
# torchvision / triton, deliberately skipped: Unsloth backend
# already pins a torch and the +cu* / +cpu local-version tags
# trip up the PyPI resolver in `-r` mode.
run: |
@ -251,7 +253,7 @@ jobs:
# `-r requirements.txt` resolves the requirements through pip's
# dependency resolver against PyPI metadata and audits the
# resolved tree without ever executing setup.py / install
# hooks. Way faster than installing the full Studio runtime
# hooks. Way faster than installing the full Unsloth runtime
# and -- critically -- safer: an attacker who has compromised
# a transitive dep cannot run code in this job.
#
@ -324,9 +326,9 @@ jobs:
} >> "$GITHUB_STEP_SUMMARY"
# ─────────────────────────────────────────────────────────────
# npm: Studio frontend
# npm: Unsloth frontend
# ─────────────────────────────────────────────────────────────
- name: npm audit (Studio frontend)
- name: npm audit (Unsloth frontend)
# `npm audit` resolves the lockfile through the npmjs.com
# advisory DB. `--audit-level=high` filters the noise floor
# to only HIGH and CRITICAL. We do NOT pass --omit=dev: a
@ -340,7 +342,7 @@ jobs:
# Always also write the full JSON for grep-ability.
npm audit --json > ../../logs-npm-audit.json || true
{
echo "## npm audit (Studio frontend)"
echo "## npm audit (Unsloth frontend)"
echo
echo '```'
tail -200 ../../logs-npm-audit.txt
@ -348,9 +350,9 @@ jobs:
} >> "$GITHUB_STEP_SUMMARY"
# ─────────────────────────────────────────────────────────────
# cargo: Studio Tauri shell
# cargo: Unsloth Tauri shell
# ─────────────────────────────────────────────────────────────
- name: cargo audit (Studio Tauri)
- name: cargo audit (Unsloth Tauri)
# `--deny warnings` would make the job fail on any advisory.
# Keep non-blocking initially; drop continue-on-error after
# the baseline closes.
@ -360,7 +362,7 @@ jobs:
set +e
cargo audit | tee ../../logs-cargo-audit.txt
{
echo "## cargo audit (Studio Tauri)"
echo "## cargo audit (Unsloth Tauri)"
echo
echo '```'
tail -200 ../../logs-cargo-audit.txt
@ -557,7 +559,7 @@ jobs:
# ─────────────────────────────────────────────────────────────
# CycloneDX SBOM. Lets downstream consumers audit what's
# actually shipped in unsloth wheels and the Studio backend
# actually shipped in unsloth wheels and the Unsloth backend
# runtime. Generates one JSON file per requirements input plus
# a combined SBOM keyed off pyproject.toml; uploads as a build
# artifact (and a future step can attest it via SLSA).
@ -738,7 +740,7 @@ jobs:
# `--with-deps` makes the scan transitive: every package the
# declared set resolves to gets fetched and pattern-scanned, not
# just the top-level pins. Resolving the full transitive closure
# of the unsloth + Studio dep tree downloads several hundred
# of the unsloth + Unsloth dep tree downloads several hundred
# archives, hence the longer timeout.
#
# Sharded across runners for wall-clock parallelism. Each shard
@ -747,7 +749,7 @@ jobs:
# composition tries to balance load:
# - hf-stack: pyproject extras + no-torch-runtime
# (~150 archives, transformers/peft/accelerate/...)
# - studio: FastAPI/Studio backend + overrides + extras-no-deps
# - studio: FastAPI/Unsloth backend + overrides + extras-no-deps
# (~150 archives, smaller scientific stack)
# - extras: the heavy openai-whisper / scikit-learn / librosa
# stack (~250 archives, dominant cost)
@ -962,7 +964,7 @@ jobs:
# documented at scripts/scan_npm_packages.py top-of-file. The
# script is stdlib-only so adding it does not increase the
# transitive supply-chain surface.
name: npm scan-packages (Studio frontend tarballs)
name: npm scan-packages (Unsloth frontend tarballs)
runs-on: ubuntu-latest
timeout-minutes: 30
needs: []
@ -1171,7 +1173,7 @@ jobs:
with:
python-version: '3.12'
- name: Install Studio frontend deps (--ignore-scripts)
- name: Install Unsloth frontend deps (--ignore-scripts)
# `npm audit signatures` requires node_modules to be populated.
# `--ignore-scripts` is mandatory: this is exactly the lever the
# new-install-script gate below protects against, and we must

14
.github/workflows/studio-api-smoke.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.
# Studio API & Auth Tests -- HTTP-level integration tests for the
# Unsloth API & Auth Tests -- HTTP-level integration tests for the
# FastAPI surface. No Playwright, no model UI; tests/studio/test_studio_api_smoke.py
# runs ~30 s and asserts:
# - CORS hardening (no wildcard + credentials, no bootstrap leak)
@ -15,7 +15,7 @@
# Reuses the GGUF cache key from studio-ui-smoke.yml so the model
# download is one cache-hit on the second job.
name: Studio API CI
name: Unsloth API CI
on:
pull_request:
@ -40,7 +40,7 @@ permissions:
jobs:
api-smoke:
name: Studio API & Auth Tests
name: Unsloth API & Auth Tests
runs-on: ubuntu-latest
timeout-minutes: 12
env:
@ -98,7 +98,7 @@ jobs:
path: hf-cache
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2
- name: Install Studio (--local, --no-torch)
- name: Install Unsloth (--local, --no-torch)
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
# Withheld on PR: this step runs checked-out PR code; public GGUF still downloads.
@ -111,7 +111,7 @@ jobs:
- name: Install pyjwt for the JWT-expiry forge test
run: pip install 'pyjwt>=2.6'
- name: Reset auth + boot Studio (API-only)
- name: Reset auth + boot Unsloth (API-only)
run: |
unsloth studio reset-password
mkdir -p logs
@ -144,7 +144,7 @@ jobs:
echo "STUDIO_NEW_PW=$NEW" >> "$GITHUB_ENV"
echo "STUDIO_NEW2_PW=$NEW2" >> "$GITHUB_ENV"
- name: Run Studio API & Auth tests
- name: Run Unsloth API & Auth tests
# The script is named WITHOUT a `test_` prefix so it isn't
# auto-collected by pytest in Backend CI's `tests/` walk
# (which doesn't set BASE_URL and would crash at import).
@ -153,7 +153,7 @@ jobs:
STUDIO_AUTH_DIR: /home/runner/.unsloth/studio/auth
run: python tests/studio/studio_api_smoke.py
- name: Stop Studio
- name: Stop Unsloth
if: always()
run: |
kill "${STUDIO_PID}" 2>/dev/null || true

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

@ -64,7 +64,7 @@ jobs:
- name: Install backend test dependencies (CPU only)
run: |
python -m pip install --upgrade pip
# Studio's declared backend deps:
# Unsloth's declared backend deps:
pip install -r studio/backend/requirements/studio.txt
# Extras that studio.txt does not list but the import chain needs
# (python-multipart for FastAPI form/file uploads, sqlalchemy/cryptography

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

@ -9,7 +9,7 @@
# export is validated separately. No GPU / model / llama.cpp: the tests mock the probes and block
# torch/unsloth, so the job installs only a CPU PyTorch plus import deps.
name: Studio export capability
name: Unsloth export capability
on:
pull_request:

4
.github/workflows/studio-frontend-ci.yml vendored
View file

@ -136,7 +136,7 @@ jobs:
- name: Build
run: npm run build
- name: Built bundle must not contain Studio's unstable_Provider call site
- name: Built bundle must not contain Unsloth's unstable_Provider call site
run: |
set -e
JS=$(ls dist/assets/index-*.js | head -1)
@ -144,7 +144,7 @@ jobs:
echo "main bundle: $JS"
echo "unstable_Provider: hits=$HITS (assistant-ui internals contribute up to 3)"
if [ "$HITS" -gt 3 ]; then
echo "::error file=studio/frontend/src/features/chat/runtime-provider.tsx::Studio bundle still passes unstable_Provider through useRemoteThreadListRuntime; this is the 2026.5.1 chat-history regression. Pass adapters directly into useLocalRuntime instead."
echo "::error file=studio/frontend/src/features/chat/runtime-provider.tsx::Unsloth bundle still passes unstable_Provider through useRemoteThreadListRuntime; this is the 2026.5.1 chat-history regression. Pass adapters directly into useLocalRuntime instead."
exit 1
fi

237
.github/workflows/studio-inference-smoke.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.
# Three end-to-end smoke jobs that boot a freshly-installed Studio and
# Three end-to-end smoke jobs that boot a freshly-installed Unsloth and
# exercise the surfaces real users hit through the OpenAI / Anthropic
# SDKs and curl. Each job picks the smallest model that exercises the
# behaviour under test, primes HF_HOME via actions/cache, and shares
@ -27,7 +27,7 @@
# All three jobs run in parallel. Total wall time is dominated by job 3
# on a cold cache; warm cache cuts that to ~3 min.
name: Studio GGUF CI
name: Unsloth GGUF CI
on:
pull_request:
@ -112,7 +112,7 @@ jobs:
path: hf-cache
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2
- name: Install Studio (--local, --no-torch)
- name: Install Unsloth (--local, --no-torch)
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
# Withheld on PR: this step runs checked-out PR code; public GGUF still downloads.
@ -125,7 +125,7 @@ jobs:
- name: Install OpenAI + Anthropic Python SDKs
run: pip install 'openai>=1.50' 'anthropic>=0.40'
- name: Reset auth + boot Studio (API-only)
- name: Reset auth + boot Unsloth (API-only)
run: |
unsloth studio reset-password
mkdir -p logs
@ -142,7 +142,7 @@ jobs:
fi
sleep 1
done
echo "Studio did not become healthy in 180s"
echo "Unsloth did not become healthy in 180s"
tail -200 logs/studio.log
exit 1
@ -229,11 +229,11 @@ jobs:
return replies
def run_anthropic():
# Two SDK quirks vs. Studio:
# Two SDK quirks vs. Unsloth:
# 1. base_url must NOT include /v1 -- the SDK appends
# /v1/messages itself; otherwise the request hits
# /v1/v1/messages and 405s.
# 2. The SDK sends `x-api-key` by default, but Studio's
# 2. The SDK sends `x-api-key` by default, but Unsloth's
# auth layer is HTTPBearer-only. Override via
# default_headers so Authorization: Bearer ... is
# sent instead.
@ -276,7 +276,7 @@ jobs:
print(
f"[{label}] WARN non-determinism at temperature=0.0 across "
f"{len(determinism_failures)} of {len(first)} turn(s); "
f"small-quant model drift, not a Studio regression. "
f"small-quant model drift, not an Unsloth regression. "
f"Details: " + " | ".join(determinism_failures)
)
# Sanity: turn-2 reply should mention the earlier question, and
@ -290,7 +290,7 @@ jobs:
print(f"[{label}] {status_word} -- 4 turns, history grounded ('paris' present)")
PY
- name: Stop Studio
- name: Stop Unsloth
if: always()
run: |
kill "${STUDIO_PID}" 2>/dev/null || true
@ -323,7 +323,7 @@ jobs:
# store xet chunks + blobs + snapshots = ~4 GiB compressed --
# 4-5x file-size inflation, dominated by xet chunks. Use main's
# `--local-dir gguf-cache` pattern to cache the flat .gguf only.
# Studio's /api/inference/load accepts either a HF repo (which
# Unsloth's /api/inference/load accepts either a HF repo (which
# uses HF_HOME) or an absolute file path; passing the absolute
# path keeps the test off HF_HOME entirely so the cache size
# tracks the GGUF file 1:1. The OpenAI/Anth and JSON+images
@ -380,7 +380,7 @@ jobs:
path: gguf-cache
key: ${{ runner.os }}-gguf-${{ env.GGUF_REPO }}-${{ env.GGUF_FILE }}-v1
- name: Install Studio (--local, --no-torch)
- name: Install Unsloth (--local, --no-torch)
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
# Withheld on PR: this step runs checked-out PR code; public GGUF still downloads.
@ -390,7 +390,7 @@ jobs:
set -o pipefail
bash install.sh --local --no-torch 2>&1 | tee logs/install.log
- name: Reset auth + boot Studio (API-only, default tool policy)
- name: Reset auth + boot Unsloth (API-only, default tool policy)
# We deliberately use the API-only mode rather than
# `unsloth studio run` because the latter calls
# `set_tool_policy(...)` with a resolved bool: on loopback the
@ -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 Unsloth
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",
@ -522,11 +575,11 @@ jobs:
def _tool_invoked(events):
"""Structural check: True iff some SSE payload is a real
tool envelope (Studio tool_start/tool_end, Anthropic
tool envelope (Unsloth tool_start/tool_end, Anthropic
tool_use/tool_result, OpenAI non-empty delta.tool_calls /
message.tool_calls / finish_reason='tool_calls' /
role:'tool' / function_call). tool_status is NOT
evidence: Studio emits empty tool_status events on
evidence: Unsloth emits empty tool_status events on
iteration boundaries even when no tool ran.
"""
for raw in events:
@ -645,23 +698,61 @@ jobs:
attempt has structural invocation evidence. WARN (not
FAIL) if invoked but no attempt produces the expected
literal in tool_end.result -- small-quant Qwen3.5-2B can
emit OpenAI tool_calls deltas without Studio's GGUF
emit OpenAI tool_calls deltas without Unsloth's GGUF
agentic loop intercepting them, and that GGUF-vs-OpenAI
format mismatch is out of scope for #5642.
"""
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({
@ -720,17 +811,21 @@ jobs:
# because (a) the search may legitimately return no results,
# and (b) DuckDuckGo upstream blocks GHA IP ranges often
# enough that requiring a tool_call marker would create
# red-herring failures from infra rather than from Studio.
# red-herring failures from infra rather than from Unsloth.
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)"
@ -739,7 +834,7 @@ jobs:
print(f"[tools] WARN web_search probe failed (non-blocking): {exc}")
# ── 5. Thinking on / off ─────────────────────────────────────
# Studio strips think blocks from message.content for tools-mode
# Unsloth strips think blocks from message.content for tools-mode
# responses, so we toggle plain chat (no enable_tools) and look
# at the surfaced reasoning_content / message.thinking field.
def thinking_call(enable):
@ -753,7 +848,7 @@ jobs:
})
assert status == 200
msg = data["choices"][0]["message"]
# Studio surfaces thinking via reasoning_content (OpenAI
# Unsloth surfaces thinking via reasoning_content (OpenAI
# extension). Fall back to inline <think> markers for
# robustness across template versions.
raw = (msg.get("content") or "") + (msg.get("reasoning_content") or "")
@ -773,7 +868,7 @@ jobs:
print(f"[tools] PASS thinking on/off (on={len(on_text)} chars, off={len(off_text)} chars)")
PY
- name: Stop Studio
- name: Stop Unsloth
if: always()
run: |
kill "${STUDIO_PID}" 2>/dev/null || true
@ -865,7 +960,7 @@ jobs:
path: hf-cache
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-${{ env.MMPROJ_FILE }}-v2
- name: Install Studio (--local, --no-torch)
- name: Install Unsloth (--local, --no-torch)
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
# Withheld on PR: this step runs checked-out PR code; public GGUF still downloads.
@ -878,7 +973,7 @@ jobs:
- name: Install OpenAI + Anthropic Python SDKs
run: pip install 'openai>=1.50' 'anthropic>=0.40'
- name: Reset auth + boot Studio (API-only)
- name: Reset auth + boot Unsloth (API-only)
# See Job 2's comment: API-only mode keeps tool_policy=None so
# response_format requests aren't routed through the agentic
# tool loop.
@ -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,20 +1053,36 @@ 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
# mode: `response_format: {"type": "json_object"}` constrains
# the model to emit syntactically-valid JSON. We use raw HTTP
# rather than the OpenAI SDK so that the field shape Studio
# rather than the OpenAI SDK so that the field shape Unsloth
# forwards to llama-server is unambiguous (the SDK rewrites
# response_format depending on which variant it recognises).
# We deliberately do NOT pass a strict JSON schema -- on
# small Gemma-4 quants the GBNF-from-schema path occasionally
# produces empty output, and JSON mode is the surface we care
# about exposing through Studio.
# about exposing through Unsloth.
status, data = post("/v1/chat/completions", {
"model": "default",
"messages": [
@ -999,7 +1112,7 @@ jobs:
print(f"[json] PASS json_object -> {parsed}")
# ── 2. OpenAI image_url (data URI base64) ───────────────────
# 64x64 solid-red PNG. stb_image (used by Studio's image
# 64x64 solid-red PNG. stb_image (used by Unsloth's image
# normaliser at routes/inference.py:3410) rejects 4x4 or
# smaller PNGs as truncated, so we go up to 64x64 -- still
# tiny in token cost. The assertion is loose: any non-empty
@ -1035,9 +1148,9 @@ jobs:
print("[image/openai] PASS image_url accepted, non-empty response")
# ── 3. Anthropic source/base64 image ────────────────────────
# Two SDK quirks vs. Studio: base_url must NOT include /v1
# Two SDK quirks vs. Unsloth: base_url must NOT include /v1
# (the SDK appends it itself; otherwise /v1/v1/messages -> 405),
# and Studio's auth is HTTPBearer-only so the SDK's default
# and Unsloth's auth is HTTPBearer-only so the SDK's default
# x-api-key header is ignored -- send Authorization: Bearer
# via default_headers.
anthropic = Anthropic(
@ -1071,7 +1184,7 @@ jobs:
print("[image/anthropic] PASS source/base64 accepted, non-empty response")
PY
- name: Stop Studio
- name: Stop Unsloth
if: always()
run: |
kill "${STUDIO_PID}" 2>/dev/null || true

4
.github/workflows/studio-load-orchestrator-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.
#
# Event-loop regression test for the Studio model-load orchestrator.
# Event-loop regression test for the Unsloth model-load orchestrator.
# Pins down issue #5642 (Win10 UI freeze on model load): the /load
# route calls LlamaCppBackend.detect_audio_type synchronously, blocking
# the FastAPI event loop on a chain of sync httpx.Client.post() probes.
@ -14,7 +14,7 @@
# danielhanchen/unsloth-staging-2 (Ubuntu / macOS / Windows all
# green at PR time).
name: Studio load-orchestrator CI
name: Unsloth load-orchestrator CI
on:
pull_request:

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

@ -33,7 +33,7 @@ permissions:
jobs:
api-smoke:
name: Studio API & Auth Tests
name: Unsloth API & Auth Tests
runs-on: macos-14
timeout-minutes: 25
env:
@ -83,7 +83,7 @@ jobs:
path: hf-cache
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2
- name: Install Studio (--local, --no-torch)
- name: Install Unsloth (--local, --no-torch)
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
# Withheld on PR: this step runs checked-out PR code; public GGUF still downloads.
@ -99,7 +99,7 @@ jobs:
- name: Install pyjwt for the JWT-expiry forge test
run: pip install 'pyjwt>=2.6'
- name: Reset auth + boot Studio (API-only)
- name: Reset auth + boot Unsloth (API-only)
run: |
unsloth studio reset-password
mkdir -p logs
@ -129,13 +129,13 @@ jobs:
echo "STUDIO_NEW_PW=$NEW" >> "$GITHUB_ENV"
echo "STUDIO_NEW2_PW=$NEW2" >> "$GITHUB_ENV"
- name: Run Studio API & Auth tests
- name: Run Unsloth API & Auth tests
env:
BASE_URL: http://127.0.0.1:18895
STUDIO_AUTH_DIR: /Users/runner/.unsloth/studio/auth
run: python tests/studio/studio_api_smoke.py
- name: Stop Studio
- name: Stop Unsloth
if: always()
run: |
kill "${STUDIO_PID}" 2>/dev/null || true

184
.github/workflows/studio-mac-inference-smoke.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.
# Three end-to-end smoke jobs that boot a freshly-installed Studio and
# Three end-to-end smoke jobs that boot a freshly-installed Unsloth and
# exercise the surfaces real users hit through the OpenAI / Anthropic
# SDKs and curl. Each job picks the smallest model that exercises the
# behaviour under test, primes a model cache via actions/cache, and
@ -108,7 +108,7 @@ jobs:
path: hf-cache
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2
- name: Install Studio (--local, --no-torch)
- name: Install Unsloth (--local, --no-torch)
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
# Withheld on PR: this step runs checked-out PR code; public GGUF still downloads.
@ -124,7 +124,7 @@ jobs:
- name: Install OpenAI + Anthropic Python SDKs
run: pip install 'openai>=1.50' 'anthropic>=0.40'
- name: Reset auth + boot Studio (API-only)
- name: Reset auth + boot Unsloth (API-only)
run: |
unsloth studio reset-password
mkdir -p logs
@ -141,7 +141,7 @@ jobs:
fi
sleep 1
done
echo "Studio did not become healthy in 180s"
echo "Unsloth did not become healthy in 180s"
tail -200 logs/studio.log
exit 1
@ -228,11 +228,11 @@ jobs:
return replies
def run_anthropic():
# Two SDK quirks vs. Studio:
# Two SDK quirks vs. Unsloth:
# 1. base_url must NOT include /v1 -- the SDK appends
# /v1/messages itself; otherwise the request hits
# /v1/v1/messages and 405s.
# 2. The SDK sends `x-api-key` by default, but Studio's
# 2. The SDK sends `x-api-key` by default, but Unsloth's
# auth layer is HTTPBearer-only. Override via
# default_headers so Authorization: Bearer ... is
# sent instead.
@ -283,7 +283,7 @@ jobs:
print(f"[{label}] OK -- 4 turns, run1 == run2, history grounded")
PY
- name: Stop Studio
- name: Stop Unsloth
if: always()
run: |
kill "${STUDIO_PID}" 2>/dev/null || true
@ -363,7 +363,7 @@ jobs:
path: gguf-cache
key: ${{ runner.os }}-gguf-${{ env.GGUF_REPO }}-${{ env.GGUF_FILE }}-v1
- name: Install Studio (--local, --no-torch)
- name: Install Unsloth (--local, --no-torch)
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
# Withheld on PR: this step runs checked-out PR code; public GGUF still downloads.
@ -376,7 +376,7 @@ jobs:
- name: Assert llama.cpp loads on this macOS
run: bash .github/scripts/assert-llama-loads.sh
- name: Reset auth + boot Studio (API-only, default tool policy)
- name: Reset auth + boot Unsloth (API-only, default tool policy)
# We deliberately use the API-only mode rather than
# `unsloth studio run` because the latter calls
# `set_tool_policy(...)` with a resolved bool: on loopback the
@ -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 Unsloth
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 = {
@ -526,11 +574,11 @@ jobs:
assert status == 200, f"tool call status {status}: {data}"
choice = data["choices"][0]
tool_calls = (choice.get("message") or {}).get("tool_calls") or []
# Studio's contract: when tool_choice='required', llama.cpp's
# Unsloth's contract: when tool_choice='required', llama.cpp's
# grammar should force a tool_calls payload. On Mac that
# contract is sometimes broken by the underlying quant; the
# PASS path is "tool_calls present + correct schema", the
# WARN path documents Studio still returned 200 with a
# WARN path documents Unsloth still returned 200 with a
# well-formed choices[] envelope.
if tool_calls:
tc = tool_calls[0]
@ -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,18 +648,19 @@ 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}")
# ── 4. Thinking on / off ─────────────────────────────────────
# Studio strips think blocks from message.content for tools-mode
# Unsloth strips think blocks from message.content for tools-mode
# responses, so we toggle plain chat (no enable_tools) and look
# at the surfaced reasoning_content / message.thinking field.
def thinking_call(enable):
@ -622,7 +678,7 @@ jobs:
}, timeout = 180)
assert status == 200
msg = data["choices"][0]["message"]
# Studio surfaces thinking via reasoning_content (OpenAI
# Unsloth surfaces thinking via reasoning_content (OpenAI
# extension). Fall back to inline <think> markers for
# robustness across template versions.
raw = (msg.get("content") or "") + (msg.get("reasoning_content") or "")
@ -648,7 +704,7 @@ jobs:
print(f"[tools] PASS thinking on/off (on={len(on_text)} chars, off={len(off_text)} chars)")
PY
- name: Stop Studio
- name: Stop Unsloth
if: always()
run: |
kill "${STUDIO_PID}" 2>/dev/null || true
@ -754,7 +810,7 @@ jobs:
path: gguf-cache
key: ${{ runner.os }}-gguf-${{ env.GGUF_REPO }}-${{ env.GGUF_FILE }}-${{ env.MMPROJ_FILE }}-v2
- name: Install Studio (--local, --no-torch)
- name: Install Unsloth (--local, --no-torch)
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
# Withheld on PR: this step runs checked-out PR code; public GGUF still downloads.
@ -770,7 +826,7 @@ jobs:
- name: Install OpenAI + Anthropic Python SDKs
run: pip install 'openai>=1.50' 'anthropic>=0.40'
- name: Reset auth + boot Studio (API-only)
- name: Reset auth + boot Unsloth (API-only)
# See Job 2's comment: API-only mode keeps tool_policy=None so
# response_format requests aren't routed through the agentic
# tool loop.
@ -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,20 +906,36 @@ 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
# mode: `response_format: {"type": "json_object"}` constrains
# the model to emit syntactically-valid JSON. We use raw HTTP
# rather than the OpenAI SDK so that the field shape Studio
# rather than the OpenAI SDK so that the field shape Unsloth
# forwards to llama-server is unambiguous (the SDK rewrites
# response_format depending on which variant it recognises).
# We deliberately do NOT pass a strict JSON schema -- on
# small Gemma-4 quants the GBNF-from-schema path occasionally
# produces empty output, and JSON mode is the surface we care
# about exposing through Studio.
# about exposing through Unsloth.
status, data = post("/v1/chat/completions", {
"model": "default",
"messages": [
@ -933,7 +1007,7 @@ jobs:
)
# ── 2. OpenAI image_url (data URI base64) ───────────────────
# 64x64 solid-red PNG. stb_image (used by Studio's image
# 64x64 solid-red PNG. stb_image (used by Unsloth's image
# normaliser at routes/inference.py:3410) rejects 4x4 or
# smaller PNGs as truncated, so we go up to 64x64 -- still
# tiny in token cost. The assertion is loose: any non-empty
@ -949,11 +1023,11 @@ jobs:
# The Mac prebuilt llama.cpp server has a known crash when
# processing image inputs alongside the gemma-4-E2B mmproj
# (server disconnects mid-completion). This is upstream
# llama.cpp behaviour, not Studio. Wrap both SDK calls in
# llama.cpp behaviour, not Unsloth. Wrap both SDK calls in
# try/except so an upstream crash registers as a WARN rather
# than failing the whole job. Studio's contract (OpenAI/
# than failing the whole job. Unsloth's contract (OpenAI/
# Anthropic image fields are accepted and forwarded) is
# validated by the request body Studio constructs, not by
# validated by the request body Unsloth constructs, not by
# whether llama.cpp can decode it on Mac Metal.
client = OpenAI(base_url = f"{BASE}/v1", api_key = KEY)
try:
@ -979,14 +1053,14 @@ jobs:
except Exception as exc:
print(
f"[image/openai] WARN image_url SDK call raised: {type(exc).__name__}: "
f"{exc}. Likely upstream llama.cpp Mac+vision crash, NOT a Studio "
f"regression. Studio successfully forwarded the request."
f"{exc}. Likely upstream llama.cpp Mac+vision crash, NOT an Unsloth "
f"regression. Unsloth successfully forwarded the request."
)
# ── 3. Anthropic source/base64 image ────────────────────────
# Two SDK quirks vs. Studio: base_url must NOT include /v1
# Two SDK quirks vs. Unsloth: base_url must NOT include /v1
# (the SDK appends it itself; otherwise /v1/v1/messages -> 405),
# and Studio's auth is HTTPBearer-only so the SDK's default
# and Unsloth's auth is HTTPBearer-only so the SDK's default
# x-api-key header is ignored -- send Authorization: Bearer
# via default_headers.
anthropic = Anthropic(
@ -1025,11 +1099,11 @@ jobs:
print(
f"[image/anthropic] WARN anthropic image SDK call raised: "
f"{type(exc).__name__}: {exc}. Likely upstream llama.cpp Mac+vision "
f"crash, NOT a Studio regression."
f"crash, NOT an Unsloth regression."
)
PY
- name: Stop Studio
- name: Stop Unsloth
if: always()
run: |
kill "${STUDIO_PID}" 2>/dev/null || true

4
.github/workflows/studio-mac-install-matrix.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.
# Proves Studio's llama.cpp install loads on every supported macOS. The heavy
# Proves Unsloth's llama.cpp install loads on every supported macOS. The heavy
# app smokes stay single-OS; this matrix covers the OS-version dimension cheaply
# (install.sh + binary-load assert). Regression guard for the macOS-version
# selection in studio/install_llama_prebuilt.py.
@ -60,7 +60,7 @@ jobs:
with:
python-version: '3.12'
- name: Install Studio (--local, --no-torch)
- name: Install Unsloth (--local, --no-torch)
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
# Withheld on PR: this step runs checked-out PR code; public GGUF still downloads.

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

@ -19,6 +19,7 @@ on:
- 'install.sh'
- 'pyproject.toml'
- 'tests/studio/**'
- '.github/scripts/run-studio-permission-browser.sh'
- '.github/workflows/studio-mac-ui-smoke.yml'
push:
branches: [main, pip]
@ -83,7 +84,7 @@ jobs:
path: hf-cache
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2
- name: Install Studio (--local, --no-torch)
- name: Install Unsloth (--local, --no-torch)
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
# Withheld on PR: this step runs checked-out PR code; public GGUF still downloads.
@ -96,7 +97,7 @@ jobs:
- name: Assert llama.cpp loads on this macOS
run: bash .github/scripts/assert-llama-loads.sh
- name: Install Playwright + Chromium
- name: Install Playwright browsers
# No --with-deps on Mac: that flag installs Linux apt packages.
# GitHub-hosted macos-14 ships the system frameworks Chromium
# needs already.
@ -112,7 +113,7 @@ jobs:
# in-script retry recover from any residual flakes.
run: |
pip install 'playwright>=1.55,<1.58'
python -m playwright install chromium
python -m playwright install chromium webkit
- name: Patch Playwright pipeTransport.js to tolerate malformed JSON
# In Playwright 1.55-1.58, pipeTransport.js does
@ -143,7 +144,7 @@ jobs:
print(f"pipeTransport.js: patched JSON.parse calls in {path}")
PY
- name: Reset auth + boot Studio
- name: Reset auth + boot Unsloth
run: |
unsloth studio reset-password
mkdir -p logs
@ -188,7 +189,7 @@ jobs:
# dies mid-test, (2) Chromium net::ERR_NO_BUFFER_SPACE when the
# runner's kernel briefly runs out of socket buffers, and (3) a
# goto 'interrupted by another navigation' when the SPA auth
# guard redirects mid-navigation. The retry FULLY resets Studio
# guard redirects mid-navigation. The retry FULLY resets Unsloth
# (kill, reset-password, reboot, wait /api/health, re-export
# bootstrap pw) before re-running the script. A real test failure
# (assertion / timeout) does NOT match any pattern so it bypasses
@ -209,7 +210,7 @@ jobs:
|| grep -q "ERR_NO_BUFFER_SPACE" logs/playwright_attempt_${attempt}.log \
|| grep -q "interrupted by another navigation" logs/playwright_attempt_${attempt}.log; } \
&& [ "$attempt" -lt "$max_attempts" ]; then
echo "::warning::Playwright flake on attempt ${attempt}; resetting Studio and retrying..."
echo "::warning::Playwright flake on attempt ${attempt}; resetting Unsloth and retrying..."
kill "${STUDIO_PID}" 2>/dev/null || true
sleep 2
unsloth studio reset-password
@ -238,13 +239,17 @@ jobs:
exit "$rc"
done
- name: Stop Studio (chat-ui ends with Shutdown click; this is belt-and-suspenders)
- name: Stop Unsloth (chat-ui ends with Shutdown click; this is belt-and-suspenders)
if: always()
run: |
kill "${STUDIO_PID}" 2>/dev/null || true
sleep 2
- name: Reset auth + boot Studio for extra UI tests (port 18897)
- name: Cross-browser permission controls
run: |
bash .github/scripts/run-studio-permission-browser.sh 18895 webkit
- name: Reset auth + boot Unsloth for extra UI tests (port 18897)
run: |
unsloth studio reset-password
mkdir -p logs
@ -271,7 +276,7 @@ jobs:
echo "STUDIO_EXTRA_OLD_PW=$OLD" >> "$GITHUB_ENV"
echo "STUDIO_EXTRA_NEW_PW=$NEW" >> "$GITHUB_ENV"
- name: Drive Compare/Recipes/Export/Studio/Settings with Playwright
- name: Drive Compare/Recipes/Export/Unsloth/Settings with Playwright
env:
BASE_URL: http://127.0.0.1:18897
STUDIO_OLD_PW: ${{ env.STUDIO_EXTRA_OLD_PW }}
@ -300,7 +305,7 @@ jobs:
|| grep -q "ERR_NO_BUFFER_SPACE" logs/playwright_extra_attempt_${attempt}.log \
|| grep -q "interrupted by another navigation" logs/playwright_extra_attempt_${attempt}.log; } \
&& [ "$attempt" -lt "$max_attempts" ]; then
echo "::warning::Playwright flake on attempt ${attempt}; resetting Studio and retrying..."
echo "::warning::Playwright flake on attempt ${attempt}; resetting Unsloth and retrying..."
kill "${STUDIO_EXTRA_PID}" 2>/dev/null || true
sleep 2
unsloth studio reset-password
@ -327,7 +332,7 @@ jobs:
exit "$rc"
done
- name: Stop second Studio
- name: Stop second Unsloth
if: always()
run: |
kill "${STUDIO_EXTRA_PID}" 2>/dev/null || true
@ -343,5 +348,7 @@ jobs:
logs/studio_extra.log
logs/install.log
logs/playwright
logs/playwright-permissions-*
logs/playwright_extra
logs/studio-permissions-*.log
retention-days: 7

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

@ -4,15 +4,15 @@
# Mac counterpart to studio-update-smoke.yml. Verifies that on a real
# Apple Silicon (macos-14, M1) runner:
#
# 1. install.sh --local --no-torch installs Studio AND auto-fetches
# 1. install.sh --local --no-torch installs Unsloth AND auto-fetches
# the prebuilt llama.cpp Mac binary (llama-bNNNN-bin-macos-arm64
# from ggml-org/llama.cpp). Hitting the source-build fallback is
# treated as an Unsloth bug -- Studio must always pick the
# treated as an Unsloth bug -- Unsloth must always pick the
# prebuilt on Mac.
# 2. unsloth studio update --local is idempotent. Two consecutive
# runs both report "prebuilt up to date and validated", no
# source-build fallback.
# 3. The installed Studio still boots and /api/health returns
# 3. The installed Unsloth still boots and /api/health returns
# healthy after the update path.
name: Mac Studio Update CI
@ -42,7 +42,7 @@ permissions:
jobs:
update-idempotency:
name: Studio Updating Tests
name: Unsloth Updating Tests
runs-on: macos-14
timeout-minutes: 30
steps:
@ -59,7 +59,7 @@ jobs:
python-version: '3.12'
cache: 'pip'
- name: Install Studio (--local, --no-torch)
- name: Install Unsloth (--local, --no-torch)
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
# Withheld on PR: this step runs checked-out PR code; public GGUF still downloads.
@ -106,7 +106,7 @@ jobs:
grep -qE "prebuilt up to date and validated|prebuilt installed and validated" logs/update2.log
echo "second update was clean"
- name: Boot Studio briefly to confirm the install is still usable
- name: Boot Unsloth briefly to confirm the install is still usable
run: |
mkdir -p logs
UNSLOTH_API_ONLY=1 unsloth studio -H 127.0.0.1 -p 18891 \
@ -123,13 +123,13 @@ jobs:
sleep 1
done
if [ -z "$HEALTHY" ]; then
echo "Studio failed to come up after \`update\`"
echo "Unsloth failed to come up after \`update\`"
tail -200 logs/studio.log
kill "$PID" 2>/dev/null || true
exit 1
fi
kill "$PID" 2>/dev/null || true
echo "post-update Studio /api/health OK"
echo "post-update Unsloth /api/health OK"
- name: Uninstall and verify clean
# Round-trip through scripts/uninstall.sh on real macOS. As a side

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

@ -12,7 +12,7 @@
# stay in release-desktop.yml (manual `workflow_dispatch`) because they need
# code-signing secrets and ~30 min of runner time each.
name: Studio Tauri CI
name: Unsloth Tauri CI
on:
pull_request:

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

@ -1,8 +1,8 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
# End-to-end Studio chat UI smoke via Playwright + Chromium against a
# headless Linux runner. Boots Studio with the smallest GGUF
# End-to-end Unsloth chat UI smoke via Playwright + Chromium against a
# headless Linux runner. Boots Unsloth with the smallest GGUF
# (gemma-3-270m-it UD-Q4_K_XL, ~254 MiB), drives the actual frontend
# bundle, and asserts the full bootstrap-password / change-password /
# send-message / persist-on-reload journey works end to end.
@ -14,7 +14,7 @@
# frontend-only CI happily pass while the actual user-visible UI is
# broken (cf. the 2026.5.1 chat-history release).
name: Studio UI CI
name: Unsloth UI CI
on:
pull_request:
@ -27,6 +27,7 @@ on:
# The Playwright test files themselves -- a PR that ONLY edits
# the test must still trigger UI CI.
- 'tests/studio/**'
- '.github/scripts/run-studio-permission-browser.sh'
- '.github/workflows/studio-ui-smoke.yml'
push:
branches: [main, pip]
@ -97,7 +98,7 @@ jobs:
path: hf-cache
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2
- name: Install Studio (--local, --no-torch)
- name: Install Unsloth (--local, --no-torch)
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
# Withheld on PR: this step runs checked-out PR code; public GGUF still downloads.
@ -107,15 +108,12 @@ jobs:
set -o pipefail
bash install.sh --local --no-torch 2>&1 | tee logs/install.log
- name: Install Playwright + Chromium
- name: Install Playwright browsers
run: |
pip install 'playwright>=1.45'
# --with-deps installs the OS-level runtime libs Chromium
# needs (libnss3, libxkbcommon, etc.). About 30 s on a
# warm runner.
python -m playwright install --with-deps chromium
python -m playwright install --with-deps chromium firefox webkit
- name: Reset auth + boot Studio
- name: Reset auth + boot Unsloth
run: |
unsloth studio reset-password
mkdir -p logs
@ -147,7 +145,7 @@ jobs:
# NEW + NEW2 are generated freshly per CI run via secrets.token_urlsafe
# rather than hardcoded. If a workflow gets compromised, the
# attacker can't replay a known-good rotated password against
# any future / parallel Studio install -- the rotated value
# any future / parallel Unsloth install -- the rotated value
# only ever exists for the lifetime of this single job, masked
# in the log via ::add-mask::.
run: |
@ -165,29 +163,35 @@ jobs:
env:
BASE_URL: http://127.0.0.1:18892
# The test file lives in the repo so it can be run locally
# against a freshly-installed Studio (BASE_URL=...; STUDIO_OLD_PW=
# against a freshly-installed Unsloth (BASE_URL=...; STUDIO_OLD_PW=
# $(cat ~/.unsloth/studio/auth/.bootstrap_password); python ...).
PW_ART_DIR: logs/playwright
# Strict mode: in CI a missing button / nav / dialog must
# FAIL the test. Locally the test still runs against partial
# Studio installs without STUDIO_UI_STRICT.
# Unsloth installs without STUDIO_UI_STRICT.
STUDIO_UI_STRICT: '1'
run: |
mkdir -p logs/playwright
python tests/studio/playwright_chat_ui.py
- name: Stop Studio (chat-ui ends with Shutdown click; this is belt-and-suspenders)
- name: Stop Unsloth (chat-ui ends with Shutdown click; this is belt-and-suspenders)
if: always()
run: |
kill "${STUDIO_PID}" 2>/dev/null || true
sleep 2
- name: Cross-browser permission controls
run: |
bash .github/scripts/run-studio-permission-browser.sh 18893 firefox
bash .github/scripts/run-studio-permission-browser.sh 18893 webkit
bash .github/scripts/run-studio-permission-browser.sh 18893 chromium chrome
# The chat UI test ends by clicking the Shutdown menuitem, which
# leaves the server dead. The extra UI test (Compare / Recipes /
# Export / Studio / Settings) needs a fresh Studio, so we boot a
# Export / Unsloth / Settings) needs a fresh Unsloth, so we boot a
# second one on a different port. Boot is fast (~3-5s on the
# warm install we already did) so this adds little wall time.
- name: Reset auth + boot Studio for extra UI tests (port 18894)
- name: Reset auth + boot Unsloth for extra UI tests (port 18894)
run: |
unsloth studio reset-password
mkdir -p logs
@ -214,7 +218,7 @@ jobs:
echo "STUDIO_EXTRA_OLD_PW=$OLD" >> "$GITHUB_ENV"
echo "STUDIO_EXTRA_NEW_PW=$NEW" >> "$GITHUB_ENV"
- name: Drive Compare/Recipes/Export/Studio/Settings with Playwright
- name: Drive Compare/Recipes/Export/Unsloth/Settings with Playwright
env:
BASE_URL: http://127.0.0.1:18894
STUDIO_OLD_PW: ${{ env.STUDIO_EXTRA_OLD_PW }}
@ -227,16 +231,73 @@ jobs:
mkdir -p logs/playwright_extra
python tests/studio/playwright_extra_ui.py
- name: Stop second Studio
- name: UI font size scaling regression (Playwright)
env:
BASE_URL: http://127.0.0.1:18894
STUDIO_PW: ${{ env.STUDIO_EXTRA_NEW_PW }}
PW_ART_DIR: logs/playwright_fontscale
run: |
mkdir -p logs/playwright_fontscale
python tests/studio/playwright_ui_font_scale.py
- name: Stop second Unsloth
if: always()
run: |
kill "${STUDIO_EXTRA_PID}" 2>/dev/null || true
sleep 2
# Model-picker per-model-config regression (PR #7207 re-land of #6647).
# Fourth Unsloth on its own port; loads the tiny GGUF and drives the
# picker's run-settings surface: Context Length persists across a reload,
# Reset clears the stored override (never pins it), and the infra models
# (RAG embedder + llama.cpp probe) stay hidden from the picker.
- name: Reset auth + boot Unsloth for model-config tests (port 18898)
run: |
unsloth studio reset-password
mkdir -p logs
UNSLOTH_API_ONLY=1 unsloth studio -H 127.0.0.1 -p 18898 \
> logs/studio_modelcfg.log 2>&1 &
echo "STUDIO_MODELCFG_PID=$!" >> "$GITHUB_ENV"
- name: Wait for /api/health on 18898
run: |
for i in $(seq 1 180); do
if curl -fs "http://127.0.0.1:18898/api/health" > /tmp/health4.json; then
jq -e '.status == "healthy"' /tmp/health4.json && break
fi
sleep 1
done
jq -e '.status == "healthy"' /tmp/health4.json
- name: Pass bootstrap pw for model-config test
run: |
NEW="CIModelCfg-$(python -c 'import secrets; print(secrets.token_urlsafe(16))')"
echo "::add-mask::$NEW"
echo "STUDIO_MODELCFG_NEW_PW=$NEW" >> "$GITHUB_ENV"
- name: Drive model-picker per-model-config with Playwright
env:
BASE_URL: http://127.0.0.1:18898
STUDIO_NEW_PW: ${{ env.STUDIO_MODELCFG_NEW_PW }}
PW_ART_DIR: logs/playwright_modelcfg
STUDIO_UI_STRICT: '1'
GGUF_REPO: ${{ env.GGUF_REPO }}
GGUF_VARIANT: ${{ env.GGUF_VARIANT }}
STUDIO_MODEL_HINT: gemma-3-270m
run: |
mkdir -p logs/playwright_modelcfg
python tests/studio/playwright_model_config.py
- name: Stop fourth Unsloth
if: always()
run: |
kill "${STUDIO_MODELCFG_PID}" 2>/dev/null || true
sleep 2
# IME + multilingual paste regression (issue #5318 / PR #5327).
# Third Studio on its own port so a hang here cannot poison the
# Third Unsloth on its own port so a hang here cannot poison the
# earlier UI tests. No GGUF -- the bug surface is the composer.
- name: Reset auth + boot Studio for IME / i18n tests (port 18896)
- name: Reset auth + boot Unsloth for IME / i18n tests (port 18896)
run: |
unsloth studio reset-password
mkdir -p logs
@ -256,7 +317,7 @@ jobs:
- name: Pass bootstrap pw for IME / i18n test
# IME smoke does the change-password against the bootstrap that
# Studio's frontend injects into the page, so it only needs the
# Unsloth's frontend injects into the page, so it only needs the
# NEW password.
run: |
NEW="CIIme-$(python -c 'import secrets; print(secrets.token_urlsafe(16))')"
@ -273,7 +334,7 @@ jobs:
mkdir -p logs/playwright_ime
python tests/studio/playwright_chat_ime_i18n.py
- name: Stop third Studio
- name: Stop third Unsloth
if: always()
run: |
kill "${STUDIO_IME_PID}" 2>/dev/null || true
@ -293,10 +354,15 @@ jobs:
path: |
logs/studio.log
logs/studio_extra.log
logs/studio_modelcfg.log
logs/studio_ime.log
logs/install.log
logs/server-logs/
logs/playwright
logs/playwright-permissions-*
logs/playwright_extra
logs/playwright_fontscale
logs/playwright_modelcfg
logs/playwright_ime
logs/studio-permissions-*.log
retention-days: 7

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

@ -9,7 +9,7 @@
# This catches regressions in setup.sh's update path that the existing
# GGUF / wheel jobs would miss because they only invoke install.sh once.
name: Studio Update CI
name: Unsloth Update CI
on:
pull_request:
@ -36,7 +36,7 @@ permissions:
jobs:
update-idempotency:
name: Studio Updating Tests
name: Unsloth Updating Tests
runs-on: ubuntu-latest
timeout-minutes: 15
steps:
@ -63,7 +63,7 @@ jobs:
# post-step then fatal-errors with "Cache folder path is
# retrieved for pip but doesn't exist on disk".
- name: Install Studio (--local, --no-torch)
- name: Install Unsloth (--local, --no-torch)
# Pass the workflow token so the llama.cpp prebuilt installer's
# GitHub-API call to list releases isn't rate-limited (60/hr
# unauthenticated). Without this, three consecutive install +
@ -122,7 +122,7 @@ jobs:
grep -qE "prebuilt up to date and validated|prebuilt installed and validated" logs/update2.log
echo "second update was clean"
- name: Boot Studio briefly to confirm the install is still usable
- name: Boot Unsloth briefly to confirm the install is still usable
# If `update --local` accidentally broke the venv or wiped the
# llama-server binary, the server would fail to start here.
run: |
@ -138,13 +138,13 @@ jobs:
sleep 1
done
if ! jq -e '.status == "healthy"' /tmp/health.json 2>/dev/null; then
echo "Studio failed to come up after `update`"
echo "Unsloth failed to come up after `update`"
tail -200 logs/studio.log
kill "$PID" 2>/dev/null || true
exit 1
fi
kill "$PID" 2>/dev/null || true
echo "post-update Studio /api/health OK"
echo "post-update Unsloth /api/health OK"
- name: Uninstall and verify clean
# Round-trip the installer through scripts/uninstall.sh: confirms the

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

@ -9,7 +9,7 @@
# (Section 6) is Linux-only and short-circuits on non-POSIX; the rest
# is platform-portable.
name: Windows Studio API CI
name: Windows Unsloth API CI
on:
pull_request:
@ -34,7 +34,7 @@ permissions:
jobs:
api-smoke:
name: Studio API & Auth Tests
name: Unsloth API & Auth Tests
runs-on: windows-latest
timeout-minutes: 30
defaults:
@ -105,7 +105,7 @@ jobs:
# studio-windows-update-smoke.yml for the full rationale --
# creating an empty studio/frontend/dist trips setup.ps1's
# mtime-based staleness check into "frontend up to date, skip
# rebuild" and Studio boots with an empty dist directory.
# rebuild" and Unsloth boots with an empty dist directory.
# Add-MpPreference accepts paths that do not yet exist.
foreach ($p in @(
"$env:USERPROFILE\.unsloth",
@ -121,7 +121,7 @@ jobs:
}
}
- name: Install Studio (--local, --no-torch)
- name: Install Unsloth (--local, --no-torch)
shell: pwsh
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
@ -161,7 +161,7 @@ jobs:
echo "install.ps1 installed the Windows prebuilt llama.cpp:"
cat "$INFO"
- name: Add Studio shim to GITHUB_PATH
- name: Add Unsloth shim to GITHUB_PATH
# install.ps1's User-PATH update doesn't propagate to a
# running Git Bash session; export the shim dir so the
# next `unsloth ...` invocation finds it.
@ -177,7 +177,7 @@ jobs:
- name: Install pyjwt for the JWT-expiry forge test
run: python -m pip install 'pyjwt>=2.6'
- name: Reset auth + boot Studio (API-only)
- name: Reset auth + boot Unsloth (API-only)
run: |
unsloth studio reset-password
mkdir -p logs
@ -207,7 +207,7 @@ jobs:
echo "STUDIO_NEW_PW=$NEW" >> "$GITHUB_ENV"
echo "STUDIO_NEW2_PW=$NEW2" >> "$GITHUB_ENV"
- name: Run Studio API & Auth tests
- name: Run Unsloth API & Auth tests
# Do NOT pin STUDIO_AUTH_DIR here. The Mac/Linux mirrors
# hardcode runner-specific paths (/Users/runner/...,
# /home/runner/...), but on Windows the path is
@ -219,7 +219,7 @@ jobs:
BASE_URL: http://127.0.0.1:18895
run: python tests/studio/studio_api_smoke.py
- name: Stop Studio
- name: Stop Unsloth
if: always()
run: |
kill "${STUDIO_PID}" 2>/dev/null || true

394
.github/workflows/studio-windows-inference-smoke.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.
# Three end-to-end smoke jobs that boot a freshly-installed Studio and
# Three end-to-end smoke jobs that boot a freshly-installed Unsloth and
# exercise the surfaces real users hit through the OpenAI / Anthropic
# SDKs and curl, on the FREE windows-latest runner. Each job picks the
# smallest model that exercises the behaviour under test, primes
@ -16,7 +16,7 @@
# Qwen3-VL-2B-Instruct UD-IQ2_XXS + mmproj-F16 (~1.4 GiB total).
# Within the 14 GB windows-latest SSD budget.
name: Windows Studio GGUF CI
name: Windows Unsloth GGUF CI
on:
pull_request:
@ -57,7 +57,7 @@ jobs:
STUDIO_PORT: '18888'
HF_HOME: ${{ github.workspace }}/hf-cache
# Force UTF-8 for stdio (Windows defaults to cp1252; hf
# download / Studio CLI print "✓" checkmarks and crash
# download / Unsloth CLI print "✓" checkmarks and crash
# otherwise).
PYTHONIOENCODING: utf-8
PYTHONUTF8: '1'
@ -160,7 +160,7 @@ jobs:
# studio-windows-update-smoke.yml for the full rationale --
# creating an empty studio/frontend/dist trips setup.ps1's
# mtime-based staleness check into "frontend up to date, skip
# rebuild" and Studio boots with an empty dist directory.
# rebuild" and Unsloth boots with an empty dist directory.
# Add-MpPreference accepts paths that do not yet exist.
foreach ($p in @(
"$env:USERPROFILE\.unsloth",
@ -176,7 +176,7 @@ jobs:
}
}
- name: Install Studio (--local, --no-torch)
- name: Install Unsloth (--local, --no-torch)
shell: pwsh
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
@ -214,7 +214,7 @@ jobs:
echo "install.ps1 installed the Windows prebuilt llama.cpp:"
cat "$INFO"
- name: Add Studio shim to GITHUB_PATH
- name: Add Unsloth shim to GITHUB_PATH
run: |
SHIM_DIR=~/.unsloth/studio/bin
if [ ! -f "$SHIM_DIR/unsloth.exe" ]; then
@ -227,7 +227,7 @@ jobs:
- name: Install OpenAI + Anthropic Python SDKs
run: python -m pip install 'openai>=1.50' 'anthropic>=0.40'
- name: Reset auth + boot Studio (API-only)
- name: Reset auth + boot Unsloth (API-only)
run: |
unsloth studio reset-password
mkdir -p logs
@ -244,7 +244,7 @@ jobs:
fi
sleep 1
done
echo "Studio did not become healthy in 180s"
echo "Unsloth did not become healthy in 180s"
tail -200 logs/studio.log
exit 1
@ -281,7 +281,7 @@ jobs:
# Retry the load step a few times so a transient TCP RST during
# llama-server warm-up (Windows runner image churn,
# windows-latest -> windows-2025-vs2026 rollout) doesn't fail
# the whole job. The Studio backend's _wait_for_health now
# the whole job. The Unsloth backend's _wait_for_health now
# catches httpx.ReadError too; this retry layer covers the
# cases the backend can't recover from on its own.
LOAD_OK=0
@ -382,15 +382,15 @@ jobs:
print(f"[{label}] OK -- 4 turns, run1 == run2, history grounded")
PY
- name: Stop Studio
- name: Stop Unsloth
if: always()
# Run as cmd so we are not running through the Git Bash shell;
# Git Bash on windows-latest has been observed to exit 143
# (SIGTERM) from any inline kill/sleep block, masking a green
# test run. The runner reclaims the Studio child process at
# test run. The runner reclaims the Unsloth child process at
# job end either way, so just emit a marker and exit 0.
shell: cmd
run: echo Stop Studio (no-op; runner reclaims STUDIO_PID=%STUDIO_PID% at job end)
run: echo Stop Unsloth (no-op; runner reclaims STUDIO_PID=%STUDIO_PID% at job end)
- name: Collect llama-server logs
if: always()
@ -398,10 +398,10 @@ jobs:
# copy must not fail an otherwise-green job.
continue-on-error: true
shell: bash
# Copy llama-server's own stdout/stderr (teed by Studio under
# Copy llama-server's own stdout/stderr (teed by Unsloth under
# ~/.unsloth/studio/logs/llama-server/) into the workspace so
# upload-artifact can pick it up. Crucial for diagnosing a
# subprocess crash where Studio's traceback only shows the
# subprocess crash where Unsloth's traceback only shows the
# symptom (httpx ReadError) but not the cause.
run: |
mkdir -p logs/llama-server
@ -439,14 +439,14 @@ jobs:
# (211 s on first run; subsequent runs hit the cache, but the
# one-time cost recurs every time the cache key bumps). Use
# main's `--local-dir gguf-cache` pattern: cache the flat .gguf
# only, pass an absolute path to Studio's /api/inference/load.
# only, pass an absolute path to Unsloth's /api/inference/load.
# The OpenAI/Anth and JSON+images jobs still cover the
# gguf_variant resolution path.
GGUF_REPO: unsloth/Qwen3.5-2B-GGUF
GGUF_FILE: Qwen3.5-2B-UD-Q4_K_XL.gguf
STUDIO_PORT: '18898'
# Force UTF-8 for stdio (Windows defaults to cp1252; hf
# download / Studio CLI print "✓" checkmarks and crash
# download / Unsloth CLI print "✓" checkmarks and crash
# otherwise).
PYTHONIOENCODING: utf-8
PYTHONUTF8: '1'
@ -507,7 +507,7 @@ jobs:
# studio-windows-update-smoke.yml for the full rationale --
# creating an empty studio/frontend/dist trips setup.ps1's
# mtime-based staleness check into "frontend up to date, skip
# rebuild" and Studio boots with an empty dist directory.
# rebuild" and Unsloth boots with an empty dist directory.
# Add-MpPreference accepts paths that do not yet exist.
foreach ($p in @(
"$env:USERPROFILE\.unsloth",
@ -523,7 +523,7 @@ jobs:
}
}
- name: Install Studio (--local, --no-torch)
- name: Install Unsloth (--local, --no-torch)
shell: pwsh
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
@ -561,7 +561,7 @@ jobs:
echo "install.ps1 installed the Windows prebuilt llama.cpp:"
cat "$INFO"
- name: Add Studio shim to GITHUB_PATH
- name: Add Unsloth shim to GITHUB_PATH
run: |
SHIM_DIR=~/.unsloth/studio/bin
if [ ! -f "$SHIM_DIR/unsloth.exe" ]; then
@ -571,7 +571,7 @@ jobs:
fi
cygpath -w "$SHIM_DIR" >> "$GITHUB_PATH"
- name: Reset auth + boot Studio (API-only, default tool policy)
- name: Reset auth + boot Unsloth (API-only, default tool policy)
run: |
unsloth studio reset-password
mkdir -p logs
@ -607,7 +607,7 @@ jobs:
# raw string, but we cannot embed `\a` etc. in JSON without
# JSON-string-escaping every backslash. Replace `\` with `/`
# via bash parameter expansion -- pathlib.Path on Windows
# accepts forward slashes natively, so Studio's loader sees
# accepts forward slashes natively, so Unsloth's loader sees
# a normal path.
GGUF_PATH="${GITHUB_WORKSPACE//\\//}/gguf-cache/${GGUF_FILE}"
ls -lh "$GGUF_PATH"
@ -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 Unsloth
# 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}")
@ -818,15 +882,15 @@ jobs:
print(f"[tools] PASS thinking on/off (on={len(on_text)} chars, off={len(off_text)} chars)")
PY
- name: Stop Studio
- name: Stop Unsloth
if: always()
# Run as cmd so we are not running through the Git Bash shell;
# Git Bash on windows-latest has been observed to exit 143
# (SIGTERM) from any inline kill/sleep block, masking a green
# test run. The runner reclaims the Studio child process at
# test run. The runner reclaims the Unsloth child process at
# job end either way, so just emit a marker and exit 0.
shell: cmd
run: echo Stop Studio (no-op; runner reclaims STUDIO_PID=%STUDIO_PID% at job end)
run: echo Stop Unsloth (no-op; runner reclaims STUDIO_PID=%STUDIO_PID% at job end)
- name: Collect llama-server logs
if: always()
@ -834,10 +898,10 @@ jobs:
# copy must not fail an otherwise-green job.
continue-on-error: true
shell: bash
# Copy llama-server's own stdout/stderr (teed by Studio under
# Copy llama-server's own stdout/stderr (teed by Unsloth under
# ~/.unsloth/studio/logs/llama-server/) into the workspace so
# upload-artifact can pick it up. Crucial for diagnosing a
# subprocess crash where Studio's traceback only shows the
# subprocess crash where Unsloth's traceback only shows the
# symptom (httpx ReadError) but not the cause.
run: |
mkdir -p logs/llama-server
@ -875,7 +939,7 @@ jobs:
STUDIO_PORT: '18899'
HF_HOME: ${{ github.workspace }}/hf-cache
# Force UTF-8 for stdio (Windows defaults to cp1252; hf
# download / Studio CLI print "✓" checkmarks and crash
# download / Unsloth CLI print "✓" checkmarks and crash
# otherwise).
PYTHONIOENCODING: utf-8
PYTHONUTF8: '1'
@ -941,7 +1005,7 @@ jobs:
# studio-windows-update-smoke.yml for the full rationale --
# creating an empty studio/frontend/dist trips setup.ps1's
# mtime-based staleness check into "frontend up to date, skip
# rebuild" and Studio boots with an empty dist directory.
# rebuild" and Unsloth boots with an empty dist directory.
# Add-MpPreference accepts paths that do not yet exist.
foreach ($p in @(
"$env:USERPROFILE\.unsloth",
@ -957,7 +1021,7 @@ jobs:
}
}
- name: Install Studio (--local, --no-torch)
- name: Install Unsloth (--local, --no-torch)
shell: pwsh
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
@ -995,7 +1059,7 @@ jobs:
echo "install.ps1 installed the Windows prebuilt llama.cpp:"
cat "$INFO"
- name: Add Studio shim to GITHUB_PATH
- name: Add Unsloth shim to GITHUB_PATH
run: |
SHIM_DIR=~/.unsloth/studio/bin
if [ ! -f "$SHIM_DIR/unsloth.exe" ]; then
@ -1008,7 +1072,7 @@ jobs:
- name: Install OpenAI + Anthropic Python SDKs
run: python -m pip install 'openai>=1.50' 'anthropic>=0.40'
- name: Reset auth + boot Studio (API-only)
- name: Reset auth + boot Unsloth (API-only)
run: |
unsloth studio reset-password
mkdir -p logs
@ -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", {
@ -1180,7 +1262,7 @@ jobs:
except Exception as exc:
print(
f"[image/openai] WARN image_url SDK call raised: {type(exc).__name__}: "
f"{exc}. Studio successfully forwarded the request; failure here is "
f"{exc}. Unsloth successfully forwarded the request; failure here is "
f"upstream llama.cpp vision behaviour."
)
@ -1221,19 +1303,19 @@ jobs:
print(
f"[image/anthropic] WARN anthropic image SDK call raised: "
f"{type(exc).__name__}: {exc}. Likely upstream llama.cpp vision "
f"behaviour, NOT a Studio regression."
f"behaviour, NOT an Unsloth regression."
)
PY
- name: Stop Studio
- name: Stop Unsloth
if: always()
# Run as cmd so we are not running through the Git Bash shell;
# Git Bash on windows-latest has been observed to exit 143
# (SIGTERM) from any inline kill/sleep block, masking a green
# test run. The runner reclaims the Studio child process at
# test run. The runner reclaims the Unsloth child process at
# job end either way, so just emit a marker and exit 0.
shell: cmd
run: echo Stop Studio (no-op; runner reclaims STUDIO_PID=%STUDIO_PID% at job end)
run: echo Stop Unsloth (no-op; runner reclaims STUDIO_PID=%STUDIO_PID% at job end)
- name: Collect llama-server logs
if: always()
@ -1241,10 +1323,10 @@ jobs:
# copy must not fail an otherwise-green job.
continue-on-error: true
shell: bash
# Copy llama-server's own stdout/stderr (teed by Studio under
# Copy llama-server's own stdout/stderr (teed by Unsloth under
# ~/.unsloth/studio/logs/llama-server/) into the workspace so
# upload-artifact can pick it up. Crucial for diagnosing a
# subprocess crash where Studio's traceback only shows the
# subprocess crash where Unsloth's traceback only shows the
# symptom (httpx ReadError) but not the cause.
run: |
mkdir -p logs/llama-server
@ -1266,7 +1348,7 @@ jobs:
# ── folded from studio-windows-no-vs-smoke.yml: install + run with no Visual Studio ──
no-vs-cpu:
name: Studio install + inference without Visual Studio
name: Unsloth install + inference without Visual Studio
runs-on: windows-latest
timeout-minutes: 35
defaults:
@ -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')) {
@ -1387,13 +1502,17 @@ jobs:
python -m pip install torch --index-url https://download.pytorch.org/whl/cpu --extra-index-url https://pypi.org/simple
python -c "import torch; print('torch', torch.__version__, 'cuda?', torch.cuda.is_available())"
- name: Install Studio (--local, --no-torch) with no build tools present
- name: Install Unsloth (--local, --no-torch) with no build tools present
shell: pwsh
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
# Withheld on PR: this step runs checked-out PR code; public GGUF still downloads.
HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }}
run: |
# Set 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
@ -1419,13 +1538,13 @@ jobs:
echo "Prebuilt installed with no build tools:"
cat "$INFO"
- name: Add Studio shim to GITHUB_PATH
- name: Add Unsloth shim to GITHUB_PATH
run: |
SHIM_DIR=~/.unsloth/studio/bin
[ -f "$SHIM_DIR/unsloth.exe" ] || { echo "::error::unsloth.exe shim not found"; ls -la ~/.unsloth/studio/ || true; exit 1; }
cygpath -w "$SHIM_DIR" >> "$GITHUB_PATH"
- name: Reset auth + boot Studio (API-only)
- name: Reset auth + boot Unsloth (API-only)
run: |
unsloth studio reset-password
mkdir -p logs
@ -1480,24 +1599,24 @@ 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
- name: Stop Unsloth
if: always()
shell: cmd
run: echo Stop Studio (no-op; runner reclaims STUDIO_PID=%STUDIO_PID% at job end)
run: echo Stop Unsloth (no-op; runner reclaims STUDIO_PID=%STUDIO_PID% at job end)
- name: Collect llama-server logs
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:
@ -1610,6 +1751,13 @@ jobs:
- name: Install Pester v5
shell: pwsh
run: |
# PSGallery is intermittently absent from the repository list on GitHub's Windows
# runners, which makes `Set-PSRepository PSGallery` fail with "No repository with the
# name 'PSGallery' was found." Re-register the default gallery first so the policy
# change and module install below always have a repository to target.
if (-not (Get-PSRepository -Name PSGallery -ErrorAction SilentlyContinue)) {
Register-PSRepository -Default -ErrorAction SilentlyContinue
}
Set-PSRepository PSGallery -InstallationPolicy Trusted
Install-Module Pester -MinimumVersion 5.5.0 -Force -SkipPublisherCheck -Scope CurrentUser
Import-Module Pester -MinimumVersion 5.5.0

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

@ -4,11 +4,11 @@
# Windows counterpart to studio-ui-smoke.yml / studio-mac-ui-smoke.yml.
# Same Playwright + Chromium end-to-end chat UI flow + extra UI flow,
# but on the FREE windows-latest runner so we catch Windows-specific
# regressions in the install path (install.ps1), the Studio CLI's
# regressions in the install path (install.ps1), the Unsloth CLI's
# Windows process-management branches, and the llama.cpp prebuilt's
# Windows HTTP layer.
name: Windows Studio UI CI
name: Windows Unsloth UI CI
on:
pull_request:
@ -19,6 +19,7 @@ on:
- 'install.ps1'
- 'pyproject.toml'
- 'tests/studio/**'
- '.github/scripts/run-studio-permission-browser.sh'
- '.github/workflows/studio-windows-ui-smoke.yml'
push:
branches: [main, pip]
@ -49,7 +50,7 @@ jobs:
GGUF_FILE: gemma-3-270m-it-UD-Q4_K_XL.gguf
STUDIO_PORT: '18896'
HF_HOME: ${{ github.workspace }}/hf-cache
# Force UTF-8 for stdio so Python tools (hf download, Studio
# Force UTF-8 for stdio so Python tools (hf download, Unsloth
# CLI, etc.) can print Unicode characters like the success
# checkmark "✓". Windows defaults to cp1252 / charmap and
# any tool that prints "OK ✓" hits a UnicodeEncodeError.
@ -121,7 +122,7 @@ jobs:
# studio-windows-update-smoke.yml for the full rationale --
# creating an empty studio/frontend/dist trips setup.ps1's
# mtime-based staleness check into "frontend up to date, skip
# rebuild" and Studio boots with an empty dist directory.
# rebuild" and Unsloth boots with an empty dist directory.
# Add-MpPreference accepts paths that do not yet exist.
foreach ($p in @(
"$env:USERPROFILE\.unsloth",
@ -148,7 +149,7 @@ jobs:
Set-Content -LiteralPath (Join-Path $appDir 'launch-studio.vbs') -Value 'WScript.Echo "legacy"' -Encoding Unicode
Write-Host "seeded legacy launch-studio.vbs at $appDir"
- name: Install Studio (--local, --no-torch)
- name: Install Unsloth (--local, --no-torch)
# install.ps1 is the supported Windows installer. install.sh
# has no Windows branch (apt-get / brew calls). The PS1
# script's `Install-UnslothStudio @args` line at the bottom
@ -205,7 +206,7 @@ jobs:
echo "install.ps1 installed the Windows prebuilt llama.cpp:"
cat "$INFO"
- name: Assert Studio launcher chain (no VBS, hidden PowerShell shortcut)
- name: Assert Unsloth launcher chain (no VBS, hidden PowerShell shortcut)
# The shortcut launch path is otherwise untested here (the steps below
# boot `unsloth studio` directly). Guard against re-introducing the VBS
# that tripped Kaspersky HEUR:Trojan.VBS.Agent.gen and against the .lnk
@ -234,7 +235,7 @@ jobs:
}
Write-Host "launcher chain OK (no VBS; hidden powershell over launch-studio.ps1)"
- name: Launch Studio via the shortcut and assert health
- name: Launch Unsloth via the shortcut and assert health
# Run the exact command the .lnk stores (hidden PowerShell over
# launch-studio.ps1) and confirm it brings the backend up. This is the
# only step that proves the shortcut launch is not silently broken.
@ -265,10 +266,10 @@ jobs:
$owner = (Get-NetTCPConnection -LocalPort $foundPort -State Listen -ErrorAction Stop | Select-Object -First 1).OwningProcess
if ($owner) { taskkill /PID $owner /T /F 2>$null | Out-Null }
} catch {}
if (-not $foundPort) { throw "Studio did not become healthy when launched via the shortcut" }
Write-Host "Studio healthy on port $foundPort (launched via the shortcut)"
if (-not $foundPort) { throw "Unsloth did not become healthy when launched via the shortcut" }
Write-Host "Unsloth healthy on port $foundPort (launched via the shortcut)"
- name: Add Studio shim to GITHUB_PATH
- name: Add Unsloth shim to GITHUB_PATH
# install.ps1 puts unsloth.exe at $StudioHome\bin\unsloth.exe
# and adds that dir to the User PATH via the Windows registry.
# Registry-level PATH updates don't propagate to a running
@ -284,7 +285,7 @@ jobs:
fi
# GITHUB_PATH wants Windows-style paths; convert via cygpath.
cygpath -w "$SHIM_DIR" >> "$GITHUB_PATH"
echo "Added Studio shim dir to PATH: $(cygpath -w "$SHIM_DIR")"
echo "Added Unsloth shim dir to PATH: $(cygpath -w "$SHIM_DIR")"
- name: Install Playwright + Chromium
# No --with-deps on Windows: that flag installs Linux apt
@ -294,7 +295,7 @@ jobs:
python -m pip install 'playwright>=1.45'
python -m playwright install chromium
- name: Reset auth + boot Studio
- name: Reset auth + boot Unsloth
run: |
unsloth studio reset-password
mkdir -p logs
@ -339,13 +340,17 @@ jobs:
mkdir -p logs/playwright
python tests/studio/playwright_chat_ui.py
- name: Stop Studio (chat-ui ends with Shutdown click; this is belt-and-suspenders)
- name: Stop Unsloth (chat-ui ends with Shutdown click; this is belt-and-suspenders)
if: always()
run: |
kill "${STUDIO_PID}" 2>/dev/null || true
sleep 2
- name: Reset auth + boot Studio for extra UI tests (port 18897)
- name: Edge permission controls
run: |
bash .github/scripts/run-studio-permission-browser.sh 18895 chromium msedge
- name: Reset auth + boot Unsloth for extra UI tests (port 18897)
run: |
unsloth studio reset-password
mkdir -p logs
@ -372,7 +377,7 @@ jobs:
echo "STUDIO_EXTRA_OLD_PW=$OLD" >> "$GITHUB_ENV"
echo "STUDIO_EXTRA_NEW_PW=$NEW" >> "$GITHUB_ENV"
- name: Drive Compare/Recipes/Export/Studio/Settings with Playwright
- name: Drive Compare/Recipes/Export/Unsloth/Settings with Playwright
env:
BASE_URL: http://127.0.0.1:18897
STUDIO_OLD_PW: ${{ env.STUDIO_EXTRA_OLD_PW }}
@ -386,7 +391,7 @@ jobs:
mkdir -p logs/playwright_extra
python tests/studio/playwright_extra_ui.py
- name: Stop second Studio
- name: Stop second Unsloth
if: always()
run: |
kill "${STUDIO_EXTRA_PID}" 2>/dev/null || true
@ -402,5 +407,7 @@ jobs:
logs/studio_extra.log
logs/install.log
logs/playwright
logs/playwright-permissions-*
logs/playwright_extra
logs/studio-permissions-*.log
retention-days: 7

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

@ -5,19 +5,19 @@
# studio-mac-update-smoke.yml. Verifies that on the FREE
# windows-latest runner:
#
# 1. install.ps1 --local --no-torch installs Studio AND auto-fetches
# the prebuilt llama.cpp Windows binary (llama-bNNNN-bin-win-cpu-
# x64 from ggml-org/llama.cpp). Hitting the source-build fallback
# is treated as an Unsloth bug -- Studio must always pick the
# 1. install.ps1 --local --no-torch installs Unsloth AND auto-fetches
# the prebuilt llama.cpp Windows binary (app-<tag>-windows-x64-cpu
# from unslothai/llama.cpp). Hitting the source-build fallback is
# treated as an Unsloth bug -- Unsloth must always pick the
# prebuilt on Windows.
# 2. unsloth studio update --local is idempotent. Two consecutive
# runs both report "prebuilt up to date and validated", no
# source-build fallback. The CLI's _find_setup_script picks
# setup.ps1 on Windows automatically.
# 3. The installed Studio still boots and /api/health returns
# 3. The installed Unsloth still boots and /api/health returns
# healthy after the update path.
name: Windows Studio Update CI
name: Windows Unsloth Update CI
on:
pull_request:
@ -45,7 +45,7 @@ permissions:
jobs:
update-idempotency:
name: Studio Updating Tests
name: Unsloth Updating Tests
runs-on: windows-latest
timeout-minutes: 30
defaults:
@ -53,7 +53,7 @@ jobs:
shell: bash
env:
# Force UTF-8 for stdio (Windows defaults to cp1252; hf
# download / Studio CLI print "✓" checkmarks and crash
# download / Unsloth CLI print "✓" checkmarks and crash
# otherwise).
PYTHONIOENCODING: utf-8
PYTHONUTF8: '1'
@ -90,7 +90,7 @@ jobs:
# reuses the existing Node with no download.
#
# (2) Defender. windows-latest's real-time scan opens / hashes
# every file Studio writes during install (Vite output =
# every file Unsloth writes during install (Vite output =
# thousands of small chunks, uv pip = wheel-extraction =
# thousands of small files). The latency dominates the
# 200 s frontend build and the 90 s deps install. Adding
@ -109,7 +109,7 @@ jobs:
# setup.ps1 line 1281-1296's mtime-based "is the frontend
# stale?" check into "up to date, skip rebuild", because the
# newly-created dist's mtime is younger than every source
# file. Studio then boots with an empty dist and 500s on
# file. Unsloth then boots with an empty dist and 500s on
# GET / with FileNotFoundError: dist\index.html. See run
# 25546676715 / job 74984469728.
# Add-MpPreference accepts paths that do not yet exist; the
@ -129,7 +129,7 @@ jobs:
}
}
- name: Install Studio (--local, --no-torch)
- name: Install Unsloth (--local, --no-torch)
shell: pwsh
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
@ -168,7 +168,7 @@ jobs:
echo "install.ps1 installed the Windows prebuilt llama.cpp:"
cat "$INFO"
- name: Add Studio shim to GITHUB_PATH
- name: Add Unsloth shim to GITHUB_PATH
run: |
SHIM_DIR=~/.unsloth/studio/bin
if [ ! -f "$SHIM_DIR/unsloth.exe" ]; then
@ -212,7 +212,7 @@ jobs:
grep -qE "prebuilt up to date and validated|prebuilt installed and validated" logs/update2.log
echo "second update was clean"
- name: Boot Studio briefly to confirm the install is still usable
- name: Boot Unsloth briefly to confirm the install is still usable
run: |
mkdir -p logs
UNSLOTH_API_ONLY=1 unsloth studio -H 127.0.0.1 -p 18891 \
@ -239,13 +239,13 @@ jobs:
sleep 1
done
if [ -z "$HEALTHY" ]; then
echo "Studio failed to come up after \`update\`"
echo "Unsloth failed to come up after \`update\`"
tail -200 logs/studio.log
kill "$PID" 2>/dev/null || true
exit 1
fi
kill "$PID" 2>/dev/null || true
echo "post-update Studio /api/health OK"
echo "post-update Unsloth /api/health OK"
- name: Uninstall and verify clean
# Round-trip through scripts/uninstall.ps1 against the default

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

@ -285,6 +285,92 @@ jobs:
tests/vllm_compat/test_extended_module_imports.py \
-v --tb=short
# Fake-CUDA GRPO/SFT/DPO patch run against REAL TRL (latest + main). Unlike
# the static symbol/source greps above, this drives unsloth's actual
# source-transform patchers (models/rl.py + rl_replacements.py) on a CPU-only
# runner under the tests/conftest.py spoof harness -- no GPU, no training.
# Catches structural TRL drift the greps miss (e.g. TRL 1.7.0's 2->3-tuple
# per-token-logps return, restructured PEFT ref-adapter block) by asserting
# the generated Unsloth trainer still satisfies the transform contracts.
grpo-fake-run:
name: GRPO fake-run (latest + main TRL, CPU spoof)
runs-on: ubuntu-latest
timeout-minutes: 18
steps:
- uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
with:
persist-credentials: false
path: unsloth
- name: Clone unsloth-zoo @ main
run: |
for attempt in 1 2 3; do
rm -rf "$RUNNER_TEMP/unsloth-zoo"
if git clone --depth=1 https://github.com/unslothai/unsloth-zoo \
"$RUNNER_TEMP/unsloth-zoo"; then
break
fi
if [ "$attempt" -eq 3 ]; then
echo "::error::git clone unsloth-zoo failed after 3 attempts"
exit 1
fi
delay=$((5 * attempt))
echo "::warning::clone failed (attempt $attempt/3), retrying in ${delay}s..."
sleep "$delay"
done
- uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0
with:
python-version: '3.12'
cache: 'pip'
- name: Install CPU torch + ecosystem + TRL latest
run: |
python -m pip install --upgrade pip
pip install --index-url https://download.pytorch.org/whl/cpu --extra-index-url https://pypi.org/simple \
'torch>=2.4,<2.11' 'torchvision<0.26' 'torchcodec<0.10'
# Ecosystem floors unsloth needs; TRL itself is installed last so it
# can pull the transformers/peft it requires.
pip install \
'transformers>=4.57' 'peft>=0.18.0' 'accelerate>=1.0' 'datasets>=3.4,<5' \
'bitsandbytes>=0.45.5' sentencepiece protobuf safetensors numpy 'pytest>=8' \
'huggingface_hub>=0.34' tqdm packaging psutil triton Pillow
pip install --upgrade trl
pip install --no-deps -e "$RUNNER_TEMP/unsloth-zoo"
pip install --no-deps -e ./unsloth
- name: Fake-run vs TRL latest
env:
UNSLOTH_IS_PRESENT: '1'
UNSLOTH_COMPILE_DISABLE: '1'
# Disable dynamo/inductor at the process level, before conftest.py's early
# `import unsloth`, so the GRPO hot path never compiles on the GPU-less runner
# (defense in depth; the CPU fake-train also flips this at runtime).
TORCHDYNAMO_DISABLE: '1'
TORCH_COMPILE_DISABLE: '1'
PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION: python
run: |
cd unsloth
python -c "import trl; print('Resolved TRL', trl.__version__)"
PYTHONPATH=. python -m pytest \
tests/version_compat/test_trl_grpo_fake_run.py \
tests/version_compat/test_trl_fake_train_cpu.py \
-v --tb=short
# `main` is scheduled/dispatch-only so PR jobs stay fast and a bleeding-edge
# TRL break does not red every PR. github.event_name is valid in a step if.
- name: Fake-run vs TRL main (scheduled / dispatch only)
if: ${{ github.event_name != 'pull_request' }}
env:
UNSLOTH_IS_PRESENT: '1'
UNSLOTH_COMPILE_DISABLE: '1'
TORCHDYNAMO_DISABLE: '1'
TORCH_COMPILE_DISABLE: '1'
PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION: python
run: |
pip install --upgrade "git+https://github.com/huggingface/trl"
cd unsloth
python -c "import trl; print('Resolved TRL', trl.__version__)"
PYTHONPATH=. python -m pytest \
tests/version_compat/test_trl_grpo_fake_run.py \
tests/version_compat/test_trl_fake_train_cpu.py \
-v --tb=short
# Daily-only: same suites but with --strict on importable upstream
# tags. Schedule-only so PR jobs stay fast; cron tolerates a flake.
daily-fresh-fetch:

10
.github/workflows/wheel-smoke.yml vendored
View file

@ -3,7 +3,7 @@
# Builds the PyPI wheel from the PR branch, then verifies the built wheel
# actually contains what we expect to ship and does NOT contain the broken
# Studio bundle that 2026.5.1 published. This is the single workflow that
# Unsloth bundle that 2026.5.1 published. This is the single workflow that
# would have blocked the 2026.5.1 release before twine upload.
#
# Verified locally end-to-end against this branch:
@ -12,7 +12,7 @@
# lockfile shipped, frontend dist shipped,
# no node_modules in wheel, no bun.lock in wheel,
# main bundle has unstable_Provider hits=1 (assistant-ui internals only).
# - Studio backend imports cleanly from the installed wheel with the
# - Unsloth backend imports cleanly from the installed wheel with the
# lightweight dep set below.
name: Wheel CI
@ -101,7 +101,7 @@ jobs:
hits = data.count("unstable_Provider:")
print(f"main bundle: {js[0]}")
print(f"unstable_Provider hits: {hits} (>=4 indicates 2026.5.1 regression)")
checks["bundle has no Studio unstable_Provider call site"] = (hits < 4)
checks["bundle has no Unsloth unstable_Provider call site"] = (hits < 4)
print()
for k, v in checks.items():
@ -109,7 +109,7 @@ jobs:
sys.exit(0 if all(checks.values()) else 1)
PY
- name: Studio backend import smoke
- name: Unsloth backend import smoke
# Imports `studio.backend.main:app` from the freshly-installed wheel in
# a clean venv. This catches the class of bug that 2026.5.1 shipped with:
# frontend dist missing, package-lock.json missing, or the wheel's Python
@ -125,7 +125,7 @@ jobs:
/tmp/v/bin/pip install --no-deps dist/unsloth-*.whl
# Run from /tmp so Python imports the installed package, not the source tree.
cd /tmp
/tmp/v/bin/python -c "from studio.backend.main import app; print('Studio backend OK:', app.title)"
/tmp/v/bin/python -c "from studio.backend.main import app; print('Unsloth backend OK:', app.title)"
- name: Upload wheel on failure
if: failure()

2
.gitignore vendored
View file

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

104
README.md
View file

@ -11,6 +11,7 @@ Unsloth Studio lets you run and train models locally.
<p align="center">
<a href="#-features">Features</a> •
<a href="#-unsloth-news">News</a> •
<a href="#-install">Quickstart</a> •
<a href="#-free-notebooks">Notebooks</a> •
<a href="https://unsloth.ai/docs">Documentation</a>
@ -47,15 +48,51 @@ Unsloth Studio (Beta) lets you run and train text, [audio](https://unsloth.ai/do
* [Auto set inference settings](https://unsloth.ai/docs/new/studio/chat#auto-parameter-tuning) and customize chat templates.
* We work directly with teams behind [gpt-oss](https://docs.unsloth.ai/new/gpt-oss-how-to-run-and-fine-tune#unsloth-fixes-for-gpt-oss), [Qwen3](https://www.reddit.com/r/LocalLLaMA/comments/1kaodxu/qwen3_unsloth_dynamic_ggufs_128k_context_bug_fixes/), [Llama 4](https://github.com/ggml-org/llama.cpp/pull/12889), [Mistral](https://huggingface.co/mistralai/Mistral-Medium-3.5-128B/discussions/18), [Gemma 1-3](https://news.ycombinator.com/item?id=39671146), and [Phi-4](https://unsloth.ai/blog/phi4), where we’ve fixed bugs that improve model accuracy.
* Chat with images, audio, PDFs, code, DOCX and more. [Connect API providers](https://unsloth.ai/docs/integrations/connections) (OpenAI, Anthropic) or servers (vLLM, Ollama).
* [**Compare any two models**](https://unsloth.ai/docs/new/studio/chat#model-arena) side by side with the same prompt.
* **OpenAI/Anthropic-compatible APIs**: Serve local models through `/v1/chat/completions`, `/v1/responses` and `/v1/messages`.
* **Connect local models to agents**: Use `unsloth start` with Claude Code, Codex, Hermes and more.
* **Web/PDF search** can read PDF papers, manuals and other PDF results.
* **GGUF hardware controls**: Choose GPUs/layers, offload MoE experts, use multi-GPU or Tensor Parallelism.
* The opt-in **MCP control endpoint** lets AI clients manage models, training, recipes and exports.
### Training
* Train and RL **500+ models** up to **2x faster** with up to **70% less VRAM**, with no accuracy loss.
* Custom Triton and mathematical **kernels**. See some collabs we did with [PyTorch](https://unsloth.ai/docs/get-started/reinforcement-learning-rl-guide/fp8-reinforcement-learning) and [Hugging Face](https://unsloth.ai/docs/new/faster-moe).
* Train and RL **500+ models** up to **2x faster** with **70% less VRAM**; MoE up to **12x faster**.
* Train and run RL on [AMD GPUs](https://unsloth.ai/docs/basics/amd) across Windows, WSL and Linux.
* **Data Recipes**: [Auto-create datasets](https://unsloth.ai/docs/new/studio/data-recipe) from **PDF, CSV, DOCX** etc. Edit data in a visual-node workflow.
* **[Reinforcement Learning](https://unsloth.ai/docs/get-started/reinforcement-learning-rl-guide)** (RL): The most efficient [RL](https://unsloth.ai/docs/get-started/reinforcement-learning-rl-guide) library, using **80% less VRAM** for GRPO, [FP8](https://unsloth.ai/docs/get-started/reinforcement-learning-rl-guide/fp8-reinforcement-learning) etc.
* Supports full fine-tuning, RL, pretraining, 4-bit, 16-bit and, FP8 training.
* **[Reinforcement Learning](https://unsloth.ai/docs/get-started/reinforcement-learning-rl-guide)** uses **80% less VRAM** for GRPO, FP8 and vision RL, with 7x longer contexts.
* [**Long-context training**](https://unsloth.ai/docs/new/3x-faster-training-packing): **3x faster**, 30% less VRAM and 500K+ context.
* Supports LoRA/QLoRA, full fine-tuning, RL, pretraining, 4-bit, 16-bit and FP8.
* Custom Triton and mathematical **kernels** built with PyTorch and Hugging Face.
* **Observability**: Monitor training live, track loss and GPU usage and customize graphs.
* [Multi-GPU](https://unsloth.ai/docs/basics/multi-gpu-training-with-unsloth) training is supported, with major improvements coming soon.
## 🚀 Unsloth Start
[Unsloth Start](https://unsloth.ai/docs/integrations/unsloth-start) connects [Claude Code](https://unsloth.ai/docs/basics/claude-code), [Codex](https://unsloth.ai/docs/basics/codex) and other agents to local models with one command.
Start Unsloth, load a model, open your project folder, then run:
```bash
unsloth start claude
```
Replace `claude` with any supported agent:
| Agent | Command |
| --- | --- |
| Claude Code | `unsloth start claude` |
| OpenAI Codex | `unsloth start codex` |
| Hermes Agent | `unsloth start hermes` |
| OpenClaw | `unsloth start openclaw` |
| OpenCode | `unsloth start opencode` |
| Pi Coding Agent | `unsloth start pi` |
Claude Code, Codex, OpenCode and Pi can keep their current model and use Unsloth as a local
subagent:
```bash
unsloth start claude --as-subagent --model unsloth/model-GGUF:quant
```
## 📥 Install
Unsloth can be used in two ways: through **[Unsloth Studio](https://unsloth.ai/docs/new/studio/)**, the web UI, or through **Unsloth Core**, the code-based version. Each has different requirements.
@ -65,7 +102,8 @@ Unsloth Studio (Beta) works on **Windows, Linux, WSL** and **macOS**.
* **CPU:** Supported for Chat and Data Recipes currently
* **NVIDIA:** Training works on RTX 30/40/50, Blackwell, DGX Spark, Station and more
* **macOS:** Training, MLX and GGUF inference are ALL supported.
* **AMD:** Chat + Data works. Train with [Unsloth Core](#unsloth-core-code-based). Studio support is out soon.
* **AMD:** Training, RL, chat and deployment work on Windows, WSL and Linux. [Read the AMD guide](https://unsloth.ai/docs/basics/amd).
* **Vulkan:** GGUF inference is supported on [compatible GPUs, including Intel GPUs](https://github.com/unslothai/unsloth/pull/5819).
* **Multi-GPU:** Available now, with a major upgrade on the way
#### macOS, Linux, WSL:
@ -84,9 +122,9 @@ 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.
To reach Studio over HTTPS, use `unsloth studio --secure`. Studio stays bound to localhost and is reached only through a free Cloudflare tunnel, which publishes it at a public `https://*.trycloudflare.com` URL (it fails closed if the tunnel can't start, so the raw port is never exposed). This makes Studio reachable from the internet, so anyone with the link and API key can use it and run code: keep your API key private (see Remote access below).
To reach Unsloth over HTTPS, use `unsloth studio --secure`. Unsloth stays bound to localhost and is reached only through a free Cloudflare tunnel, which publishes it at a public `https://*.trycloudflare.com` URL (it fails closed if the tunnel can't start, so the raw port is never exposed). This makes Unsloth reachable from the internet, so anyone with the link and API key can use it and run code: keep your API key private (see Remote access below).
#### Docker
Use our [Docker image](https://hub.docker.com/r/unsloth/unsloth) ```unsloth/unsloth``` container. Run:
@ -122,7 +160,7 @@ You can use the same Docker image as Unsloth Studio.
#### AMD, Intel:
For RTX 50x, B200, 6000 GPUs: `uv pip install unsloth --torch-backend=auto`. Read our guides for: [Blackwell](https://unsloth.ai/docs/blog/fine-tuning-llms-with-blackwell-rtx-50-series-and-unsloth) and [DGX Spark](https://unsloth.ai/docs/blog/fine-tuning-llms-with-nvidia-dgx-spark-and-unsloth). <br>
To install Unsloth on **AMD** and **Intel** GPUs, follow our [AMD Guide](https://unsloth.ai/docs/get-started/install/amd) and [Intel Guide](https://unsloth.ai/docs/get-started/install/intel).
To install Unsloth on **AMD** and **Intel** GPUs, follow our [AMD Guide](https://unsloth.ai/docs/basics/amd) and [Intel Guide](https://unsloth.ai/docs/get-started/install/intel).
## 📒 Free Notebooks
@ -148,13 +186,20 @@ Read our [guide](https://unsloth.ai/docs/get-started/fine-tuning-llms-guide). Ad
- See detailed documentation for Unsloth [here](https://unsloth.ai/docs)
## 🦥 Unsloth News
- **Connections**: Connect any API provider (OpenAI, Anthropic) or server (vLLM, Ollama). [Guide](https://unsloth.ai/docs/integrations/connections)
- **MTP**: Run Qwen3.6 MTP in Unsloth. MTP settings are autoset specific to your hardware. [Guide](https://unsloth.ai/docs/models/qwen3.6#mtp-guide)
- **API inference endpoint**: Deploy and run local LLMs in Claude Code, Codex tools. [Guide](https://unsloth.ai/docs/basics/api)
- **Qwen3.6**: Qwen3.6-35B-A3B can now be trained and run in Unsloth Studio. [Blog](https://unsloth.ai/docs/models/qwen3.6)
- **Gemma 4**: Run and train Google’s new models directly in Unsloth. [Blog](https://unsloth.ai/docs/models/gemma-4)
- **AMD training**: Train, run RL, chat and deploy on AMD GPUs across Windows, WSL and Linux. [Guide](https://unsloth.ai/docs/basics/amd)
- **GGUF hardware controls**: Choose GPU/layer placement, offload MoE experts and use multi-GPU or Tensor Parallelism. [#6414](https://github.com/unslothai/unsloth/pull/6414)
- **Local models for any agent**: Use `unsloth start` with Claude Code, Codex, Hermes, OpenCode, OpenClaw, Pi and more through Unsloth's OpenAI- and Anthropic-compatible APIs. [Guide](https://unsloth.ai/docs/basics/api)
- **MCP control endpoint**: Let compatible clients manage models, training, recipes, checkpoints and exports. [#7191](https://github.com/unslothai/unsloth/pull/7191)
- **Local inference reliability**: Resume long chats faster, recover stalled downloads and reuse existing GGUF files. [#7204](https://github.com/unslothai/unsloth/pull/7204) • [#6858](https://github.com/unslothai/unsloth/pull/6858) • [#7209](https://github.com/unslothai/unsloth/pull/7209)
- **New models**: [Qwen-AgentWorld](https://huggingface.co/unsloth/Qwen-AgentWorld-35B-A3B-GGUF), [Ornith](https://huggingface.co/unsloth/models?search=ornith), [Kimi K2.7 Code](https://unsloth.ai/docs/models/kimi-k2.7-code) and [MiniMax M3](https://unsloth.ai/docs/models/minimax-m3)
- **GLM-5.2**: Run Z.ai's 744B-parameter, 1M-context open model locally with Unsloth Dynamic GGUFs. [Guide](https://unsloth.ai/docs/models/glm-5.2)
- **DeepSeek-V4**: Run DeepSeek-V4-Flash locally with corrected multi-turn and tool-calling behavior. [Guide](https://unsloth.ai/docs/models/deepseek-v4)
- **DiffusionGemma**: Run and fine-tune Google's diffusion language model with 1.8x faster inference in Unsloth Studio. [Guide](https://unsloth.ai/docs/models/diffusiongemma)
- **Qwen3.6**: Run and train Qwen3.6 with MTP for 1.4-2.2x faster inference and NVFP4 quants for supported GPUs. [Guide](https://unsloth.ai/docs/models/qwen3.6)
- **Gemma 4**: Run and train Gemma 4 text, image and audio models with QAT, MTP, GGUF and MLX support. [Guide](https://unsloth.ai/docs/models/gemma-4)
- **MCP servers**: Connect local models to files, apps, databases and external tools through Model Context Protocol. [Guide](https://unsloth.ai/docs/basics/mcp)
- **Connections**: Mix local models with API providers (OpenAI, Anthropic) or servers (vLLM, Ollama) in the same interface. [Guide](https://unsloth.ai/docs/integrations/connections)
- **Introducing Unsloth Studio**: our new web UI for running and training LLMs. [Blog](https://unsloth.ai/docs/new/studio)
- **Qwen3.5** - 0.8B, 2B, 4B, 9B, 27B, 35-A3B, 112B-A10B are now supported. [Guide + notebooks](https://unsloth.ai/docs/models/qwen3.5/fine-tune)
- Train **MoE LLMs 12x faster** with 35% less VRAM - DeepSeek, GLM, Qwen and gpt-oss. [Blog](https://unsloth.ai/docs/new/faster-moe)
- **Embedding models**: Unsloth now supports ~1.8-3.3x faster embedding fine-tuning. [Blog](https://unsloth.ai/docs/new/embedding-finetuning) • [Notebooks](https://unsloth.ai/docs/get-started/unsloth-notebooks#embedding-models)
- New **7x longer context RL** vs. all other setups, via our new batching algorithms. [Blog](https://unsloth.ai/docs/new/grpo-long-context)
@ -208,16 +253,29 @@ unsloth studio -p 8888
#### Remote access: `--secure` (HTTPS tunnel) vs raw port
By default `unsloth studio` binds to `127.0.0.1` (this machine only). To reach it from another device, pick one of:
- `--secure` (recommended): serve **only** through a free Cloudflare HTTPS link. Studio stays bound to localhost and the tunnel provides the public URL; it fails closed (does not start) if the tunnel can't come up, so the raw port is never exposed.
- `--secure` (recommended): serve **only** through a free Cloudflare HTTPS link. Unsloth stays bound to localhost and the tunnel provides the public URL; it fails closed (does not start) if the tunnel can't come up, so the raw port is never exposed.
```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.
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.
The first time Unsloth 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: Unsloth 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 Unsloth.
#### Advanced launch options
Installer options can be passed as environment variables. On macOS, Linux and WSL place the variable after the pipe so the shell passes it to `sh`; on Windows set it with `$env:` before piping to `iex`.
@ -230,6 +288,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 Unsloth (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
@ -258,9 +324,9 @@ UNSLOTH_NPM_REGISTRY=https://artifactory.example.com/api/npm/npm/ ./install.sh -
```powershell
$env:UNSLOTH_NPM_REGISTRY='https://artifactory.example.com/api/npm/npm/'; .\install.ps1 --local
```
It is threaded as `--registry` into the Studio frontend `npm`/`bun` installs; the supply-chain locks (7-day `min-release-age`, exact version pins) stay in force.
It is threaded as `--registry` into the Unsloth frontend `npm`/`bun` installs; the supply-chain locks (7-day `min-release-age`, exact version pins) stay in force.
Cap Studio's native CPU thread pools on high-core hosts: `UNSLOTH_CPU_THREADS=8 unsloth studio -p 8888`.
Cap Unsloth's native CPU thread pools on high-core hosts: `UNSLOTH_CPU_THREADS=8 unsloth studio -p 8888`.
#### Uninstall
The recommended way to fully remove Unsloth Studio is the matching uninstall script for your OS. It stops any running servers, removes the install dir, the launcher data dir, the desktop shortcut, and any platform-specific entries (macOS `.app` bundle + Launch Services on Mac; Start Menu, `HKCU\Software\Unsloth` registry key and user `PATH` entries on Windows):

8
build.sh
View file

@ -4,9 +4,9 @@
set -euo pipefail
# PyPI/Studio release publishing must use `./build.sh publish` (or an
# equivalent stamp -> build -> verify-dist -> upload flow) so packaged Studio
# artifacts include the display-only Studio release version.
# PyPI/Unsloth release publishing must use `./build.sh publish` (or an
# equivalent stamp -> build -> verify-dist -> upload flow) so packaged Unsloth
# artifacts include the display-only Unsloth release version.
# 1. Build frontend (Vite outputs to dist/)
cd studio/frontend
@ -87,7 +87,7 @@ cd ../..
# 2. Clean old artifacts
rm -rf build dist *.egg-info
# 3. Stamp display-only Studio release metadata for packaged builds.
# 3. Stamp display-only Unsloth release metadata for packaged builds.
_STUDIO_BUILD_INFO="studio/backend/utils/_studio_release_build.py"
_STUDIO_BUILD_INFO_BACKUP="$(mktemp)"
cp "$_STUDIO_BUILD_INFO" "$_STUDIO_BUILD_INFO_BACKUP"

327
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
@ -52,7 +53,8 @@ function Install-UnslothStudio {
param([string]$TorchIndexUrl)
if ($SkipTorch) { return "none" }
if ([string]::IsNullOrWhiteSpace($TorchIndexUrl)) { return "none" }
$leaf = ($TorchIndexUrl.TrimEnd('/') -split '/')[-1].ToLowerInvariant()
# Drop query/fragment first so a token-authenticated pin classifies by family.
$leaf = (($TorchIndexUrl -split '[?#]', 2)[0].TrimEnd('/') -split '/')[-1].ToLowerInvariant()
if (@("cpu", "cu118", "cu124", "cu126", "cu128", "cu130") -contains $leaf) { return $leaf }
if ($leaf -match '^rocm[0-9]+\.[0-9]+$') { return $leaf }
return "auto"
@ -61,7 +63,8 @@ function Install-UnslothStudio {
function Get-TauriGpuBranch {
param([string]$TorchIndexFamily)
if ($SkipTorch) { return "no_torch" }
if ($TorchIndexFamily -like "cu*") { return "cuda" }
# Require a digit after "cu" so /current or /custom isn't branded CUDA (parity ^cu[0-9]).
if ($TorchIndexFamily -match '^cu[0-9]') { return "cuda" }
if ($TorchIndexFamily -like "rocm*") { return "rocm" }
if ($TorchIndexFamily -eq "cpu") { return "cpu" }
return "unknown"
@ -90,6 +93,7 @@ function Install-UnslothStudio {
if ($TauriMode) {
exit $Code
}
throw $Message
}
# ── Parse flags ──
@ -98,6 +102,7 @@ function Install-UnslothStudio {
$RepoRoot = ""
$TauriMode = $false
$SkipTorch = $false
$SkipAutostart = $false
$ShortcutsOnly = $false
$WithLlamaCppDir = ""
$argList = $args
@ -130,6 +135,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.
@ -172,7 +178,7 @@ function Install-UnslothStudio {
$envOverride = $env:STUDIO_HOME.Trim()
}
# Custom Studio roots are not supported with --tauri (desktop app still
# Custom Unsloth roots are not supported with --tauri (desktop app still
# resolves %USERPROFILE%\.unsloth\studio). Pass through if override == legacy.
if ($TauriMode -and $envOverride) {
$_tauriOverride = $envOverride
@ -463,12 +469,36 @@ function Install-UnslothStudio {
}
}
# Redact index-URL credentials (userinfo + ?query= + #fragment) from captured installer
# output before printing on failure; uv/pip errors echo the failing --index-url verbatim.
# Mirrors the other installers. Verbose mode streams uncaptured, so it isn't redacted.
function Redact-InstallOutput {
param([string]$Text)
if (-not $Text) { return $Text }
$Text = $Text -replace '(https?://)[^/@\s`]+@', '$1<redacted>@'
$Text = $Text -replace '([?&][^=\s&`]+)=[^&#\s`]+', '$1=<redacted>'
# A #token=... fragment is as sensitive as a query; URL-anchored.
return $Text -replace '(https?://[^\s`#]+)#[^\s`]+', '$1#<redacted>'
}
# Run native commands quietly by default to match install.sh behavior.
# Full command output is shown only when --verbose / UNSLOTH_VERBOSE=1.
function Invoke-InstallCommand {
param(
[Parameter(Mandatory = $true)][ScriptBlock]$Command
)
# Installer-pinned index installs (torch) must beat an inherited uv mirror (#6898):
# for --default-index, clear the uv index env vars (restore in finally) and set
# UV_NO_CONFIG=1 so a uv.toml/pyproject index can't outrank the CLI pin (uv 0.10).
$savedUvIndex = $null
if ($Command.ToString() -match '--default-index') {
$savedUvIndex = @{}
foreach ($n in 'UV_DEFAULT_INDEX', 'UV_INDEX_URL', 'UV_INDEX', 'UV_EXTRA_INDEX_URL', 'UV_TORCH_BACKEND', 'UV_FIND_LINKS', 'UV_CONFIG_FILE', 'UV_NO_CONFIG') {
$savedUvIndex[$n] = [Environment]::GetEnvironmentVariable($n)
Remove-Item "Env:$n" -ErrorAction SilentlyContinue
}
$env:UV_NO_CONFIG = '1'
}
$prevEap = $ErrorActionPreference
$ErrorActionPreference = "Continue"
try {
@ -478,16 +508,23 @@ function Install-UnslothStudio {
# Merge stderr into stdout so progress/warning output stays visible
# without flipping $? on successful native commands (PS 5.1 treats
# stderr records as errors that set $? = $false even on exit code 0).
& $Command 2>&1 | Out-Host
# Redact per record: uv echoes index URLs (credentials and all) in
# its errors, and verbose mode must not bypass the quiet path's
# redaction. ForEach-Object/Out-Host leave $LASTEXITCODE untouched.
& $Command 2>&1 | ForEach-Object { Redact-InstallOutput "$_" } | Out-Host
} else {
$output = & $Command 2>&1 | Out-String
if ($LASTEXITCODE -ne 0) {
Write-Host $output -ForegroundColor Red
Write-Host (Redact-InstallOutput $output) -ForegroundColor Red
}
}
return [int]$LASTEXITCODE
} finally {
$ErrorActionPreference = $prevEap
if ($savedUvIndex) {
Remove-Item "Env:UV_NO_CONFIG" -ErrorAction SilentlyContinue
foreach ($n in $savedUvIndex.Keys) { if ($null -ne $savedUvIndex[$n]) { Set-Item "Env:$n" $savedUvIndex[$n] } }
}
}
}
@ -740,7 +777,7 @@ function Find-FreeLaunchPort {
return `$null
}
# If Studio is already healthy on any expected port, just open it and exit.
# If Unsloth is already healthy on any expected port, just open it and exit.
`$existingPort = Find-HealthyStudioPort
if (`$existingPort) {
Start-Process "http://localhost:`$existingPort"
@ -756,7 +793,7 @@ try {
`$haveMutex = `$true
}
if (-not `$haveMutex) {
# Another launcher is already running; wait for it to bring Studio up
# Another launcher is already running; wait for it to bring Unsloth up
`$deadline = (Get-Date).AddSeconds(`$timeoutSec)
while ((Get-Date) -lt `$deadline) {
`$port = Find-HealthyStudioPort
@ -1379,13 +1416,82 @@ exit 0
$suffix++
$candidate = Join-Path $StudioHome "unsloth_studio.rollback.$stamp.$PID.$suffix"
}
Move-Item -LiteralPath $ExistingDir -Destination $candidate -ErrorAction Stop
$script:StudioVenvRollbackDir = $candidate
$script:StudioVenvRollbackTarget = $ExistingDir
$script:StudioVenvRollbackActive = $true
# Publish the rollback state before the atomic rename so interruption
# cannot land after Move-Item but before cleanup knows where the old venv went.
try {
Move-Item -LiteralPath $ExistingDir -Destination $candidate -ErrorAction Stop
} catch {
# A collision or ordinary rename failure leaves the original in place.
# Keep state active only when the rename happened before interruption.
if (Test-Path -LiteralPath $ExistingDir) {
$script:StudioVenvRollbackActive = $false
$script:StudioVenvRollbackDir = $null
}
throw
}
substep "previous environment preserved for rollback"
}
function Remove-StudioVenvTreeWithRetry {
param(
[Parameter(Mandatory = $true)][string]$Path,
[Parameter(Mandatory = $true)][string]$Label
)
$lastError = $null
for ($attempt = 1; $attempt -le 3; $attempt++) {
try {
Remove-Item -LiteralPath $Path -Recurse -Force -ErrorAction Stop
} catch {
$lastError = $_.Exception.Message
}
if (-not (Test-Path -LiteralPath $Path)) { return $true }
if ($attempt -lt 3) { Start-Sleep -Milliseconds (250 * $attempt) }
}
Write-Host "[WARN] Could not remove $Label at $Path" -ForegroundColor Yellow
if ($lastError) { Write-Host " $lastError" -ForegroundColor Yellow }
return $false
}
function Test-StudioVenvRollbackMustBePreserved {
param([Parameter(Mandatory = $true)][System.IO.FileSystemInfo]$Rollback)
# Preserve anything outside the installer's timestamp.PID[.suffix] format.
if ($Rollback.Name -notmatch '^unsloth_studio\.rollback\.[0-9]{14}\.([0-9]+)(?:\.[0-9]+)?$') {
return $true
}
$ownerPid = 0
if (-not [int]::TryParse($Matches[1], [ref]$ownerPid)) { return $true }
if ($ownerPid -eq $PID) { return $true }
return $null -ne (Get-Process -Id $ownerPid -ErrorAction SilentlyContinue)
}
function Remove-StaleStudioVenvRollbacks {
try {
$rollbacks = @(
Get-ChildItem -LiteralPath $StudioHome -Directory -Force -ErrorAction Stop |
Where-Object { $_.Name -like 'unsloth_studio.rollback.*' }
)
} catch {
Write-Host "[WARN] Could not inspect stale environment rollbacks in $StudioHome" -ForegroundColor Yellow
Write-Host " $($_.Exception.Message)" -ForegroundColor Yellow
return
}
foreach ($rollback in $rollbacks) {
if (($rollback.Attributes -band [System.IO.FileAttributes]::ReparsePoint) -ne 0) {
Write-Host "[WARN] Refusing to remove rollback reparse point $($rollback.FullName)" -ForegroundColor Yellow
continue
}
# A concurrent installer may have moved its live venv aside. The PID
# in the generated name keeps this run from deleting its rescue copy.
if (Test-StudioVenvRollbackMustBePreserved -Rollback $rollback) { continue }
if (Remove-StudioVenvTreeWithRetry -Path $rollback.FullName -Label "stale environment rollback") {
substep "removed stale environment rollback $($rollback.Name)"
}
}
}
function Restore-StudioVenvRollback {
if (-not $script:StudioVenvRollbackActive) { return }
$backup = $script:StudioVenvRollbackDir
@ -1397,7 +1503,9 @@ exit 0
substep "restoring previous environment after failed install..." "Yellow"
try {
if (Test-Path -LiteralPath $target) {
Remove-Item -LiteralPath $target -Recurse -Force -ErrorAction SilentlyContinue
if (-not (Remove-StudioVenvTreeWithRetry -Path $target -Label "incomplete environment")) {
throw "Could not remove incomplete environment at $target"
}
}
Move-Item -LiteralPath $backup -Destination $target -Force -ErrorAction Stop
substep "restored previous environment"
@ -1412,17 +1520,21 @@ exit 0
function Complete-StudioVenvRollback {
if (-not $script:StudioVenvRollbackActive) { return }
$backup = $script:StudioVenvRollbackDir
if ($backup -and (Test-Path -LiteralPath $backup)) {
Remove-Item -LiteralPath $backup -Recurse -Force -ErrorAction SilentlyContinue
}
# The replacement is committed. Disable restoration before deleting the
# backup so interruption cannot restore a partially deleted environment.
$script:StudioVenvRollbackActive = $false
$script:StudioVenvRollbackDir = $null
if ($backup -and (Test-Path -LiteralPath $backup)) {
Remove-StudioVenvTreeWithRetry -Path $backup -Label "environment rollback" | Out-Null
}
}
$studioVenvReplacementCommitted = $false
try {
if (Test-Path -LiteralPath $VenvPython) {
# why: matching guard to the .venv branch below -- in env-mode
# $StudioHome is a user-chosen workspace, so refuse to nuke an
# existing $StudioHome\unsloth_studio that lacks Studio sentinels.
# existing $StudioHome\unsloth_studio that lacks Unsloth sentinels.
# -PathType Leaf rejects a directory at the sentinel path. Accept the
# in-VENV ownership marker so partial-install retries are not blocked.
if (
@ -1433,7 +1545,7 @@ exit 0
) {
Write-Host "[ERROR] $VenvDir already exists but does not look like an Unsloth Studio install." -ForegroundColor Red
Write-Host " Move it aside or choose an empty UNSLOTH_STUDIO_HOME." -ForegroundColor Yellow
throw "Refusing to delete non-Studio venv at $VenvDir"
throw "Refusing to delete non-Unsloth venv at $VenvDir"
}
# New layout already exists -- replace only after preserving rollback copy.
substep "preserving existing environment for rollback..."
@ -1452,7 +1564,7 @@ exit 0
# workspace root (e.g. user's existing project Python venv).
$OldVenv = Join-Path $StudioHome ".venv"
$OldPy = Join-Path $OldVenv "Scripts\python.exe"
substep "found legacy Studio environment, validating..."
substep "found legacy Unsloth environment, validating..."
$prevEAP2 = $ErrorActionPreference
$ErrorActionPreference = "Continue"
try {
@ -1482,7 +1594,7 @@ exit 0
# Skip in env-mode so we don't relocate the default-install venv into
# the workspace root.
$CwdVenv = Join-Path $env:USERPROFILE "unsloth_studio"
substep "found CWD-relative Studio environment, migrating to $VenvDir..."
substep "found CWD-relative Unsloth environment, migrating to $VenvDir..."
Move-Item -LiteralPath $CwdVenv -Destination $VenvDir -Force
substep "moved ~/unsloth_studio -> ~/.unsloth/studio/unsloth_studio"
$_Migrated = $true
@ -1501,7 +1613,7 @@ exit 0
substep "$VenvDir"
}
# Mark the freshly-created venv as Studio-owned so a partial install can be
# Mark the freshly-created venv as Unsloth-owned so a partial install can be
# repaired by re-running install.ps1; the env-mode deletion guard above
# accepts this marker as the primary sentinel.
if (Test-Path -LiteralPath $VenvDir -PathType Container) {
@ -1510,7 +1622,7 @@ exit 0
# ── Helper: run amd-smi without triggering a UAC elevation prompt ──
# amd-smi on Windows auto-elevates to read GPU/APU memory, surfacing a confusing
# DiskPart UAC prompt mid-install (Studio backend amd.py hits the same).
# DiskPart UAC prompt mid-install (Unsloth backend amd.py hits the same).
# __COMPAT_LAYER=RunAsInvoker forces it (and helpers it spawns) to run
# un-elevated; on failure the WMI name -> gfx fallback still resolves the arch.
function Invoke-AmdSmiNoElevate {
@ -1637,7 +1749,7 @@ exit 0
function Test-HipinfoIsVenvInternal {
param([AllowNull()][string]$HipinfoPath)
if ([string]::IsNullOrWhiteSpace($HipinfoPath)) { return $false }
# Also derive the venv from the setup python + default Studio home, so
# Also derive the venv from the setup python + default Unsloth home, so
# the venv hipInfo is caught when VenvDir/VIRTUAL_ENV are unset.
$venvRoots = @()
if ($env:VIRTUAL_ENV) { $venvRoots += $env:VIRTUAL_ENV }
@ -1647,7 +1759,7 @@ exit 0
try { $venvRoots += (Split-Path -Parent (Split-Path -Parent $env:UNSLOTH_SETUP_PYTHON)) } catch {}
}
if ($env:USERPROFILE) { $venvRoots += (Join-Path $env:USERPROFILE ".unsloth\studio\unsloth_studio") }
# A custom Studio home (UNSLOTH_STUDIO_HOME / STUDIO_HOME alias) moves the
# A custom Unsloth home (UNSLOTH_STUDIO_HOME / STUDIO_HOME alias) moves the
# venv off the default path; seed it too or its hipInfo escapes the filter.
$studioHomeEnv = if (-not [string]::IsNullOrWhiteSpace($env:UNSLOTH_STUDIO_HOME)) { $env:UNSLOTH_STUDIO_HOME.Trim() } elseif (-not [string]::IsNullOrWhiteSpace($env:STUDIO_HOME)) { $env:STUDIO_HOME.Trim() } else { $null }
if ($studioHomeEnv) {
@ -1807,7 +1919,7 @@ exit 0
$nameArchTable = @(
@{ P = "9070 XT|9080"; A = "gfx1201" } # RDNA 4 (RX 9070 XT / 9080)
@{ P = "9070|9060"; A = "gfx1200" } # RDNA 4 (RX 9070 / 9060)
@{ P = "8060S|8050S|8040S|Strix Halo|Ryzen AI Max|AI Max"; A = "gfx1151" } # RDNA 3.5 (Strix Halo: Radeon 8060S/8050S/8040S iGPU, Ryzen AI Max+)
@{ P = "8065S|8060S|8050S|8040S|Strix Halo|Ryzen AI Max|AI Max"; A = "gfx1151" } # RDNA 3.5 (Strix Halo + Gorgon Halo: Radeon 8065S/8060S/8050S/8040S iGPU, Ryzen AI Max / Max+)
@{ P = "890M|880M|860M|840M|Strix Point|Krackan|HX 37[05]|AI 9 HX|AI 9 36[05]|AI 7 35[05]|AI 5 34[05]|AI 7 PRO 35|AI 5 33"; A = "gfx1150" } # RDNA 3.5 (Strix/Krackan Point: Radeon 890M/880M iGPU, Ryzen AI 9 HX 370/375)
@{ P = "RX 7900|RX 7800|RX 7700(?!S)|PRO W7900|PRO W7800|PRO W7700"; A = "gfx1100" } # RDNA 3 desktop/workstation (Navi 31)
@{ P = "RX 7600|RX 7700S|RX 7650|PRO W7600|PRO W7500|PRO V710"; A = "gfx1102" } # RDNA 3 (Navi 33)
@ -1926,7 +2038,7 @@ exit 0
substep " Ensure the ROCm compute driver is installed alongside the display driver:" "Yellow"
substep " https://rocm.docs.amd.com/en/latest/deploy/windows/index.html" "Yellow"
} elseif ($ROCmGfxArch) {
# Known arch: Studio setup installs AMD's bundled-runtime ROCm PyTorch wheels
# Known arch: Unsloth setup installs AMD's bundled-runtime ROCm PyTorch wheels
# (repo.amd.com), which ship their own runtime -- HIP SDK optional.
step "gpu" "AMD ROCm ($ROCmGfxArch)" "Cyan"
substep "Detected: $ROCmGpuLabel" "Cyan"
@ -1944,10 +2056,31 @@ exit 0
# On an AMD GPU (no NVIDIA), surface the optional WSL-ROCm driver hint.
if (-not $HasNvidiaSmi -and ($ROCmGfxArch -or $ROCmGpuLabel)) { Show-AmdWslDriverHint }
# Trim trailing slashes from the URL PATH only, preserving ?query / #fragment: a whole-URL
# TrimEnd corrupts a token ending in "/", a single strip leaves .../cu128// empty. Shared.
function Trim-IndexPathSlashes {
param([string]$Url)
$value = $Url.Trim()
$idx = $value.IndexOfAny([char[]]@('?', '#'))
if ($idx -lt 0) {
return $value.TrimEnd('/')
}
return $value.Substring(0, $idx).TrimEnd('/') + $value.Substring($idx)
}
# ── Choose the correct PyTorch index URL based on driver CUDA version ──
# Mirrors Get-PytorchCudaTag in setup.ps1.
function Get-TorchIndexUrl {
$baseUrl = if ($env:UNSLOTH_PYTORCH_MIRROR) { $env:UNSLOTH_PYTORCH_MIRROR.TrimEnd('/') } else { "https://download.pytorch.org/whl" }
# Explicit pin -- skip ALL GPU probing (headless / CI / cross-install).
# UNSLOTH_TORCH_INDEX_URL wins (full URL, verbatim); _FAMILY is the leaf appended
# to the mirror base. Matches install.sh / install_python_stack.py.
if (-not [string]::IsNullOrWhiteSpace($env:UNSLOTH_TORCH_INDEX_URL)) {
return (Trim-IndexPathSlashes $env:UNSLOTH_TORCH_INDEX_URL)
}
if (-not [string]::IsNullOrWhiteSpace($env:UNSLOTH_TORCH_INDEX_FAMILY)) {
return "$baseUrl/$($env:UNSLOTH_TORCH_INDEX_FAMILY.Trim().Trim('/'))"
}
if (-not $NvidiaSmiExe) { return "$baseUrl/cpu" }
try {
$output = Invoke-NvidiaSmiBounded $NvidiaSmiExe
@ -1968,6 +2101,27 @@ exit 0
return "$baseUrl/cu126"
}
# Strip userinfo AND query/fragment so an authenticated pin never leaks. Shared with
# _strip_index_url_credentials (install.sh / py / setup.ps1).
function Remove-IndexUrlCredentials {
param([string]$Url)
# Ordinal, not culture-aware: on non-English locales (e.g. th-TH) linguistic
# IndexOf treats "://" as ignorable, mis-locates it, and crashes Substring (issue #7279).
$sep = $Url.IndexOf('://', [System.StringComparison]::Ordinal)
if ($sep -lt 0) { return $Url }
$scheme = $Url.Substring(0, $sep)
$rest = $Url.Substring($sep + 3)
# Drop query / fragment (may hold auth tokens).
$q = $rest.IndexOfAny([char[]]('?', '#'))
if ($q -ge 0) { $rest = $rest.Substring(0, $q) }
$slash = $rest.IndexOf('/', [System.StringComparison]::Ordinal)
$authority = if ($slash -ge 0) { $rest.Substring(0, $slash) } else { $rest }
$at = $authority.LastIndexOf('@', [System.StringComparison]::Ordinal)
$host_ = if ($at -ge 0) { $authority.Substring($at + 1) } else { $authority }
if ($slash -ge 0) { return "${scheme}://${host_}$($rest.Substring($slash))" }
return "${scheme}://${host_}"
}
# ── Torch flavor helpers (to repair a stale CPU / wrong-CUDA wheel) ──
# torch.__version__ -> flavor tag (cuXXX / rocm / cpu); untagged wheel = cpu,
# matching setup.ps1's stale-venv parse.
@ -1986,11 +2140,13 @@ exit 0
param([string]$TorchIndexUrl, [string]$ROCmIndexUrl)
if (-not [string]::IsNullOrWhiteSpace($ROCmIndexUrl)) { return 'rocm' }
if ([string]::IsNullOrWhiteSpace($TorchIndexUrl)) { return $null }
$leaf = ($TorchIndexUrl.TrimEnd('/') -split '/')[-1].ToLowerInvariant()
# Drop query/fragment first so .../cu128?token=x classifies as cu128 (else it reinstalls every run).
$leaf = (($TorchIndexUrl -split '[?#]', 2)[0].TrimEnd('/') -split '/')[-1].ToLowerInvariant()
if ($leaf -match '^cu\d+$') { return $leaf }
if ($leaf -eq 'cpu') { return 'cpu' }
if ($leaf -match '^rocm') { return 'rocm' }
if ($leaf -match '^gfx') { return 'rocm' }
# gfx must be followed by a digit (an architecture leaf); gfx-private is custom.
if ($leaf -match '^gfx[0-9]') { return 'rocm' }
return $null
}
@ -2025,6 +2181,10 @@ exit 0
} catch { return $null }
}
# An explicit pin is authoritative: the AMD ROCm reroute below must not rewrite it
# (e.g. a deliberate cpu pin on an AMD host).
$TorchIndexPinned = (-not [string]::IsNullOrWhiteSpace($env:UNSLOTH_TORCH_INDEX_URL)) -or `
(-not [string]::IsNullOrWhiteSpace($env:UNSLOTH_TORCH_INDEX_FAMILY))
$TorchIndexUrl = Get-TorchIndexUrl
# ── GPU arch → newest compatible Windows ROCm wheel release ──
@ -2036,13 +2196,19 @@ exit 0
# Override with UNSLOTH_ROCM_WINDOWS_MIRROR for air-gapped / mirror installs.
$ROCmIndexUrl = $null
$ROCmTorchFloor = $null
if (($HasROCm -or $ROCmGfxArch) -and $TorchIndexUrl -like "*/cpu" -and -not $SkipTorch) {
$PinnedRocmVisionSpec = $null
$PinnedRocmAudioSpec = $null
if (-not $TorchIndexPinned -and ($HasROCm -or $ROCmGfxArch) -and $TorchIndexUrl -like "*/cpu" -and -not $SkipTorch) {
$amdIndexBase = if ($env:UNSLOTH_ROCM_WINDOWS_MIRROR) { $env:UNSLOTH_ROCM_WINDOWS_MIRROR.TrimEnd('/') } else { "https://repo.amd.com/rocm/whl" }
$archFamilyMap = @{
"gfx1201" = "gfx120X-all"; "gfx1200" = "gfx120X-all" # RDNA 4
"gfx1151" = "gfx1151"; "gfx1150" = "gfx1150" # RDNA 3.5 (Strix Halo/Point)
"gfx1103" = "gfx110X-all"; "gfx1102" = "gfx110X-all" # RDNA 3
"gfx1101" = "gfx110X-all"; "gfx1100" = "gfx110X-all"
"gfx1036" = "gfx103X-all"; "gfx1035" = "gfx103X-all" # RDNA 2 (RX 6000)
"gfx1034" = "gfx103X-all"; "gfx1033" = "gfx103X-all"
"gfx1032" = "gfx103X-all"; "gfx1031" = "gfx103X-all"
"gfx1030" = "gfx103X-all"
"gfx90a" = "gfx90a"; "gfx908" = "gfx908" # MI200/MI100
}
# gfx120X (RDNA 4) and gfx1151/gfx1150 (Strix) have a null-pointer bug in
@ -2086,6 +2252,32 @@ exit 0
}
}
# A gfx*/rocm pin skips the auto-reroute above, but the generic CPU/CUDA install below
# would use torch>=2.4,<2.11 and pull a known-bad wheel on the gfx115x/gfx120x/rocm>=7.2
# indexes (the _grouped_mm bug). Route a pinned ROCm index through the ROCm path.
if ($TorchIndexPinned -and -not $ROCmIndexUrl -and -not $SkipTorch) {
$_pinLeaf = (($TorchIndexUrl -split '[?#]', 2)[0].TrimEnd('/') -split '/')[-1].ToLower()
$_pinRocm211 = $false
# Anchor ($) so a suffixed custom leaf (rocm7.2-private) falls through to verbatim.
if ($_pinLeaf -match '^rocm(\d+)\.(\d+)$') {
# Only KNOWN-2.11 rocm (rocm7.2) gets the floor. Matches Test-RocmKnown211Version.
$_pinRocm211 = ([int]$Matches[1] -eq 7 -and [int]$Matches[2] -eq 2)
}
# Only the 2.11-allowlist gfx arches need the floor; others publish <2.11 and stay bare.
$_pinGfx211 = @('gfx120x-all', 'gfx1151', 'gfx1150') -contains $_pinLeaf
if ($_pinGfx211 -or $_pinRocm211) {
$ROCmIndexUrl = $TorchIndexUrl
$ROCmTorchFloor = "torch>=2.11.0,<2.12.0"
$PinnedRocmVisionSpec = "torchvision>=0.26.0,<0.27.0"
$PinnedRocmAudioSpec = "torchaudio>=2.11.0,<2.12.0"
substep "pinned ROCm index ($_pinLeaf) -- enforcing $ROCmTorchFloor" "Cyan"
} elseif ($_pinLeaf -match '^gfx[0-9]' -or $_pinLeaf -match '^rocm[0-9]+(\.[0-9]+)?$') {
# Other gfx / older rocm (<=7.1) ship torch <2.11; route via the ROCm path with
# bare specs. Only EXACT rocm<digits>/gfx* are families; a suffixed leaf is verbatim.
$ROCmIndexUrl = $TorchIndexUrl
}
}
if ($ROCmIndexUrl) {
$TorchIndexFamily = "rocm"
} else {
@ -2148,14 +2340,14 @@ exit 0
}
if ($_Migrated) {
# Migrated env: force-reinstall unsloth+unsloth-zoo to ensure clean state
# in the new venv location, while preserving existing torch/CUDA
# Migrated env: force-reinstall unsloth+unsloth-zoo for a clean state, preserving
# existing torch/CUDA unless the flavor repair below re-lands it.
Write-TauriLog "STEP" "Installing unsloth"
substep "upgrading unsloth in migrated environment..."
if ($SkipTorch) {
# No-torch: install unsloth + unsloth-zoo with --no-deps, then
# runtime deps (typer, safetensors, transformers, etc.) with --no-deps.
$baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated no-torch)" { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.6.9" "unsloth-zoo>=2026.6.7" }
$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.5" "unsloth-zoo>=2026.7.6" }
if ($baseInstallExit -eq 0) {
# Resolve pydantic WITH deps so pip pins pydantic-core
# to the matching version (no-torch-runtime.txt below
@ -2169,7 +2361,7 @@ exit 0
}
}
} else {
$baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated)" { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.6.9" "unsloth-zoo>=2026.6.7" }
$baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated)" { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.7.5" "unsloth-zoo>=2026.7.6" }
}
if ($baseInstallExit -ne 0) {
Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red
@ -2194,22 +2386,24 @@ exit 0
substep "skipping PyTorch (--no-torch flag set)." "Yellow"
} elseif ($ROCmIndexUrl) {
Write-TauriLog "STEP" "Installing PyTorch (AMD ROCm Windows)"
substep "installing PyTorch from $ROCmIndexUrl..."
substep "installing PyTorch from $(Remove-IndexUrlCredentials $ROCmIndexUrl)..."
$torchSpec = if ($ROCmTorchFloor) { $ROCmTorchFloor } else { "torch" }
# Pin the companions to match $torchSpec; bare names can resolve an
# 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 }
$visionSpec = if ($PinnedRocmVisionSpec) { $PinnedRocmVisionSpec } elseif ($ROCmGfxArch -and $torchvisionFloorMap -and $torchvisionFloorMap.ContainsKey($ROCmGfxArch)) { $torchvisionFloorMap[$ROCmGfxArch] } else { "torchvision" }
$audioSpec = if ($PinnedRocmAudioSpec) { $PinnedRocmAudioSpec } elseif ($ROCmGfxArch -and $torchaudioFloorMap -and $torchaudioFloorMap.ContainsKey($ROCmGfxArch)) { $torchaudioFloorMap[$ROCmGfxArch] } else { "torchaudio" }
$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.
substep "ROCm PyTorch install failed (exit $torchInstallExit); using a CPU base, Studio setup retries ROCm." "Yellow"
# Transient AMD-index failure: fall back to a CPU base (Unsloth setup retries
# ROCm). Use an explicit CPU index -- for a pinned ROCm index $TorchIndexUrl IS
# the ROCm mirror, so reusing it would just retry it.
$CpuFallbackIndexUrl = if ($env:UNSLOTH_PYTORCH_MIRROR) { "$($env:UNSLOTH_PYTORCH_MIRROR.TrimEnd('/'))/cpu" } else { "https://download.pytorch.org/whl/cpu" }
substep "ROCm PyTorch install failed (exit $torchInstallExit); using a CPU base, Unsloth setup retries ROCm." "Yellow"
# --force-reinstall: a failed ROCm install can leave an unpinned ROCm
# 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>=0.19,<0.26.0" "torchaudio>=2.4,<2.11.0" --default-index $CpuFallbackIndexUrl }
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)
@ -2222,8 +2416,14 @@ 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 }
substep "installing PyTorch ($(Remove-IndexUrlCredentials $TorchIndexUrl))..."
# Bound the companions to the capped torch on EVERY index, cu<digits>
# families included: torchaudio 2.11 dropped its exact torch pin from
# the wheel metadata, so a bare companion next to torch<2.11 can
# resolve a mismatched 2.11.0 build. Mirrors install.sh.
$_pinVisionSpec = "torchvision>=0.19,<0.26.0"
$_pinAudioSpec = "torchaudio>=2.4,<2.11.0"
$torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch" { uv pip install --python $VenvPython "torch>=2.4,<2.11.0" $_pinVisionSpec $_pinAudioSpec --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)
@ -2235,7 +2435,7 @@ exit 0
if ($SkipTorch) {
# No-torch: install unsloth + unsloth-zoo with --no-deps, then
# runtime deps (typer, safetensors, transformers, etc.) with --no-deps.
$baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (no-torch)" { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.6.9" "unsloth-zoo>=2026.6.7" }
$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.5" "unsloth-zoo>=2026.7.6" }
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 }
@ -2247,7 +2447,7 @@ exit 0
}
}
} elseif ($StudioLocalInstall) {
$baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (local)" { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.6.9" "unsloth-zoo>=2026.6.7" }
$baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (local)" { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.7.5" "unsloth-zoo>=2026.7.6" }
} else {
$baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth" { uv pip install --python $VenvPython --upgrade-package unsloth -- "$PackageName" }
}
@ -2275,7 +2475,7 @@ exit 0
Write-TauriLog "STEP" "Installing unsloth"
substep "installing unsloth (this may take a few minutes)..."
if ($StudioLocalInstall) {
$baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (auto torch backend)" { uv pip install --python $VenvPython "unsloth-zoo>=2026.6.7" "unsloth>=2026.6.9" --torch-backend=auto }
$baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (auto torch backend)" { uv pip install --python $VenvPython "unsloth-zoo>=2026.7.6" "unsloth>=2026.7.5" --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)
@ -2301,12 +2501,19 @@ exit 0
}
}
$installedPackageVersion = (& $VenvPython -c "from importlib.metadata import version; import sys; print(version(sys.argv[1]))" $PackageName 2>$null | Out-String).Trim()
if ($LASTEXITCODE -eq 0 -and $installedPackageVersion) {
step $PackageName "$installedPackageVersion installed"
} else {
substep "[WARN] installed $PackageName version could not be determined" "Yellow"
}
# ── Enforce the installed torch flavor matches the detected GPU build ──
# PEP 440 ignores the +cpu/+cuXXX/+rocm local label in a version range, so uv
# 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
@ -2319,10 +2526,10 @@ exit 0
$rocmSpec = if ($ROCmTorchFloor) { $ROCmTorchFloor } else { "torch" }
# Pin companions like the fresh ROCm path (bare names can pull an
# ABI-incompatible torchvision/torchaudio from the 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" }
$visionSpec = if ($PinnedRocmVisionSpec) { $PinnedRocmVisionSpec } elseif ($ROCmGfxArch -and $torchvisionFloorMap -and $torchvisionFloorMap.ContainsKey($ROCmGfxArch)) { $torchvisionFloorMap[$ROCmGfxArch] } else { "torchvision" }
$audioSpec = if ($PinnedRocmAudioSpec) { $PinnedRocmAudioSpec } elseif ($ROCmGfxArch -and $torchaudioFloorMap -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)
@ -2331,7 +2538,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>=0.19,<0.26.0" "torchaudio>=2.4,<2.11.0" --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)
@ -2406,7 +2613,7 @@ exit 0
Write-TauriLog "ERROR" "unsloth CLI was not installed correctly"
Write-Host "[ERROR] unsloth CLI was not installed correctly." -ForegroundColor Red
Write-Host " Expected: $UnslothExe" -ForegroundColor Yellow
Write-Host " This usually means an older unsloth version was installed that does not include the Studio CLI." -ForegroundColor Yellow
Write-Host " This usually means an older unsloth version was installed that does not include the Unsloth CLI." -ForegroundColor Yellow
Write-Host " Try re-running the installer or see: https://github.com/unslothai/unsloth?tab=readme-ov-file#-quickstart" -ForegroundColor Yellow
return (Exit-InstallFailure "unsloth CLI was not installed correctly")
}
@ -2517,7 +2724,7 @@ exit 0
Write-Host " Move or remove it manually, then re-run the installer." -ForegroundColor Yellow
throw "Cannot create unsloth launcher: $ShimExe is a directory."
}
# try/catch: if unsloth.exe is locked (Studio running), keep the old shim.
# try/catch: if unsloth.exe is locked (Unsloth running), keep the old shim.
$shimUpdated = $false
try {
if (Test-Path -LiteralPath $ShimExe) { Remove-Item -LiteralPath $ShimExe -Force -ErrorAction Stop }
@ -2535,7 +2742,7 @@ exit 0
if (Test-Path -LiteralPath $ShimExe) {
Write-Host "[WARN] Could not refresh unsloth launcher at $ShimExe." -ForegroundColor Yellow
Write-Host " This usually means a running 'unsloth studio' process still holds the file open." -ForegroundColor Yellow
Write-Host " Close Studio and re-run the installer to pick up the latest launcher." -ForegroundColor Yellow
Write-Host " Close Unsloth and re-run the installer to pick up the latest launcher." -ForegroundColor Yellow
Write-Host " Continuing with the existing launcher." -ForegroundColor Yellow
} else {
Write-Host "[WARN] Could not create unsloth launcher at $ShimExe" -ForegroundColor Yellow
@ -2556,6 +2763,13 @@ exit 0
}
Refresh-SessionPath # sync current session with registry
Complete-StudioVenvRollback
$studioVenvReplacementCommitted = $true
Remove-StaleStudioVenvRollbacks
} finally {
if (-not $studioVenvReplacementCommitted) {
Restore-StudioVenvRollback
}
}
# Env-mode session export AFTER Refresh-SessionPath; otherwise a legacy
# User PATH entry (Machine > User > current $env:Path) would win.
@ -2600,9 +2814,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 Unsloth 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)
if ($IsInteractive) {
Write-Host ""
$reply = Read-Host " Start Unsloth Studio now? [Y/n]"
@ -2611,8 +2826,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 {
@ -2632,8 +2847,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 ""
}
}

1262
install.sh
View file

File diff suppressed because it is too large Load diff

9
pyproject.toml
View file

@ -25,7 +25,7 @@ classifiers = [
"Topic :: Scientific/Engineering :: Artificial Intelligence",
]
dependencies = [
"typer",
"typer>=0.12.0",
"rich",
"pydantic",
"pyyaml",
@ -42,6 +42,7 @@ version = {attr = "unsloth.models._utils.__version__"}
include-package-data = true
[tool.setuptools.package-data]
unsloth_cli = ["codex_fallback_prompt.md", "pi_subagent.ts"]
studio = [
"*.sh",
"*.ps1",
@ -73,7 +74,7 @@ triton = [
]
huggingfacenotorch = [
"unsloth_zoo>=2026.6.7",
"unsloth_zoo>=2026.7.6",
"wheel>=0.42.0",
"packaging",
"numpy",
@ -94,7 +95,7 @@ huggingfacenotorch = [
]
huggingface = [
"unsloth[huggingfacenotorch]",
"unsloth_zoo>=2026.6.7",
"unsloth_zoo>=2026.7.6",
"torchvision",
"unsloth[triton]",
]
@ -579,7 +580,7 @@ colab-ampere-torch220 = [
"flash-attn>=2.6.3 ; ('linux' in sys_platform)",
]
colab-new = [
"unsloth_zoo>=2026.6.7",
"unsloth_zoo>=2026.7.6",
"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",

71
scripts/build_whisper_cpp.sh Executable file
View file

@ -0,0 +1,71 @@
#!/bin/sh
# Build whisper.cpp's whisper-server for Studio's GGUF dictation engine.
#
# Installs into the managed Studio home so the backend's binary discovery
# (core/inference/stt_ggml_sidecar.py::find_whisper_server_binary) picks it up:
# <UNSLOTH_STUDIO_HOME>/whisper.cpp/build/bin/whisper-server (custom home)
# ~/.unsloth/whisper.cpp/build/bin/whisper-server (default)
#
# Usage:
# ./scripts/build_whisper_cpp.sh # build the pinned tag
# WHISPER_CPP_TAG=v1.9.0 ./scripts/build_whisper_cpp.sh
#
# Requires: git, cmake, a C/C++ toolchain (the same prerequisites as a
# llama.cpp source build). GPU backends are auto-detected by whisper.cpp's
# CMake (Metal on macOS; set GGML_CUDA=1 to force a CUDA build on Linux).
set -eu
WHISPER_CPP_SOURCE="${WHISPER_CPP_SOURCE:-https://github.com/ggml-org/whisper.cpp}"
WHISPER_CPP_TAG="${WHISPER_CPP_TAG:-v1.9.1}"
STUDIO_HOME="${UNSLOTH_STUDIO_HOME:-${STUDIO_HOME:-}}"
CUSTOM_STUDIO_HOME=false
if [ -n "$STUDIO_HOME" ]; then
CUSTOM_STUDIO_HOME=true
INSTALL_DIR="$STUDIO_HOME/whisper.cpp"
else
INSTALL_DIR="$HOME/.unsloth/whisper.cpp"
fi
command -v git >/dev/null 2>&1 || { echo "ERROR: git is required" >&2; exit 1; }
command -v cmake >/dev/null 2>&1 || { echo "ERROR: cmake is required" >&2; exit 1; }
# Same policy as studio/setup.sh's _assert_studio_owned_or_absent: never delete
# a directory under a custom Studio home unless Studio itself created it (the
# marker file below). Protects a user-managed whisper.cpp/src from rm -rf.
STUDIO_OWNED_MARKER=".unsloth-studio-owned"
if [ "$CUSTOM_STUDIO_HOME" = true ] && [ -e "$INSTALL_DIR" ] && \
[ ! -f "$INSTALL_DIR/$STUDIO_OWNED_MARKER" ]; then
echo "ERROR: $INSTALL_DIR already exists and is not marked as an Unsloth-owned whisper.cpp build tree." >&2
echo " Move it aside or choose an empty UNSLOTH_STUDIO_HOME before re-running." >&2
exit 1
fi
echo "==> Building whisper.cpp ($WHISPER_CPP_TAG) into $INSTALL_DIR"
mkdir -p "$INSTALL_DIR"
: > "$INSTALL_DIR/$STUDIO_OWNED_MARKER"
if [ ! -d "$INSTALL_DIR/src/.git" ]; then
rm -rf "$INSTALL_DIR/src"
git clone --depth 1 --branch "$WHISPER_CPP_TAG" "$WHISPER_CPP_SOURCE" "$INSTALL_DIR/src"
else
git -C "$INSTALL_DIR/src" fetch --depth 1 origin "$WHISPER_CPP_TAG"
git -C "$INSTALL_DIR/src" checkout FETCH_HEAD
fi
CMAKE_FLAGS="-DCMAKE_BUILD_TYPE=Release -DBUILD_SHARED_LIBS=OFF"
if [ "${GGML_CUDA:-0}" = "1" ]; then
CMAKE_FLAGS="$CMAKE_FLAGS -DGGML_CUDA=ON"
fi
# shellcheck disable=SC2086
cmake -S "$INSTALL_DIR/src" -B "$INSTALL_DIR/src/build" $CMAKE_FLAGS
NCPU="$(getconf _NPROCESSORS_ONLN 2>/dev/null || echo 4)"
cmake --build "$INSTALL_DIR/src/build" --config Release --target whisper-server -j"$NCPU"
mkdir -p "$INSTALL_DIR/build/bin"
cp "$INSTALL_DIR/src/build/bin/whisper-server" "$INSTALL_DIR/build/bin/whisper-server"
echo "==> Installed $INSTALL_DIR/build/bin/whisper-server"
"$INSTALL_DIR/build/bin/whisper-server" --help >/dev/null 2>&1 && echo "==> Binary runs OK"

59
scripts/install_rocm_wsl_strixhalo.sh
View file

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

4
scripts/lockfile_supply_chain_audit.py
View file

@ -2,7 +2,7 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
"""Lockfile supply-chain audit for the Studio frontend and Tauri shell.
"""Lockfile supply-chain audit for the Unsloth frontend and Tauri shell.
Runs BEFORE `npm ci` / `cargo fetch` in CI. Refuses to proceed when a
lockfile contains patterns indicating supply-chain injection (npm
@ -294,7 +294,7 @@ CARGO_REGISTRY_SOURCE = "registry+https://github.com/rust-lang/crates.io-index"
# Cargo non-registry source allowlist: `(crate_name, exact_source_string)`.
# Both must match verbatim; bumping the pinned SHA forces a re-review.
# Studio's Tauri shell pulls `fix-path-env` from git because it is not
# Unsloth's Tauri shell pulls `fix-path-env` from git because it is not
# published to crates.io; commit c4c45d5 was reviewed when it landed.
CARGO_SOURCE_ALLOWLIST: tuple[tuple[str, str], ...] = (
(

2
scripts/scan_npm_packages.py
View file

@ -62,7 +62,7 @@ REPO_ROOT = Path(__file__).resolve().parents[1]
# Hard caps (deliberately conservative; npm tarballs in this repo are
# all well under these limits, so a packaging spike is noticeable).
# ─────────────────────────────────────────────────────────────────────
# Caps calibrated against the real Studio frontend transitive closure:
# Caps calibrated against the real Unsloth frontend transitive closure:
# - typescript.js is 9.1 MB (TS compiler bundled into one file)
# - mermaid 11.x dist/mermaid.js.map is ~12 MB (sourcemap)
# - lightningcss-linux-x64-{gnu,musl}.node is 10 MB

296
scripts/scan_packages_baseline.json
View file

File diff suppressed because one or more lines are too long

20
scripts/stamp_studio_release.py
View file

@ -2,7 +2,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
"""Stamp and verify display-only Studio release metadata for builds."""
"""Stamp and verify display-only Unsloth release metadata for builds."""
from __future__ import annotations
@ -50,7 +50,7 @@ MAX_VERSION_LENGTH = 64
PLACEHOLDER = """# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
\"\"\"Build-stamped Studio release metadata.
\"\"\"Build-stamped Unsloth release metadata.
Release builds may rewrite this module in the build workspace before creating
Python artifacts. Keep the committed value neutral so source checkouts do not
@ -145,7 +145,7 @@ def build_info_source(version: str | None) -> str:
return f'''# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Build-stamped Studio release metadata."""
"""Build-stamped Unsloth release metadata."""
STUDIO_RELEASE_VERSION = {literal}
'''
@ -168,7 +168,7 @@ def stamp(require_release: bool) -> int:
version, source = resolve_version()
if version is not None and not is_valid_version(version):
print(
f"Invalid Studio release version from {source}: {version!r}",
f"Invalid Unsloth release version from {source}: {version!r}",
file = sys.stderr,
)
return 2
@ -196,9 +196,9 @@ def stamp(require_release: bool) -> int:
if version is None:
if require_release:
print(
"No Studio release version available. Set "
"No Unsloth release version available. Set "
"UNSLOTH_STUDIO_RELEASE_VERSION, build from a GitHub tag, "
"or run from an exact local Studio release tag.",
"or run from an exact local Unsloth release tag.",
file = sys.stderr,
)
return 2
@ -207,7 +207,7 @@ def stamp(require_release: bool) -> int:
return 0
_atomic_write_text(BUILD_INFO_PATH, build_info_source(version), encoding = "utf-8")
print(f"Stamping Studio release version {version} from {source}", file = sys.stderr)
print(f"Stamping Unsloth release version {version} from {source}", file = sys.stderr)
print(version)
return 0
@ -233,7 +233,7 @@ def _read_sdist_member(path: Path) -> str | None:
def verify_dist(expected: str, dist_dir: Path) -> int:
if not is_valid_version(expected):
print(f"Invalid expected Studio release version: {expected!r}", file = sys.stderr)
print(f"Invalid expected Unsloth release version: {expected!r}", file = sys.stderr)
return 2
artifacts = list(dist_dir.glob("*.whl")) + list(dist_dir.glob("*.tar.gz"))
@ -251,14 +251,14 @@ def verify_dist(expected: str, dist_dir: Path) -> int:
if content is None:
failures.append(f"{artifact.name}: missing {BUILD_INFO_SUFFIX}")
elif expected_line not in content:
failures.append(f"{artifact.name}: Studio release version mismatch")
failures.append(f"{artifact.name}: Unsloth release version mismatch")
if failures:
for failure in failures:
print(failure, file = sys.stderr)
return 2
print(f"Verified Studio release version {expected} in {len(artifacts)} artifact(s)")
print(f"Verified Unsloth release version {expected} in {len(artifacts)} artifact(s)")
return 0

14
scripts/uninstall.ps1
View file

@ -83,7 +83,7 @@ function Uninstall-UnslothStudio {
}
}
# A path is a Studio-owned root iff one of install.ps1's sentinels exists:
# A path is an Unsloth-owned root iff one of install.ps1's sentinels exists:
# <root>\share\studio.conf, <root>\unsloth_studio\.unsloth-studio-owned,
# or <root>\bin\unsloth.exe.
function _IsStudioRoot {
@ -164,7 +164,7 @@ function Uninstall-UnslothStudio {
return $p
}
# Discover non-default Studio roots from env vars + studio.conf files.
# Discover non-default Unsloth roots from env vars + studio.conf files.
# Mirrors install.ps1's precedence: UNSLOTH_STUDIO_HOME wins, STUDIO_HOME
# is ignored when both are set, so uninstalling install A doesn't also
# delete install B if the user has a stale STUDIO_HOME pointing at B.
@ -207,7 +207,7 @@ function Uninstall-UnslothStudio {
# Return $true iff the PID's image path lives under one of $KnownRoots.
# Prevents killing an unrelated process that happens to listen on a stale
# Studio port.
# Unsloth port.
function _PidUnderKnownRoot {
param([int]$Pid_, [string[]]$KnownRoots)
if (-not $KnownRoots -or $KnownRoots.Count -eq 0) { return $false }
@ -223,8 +223,8 @@ function Uninstall-UnslothStudio {
return $false
}
# Stop a Studio backend whose port is recorded in <DataDir>\studio.port.
# Only kills if the listening PID's exe path is under a known Studio root.
# Stop an Unsloth backend whose port is recorded in <DataDir>\studio.port.
# Only kills if the listening PID's exe path is under a known Unsloth root.
function _StopByPortFile {
param([string]$PortFile, [string[]]$KnownRoots)
if (-not (Test-Path -LiteralPath $PortFile -PathType Leaf)) { return }
@ -372,7 +372,7 @@ function Uninstall-UnslothStudio {
continue
}
if (-not (_IsStudioRoot $r)) {
_Substep "refusing to remove non-Studio path: $r" "Yellow"
_Substep "refusing to remove non-Unsloth path: $r" "Yellow"
continue
}
_RemovePath $r
@ -436,7 +436,7 @@ function Uninstall-UnslothStudio {
$entries = $rawPath -split ';'
$kept = New-Object System.Collections.ArrayList
$removedAny = $false
# Only remove PATH entries that live inside a Studio root we
# Only remove PATH entries that live inside an Unsloth root we
# actually own (default or env-mode). A literal substring
# match on `unsloth_studio` would clobber unrelated user
# virtualenvs that happen to share the name.

14
scripts/uninstall.sh
View file

@ -12,7 +12,7 @@
set -e
# Stop a Studio server via its PID file (written by install.sh's _spawn_terminal).
# Stop an Unsloth server via its PID file (written by install.sh's _spawn_terminal).
_kill_pid_file() {
_pid_file="$1"
[ -f "$_pid_file" ] || return 0
@ -47,7 +47,7 @@ _pkill_studio() {
command -v pkill >/dev/null 2>&1 || return 0
# Scope fallback patterns to the install roots we are removing so a
# different Studio install (different UNSLOTH_STUDIO_HOME) is not touched.
# different Unsloth install (different UNSLOTH_STUDIO_HOME) is not touched.
_kill_roots="$HOME/.unsloth/studio"
_roots_from_conf=$(_custom_studio_roots 2>/dev/null || true)
[ -n "$_roots_from_conf" ] && _kill_roots="$_kill_roots
@ -89,7 +89,7 @@ _remove_path() {
fi
}
# Accept as Studio root only if Studio sentinels exist (matches install.sh's
# Accept as Unsloth root only if Unsloth sentinels exist (matches install.sh's
# env-mode ownership guard at install.sh:1358-1361). A bare unsloth_studio/
# directory is NOT enough -- require the install-time owner marker so a user
# directory that happens to contain a folder named "unsloth_studio" is safe.
@ -175,8 +175,8 @@ _custom_studio_roots() {
_from_conf "$HOME/.local/share/unsloth/studio.conf"
}
# Remove $HOME/.local/bin/unsloth only if it's a Studio-managed symlink.
# Studio's install.sh writes this as a symlink into the studio venv
# Remove $HOME/.local/bin/unsloth only if it's an Unsloth-managed symlink.
# Unsloth's install.sh writes this as a symlink into the studio venv
# (install.sh: `ln -sfn "$VENV_DIR/bin/unsloth" "$_shim_path"`). A
# pip-installed `unsloth` CLI is a regular file — leave it alone to avoid
# wiping an unrelated install.
@ -206,7 +206,7 @@ _custom_studio_roots | while IFS= read -r _custom_root; do
continue
fi
if ! _is_studio_root "$_custom_root"; then
echo " refusing to remove non-Studio path: $_custom_root" >&2
echo " refusing to remove non-Unsloth path: $_custom_root" >&2
continue
fi
_remove_path "$_custom_root"
@ -234,7 +234,7 @@ _remove_path "$HOME/.unsloth/rocm-smoketest"
# Drop ~/.unsloth only if now empty (rmdir refuses non-empty, so user content is kept).
rmdir "$HOME/.unsloth" 2>/dev/null || true
_remove_path "$HOME/.local/share/unsloth"
# CLI shim: only the symlink Studio created, never a pip-installed file.
# CLI shim: only the symlink Unsloth created, never a pip-installed file.
_remove_cli_shim
echo "Removing desktop shortcut and launcher lock..."

20
scripts/verify_import_hoist.py
View file

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

34
studio/MCP.md Normal file
View file

@ -0,0 +1,34 @@
# Unsloth Studio MCP server
Unsloth 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 Unsloth 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 Unsloth uses its default port
(a request to `/mcp` redirects to the canonical `/mcp/`). Use the actual Unsloth
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 Unsloth `TrainingStartRequest`.
The request is validated by the existing Pydantic model before a subprocess is
started. Export paths use the existing Unsloth 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.

2
studio/Unsloth_Studio_Colab.ipynb
View file

@ -33,7 +33,7 @@
"\n",
"We are actively working on making Unsloth Studio install on Colab T4 GPUs faster.\n",
"\n",
"[Features](https://unsloth.ai/docs/new/unsloth-studio#features) • [Quickstart](https://unsloth.ai/docs/new/unsloth-studio/start) • [Data Recipes](https://unsloth.ai/docs/new/unsloth-studio/data-recipe) • [Studio Chat](https://unsloth.ai/docs/new/unsloth-studio/chat) • [Export](https://unsloth.ai/docs/new/unsloth-studio/export)"
"[Features](https://unsloth.ai/docs/new/unsloth-studio#features) • [Quickstart](https://unsloth.ai/docs/new/unsloth-studio/start) • [Data Recipes](https://unsloth.ai/docs/new/unsloth-studio/data-recipe) • [Unsloth Chat](https://unsloth.ai/docs/new/unsloth-studio/chat) • [Export](https://unsloth.ai/docs/new/unsloth-studio/export)"
]
},
{

2
studio/backend/assets/chat_templates/gemma-4-edge.jinja
View file

@ -3,7 +3,7 @@
Source: google/gemma-4-31B-it HF discussion/PR #118 (adds the preserve_thinking
flag plus null-rendering, string-arguments validation, balanced turn tags, empty
messages handling, and OpenAI image_url/input_audio aliases).
Studio-local changes vs PR #118:
Unsloth-local changes vs PR #118:
1. preserve_thinking defaults to false (see SETUP block below).
2. The empty "<|channel>thought\n<channel|>" block on enable_thinking=false is
NOT emitted. Google ships a distinct template for E2B/E4B (google/gemma-4-E2B-it,

2
studio/backend/assets/chat_templates/gemma-4.jinja
View file

@ -3,7 +3,7 @@
Source: google/gemma-4-31B-it HF discussion/PR #118 (adds the preserve_thinking
flag plus null-rendering, string-arguments validation, balanced turn tags, empty
messages handling, and OpenAI image_url/input_audio aliases).
Studio-local change: preserve_thinking defaults to false (see SETUP block below).
Unsloth-local change: preserve_thinking defaults to false (see SETUP block below).
Applied to unsloth/gemma-4-*-GGUF models so the embedded GGUF template does not
need re-downloading. Keep in sync with upstream if PR #118 changes.
-#}

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

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

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

@ -148,7 +148,7 @@ async def authenticated_via_api_key(
) -> bool:
"""True when the caller used an sk-unsloth API key, not a UI session JWT.
Lets routes treat programmatic API callers differently from the Studio UI
Lets routes treat programmatic API callers differently from the Unsloth UI
(e.g. refuse a teardown the UI would allow).
"""
return bool(credentials and credentials.credentials.startswith(API_KEY_PREFIX))

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

@ -1,13 +1,13 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Auto-shutdown for an exposed first-run Studio whose admin password is unchanged.
"""Auto-shutdown for an exposed first-run Unsloth whose admin password is unchanged.
On a fresh install the seeded bootstrap admin password stays a valid login
credential until first login changes it. When the web UI is put on the network
(``--secure`` / ``0.0.0.0``) and nobody completes that first-login change within
a deadline, tear Studio down so a fresh, unconfigured instance does not stay
publicly reachable indefinitely. If the password was changed, Studio keeps
a deadline, tear Unsloth down so a fresh, unconfigured instance does not stay
publicly reachable indefinitely. If the password was changed, Unsloth keeps
running.
Scope: web UI launches only (never ``--api-only``, which authenticates by API
@ -98,7 +98,7 @@ def enforce_bootstrap_password_deadline(
) -> bool:
"""Deadline handler: shut down iff the seeded admin password is still unchanged.
Returns True if it shut Studio down, False if it left it running (the
Returns True if it shut Unsloth down, False if it left it running (the
password was changed in time).
"""
try:
@ -106,7 +106,7 @@ def enforce_bootstrap_password_deadline(
except Exception:
return False
if not still_default:
return False # password changed in time -> leave Studio running
return False # password changed in time -> leave Unsloth running
message = (
"\nUnsloth Studio was exposed on the network but its default admin "

61
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:
@ -111,7 +146,7 @@ def get_connection() -> sqlite3.Connection:
pass
conn.row_factory = sqlite3.Row
# WAL lets token reads run concurrently with refresh-token writes;
# busy_timeout bounds lock waits. Matches the other Studio SQLite stores.
# busy_timeout bounds lock waits. Matches the other Unsloth SQLite stores.
# Set busy_timeout first: switching journal_mode needs a lock, so if a
# refresh-token write already holds one, journal_mode=WAL raises SQLITE_BUSY;
# with busy_timeout already in effect it waits instead of failing and leaving
@ -270,8 +305,8 @@ def get_or_create_identity_secret() -> bytes:
def compute_identity_proof(nonce: bytes, host: str, port: int) -> str:
"""HMAC-SHA256 proof that the caller holds this install's identity secret,
bound to the loopback address and port the connection landed on. A proof
relayed from a Studio on a different address/port (a squatter proxying to the
real one, e.g. localhost resolving to ::1 while Studio is on 127.0.0.1) was
relayed from an Unsloth on a different address/port (a squatter proxying to the
real one, e.g. localhost resolving to ::1 while Unsloth is on 127.0.0.1) was
computed for that other endpoint and won't match the one the client dialed."""
try:
host = ipaddress.ip_address(host).compressed # normalise 127.0.0.1 / ::1 forms
@ -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()

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

@ -0,0 +1,286 @@
# 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
Unsloth 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 Unsloth 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 any(ch.isspace() for ch in new_password):
out.write("Password cannot contain spaces; 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 Unsloth.\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

96
studio/backend/cloudflare_tunnel.py
View file

@ -1,13 +1,13 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Free Cloudflare quick tunnel for Studio's 0.0.0.0 launches.
"""Free Cloudflare quick tunnel for Unsloth's 0.0.0.0 launches.
The raw http://<ip>:<port> is often unreachable (https-vs-http, blocked ports,
closed security groups); a cloudflared quick tunnel gives a free
https://*.trycloudflare.com URL that works anywhere, with no account or domain.
Best-effort throughout: any failure collapses to "no URL" and Studio keeps
Best-effort throughout: any failure collapses to "no URL" and Unsloth keeps
running. Stdlib only (back-end imports are lazy) so it is safe to import early.
"""
@ -20,6 +20,7 @@ import shutil
import subprocess
import sys
import threading
import time
from pathlib import Path
from typing import Optional, Tuple
@ -40,6 +41,22 @@ _RELEASE_BASE = "https://github.com/cloudflare/cloudflared/releases/latest/downl
_READY_TIMEOUT = 15.0 # seconds to wait for the URL + a registered edge connection
_DOWNLOAD_TIMEOUT = 60 # urlopen timeout for the one-time binary download
# A registered edge connection does not mean the hostname resolves yet, so the
# URL is fetched once before it is advertised.
_PUBLIC_PROBE_PATH = "/api/health"
_PUBLIC_PROBE_MARKER = "Unsloth UI Backend"
# One deadline for DNS propagation + the health probe, bounding the startup stall.
_PUBLIC_PROBE_TIMEOUT = 45.0
_PUBLIC_PROBE_ATTEMPT_TIMEOUT = 5.0
_PUBLIC_PROBE_RETRY_DELAY = 1.0
# Wait for the hostname via DoH first: an early OS lookup negative-caches the
# NXDOMAIN for up to 30 min.
_DNS_POLL_DELAY = 2.0
# Retry transient DoH failures, but give up fast when DoH is blocked outright.
_DNS_MAX_DOH_ERRORS = 3
_DOH_URL = "https://cloudflare-dns.com/dns-query?name={host}&type=A"
def _windows_hidden_kwargs() -> dict:
"""Suppress a child console window on Windows; no-op elsewhere."""
@ -95,7 +112,7 @@ def _cache_path() -> Optional[Path]:
def find_cloudflared() -> Optional[str]:
"""Locate an existing cloudflared: PATH first, then the Studio bin cache."""
"""Locate an existing cloudflared: PATH first, then the Unsloth bin cache."""
on_path = shutil.which("cloudflared")
if on_path:
return on_path
@ -191,6 +208,59 @@ def ensure_cloudflared() -> Optional[str]:
return None
def _wait_for_dns(host: str, deadline: float) -> None:
import json
import urllib.request
errors = 0
while True:
answered = False
try:
req = urllib.request.Request(
_DOH_URL.format(host = host),
headers = {"Accept": "application/dns-json", "User-Agent": "unsloth-studio"},
)
with urllib.request.urlopen(req, timeout = 5) as response:
answered = bool(json.loads(response.read(65536)).get("Answer"))
errors = 0
except Exception:
errors += 1
if errors >= _DNS_MAX_DOH_ERRORS:
return
if answered:
return
remaining = deadline - time.monotonic()
if remaining <= 0:
return
time.sleep(min(_DNS_POLL_DELAY, remaining))
def verify_public_url(url: str, timeout: float = _PUBLIC_PROBE_TIMEOUT) -> bool:
import json
import urllib.request
from urllib.parse import urlsplit
deadline = time.monotonic() + timeout
host = urlsplit(url).hostname
if host:
_wait_for_dns(host, deadline)
probe_url = f"{url.rstrip('/')}{_PUBLIC_PROBE_PATH}"
while True:
try:
req = urllib.request.Request(probe_url, headers = {"User-Agent": "unsloth-studio"})
with urllib.request.urlopen(req, timeout = _PUBLIC_PROBE_ATTEMPT_TIMEOUT) as response:
body = response.read(4096)
if json.loads(body).get("service") == _PUBLIC_PROBE_MARKER:
return True
except Exception:
pass
remaining = deadline - time.monotonic()
if remaining <= 0:
return False
time.sleep(min(_PUBLIC_PROBE_RETRY_DELAY, remaining))
class CloudflareTunnel:
"""A cloudflared quick tunnel to http://localhost:<port>. Best-effort throughout.
@ -309,7 +379,7 @@ class CloudflareTunnel:
pass
# Single serving process per Studio launch, so one module-level tunnel handle is
# Single serving process per Unsloth launch, so one module-level tunnel handle is
# enough; the lock guards the start/stop/shutdown races.
_active_tunnel: Optional[CloudflareTunnel] = None
_active_lock = threading.Lock()
@ -322,11 +392,12 @@ def start_studio_tunnel(port: int, timeout: float = _READY_TIMEOUT) -> Optional[
"""Start a quick tunnel and return its public URL once it is actually
serving, or None (best-effort).
Waits for cloudflared to both mint the URL and register an edge connection
before returning, so the caller never advertises a URL that yields Cloudflare
error 1033 (HTTP 530). If a URL is minted but no connection registers within
the window (e.g. quic is blocked on this network), retries once forcing the
http2 protocol. On any failure the tunnel is stopped and None is returned.
Waits for cloudflared to both mint the URL and register an edge connection,
then fetches /api/health over the public URL, so the caller never advertises
a link that yields Cloudflare error 1033 (HTTP 530) or an unresolvable host.
If a URL is minted but no connection registers within the window (e.g. quic
is blocked on this network), retries once forcing the http2 protocol. On any
failure the tunnel is stopped and None is returned.
"""
global _active_tunnel, _shutdown_requested
binary = ensure_cloudflared()
@ -349,9 +420,13 @@ def start_studio_tunnel(port: int, timeout: float = _READY_TIMEOUT) -> Optional[
prior, _active_tunnel = _active_tunnel, tunnel
if prior is not None:
prior.stop()
registered = False
try:
tunnel.start()
url = tunnel.wait_for_ready(timeout)
registered = url is not None
if url and not verify_public_url(url):
url = None
except Exception:
url = None
if url:
@ -371,6 +446,9 @@ def start_studio_tunnel(port: int, timeout: float = _READY_TIMEOUT) -> Optional[
# http2 will not help, so do not burn another window on it.
if not saw_url:
return None
# probe failure after registering is DNS propagation; http2 would not help
if registered:
return None
return None

16
studio/backend/colab.py
View file

@ -129,7 +129,7 @@ def start_cloudflare_tunnel(port: int) -> "str | None":
logger.warning(
"Cloudflare link not started: the admin account still has its temporary "
"bootstrap password, which is exposed to anyone who can load the page. "
"Open Studio in this tab, log in and change the admin password, then re-run "
"Open Unsloth in this tab, log in and change the admin password, then re-run "
"start(cloudflare=True) to get the shareable link."
)
return None
@ -203,7 +203,7 @@ def _shareable_link_html(cloudflare_url: str) -> str:
display: flex; align-items: center; gap: 12px;">
<img src="https://github.com/unslothai/unsloth/raw/main/studio/frontend/public/unsloth-gem.png"
height="48" style="display:block;">
Shareable Studio Link is Ready!
Shareable Unsloth Link is Ready!
</h2>
<a href="{cloudflare_url}" onclick="var w=window.open(this.href,'_blank');if(!w){{return true;}}return false;"
style="display: inline-flex; align-items: center; gap: 10px; padding: 14px 28px;
@ -223,7 +223,7 @@ def _shareable_link_html(cloudflare_url: str) -> str:
def _show_and_embed(port: int, *, cloudflare_url: "str | None" = None):
"""Render the Studio header + iframe for *port*, with a shareable-link card above
"""Render the Unsloth header + iframe for *port*, with a shareable-link card above
when *cloudflare_url* is set. Falls back to serve_kernel_port_as_iframe."""
url = get_colab_url(port)
logger.info(f"🌐 Unsloth Studio URL: {url}")
@ -281,7 +281,7 @@ def start(port: int = 8888, *, cloudflare: bool = False):
Args:
port: Port to bind/serve on.
cloudflare: Opt in to a shareable Cloudflare HTTPS link reachable from any
device (default OFF). It exposes Studio's login page beyond Colab, so it
device (default OFF). It exposes Unsloth's login page beyond Colab, so it
stays an explicit opt-in; the default shows only the in-tab proxy iframe.
Usage:
@ -292,10 +292,10 @@ def start(port: int = 8888, *, cloudflare: bool = False):
logger.info("🦥 Starting Unsloth Studio...")
# Fast path: Studio already running (cell re-run). Re-launching would collide on
# Fast path: Unsloth already running (cell re-run). Re-launching would collide on
# the port, so just re-show the link and iframe.
if _is_studio_healthy(port):
logger.info(f" Studio is already running on port {port} — reusing existing server.")
logger.info(f" Unsloth is already running on port {port} — reusing existing server.")
# try/finally: tear the tunnel down even if interrupted mid-start/render.
try:
cf_url = start_cloudflare_tunnel(port) if cloudflare else None
@ -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,

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

@ -27,7 +27,6 @@ from .constants import (
)
from .parse import apply_update, coerce_event, parse_log_message
from .types import Job
from .worker import run_job_process
from loggers import get_logger
logger = get_logger(__name__)
@ -169,12 +168,18 @@ class JobManager:
native_path_secret_removed_for_child_start,
run_without_native_path_secret,
)
from utils.hf_cache_settings import child_environment_for_spawn, get_hf_cache_paths
with native_path_secret_removed_for_child_start():
cache_env = get_hf_cache_paths().child_env({})
with (
child_environment_for_spawn(cache_env),
native_path_secret_removed_for_child_start(),
):
mp_q = _CTX.Queue()
proc = _CTX.Process(
target = run_without_native_path_secret,
args = (run_job_process,),
args = ("core.data_recipe.jobs.worker", "run_job_process", cache_env),
kwargs = {"event_queue": mp_q, "recipe": recipe, "run": run_payload},
daemon = True,
)

6
studio/backend/core/data_recipe/jobs/parse.py
View file

@ -133,7 +133,7 @@ def parse_log_message(msg: str) -> ParsedUpdate | None:
source = "github",
status = "rate_limited",
retry_after_sec = seconds,
message = ("Waiting for GitHub rate limit. Studio will resume automatically."),
message = ("Waiting for GitHub rate limit. Unsloth will resume automatically."),
),
)
@ -147,7 +147,7 @@ def parse_log_message(msg: str) -> ParsedUpdate | None:
status = "rate_limited",
retry_after_sec = seconds,
message = (
"Waiting for GitHub secondary rate limit. Studio will resume automatically."
"Waiting for GitHub secondary rate limit. Unsloth will resume automatically."
),
),
)
@ -161,7 +161,7 @@ def parse_log_message(msg: str) -> ParsedUpdate | None:
source = "github",
status = "rate_limited",
retry_after_sec = seconds,
message = ("Waiting for GitHub rate limit. Studio will resume automatically."),
message = ("Waiting for GitHub rate limit. Unsloth will resume automatically."),
),
)

2
studio/backend/core/data_recipe/local_callable_validators.py
View file

@ -238,7 +238,7 @@ def _run_oxc_batch(
if not node_executable:
return _fallback_results(
len(code_values),
"Node.js not found (install Node >= 20.19, or re-run Studio setup to provision it).",
"Node.js not found (install Node >= 20.19, or re-run Unsloth setup to provision it).",
)
try:
tmp_dir = ensure_dir(oxc_validator_tmp_root())

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 Unsloth can run with
# cwd=/, so keep default callers on Unsloth'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)

105
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,20 +209,41 @@ 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,
)
from utils.hf_cache_settings import child_environment_for_spawn, get_hf_cache_paths
from .worker import run_export_process
cache_env = get_hf_cache_paths().child_env({})
with native_path_secret_removed_for_child_start():
with (
child_environment_for_spawn(cache_env),
native_path_secret_removed_for_child_start(),
):
self._cmd_queue = _CTX.Queue()
self._resp_queue = _CTX.Queue()
self._proc = _CTX.Process(
target = run_without_native_path_secret,
args = (run_export_process,),
args = ("core.export.worker", "run_export_process", cache_env),
kwargs = {
"cmd_queue": self._cmd_queue,
"resp_queue": self._resp_queue,
@ -231,11 +257,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 +297,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 +381,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 +452,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 +633,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":

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

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

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

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

@ -258,6 +258,10 @@ class AnthropicStreamEmitter:
self._open_tool_use_id: Optional[str] = None
self._open_tool_args_sent: bool = False
self._prev_text: str = ""
# Net <think> minus </think> in the text emitted to the client. Tracked
# from emitted deltas (not _prev_text, which a final bare shrink clobbers)
# so an unclosed reasoning-only block can be balanced before close.
self._open_think_tags: int = 0
self._usage: dict = {}
def start(
@ -317,6 +321,7 @@ class AnthropicStreamEmitter:
"""Close any open block and emit message_delta + message_stop."""
events = []
if self._text_block_open or self._open_tool_call_id is not None:
events.extend(self._close_open_think())
events.append(self._close_block())
self._open_tool_call_id = None
self._open_tool_use_id = None
@ -344,12 +349,33 @@ class AnthropicStreamEmitter:
)
return events
def _close_open_think(self) -> list[str]:
"""Emit a ``</think>`` delta when the streamed text left a ``<think>``
open. This emitter diffs cumulative snapshots and drops the generator's
final bare shrink, so a reasoning-only reply would otherwise end on an
unclosed tag. Mirrors the chat route's reasoning extractor, which closes
the block on finish; balances the block before it is closed."""
if not self._text_block_open or self._open_think_tags <= 0:
return []
self._open_think_tags = 0
return [
build_anthropic_sse_event(
"content_block_delta",
{
"type": "content_block_delta",
"index": self.block_index,
"delta": {"type": "text_delta", "text": "</think>"},
},
)
]
def _handle_content(self, event: dict) -> list[str]:
cumulative = event.get("text", "")
new_text = cumulative[len(self._prev_text) :]
self._prev_text = cumulative
if not new_text:
return []
self._open_think_tags += new_text.count("<think>") - new_text.count("</think>")
if not self._text_block_open:
events = self._open_text_block()
else:
@ -374,6 +400,7 @@ class AnthropicStreamEmitter:
events = []
if self._text_block_open:
events.extend(self._close_open_think())
events.append(self._close_block())
# Defensive: close a stale open tool_use block before starting another.
elif self._open_tool_call_id is not None:
@ -452,6 +479,7 @@ class AnthropicStreamEmitter:
events.extend(self._open_text_block())
# Reset text tracking for the next synthesis turn
self._prev_text = ""
self._open_think_tags = 0
return events
def _open_text_block(self) -> list[str]:
@ -511,7 +539,7 @@ class AnthropicPassthroughEmitter:
Only calls naming a tool in ``allowed_tools`` (the client's declared
tools) are promoted; everything else streams as text exactly as before.
Never enabled for Studio's own tool loop.
Never enabled for Unsloth's own tool loop.
"""
from core.inference.passthrough_healing import StreamToolCallHealer

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

@ -76,8 +76,14 @@ class AudioCodecManager:
if self._snac_model is not None:
return
from snac import SNAC
from utils.hf_cache_settings import active_hf_hub_cache
self._snac_model = SNAC.from_pretrained("hubertsiuzdak/snac_24khz").to(device).eval()
# Route weights to the selected cache; this can run in the main process.
self._snac_model = (
SNAC.from_pretrained("hubertsiuzdak/snac_24khz", cache_dir = active_hf_hub_cache())
.to(device)
.eval()
)
logger.info("Loaded SNAC codec (24kHz)")
def _load_bicodec(

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

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

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

@ -3,11 +3,328 @@
"""
Dependency-light wrapper around tokenizer.apply_chat_template with a kwarg
fallback for templates that reject reasoning/tools args.
fallback for templates that reject reasoning/tools args, plus the shared
native-chat-template fallback used by the transformers and MLX backends.
"""
import copy
import json
import logging
from 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 Unsloth'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:
"""Return the trailing open ``<think>`` prefill of a rendered prompt.
Reasoning templates (Qwen3.6, DeepSeek-R1-style) end the generation
prompt with ``<think>\\n`` so the model starts reasoning immediately.
Because that opening tag is part of the *prompt*, skip_prompt streaming
never emits it, and the frontend's ``<think>``/``</think>`` parser shows
the reasoning as plain text instead of a thinking block. (The GGUF path
is unaffected: llama-server's reasoning parser returns
``reasoning_content``, which gets re-wrapped in think tags.)
Returns the exact prompt tail to re-emit at the start of the generated
stream (e.g. ``"<think>\\n"``), or ``""`` when the prompt does not end
with an open think block, including the ``enable_thinking=False`` case
where templates prefill an already-closed ``<think>\\n\\n</think>``.
``special_tokens`` is the tokenizer's special-token list. If ``</think>``
is one, the streamer's skip_special_tokens strips the model's closing tag,
so re-emitting the open would leave an unclosed block that swallows the
answer. In that case return ``""`` and fall back to plain text.
"""
if not prompt:
return ""
open_idx = prompt.rfind(_THINK_OPEN)
if open_idx == -1:
return ""
tail = prompt[open_idx:]
if _THINK_CLOSE in tail or tail.strip() != _THINK_OPEN:
return ""
if special_tokens and _THINK_CLOSE in set(special_tokens):
return ""
return tail
logger = logging.getLogger(__name__)
def _normalize_tool_call_arguments(messages: list) -> list:
"""Coerce each assistant ``tool_calls[].function.arguments`` from a JSON
string to a dict.
The OpenAI wire format carries ``arguments`` as a JSON string, but some chat
templates (e.g. the stricter Qwen tool templates shipped with mlx-community
checkpoints) iterate ``arguments.items()`` and raise
``TypeError: Can only get item pairs from a mapping.`` on the string form
when a prior tool call is re-rendered on the next turn. A dict works on both
strict and lenient templates, so parse the string; leave non-JSON or non-dict
values untouched. Returns the original list unchanged when nothing needed
coercing (no copy)."""
mutated = False
out: list = []
for msg in messages:
tool_calls = msg.get("tool_calls") if isinstance(msg, dict) else None
if not tool_calls:
out.append(msg)
continue
new_calls = []
msg_changed = False
for call in tool_calls:
fn = call.get("function") if isinstance(call, dict) else None
args = fn.get("arguments") if isinstance(fn, dict) else None
if isinstance(args, str):
try:
parsed = json.loads(args)
except (ValueError, TypeError):
parsed = None
if isinstance(parsed, dict):
call = {**call, "function": {**fn, "arguments": parsed}}
msg_changed = True
new_calls.append(call)
if msg_changed:
out.append({**msg, "tool_calls": new_calls})
mutated = True
else:
out.append(msg)
return out if mutated else messages
def apply_chat_template_for_generation(
tokenizer,
@ -38,21 +355,241 @@ def apply_chat_template_for_generation(
attempts.append(dict(reasoning_kwargs))
attempts.append({})
last_exc: Optional[Exception] = None
for kwargs in attempts:
def _render(msgs: list) -> str:
last_exc: Optional[Exception] = None
for kwargs in attempts:
try:
return tokenizer.apply_chat_template(
msgs,
tokenize = False,
add_generation_prompt = True,
**kwargs,
)
except TypeError as e:
last_exc = e
continue
except Exception as e:
last_exc = e
break
if last_exc is not None:
raise last_exc
raise RuntimeError("apply_chat_template_for_generation: no attempt produced a result")
try:
return _render(messages)
except Exception:
# Strict tool templates reject the JSON-string ``arguments`` form via
# TypeError or a broad Jinja raise_exception, so retry with dicts coerced.
# Original messages render first, so working templates stay byte-identical.
normalized = _normalize_tool_call_arguments(messages)
if normalized is messages:
raise
return _render(normalized)
def render_native_template(
*,
model_info: dict,
active_model_name: Optional[str],
messages: list,
tools: list,
enable_thinking: Optional[bool] = None,
reasoning_effort: Optional[str] = None,
preserve_thinking: Optional[bool] = None,
apply_fn = None,
hf_token: Optional[str] = None,
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
the ``tools`` schema, so a tool-calling turn silently stops advertising tools.
The native template ships in the model repo and carries the family's
tool-calling syntax. It is loaded straight from the repo (bypassing any
override on the live tokenizer) and cached on ``model_info``. Returns the
rendered prompt only if the native template actually emits the tools (render
differs with vs without tools); otherwise ``None``. 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
fallback fails silently and keeps the override prompt that dropped tools).
``trust_remote_code`` is sourced from ``model_info`` (the value the model was
actually loaded with) rather than a call-site argument, so the native-template
reload uses exactly the consent already granted at load. A custom-code tokenizer
repo raises in ``AutoTokenizer.from_pretrained`` unless ``trust_remote_code`` is
passed, so without this the fallback fails silently and keeps the tool-dropping
prompt for a model the user already consented to run remote code for. For a LoRA
adapter the reload targets the base model, whose remote code was gated and loaded
under the same stored flag, so re-passing it executes no unconsented code.
"""
# ``apply_fn`` lets a backend inject its own render; defaults to the module helper.
if apply_fn is None:
apply_fn = apply_chat_template_for_generation
native_tpl = model_info.get("native_chat_template")
if native_tpl is None:
# A LoRA adapter's native template lives on the base model, not the adapter id.
template_source = model_info.get("base_model") or active_model_name
# Re-use the load-time trust_remote_code so a custom-code tokenizer repo can
# instantiate its class (the stored flag already covers template_source).
trust_remote_code = bool(model_info.get("trust_remote_code", False))
try:
return tokenizer.apply_chat_template(
messages,
tokenize = False,
add_generation_prompt = True,
**kwargs,
from transformers import AutoTokenizer
nt = AutoTokenizer.from_pretrained(
template_source,
token = hf_token if hf_token and hf_token.strip() else None,
trust_remote_code = trust_remote_code,
)
except TypeError as e:
last_exc = e
continue
except Exception as e:
last_exc = e
break
if last_exc is not None:
raise last_exc
raise RuntimeError("apply_chat_template_for_generation: no attempt produced a result")
native_tpl = nt.chat_template or False
except Exception as exc:
logger.warning(
"Could not load native chat template for '%s': %s",
template_source,
exc,
)
# A failed fetch is not "no template": leave the sentinel unset so the next
# call retries (caching False would pin the tool-dropping override).
return None
model_info["native_chat_template"] = native_tpl
if not native_tpl:
return None
tokenizer = model_info.get("tokenizer") or model_info.get("processor")
if tokenizer is None:
return None
tokenizer = getattr(tokenizer, "tokenizer", tokenizer)
# Render on a shallow copy: mutating the shared tokenizer.chat_template (outside the
# generation lock) races concurrent requests.
try:
render_tokenizer = copy.copy(tokenizer)
render_tokenizer.chat_template = native_tpl
except Exception as exc:
logger.warning(
"Could not clone tokenizer for native-template render of '%s': %s",
active_model_name,
exc,
)
return None
try:
with_tools = apply_fn(
render_tokenizer,
messages,
tools = tools,
enable_thinking = enable_thinking,
reasoning_effort = reasoning_effort,
preserve_thinking = preserve_thinking,
)
no_tools = apply_fn(
render_tokenizer,
messages,
tools = None,
enable_thinking = enable_thinking,
reasoning_effort = reasoning_effort,
preserve_thinking = preserve_thinking,
)
except Exception as exc:
logger.warning(
"Native-template tool render failed for '%s': %s",
active_model_name,
exc,
)
return None
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(
*,
formatted_prompt: str,
tokenizer,
model_info: dict,
active_model_name: Optional[str],
messages: list,
tools: Optional[list],
enable_thinking: Optional[bool] = None,
reasoning_effort: Optional[str] = None,
preserve_thinking: Optional[bool] = None,
apply_fn = None,
hf_token: Optional[str] = None,
return_metadata: bool = False,
):
"""Return ``formatted_prompt``, swapping in a native-template render when an
override template dropped the ``tools`` schema.
If ``tools`` were requested but the live render is identical with and without
them (detected by comparison, robust against tool names in the system prompt),
re-render with the model's native template. Shared by the transformers and MLX
backends so both advertise tools consistently. ``hf_token`` is forwarded so a
gated/private model's native template can still be fetched. 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:
# 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
# can raise here; on any error keep the valid tools prompt rather than lose it.
try:
probe_no_tools = apply_fn(
tokenizer,
messages,
tools = None,
enable_thinking = enable_thinking,
reasoning_effort = reasoning_effort,
preserve_thinking = preserve_thinking,
)
except Exception as exc:
logger.warning(
"No-tools probe failed for '%s'; keeping the existing tools prompt: %s",
active_model_name,
exc,
)
return _result(formatted_prompt)
if formatted_prompt != probe_no_tools:
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,
messages = messages,
tools = tools,
enable_thinking = enable_thinking,
reasoning_effort = reasoning_effort,
preserve_thinking = preserve_thinking,
apply_fn = apply_fn,
hf_token = hf_token,
return_metadata = return_metadata,
)
if native_prompt:
logger.info(
"Override template for '%s' dropped tool schemas; using the model's "
"native template for this tool-calling turn.",
active_model_name,
)
return native_prompt
return _result(formatted_prompt)

4
studio/backend/core/inference/chat_templates.py
View file

@ -4,13 +4,13 @@
"""Bundled chat-template selection for GGUF inference.
Some shipped GGUF quants embed an older chat template. Rather than re-cutting and
asking users to re-download every quant, Studio can override the embedded template
asking users to re-download every quant, Unsloth can override the embedded template
at llama-server launch time with a bundled, up-to-date Jinja template for known
model families. The override is wired through the existing ``chat_template_override``
-> ``--chat-template-file`` path in ``LlamaCppBackend.load_model``.
Currently this covers ``unsloth/gemma-4-*-GGUF``, which gains the upstream PR #118
``preserve_thinking`` flag (defaulted OFF here) so the Studio "Preserve thinking"
``preserve_thinking`` flag (defaulted OFF here) so the Unsloth "Preserve thinking"
toggle appears while staying disabled by default.
"""

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

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

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

@ -473,7 +473,7 @@ def _apply_mistral_reasoning_controls(
# handles every provider without storing credentials.
def _create_shared_http_client() -> httpx.AsyncClient:
# Unsupported env proxy schemes (socks:// etc) raise at construction and
# would crash Studio startup (#6090); retry ignoring env proxies instead.
# would crash Unsloth startup (#6090); retry ignoring env proxies instead.
try:
return httpx.AsyncClient()
except (ImportError, ValueError) as exc:
@ -858,7 +858,7 @@ class ExternalProviderClient:
if not self._is_openai_compatible():
# Gemini speaks its own native REST shape (contents/parts);
# `_stream_gemini` translates request/response into the OpenAI
# Chat Completions chunk format the rest of Studio expects.
# Chat Completions chunk format the rest of Unsloth expects.
# API ref: https://ai.google.dev/gemini-api/docs
if self.provider_type == "gemini":
async for line in self._stream_gemini(
@ -1706,7 +1706,7 @@ class ExternalProviderClient:
# Translate OpenAI multimodal parts -> Anthropic native shapes.
# - `image_url` -> `{type:"image", source:...}`
# - `input_document` -> `{type:"document", source:...}`
# (Studio extension; mirrors Anthropic's document block,
# (Unsloth extension; mirrors Anthropic's document block,
# which supports PDFs as base64 or URL per
# https://platform.claude.com/docs/en/build-with-claude/vision)
anthropic_parts: list[dict[str, Any]] = []
@ -1749,7 +1749,7 @@ class ExternalProviderClient:
}
)
elif part.get("type") == "input_document":
# Studio's normalised PDF/doc type (file_data data-URI or
# Unsloth's normalised PDF/doc type (file_data data-URI or
# file_url) -> Anthropic's native `document` block.
url = part.get("file_url") or ""
data_uri = part.get("file_data") or ""
@ -4704,7 +4704,7 @@ class ExternalProviderClient:
{"type": "image_generation_call", "id": call_id}
)
elif part_type == "input_document":
# Map Studio's `input_document` onto Responses' `input_file`.
# Map Unsloth's `input_document` onto Responses' `input_file`.
# https://developers.openai.com/api/docs/guides/images-vision
file_url = part.get("file_url")
file_data = part.get("file_data")
@ -6010,7 +6010,7 @@ class ExternalProviderClient:
if not models and self.provider_type == "ollama":
models = await self._list_ollama_native_models()
# Gemini's native /v1beta/models uses a different shape; repackage
# into the OpenAI-compatible one Studio expects.
# into the OpenAI-compatible one Unsloth expects.
if not models and self.provider_type == "gemini":
models = self._parse_gemini_models(data)
return models
@ -6213,7 +6213,7 @@ def _friendly_provider_error_text(
*,
model: str | None = None,
) -> str:
"""Rewrite common provider errors into actionable Studio copy."""
"""Rewrite common provider errors into actionable Unsloth copy."""
if status_code == 404 and model:
lowered = raw_message.lower()
if "not found" in lowered or "not_found" in lowered:

537
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,
@ -27,6 +28,16 @@ from utils.hardware import (
from core.inference.audio_codecs import AudioCodecManager
from core.inference.runtime_context import runtime_context_length
from core.inference.message_content import content_to_text
from core.inference.chat_eos import (
chat_eos_repair,
resolve_chat_turn_end_eos_ids_using,
)
from core.inference.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
from loggers import get_logger
@ -181,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"""
@ -210,6 +268,50 @@ class InferenceBackend:
# API uses -1 to disable top-k; transformers uses 0.
return 0 if top_k < 0 else top_k
def _resolve_chat_eos(self, model_name: str) -> None:
"""Resolve this chat model's assistant-turn-end stop tokens once at load,
cache them in model_info, and repair generation_config so every
``.generate()`` path stops at the turn boundary.
Some checkpoints (e.g. Qwen3.5 / Qwen3.6 small chat models) end turns with
``<|im_end|>`` but ship ``config.eos_token_id = <|endoftext|>`` and no
``generation_config.json``, so paths that read ``generation_config`` (the
vision path, tool loops) run past the turn and loop. Turn-end markers are
derived from the chat_template (see chat_eos.resolve_chat_turn_end_eos_ids),
so base/coder models and harmony templates are left untouched.
"""
info = self.models.get(model_name) or {}
model = info.get("model")
container = info.get("tokenizer")
tokenizer = getattr(container, "tokenizer", container) # unwrap processors
if model is None or tokenizer is None:
return
# Vision models carry the chat_template on the processor, not the inner
# tokenizer. Read markers from whichever has one, but resolve ids on the
# generation tokenizer, else the vision path misses the turn-end token.
template_source = container if getattr(container, "chat_template", None) else tokenizer
try:
turn_end_ids = resolve_chat_turn_end_eos_ids_using(template_source, tokenizer)
except Exception as e: # never block a load on eos resolution
logger.warning("Chat turn-end eos resolution failed for %s: %s", model_name, e)
return
info["chat_turn_end_eos_ids"] = turn_end_ids
gen = getattr(model, "generation_config", None)
if gen is None:
return
repaired = chat_eos_repair(gen.eos_token_id, turn_end_ids)
if repaired is None:
return
previous = gen.eos_token_id
gen.eos_token_id = repaired
logger.info(
"Repaired generation_config.eos_token_id for %s: %s -> %s",
model_name,
previous,
repaired,
)
def load_model(
self,
config: ModelConfig,
@ -221,6 +323,9 @@ class InferenceBackend:
gpu_ids: Optional[list[int]] = None,
) -> bool:
"""Load any model: base, LoRA adapter, text, or vision."""
# Keep the token so the native-template fallback can fetch a
# gated model's repo template later during generation.
self._hf_token = hf_token
# GGUF uses max_seq_length=0 as "model default"; Unsloth crashes on it.
if max_seq_length <= 0:
max_seq_length = 2048
@ -231,6 +336,8 @@ class InferenceBackend:
# Already loaded?
if model_name in self.models and self.models[model_name].get("model"):
logger.info(f"Model {model_name} already loaded")
if hf_token:
self.models[model_name]["hf_token"] = hf_token
self.active_model_name = model_name
return True
@ -246,6 +353,14 @@ class InferenceBackend:
)
self.models[model_name] = {
# Per-model token: the native-template fallback must use the
# token this model was loaded with, not whichever loaded last.
"hf_token": hf_token,
# Per-model consent: the native-template reload must re-use the
# exact trust_remote_code this model (and a LoRA's base) was loaded
# with, so a custom-code tokenizer repo can be re-fetched without
# executing any code the user did not already consent to.
"trust_remote_code": trust_remote_code,
"is_vision": config.is_vision,
"is_lora": config.is_lora,
"is_audio": config.is_audio,
@ -378,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
@ -496,6 +611,7 @@ class InferenceBackend:
max_seq_length,
)
self._resolve_chat_eos(model_name)
self._load_chat_template_info(model_name)
self.active_model_name = model_name
@ -766,9 +882,13 @@ class InferenceBackend:
preserve_thinking: Optional[bool] = None,
max_tool_iterations: int = 25,
auto_heal_tool_calls: bool = True,
nudge_tool_calls: Optional[bool] = None,
tool_call_timeout: int = 300,
session_id: Optional[str] = None,
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``.
@ -802,6 +922,7 @@ class InferenceBackend:
enable_thinking = enable_thinking,
reasoning_effort = reasoning_effort,
preserve_thinking = preserve_thinking,
presence_penalty = presence_penalty,
)
initial = list(messages)
@ -815,10 +936,13 @@ class InferenceBackend:
execute_tool = execute_tool,
cancel_event = cancel_event,
auto_heal_tool_calls = auto_heal_tool_calls,
nudge_tool_calls = nudge_tool_calls,
max_tool_iterations = max_tool_iterations,
tool_call_timeout = tool_call_timeout,
session_id = session_id,
thread_id = thread_id,
rag_scope = rag_scope,
reasoning_prefilled = reasoning_prefilled,
)
def generate_chat_response(
@ -837,12 +961,14 @@ class InferenceBackend:
enable_thinking: Optional[bool] = None,
reasoning_effort: Optional[str] = None,
preserve_thinking: Optional[bool] = None,
presence_penalty: float = 0.0,
) -> Generator[str, None, None]:
"""Generate response for text or vision models (lock held by background thread).
``tools`` / ``enable_thinking`` / ``reasoning_effort`` / ``preserve_thinking``
are forwarded into ``apply_chat_template`` so templates that understand them
(Qwen3, Llama 3.1+, gpt-oss harmony) advertise tool schemas / reasoning controls.
``presence_penalty`` matches the GGUF sampling path (0 disables it).
"""
yield from self._generate_chat_response_inner(
messages = messages,
@ -859,6 +985,7 @@ class InferenceBackend:
enable_thinking = enable_thinking,
reasoning_effort = reasoning_effort,
preserve_thinking = preserve_thinking,
presence_penalty = presence_penalty,
)
def _generate_chat_response_inner(
@ -878,6 +1005,7 @@ class InferenceBackend:
enable_thinking: Optional[bool] = None,
reasoning_effort: Optional[str] = None,
preserve_thinking: Optional[bool] = None,
presence_penalty: float = 0.0,
) -> Generator[str, None, None]:
"""Inner generation logic, called by generate_chat_response and
generate_with_adapter_control.
@ -886,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)
@ -917,6 +1044,7 @@ class InferenceBackend:
max_new_tokens,
repetition_penalty,
cancel_event = cancel_event,
presence_penalty = presence_penalty,
)
return
else:
@ -946,6 +1074,22 @@ class InferenceBackend:
tokenizer,
chat_template = template_name,
)
# The mapper installs the effective template only now, at generate
# time, so re-resolve and UNION into the load-time cache (never
# overwrite). get_chat_template can return a remapped tokenizer
# (turn-end folded onto doc-eos) while generate_stream reads the
# original, so take marker strings from the mapped template but
# resolve their ids on the original.
try:
_gen_tok = model_info.get("tokenizer") or tokenizer
refreshed = resolve_chat_turn_end_eos_ids_using(
getattr(tokenizer, "tokenizer", tokenizer),
getattr(_gen_tok, "tokenizer", _gen_tok),
)
existing = model_info.get("chat_turn_end_eos_ids") or []
model_info["chat_turn_end_eos_ids"] = sorted(set(existing) | set(refreshed))
except Exception as e:
logger.warning(f"Could not refresh chat turn-end eos after template: {e}")
else:
logger.info(
f"No registered Unsloth template for {self.active_model_name}, using tokenizer default"
@ -958,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(
@ -967,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,
@ -975,11 +1121,38 @@ class InferenceBackend:
reasoning_effort = reasoning_effort,
preserve_thinking = preserve_thinking,
)
# If tools were requested but the (possibly overridden) template ignored
# them, fall back to the model's native template (shared with MLX).
from core.inference.chat_template_helpers import (
render_with_native_template_fallback,
)
render_result = render_with_native_template_fallback(
formatted_prompt = formatted_prompt,
tokenizer = tokenizer,
model_info = model_info,
active_model_name = self.active_model_name,
messages = template_messages,
tools = tools,
enable_thinking = enable_thinking,
reasoning_effort = reasoning_effort,
preserve_thinking = preserve_thinking,
apply_fn = self._apply_chat_template_for_generation,
hf_token = model_info.get("hf_token"),
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(
@ -992,6 +1165,9 @@ class InferenceBackend:
repetition_penalty,
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(
@ -1006,6 +1182,7 @@ class InferenceBackend:
max_new_tokens,
repetition_penalty,
cancel_event = None,
presence_penalty: float = 0.0,
) -> Generator[str, None, None]:
"""Handle vision model generation with true token-by-token streaming."""
model_info = self.models[self.active_model_name]
@ -1067,21 +1244,36 @@ class InferenceBackend:
add_special_tokens = False,
return_tensors = "pt",
).to(model.device)
prompt_text = input_text
else:
# Text-only path for a vision model
formatted_prompt = self.format_chat_prompt(messages, system_prompt)
inputs = raw_tokenizer(formatted_prompt, return_tensors = "pt").to(model.device)
prompt_text = formatted_prompt
# Stream with TextIteratorStreamer + background thread
try:
from transformers import TextIteratorStreamer
# 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)
)
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(
@ -1095,6 +1287,18 @@ class InferenceBackend:
top_k = top_k,
min_p = min_p,
)
# Presence penalty (GGUF parity) for VLM chat.
_vision_input_ids = inputs.get("input_ids") if hasattr(inputs, "get") else None
if _vision_input_ids is not None:
_pp = _make_presence_penalty_processor(
presence_penalty, int(_vision_input_ids.shape[1])
)
if _pp is not None:
generation_kwargs["logits_processor"] = _pp
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] = {}
@ -1104,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:
@ -1114,14 +1320,23 @@ class InferenceBackend:
thread = threading.Thread(target = generate_fn)
thread.start()
output = ""
output = think_prefix
# Emit the prefilled <think> before the first token so the block
# renders during prompt prefill (which can take seconds).
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:
@ -1130,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,
@ -1275,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,
@ -1312,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,
@ -1323,15 +1641,18 @@ class InferenceBackend:
repetition_penalty: float = 1.0,
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).
_adapter_state: if not None, the background thread toggles adapters
before model.generate(), under _generation_lock.
``presence_penalty`` matches the GGUF sampling path via a logits processor (0 disables it).
"""
if not self.active_model_name:
yield "Error: No active model"
return
raise RuntimeError("No active model")
model_info = self.models[self.active_model_name]
model = model_info["model"]
@ -1344,33 +1665,27 @@ class InferenceBackend:
try:
inputs = tokenizer(prompt, return_tensors = "pt").to(model.device)
from transformers import TextIteratorStreamer
import threading
# 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,
)
# skip_prompt swallows an open <think> prefilled by the template;
# re-emit it so the frontend can render the thinking block.
# gpt-oss emits its own tags via HarmonyTextStreamer.
think_prefix = (
""
if self._is_gpt_oss_model()
else detect_think_prefill(prompt, getattr(tokenizer, "all_special_tokens", None))
)
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,
@ -1382,26 +1697,22 @@ class InferenceBackend:
min_p = min_p,
repetition_penalty = repetition_penalty,
do_sample = temperature > 0,
eos_token_id = tokenizer.eos_token_id,
# Resolved once at load (chat_template-derived turn-end tokens).
eos_token_id = model_info.get("chat_turn_end_eos_ids") or tokenizer.eos_token_id,
pad_token_id = tokenizer.eos_token_id
if tokenizer.pad_token_id is None
else tokenizer.pad_token_id,
)
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)]
)
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
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:
@ -1411,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:
@ -1422,14 +1735,23 @@ class InferenceBackend:
thread = threading.Thread(target = generate_fn)
thread.start()
output = ""
output = think_prefix
# Emit the prefilled <think> before the first token so the block
# renders during prompt prefill (which can take seconds).
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:
@ -1438,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
@ -1451,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 ────────────────────────────────────
@ -1949,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.
@ -1960,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()

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

File diff suppressed because it is too large Load diff

140
studio/backend/core/inference/llama_keepwarm.py
View file

@ -15,6 +15,7 @@ import asyncio
import contextlib
import threading
import time
from pathlib import Path
from loggers import get_logger
@ -30,6 +31,8 @@ _last_active = time.monotonic()
# otherwise 503 against an empty backend can reload it (set on unload, cleared on
# reload). Storing the quant means the reload restores the exact freed variant.
_last_unloaded_model = None
# Slot KV manifest saved by the idle unload; whoever pops it owns deleting its files.
_kv_resume = None
# Guards inflight bumps against the idle-check-then-unload race, and blocks new
# inference from starting mid-swap. Process-wide, not per-loop: the backend slot is
# shared across every event loop in the process, so a per-loop gate would let a
@ -59,7 +62,7 @@ _INFERENCE_SUFFIXES = (
"/messages/count_tokens", # counts via the loaded tokenizer; protect like /messages
"/embeddings",
"/responses",
"/generate/stream", # Studio's own streaming route on the same llama-server
"/generate/stream", # Unsloth's own streaming route on the same llama-server
"/audio/generate", # direct GGUF TTS; can outlive the idle TTL
)
@ -161,11 +164,17 @@ def inference_lifecycle_gate():
return _unload_gate()
def note_model_loaded() -> None:
"""Record a successful GGUF load: stamp activity and drop any reload stash so
a manual load clears it synchronously, not only on the next idle poll."""
def note_model_loaded(backend = None) -> None:
"""Stamp activity and synchronously drop any reload stash."""
_note_activity()
resume = take_kv_resume()
_set_last_unloaded(None)
if resume is None:
return
if backend is not None:
restore_kv_resume(backend, resume)
else:
_delete_resume_files(resume)
def note_model_unloaded() -> None:
@ -182,9 +191,81 @@ def get_last_unloaded_model():
def _set_last_unloaded(value) -> None:
global _last_unloaded_model
global _last_unloaded_model, _kv_resume
stale = None
with _lock:
_last_unloaded_model = value
if value is None and _kv_resume is not None:
stale, _kv_resume = _kv_resume, None
if stale:
_delete_resume_files(stale)
def _delete_resume_files(manifest) -> None:
try:
base = Path(manifest.get("dir") or "")
for entry in manifest.get("slots") or []:
with contextlib.suppress(OSError):
(base / str(entry.get("filename"))).unlink()
except Exception:
pass
def _set_kv_resume(value) -> None:
global _kv_resume
stale = None
with _lock:
if _kv_resume is not None and _kv_resume is not value:
stale = _kv_resume
_kv_resume = value
if stale:
_delete_resume_files(stale)
def take_kv_resume():
global _kv_resume
with _lock:
manifest, _kv_resume = _kv_resume, None
return manifest
def purge_kv_resume() -> None:
resume = take_kv_resume()
if resume:
_delete_resume_files(resume)
def restore_kv_resume(backend, manifest) -> None:
try:
gguf = manifest.get("gguf")
binary = manifest.get("binary")
current = getattr(backend, "_gguf_path", None)
same_gguf = bool(gguf and current) and Path(current).resolve() == Path(gguf).resolve()
if same_gguf:
# Same path is not enough: shards may have been rewritten meanwhile.
identity = getattr(backend, "_gguf_file_identity", None)
same_gguf = callable(identity) and identity(current) == manifest.get("gguf_stat")
if same_gguf:
# Nor the same file: launch overrides can invalidate KV numerics.
fingerprint = getattr(backend, "_slot_launch_fingerprint", None)
same_gguf = callable(fingerprint) and manifest.get("launch") == fingerprint()
if same_gguf and binary and binary == getattr(backend, "_slot_save_binary", None):
logger.info("Restoring saved slot KV onto the reloaded model")
backend.restore_slots_for_resume(manifest)
except Exception as exc:
logger.debug("slot restore after reload failed: %s", exc)
finally:
_delete_resume_files(manifest)
def sweep_slot_save_dir() -> None:
try:
from utils.paths.storage_roots import llama_slot_cache_root
for path in llama_slot_cache_root().glob("resume-*.bin"):
with contextlib.suppress(OSError):
path.unlink()
except Exception:
pass
class LlamaKeepWarmMiddleware:
@ -266,7 +347,10 @@ def _loaded_identity(backend):
async def idle_unload_loop(poll_seconds: float = 15.0) -> None:
"""Unload the loaded GGUF once idle past the configured TTL. Inert when off."""
from utils.openai_auto_switch_settings import get_auto_unload_idle_seconds
from utils.openai_auto_switch_settings import (
get_auto_unload_idle_seconds,
get_auto_unload_keep_kv,
)
seen_model = None
while True:
@ -281,17 +365,47 @@ async def idle_unload_loop(poll_seconds: float = 15.0) -> None:
# Track by (id, variant): a (re)loaded model -- including the same repo
# at a different quant -- counts as activity so it survives one TTL
# before its first request (loads bypass the activity middleware).
current = _loaded_identity(backend)
if current != seen_model:
seen_model = current
if current is not None:
_note_activity()
_set_last_unloaded(None) # a model is loaded; drop stale stash
async with _unload_gate():
# Purging the stash mid-reload would race the restore.
current = _loaded_identity(backend)
if current != seen_model:
seen_model = current
if current is not None:
_note_activity()
_set_last_unloaded(None) # a model is loaded; drop stale stash
if backend.is_loaded and _is_idle(ttl):
freed = _loaded_identity(backend)
await asyncio.to_thread(backend.unload_model)
manifest = None
if get_auto_unload_keep_kv():
try:
manifest = await asyncio.to_thread(
backend.save_slots_for_resume,
lambda: not _is_idle(ttl),
)
except Exception as exc:
logger.debug("slot save before idle unload failed: %s", exc)
# Re-read settings: the save can outlive a settings change.
ttl = get_auto_unload_idle_seconds()
if ttl <= 0 or not _is_idle(ttl):
if manifest:
_delete_resume_files(manifest)
continue
if manifest and not get_auto_unload_keep_kv():
_delete_resume_files(manifest)
manifest = None
try:
await asyncio.to_thread(backend.unload_model)
except Exception:
# Failed unload means nothing will stash the manifest.
if manifest:
_delete_resume_files(manifest)
raise
_set_last_unloaded(freed) # let an alias request reload it
if manifest and freed:
_set_kv_resume({"identity": freed, **manifest})
logger.info("Idle auto-unload: saved slot KV for restore on reload")
elif manifest:
_delete_resume_files(manifest)
logger.info("Idle auto-unload: freed GGUF after %ss idle", ttl)
seen_model = None
except Exception as exc:

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

@ -3,10 +3,10 @@
"""Boundary validator for user-supplied llama-server pass-through args.
Reject only flags Studio manages (model identity, auth, network, parallel
Reject only flags Unsloth manages (model identity, auth, network, parallel
slots). Everything else (sampling, ``-c``, ``-ngl``, ``--flash-attn``,
``--cache-type-*``, ``--spec-*``, ``--jinja``, ...) is appended after
Studio's auto-set flags so llama.cpp's last-wins parser lets the user override.
Unsloth's auto-set flags so llama.cpp's last-wins parser lets the user override.
Ref: https://github.com/ggml-org/llama.cpp/blob/master/tools/server/README.md
"""
@ -22,12 +22,12 @@ _DENYLIST_GROUPS: tuple[frozenset[str], ...] = (
# Parallel slots: owned by typer --parallel; a pass-through would desync
# app.state.llama_parallel_slots from llama-server.
frozenset({"-np", "--parallel", "--n-parallel"}),
# Model identity: Studio resolves it from LoadRequest; a second -m would
# load a different model than Studio thinks it loaded.
# Model identity: Unsloth resolves it from LoadRequest; a second -m would
# load a different model than Unsloth thinks it loaded.
frozenset({"-m", "--model"}),
# Public model id: Studio sets a sanitized --alias so the OpenAI API never
# Public model id: Unsloth sets a sanitized --alias so the OpenAI API never
# exposes the local .gguf path. A user-supplied alias is appended after
# Studio's and, with llama.cpp's last-wins parsing, would reintroduce the
# Unsloth's and, with llama.cpp's last-wins parsing, would reintroduce the
# path leak this is meant to prevent.
frozenset({"-a", "--alias"}),
frozenset({"-mu", "--model-url"}),
@ -39,14 +39,14 @@ _DENYLIST_GROUPS: tuple[frozenset[str], ...] = (
frozenset({"-hft", "--hf-token"}),
frozenset({"-mm", "--mmproj"}),
frozenset({"-mmu", "--mmproj-url"}),
# Networking: Studio binds + proxies; retargeting orphans the proxy.
# Networking: Unsloth binds + proxies; retargeting orphans the proxy.
frozenset({"--host"}),
frozenset({"--port"}),
frozenset({"--path"}),
frozenset({"--api-prefix"}),
frozenset({"--reuse-port"}),
# Auth / TLS: Studio terminates auth; upstream --api-key / TLS shadows
# Studio's key and breaks the proxy hop.
# Auth / TLS: Unsloth terminates auth; upstream --api-key / TLS shadows
# Unsloth's key and breaks the proxy hop.
frozenset({"--api-key"}),
frozenset({"--api-key-file"}),
frozenset({"--ssl-key-file"}),
@ -64,12 +64,14 @@ _DENYLIST_GROUPS: tuple[frozenset[str], ...] = (
frozenset({"--models-max"}),
frozenset({"--models-autoload", "--no-models-autoload"}),
# Server-mode flips: --embedding / --rerank restrict llama-server to
# those endpoints, breaking Studio's /v1/chat/completions hop.
# those endpoints, breaking Unsloth's /v1/chat/completions hop.
frozenset({"--embedding", "--embeddings"}),
frozenset({"--rerank", "--reranking"}),
# llama-server's own built-in tools flag would silently stack on top of
# Studio's --enable-tools / --disable-tools policy resolver.
# Unsloth's --enable-tools / --disable-tools policy resolver.
frozenset({"--tools"}),
# Slot-state dir: Studio owns it for KV persistence across idle unload.
frozenset({"--slot-save-path"}),
)
_DENYLIST: frozenset[str] = frozenset().union(*_DENYLIST_GROUPS)
@ -120,7 +122,7 @@ def validate_extra_args(args: Optional[Iterable[str]]) -> list[str]:
def is_managed_flag(flag: str) -> bool:
"""True if ``flag`` is Studio-managed. Normalises via ``_flag_name`` so
"""True if ``flag`` is Unsloth-managed. Normalises via ``_flag_name`` so
`-np8` / `--parallel=8` classify like the canonical tokens."""
normalised = _flag_name(flag)
return normalised is not None and normalised in _DENYLIST
@ -142,7 +144,7 @@ _SPEC_FLAGS: frozenset[str] = frozenset(
"--draft-min",
"--draft-max",
# MTP path (llama.cpp #22673). The drafter selectors (local --model-draft
# and HF --spec-draft-hf aliases) are Studio-managed since the separate-
# and HF --spec-draft-hf aliases) are Unsloth-managed since the separate-
# drafter support (Gemma 4): an inherited copy must not last-wins-override
# the auto-detected drafter. Explicit extras for the current load are never
# stripped. The per-drafter tuning knobs (--spec-draft-type-*, -ngld,
@ -179,25 +181,38 @@ _TEMPLATE_FLAGS: frozenset[str] = frozenset(
# (--split-mode tensor). Pass-through stays allowed so users keep the
# row/none/layer modes the toggle doesn't expose, but it's stripped on
# inherit and reconciled into the round-tripped tensor_parallel state.
# --tensor-split is coupled to the split mode and is stripped with it: Studio
# --tensor-split is coupled to the split mode and is stripped with it: Unsloth
# owns the tensor-mode split ratios, so an inherited/stale --tensor-split must
# not last-wins-override Studio's computed asymmetric split.
# not last-wins-override Unsloth's computed asymmetric split.
_SPLIT_MODE_FLAGS: frozenset[str] = frozenset({"-sm", "--split-mode"})
_TENSOR_SPLIT_FLAGS: frozenset[str] = frozenset({"-ts", "--tensor-split"})
_SPLIT_SHADOWING_FLAGS: frozenset[str] = _SPLIT_MODE_FLAGS | _TENSOR_SPLIT_FLAGS
# GPU-offload flags. Stripped only when the GPU Memory mode owns offload
# (manual emits --fit / --gpu-layers / --n-cpu-moe); in auto, a user's
# inherited -ngl is respected (the offload_overridden path), so this group is
# opt-in, not default. Layer flags are shared with llama_cpp's override
# detection; the MoE flags are strip-only (manual's --n-cpu-moe slider owns them).
_LAYER_OFFLOAD_FLAGS: frozenset[str] = frozenset(
{"-ngl", "--gpu-layers", "--n-gpu-layers", "-fit", "--fit"}
)
_MOE_OFFLOAD_FLAGS: frozenset[str] = frozenset({"-ncmoe", "--n-cpu-moe", "-cmoe", "--cpu-moe"})
_OFFLOAD_SHADOWING_FLAGS: frozenset[str] = _LAYER_OFFLOAD_FLAGS | _MOE_OFFLOAD_FLAGS
_SHADOWING_FLAGS: frozenset[str] = (
_CONTEXT_FLAGS | _CACHE_FLAGS | _SPEC_FLAGS | _TEMPLATE_FLAGS | _SPLIT_SHADOWING_FLAGS
)
# Shadowing flags that take no value -- strip the flag only, not the next token.
_BOOLEAN_SHADOWING_FLAGS: frozenset[str] = frozenset({"--spec-default", "--jinja", "--no-jinja"})
_BOOLEAN_SHADOWING_FLAGS: frozenset[str] = frozenset(
{"--spec-default", "--jinja", "--no-jinja", "-cmoe", "--cpu-moe"}
)
def parse_ctx_override(args: Optional[Iterable[str]]) -> Optional[int]:
"""Return the last user-supplied ``-c`` / ``--ctx-size`` value.
Mirrors llama.cpp's last-wins parsing for the one numeric knob Studio's
Mirrors llama.cpp's last-wins parsing for the one numeric knob Unsloth's
load-time fit logic needs.
"""
if not args:
@ -286,7 +301,7 @@ def parse_cache_override(args: Optional[Iterable[str]]) -> Optional[str]:
Mirrors parse_ctx_override but for cache type. Recognises both -ctk
(key) and -ctv (value). When both flags appear, returns the last-wins
value, treating key and value cache flags as the same setting because
Studio's KV estimate has a single cache_type_kv knob.
Unsloth's KV estimate has a single cache_type_kv knob.
"""
return _last_flag_value(args, _CACHE_FLAGS)
@ -341,7 +356,7 @@ def resolve_tensor_parallel(args: Optional[Iterable[str]], fallback_tensor_paral
def _env_split_mode_is_tensor(env: Optional[Mapping[str, str]] = None) -> bool:
"""True when the inherited LLAMA_ARG_SPLIT_MODE env selects tensor. Studio
"""True when the inherited LLAMA_ARG_SPLIT_MODE env selects tensor. Unsloth
emits --split-mode only on its tensor branch, so a tensor env on the layer
path would run the child tensor-parallel unbudgeted; this flips the budget
to tensor. Only tensor is heavier, so other modes are ignored."""
@ -424,14 +439,22 @@ def strip_shadowing_flags(
strip_spec: bool = True,
strip_template: bool = True,
strip_split_mode: bool = True,
strip_tensor_split: bool = False,
strip_offload: bool = False,
) -> list[str]:
"""Strip flags that shadow first-class Studio settings.
"""Strip flags that shadow first-class Unsloth settings.
Used when inheriting a previous load's ``llama_extra_args`` so an
inherited `-c 4096` can't override the current `max_seq_length`
(same for cache / spec / template / split-mode). Each ``strip_*``
toggle controls one group; the route only strips groups whose
first-class field the caller actually supplied.
``strip_split_mode`` removes both ``--split-mode`` and the coupled
``--tensor-split`` (the Tensor Parallelism toggle owns the whole split).
``strip_tensor_split`` removes ``--tensor-split`` *alone*, so manual mode can
replace an inherited per-GPU ratio while leaving the user's ``--split-mode``
row/none/layer choice intact.
"""
shadowing: set[str] = set()
if strip_context:
@ -444,6 +467,10 @@ def strip_shadowing_flags(
shadowing |= _TEMPLATE_FLAGS
if strip_split_mode:
shadowing |= _SPLIT_SHADOWING_FLAGS
if strip_tensor_split:
shadowing |= _TENSOR_SPLIT_FLAGS
if strip_offload:
shadowing |= _OFFLOAD_SHADOWING_FLAGS
tokens = [str(a) for a in (args or [])]
out: list[str] = []

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

@ -5,7 +5,7 @@
engine-stats log line (generation/prompt throughput, requests in flight).
llama-server already computes these (it needs `--metrics`); this lifts them
into Studio's structured log so the terminal shows serving health, not just
into Unsloth's structured log so the terminal shows serving health, not just
per-request access lines. Emitted only while there is activity.
"""

18
studio/backend/core/inference/local_model_resolver.py
View file

@ -130,7 +130,7 @@ def info_has_local_gguf(info) -> bool:
def _build_index() -> dict[str, _LocalGgufEntry]:
"""Map normalized id/model_id/display_name -> local GGUF entry.
Scans the same roots Studio's model picker lists (./models, the active plus
Scans the same roots Unsloth's model picker lists (./models, the active plus
legacy/default HF caches, LM Studio dirs, and user scan folders) so a named
local model is never missed and silently served as the loaded one. Ollama's
scanner is skipped: it creates symlinks as a side effect and this runs on the
@ -146,6 +146,7 @@ def _build_index() -> dict[str, _LocalGgufEntry]:
_is_hidden_model,
)
from utils.paths import legacy_hf_cache_dir, hf_default_cache_dir, lmstudio_model_dirs
from utils.hf_cache_settings import known_hf_hub_caches
index: dict[str, _LocalGgufEntry] = {}
seen_hf: set[str] = set()
@ -174,7 +175,12 @@ def _build_index() -> dict[str, _LocalGgufEntry]:
except Exception as exc:
logger.debug("auto-switch: ./models scan failed: %s", exc)
try:
for hf_dir in (_resolve_hf_cache_dir(), legacy_hf_cache_dir(), hf_default_cache_dir()):
for hf_dir in (
*known_hf_hub_caches(),
_resolve_hf_cache_dir(),
legacy_hf_cache_dir(),
hf_default_cache_dir(),
):
found += _scan_hf_once(hf_dir)
except Exception as exc:
logger.debug("auto-switch: HF cache scan failed: %s", exc)
@ -199,9 +205,13 @@ def _build_index() -> dict[str, _LocalGgufEntry]:
raw_id = getattr(info, "id", None)
if not raw_id:
continue
# Skip what Studio hides from its pickers (validation probe, RAG embed
# Skip what Unsloth hides from its pickers (validation probe, RAG embed
# weights): not chat models, so never an auto-switch target.
if _is_hidden_model(raw_id, getattr(info, "path", None)):
if _is_hidden_model(
raw_id,
getattr(info, "model_id", None),
getattr(info, "path", None),
):
continue
# Advertise a client-facing alias, not an absolute filesystem path.
loader_id = _advertised_loader_id(info)

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 an Unsloth 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}"

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

File diff suppressed because it is too large Load diff

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

File diff suppressed because it is too large Load diff

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

@ -5,7 +5,7 @@
With server-side tools disabled (``unsloth run --disable-tools``, every
``unsloth start`` coding agent), requests carrying the client's own ``tools``
bypass Studio's tool loop and are relayed to/from llama-server verbatim. Small
bypass Unsloth's tool loop and are relayed to/from llama-server verbatim. Small
GGUF models often emit their tool calls as TEXT (``<tool_call>{...}</tool_call>``,
Gemma ``<|tool_call>...``, ``<function=...>`` XML) instead of structured
``tool_calls`` -- on the passthrough that text reaches the agent as prose and
@ -18,7 +18,7 @@ promotes calls whose function name exactly matches a declared tool. Promotion
removes EXACTLY the promoted calls' markup spans (the parser reports them):
undeclared calls, unparseable blocks, and suppressed alternate formats keep
every byte and relay as text, so healing can never silently delete model
output. Responses without a tool signal, requests without tools, and Studio's
output. Responses without a tool signal, requests without tools, and Unsloth's
own enable-tools loop are untouched. Per-request opt-out:
``auto_heal_tool_calls: false``. Process kill-switch:
``UNSLOTH_DISABLE_TOOL_CALL_HEALING=1``.
@ -29,10 +29,28 @@ import os
from collections.abc import Mapping
from typing import Any, Optional
from core.inference.tool_call_parser import TOOL_XML_SIGNALS, has_tool_signal
from core.inference.tool_loop_controller import coerce_tool_arguments
from core.tool_healing import parse_tool_calls_from_text
# Only the formats this healer's parser can promote -- narrower than the loops'
# broader TOOL_XML_SIGNALS. A loop-only marker (Llama <|python_tag|>, bare
# [ARGS]) would buffer a streamed call as prose without promoting it, so keep a
# healer-aligned list. Mistral's [TOOL_CALLS] IS promotable, so it stays in.
_HEAL_SIGNALS = (
"<tool_call>",
"<|tool_call>",
"<function=",
"[TOOL_CALLS]",
# TML Inkling native call marker (leaks as text when the server-side
# parser misses a narration-then-call turn).
"<|content_invoke_tool_json|>",
)
def _has_heal_signal(text: str) -> bool:
return any(s in text for s in _HEAL_SIGNALS)
# Read once at import (same convention as the other UNSLOTH_* switches).
_HEALING_DISABLED = os.environ.get("UNSLOTH_DISABLE_TOOL_CALL_HEALING", "0") == "1"
# Nudging is OPT-IN: per-request nudge_tool_calls=true, or flip the process
@ -44,7 +62,7 @@ def nudge_enabled(request_flag: Optional[bool]) -> bool:
return _NUDGE_DEFAULT if request_flag is None else bool(request_flag)
_MAX_SIGNAL_LEN = max(len(s) for s in TOOL_XML_SIGNALS)
_MAX_SIGNAL_LEN = max(len(s) for s in _HEAL_SIGNALS)
# A suspected-but-unclosed tool block larger than this is declared a false
# alarm and flushed, bounding memory on a model rambling XML-lookalike text.
_MAX_HOLD_CHARS = 64 * 1024
@ -198,7 +216,7 @@ def heal_openai_message_events(
if not isinstance(msg, dict) or msg.get("tool_calls"):
return None
content = msg.get("content")
if not isinstance(content, str) or not has_tool_signal(content):
if not isinstance(content, str) or not _has_heal_signal(content):
return None
parsed, spans = parse_tool_calls_from_text(content, allow_incomplete = True, with_spans = True)
tool_schemas = _tool_schemas_by_name(tools) if tools is not None else None
@ -248,7 +266,7 @@ def heal_openai_message(
def _earliest_signal(buffer: str) -> int:
best = -1
for signal in TOOL_XML_SIGNALS:
for signal in _HEAL_SIGNALS:
index = buffer.find(signal)
if index >= 0 and (best < 0 or index < best):
best = index
@ -275,7 +293,7 @@ def _partial_signal_suffix(buffer: str) -> int:
"""Length of the longest buffer suffix that is a proper prefix of a signal."""
for length in range(min(len(buffer), _MAX_SIGNAL_LEN - 1), 0, -1):
tail = buffer[-length:]
if any(signal.startswith(tail) for signal in TOOL_XML_SIGNALS):
if any(signal.startswith(tail) for signal in _HEAL_SIGNALS):
return length
return 0
@ -340,9 +358,10 @@ class StreamToolCallHealer:
events.append(("text", emit))
self._buffer = self._buffer[len(self._buffer) - keep :]
return events
# HOLD: handle the FIRST complete block per pass so events keep
# document order (a later declared call must not overtake an
# earlier undeclared one flushing as text).
# HOLD: drain the first contiguous run per pass so events keep document
# order (a later declared call must not overtake an earlier undeclared one
# flushing as text). A run is one markup call OR a whole Mistral [TOOL_CALLS]
# array of contiguous spans, so later calls in it are not stranded as text.
parsed, spans = parse_tool_calls_from_text(
self._buffer,
id_offset = self._id_offset,
@ -363,26 +382,32 @@ class StreamToolCallHealer:
self._holding = False
continue
return events
start, end = spans[0]
promoted = _promote(
[parsed[0]],
self._allowed,
id_offset = self._id_offset,
tool_schemas = self._tool_schemas,
)
if promoted:
if start:
events.append(("text", self._buffer[:start]))
events.append(("tool_call", promoted[0]))
self._id_offset += 1
# Drop exactly the promoted markup span; everything else
# (leading text, later blocks) stays and is rescanned.
self._buffer = self._buffer[end:]
else:
# Undeclared or unusable name: its markup is DATA, flush it
# (and anything before it) verbatim, then rescan the rest.
events.append(("text", self._buffer[:end]))
self._buffer = self._buffer[end:]
pos = 0
run_end = spans[0][1]
for order, (call, (start, end)) in enumerate(zip(parsed, spans)):
# Stop at the first gap or incomplete trailing block: leave it for the
# next pass to re-hold and stream incrementally, not flush as text early.
if order and start != run_end:
break
promoted = _promote(
[call],
self._allowed,
id_offset = self._id_offset,
tool_schemas = self._tool_schemas,
)
if promoted:
# Flush any leading text, then drop the promoted markup span.
if self._buffer[pos:start]:
events.append(("text", self._buffer[pos:start]))
events.append(("tool_call", promoted[0]))
self._id_offset += 1
else:
# Undeclared/unusable name: markup is DATA, flush it (and prior text) verbatim.
events.append(("text", self._buffer[pos:end]))
pos = end
run_end = end
# Everything past the drained run (later blocks) stays and is rescanned.
self._buffer = self._buffer[run_end:]
self._holding = False
def finalize(self) -> list:
@ -508,7 +533,7 @@ def nudge_should_retry(
if not message or message.get("tool_calls"):
return False
text = message.get("content")
if not isinstance(text, str) or not has_tool_signal(text):
if not isinstance(text, str) or not _has_heal_signal(text):
return False
return not _heal_would_promote(text, allowed_tools, tools)

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

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

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

@ -122,12 +122,12 @@ def calculate_cost(provider: str, model: str, usage: dict[str, Any]) -> dict[str
"priced": bool(prices),
}
# Accept raw (input_tokens/output_tokens) and Studio chat-style
# Accept raw (input_tokens/output_tokens) and Unsloth chat-style
# (prompt_tokens/completion_tokens) envelopes. Cache buckets differ:
# raw Anthropic: input_tokens EXCLUDES cache buckets
# raw OpenAI: input_tokens INCLUDES cache_read
# Studio Anthropic: prompt_tokens INCLUDES cache_creation + cache_read
# Studio OpenAI: prompt_tokens == raw input_tokens
# Unsloth Anthropic: prompt_tokens INCLUDES cache_creation + cache_read
# Unsloth OpenAI: prompt_tokens == raw input_tokens
# Clamp >=0 so corrupted payloads can't produce a negative bill.
cache_creation = max(0, int(usage.get("cache_creation_input_tokens") or 0))
cache_read_native_present = (
@ -160,7 +160,7 @@ def calculate_cost(provider: str, model: str, usage: dict[str, Any]) -> dict[str
output_tokens = max(0, int(usage.get("completion_tokens") or 0))
if provider == "openai":
# Cached tokens land on input_tokens_details (raw Responses) or
# prompt_tokens_details (Studio chat-style).
# prompt_tokens_details (Unsloth chat-style).
for key in ("input_tokens_details", "prompt_tokens_details"):
details = usage.get(key) or {}
if isinstance(details, dict):

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

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

File diff suppressed because it is too large Load diff

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 Unsloth 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

876
studio/backend/core/inference/stt_ggml_sidecar.py Normal file
View file

@ -0,0 +1,876 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""whisper.cpp (GGML/GGUF) speech-to-text sidecar for Studio dictation.
Runs the same curated Whisper checkpoints as the Transformers sidecar
(stt_sidecar.py) through whisper.cpp's `whisper-server`, ~2.5x faster at
identical quality on Apple Silicon and CPU because its Metal/CPU kernels run
the weights in f16 where PyTorch MPS requires fp32.
Owns a single `whisper-server` subprocess bound to 127.0.0.1 on an ephemeral
port; the model loads on demand, stays warm between dictations, and unloads
after the same keep-alive as the Transformers sidecar. Curated GGML checkpoints
are single files from `unslothai/whisper-*-GGUF`, downloaded directly rather
than through the Model Hub (whose variant planner only handles `.gguf` chat
layouts).
Binary discovery mirrors `_find_llama_server_binary`: env override, then managed
Studio home, then PATH. With no binary the engine is unavailable and dictation
falls back to the Transformers sidecar; `scripts/build_whisper_cpp.sh` installs
the binary.
"""
from __future__ import annotations
import io
import json
import os
import re
import shutil
import socket
import subprocess
import sys
import threading
import time
import urllib.request
import uuid
import wave
from contextlib import contextmanager
from pathlib import Path
from typing import Iterator, Optional
from loggers import get_logger
from core.inference.stt_sidecar import (
STT_KEEP_ALIVE_SECONDS,
SttAudioDecodeError,
SttLanguageError,
SttLoadCancelledError,
SttModelIdError,
SttModelNotDownloadedError,
SttUnavailableError,
_decode_audio_bounded,
_known_whisper_languages,
_TARGET_SAMPLE_RATE,
_training_active,
normalize_whisper_language,
)
from utils.prebuilt.child_env import isolate_home, scrub_env, wsl_system_rocm_lib_dirs
from utils.prebuilt.runtime_libs import dedupe_existing_dirs
from utils.prebuilt.whisper_layout import lookup_marker
from utils.process_lifetime import adopt_pid, child_popen_kwargs, forget_pid
logger = get_logger(__name__)
# Curated GGML checkpoints, one repo per model. Keys match the Transformers
# sidecar's ids so the frontend reuses one picker; values are the single file
# inside each repo.
GGML_STT_REPOS: dict[str, str] = {
"tiny": "unslothai/whisper-tiny-GGUF",
"base": "unslothai/whisper-base-GGUF",
"small": "unslothai/whisper-small-GGUF",
"large-v3-turbo": "unslothai/whisper-large-v3-turbo-GGUF",
"large-v3": "unslothai/whisper-large-v3-GGUF",
}
GGML_STT_MODELS: dict[str, str] = {
"tiny": "whisper-tiny.bin",
"base": "whisper-base.bin",
"small": "whisper-small.bin",
"large-v3-turbo": "whisper-large-v3-turbo.bin",
"large-v3": "whisper-large-v3.bin",
}
DEFAULT_GGML_STT_MODEL = "small"
_SERVER_START_TIMEOUT_SECONDS = 120.0
_TRANSCRIBE_TIMEOUT_SECONDS = 600.0
class SttEngineUnavailableError(SttUnavailableError):
"""whisper-server is not installed; the GGUF dictation engine is off."""
def resolve_ggml_model_id(model: Optional[str]) -> str:
"""Validate a curated GGML model id. Custom repos are not supported here."""
if model is None or not str(model).strip():
return DEFAULT_GGML_STT_MODEL
normalized = str(model).strip()
if normalized in GGML_STT_MODELS:
return normalized
raise SttModelIdError(
f"STT model '{model}' is not a curated GGUF dictation model. "
f"Choose one of: {', '.join(GGML_STT_MODELS)}."
)
def _managed_whisper_cpp_dir() -> Path:
"""`<STUDIO_HOME>/whisper.cpp` in custom mode, else `~/.unsloth/whisper.cpp`.
Mirrors `managed_node_dir` / `_find_llama_server_binary` so managed runtimes
share one parent directory.
"""
legacy = Path.home() / ".unsloth" / "whisper.cpp"
try:
from utils.paths.storage_roots import studio_root
resolved = studio_root()
legacy_studio = Path.home() / ".unsloth" / "studio"
try:
is_legacy = resolved.resolve() == legacy_studio.resolve()
except (OSError, ValueError):
is_legacy = resolved == legacy_studio
return legacy if is_legacy else (resolved / "whisper.cpp")
except (ImportError, OSError, ValueError):
override = (
os.environ.get("UNSLOTH_STUDIO_HOME") or os.environ.get("STUDIO_HOME") or ""
).strip()
if override:
try:
return Path(override).expanduser().resolve() / "whisper.cpp"
except (OSError, ValueError):
return Path(override).expanduser() / "whisper.cpp"
return legacy
def find_whisper_server_binary() -> Optional[str]:
"""Locate the whisper-server binary.
Search order:
1. WHISPER_SERVER_PATH environment variable (direct path to binary)
2. UNSLOTH_WHISPER_CPP_PATH env var (custom whisper.cpp install dir)
3. managed dir: <STUDIO_HOME or ~/.unsloth>/whisper.cpp/{,build/bin/}whisper-server
4. whisper-server on PATH
"""
binary_name = "whisper-server.exe" if sys.platform == "win32" else "whisper-server"
def _layout_candidates(d: Path) -> list[Path]:
cands = [d / binary_name, d / "build" / "bin" / binary_name]
if sys.platform == "win32":
cands.append(d / "build" / "bin" / "Release" / binary_name)
return cands
env_path = os.environ.get("WHISPER_SERVER_PATH")
if env_path:
p = Path(env_path)
if _is_runnable(p):
return str(p)
custom_dir = os.environ.get("UNSLOTH_WHISPER_CPP_PATH")
if custom_dir:
for p in _layout_candidates(Path(custom_dir)):
if _is_runnable(p):
return str(p)
for p in _layout_candidates(_managed_whisper_cpp_dir()):
if _is_runnable(p):
return str(p)
return shutil.which(binary_name)
def _is_runnable(p: Path) -> bool:
"""A real whisper-server is an executable file. On Windows os.access(X_OK) is
effectively an existence check; on Unix it rejects a non-executable stub so a
half-written or wrong-mode file isn't mistaken for the server."""
return p.is_file() and (sys.platform == "win32" or os.access(p, os.X_OK))
def _whisper_install_marker(binary: str) -> Optional[dict]:
"""The prebuilt install marker above ``binary``, or None (source/custom builds)."""
return lookup_marker(binary).marker
def slim_runtime_intact(binary: str) -> bool:
"""True unless the marker says slim and the linked ggml runtime is missing
beside the server. New markers record the exact wired filenames
(linked_libraries), all of which must be present; legacy markers without the
field fall back to the per-OS core ggml name globs. A broken slim install
reads as engine-unavailable (reinstall via `unsloth studio update`), never a
crash at load."""
lookup = lookup_marker(binary)
marker = lookup.marker
if lookup.invalid or marker is None:
return not lookup.slim_collision
if not marker or marker.get("install_kind") != "slim":
return True
if lookup.authoritative:
valid = marker.get("component") == "whisper.cpp"
valid = valid and isinstance(marker.get("schema_version"), int)
valid = valid and all(
isinstance(marker.get(key), str) and marker[key]
for key in ("release_tag", "backend", "paired_llama_tag")
)
valid = valid and isinstance(marker.get("linked_libraries"), list)
valid = valid and bool(marker.get("linked_libraries"))
valid = valid and all(
isinstance(name, str) and name and Path(name).name == name
for name in marker["linked_libraries"]
)
if not valid:
return False
bin_dir = Path(binary).parent
linked = marker.get("linked_libraries")
if isinstance(linked, list) and linked and all(isinstance(name, str) for name in linked):
intact = all((bin_dir / name).is_file() for name in linked)
else:
if sys.platform == "win32":
required = ("ggml.dll", "ggml-base.dll")
elif sys.platform == "darwin":
required = ("libggml*.dylib", "libggml-base*.dylib")
else:
required = ("libggml.so*", "libggml-base.so*")
intact = all(any(p.is_file() for p in bin_dir.glob(pattern)) for pattern in required)
runtime_dirs = marker.get("linked_runtime_directories")
if intact and isinstance(runtime_dirs, list) and runtime_dirs:
intact = all(
isinstance(name, str)
and name
and (bin_dir / name).is_dir()
and any(path.is_file() for path in (bin_dir / name).rglob("*"))
for name in runtime_dirs
)
if intact and marker.get("backend") == "rocm":
expected_runtime_dirs = set() if sys.platform == "win32" else {"hipblaslt", "rocblas"}
intact = (
marker.get("runtime_wiring_version") == 2
and isinstance(runtime_dirs, list)
and set(runtime_dirs) == expected_runtime_dirs
)
if not intact:
logger.warning(
"slim whisper install is missing its linked ggml runtime at "
f"{bin_dir}; run `unsloth studio update` to reinstall it"
)
return intact
def is_available() -> bool:
binary = find_whisper_server_binary()
if binary is None:
return False
if not slim_runtime_intact(binary):
return False
try:
import av # noqa: F401
except Exception:
# No PyAV means every transcription 501s on decode.
return False
return True
def ensure_engine_available() -> str:
binary = find_whisper_server_binary()
if binary is None:
raise SttEngineUnavailableError(
"The local transcription runtime is not installed. Run "
"`unsloth studio update` to install it."
)
if not slim_runtime_intact(binary):
raise SttEngineUnavailableError(
"The local transcription runtime is missing its paired ggml "
"libraries. Run `unsloth studio update` to reinstall it."
)
return binary
# ---------------------------------------------------------------------------
# whisper-server child-process environment
# ---------------------------------------------------------------------------
# Build the whisper-server env: prepend the binary dir (co-located libs win, and
# a backstop where the loader ignores the rpath) and scrub secret-bearing vars the
# binary never needs. On WSL2 ROCm the system HIP libs go first, since a bundle's
# bare-metal HIP cannot drive /dev/dxg. A CUDA bundle ships libggml-cuda.so but not
# libcudart/libcublas (paired with the user's PyTorch), so add the
# CUDA-from-PyTorch runtime dirs the selection gated on, else the backend cannot
# resolve a runtime that lives only in wheels. Mirrors llama's binary_env(); the
# scrub/WSL/dedupe helpers live in utils.prebuilt.
# Module-level aliases keep the historical patch points for tests and callers.
_wsl_system_rocm_lib_dirs = wsl_system_rocm_lib_dirs
_dedupe_existing_dirs = dedupe_existing_dirs
def _whisper_server_child_env(binary: str) -> dict[str, str]:
"""Env for the whisper-server subprocess: secrets scrubbed, home/profile vars
repointed at a managed scratch dir (a downloaded binary must not see the real
home's token caches), co-located libs on the loader path, WSL system HIP first
on WSL2 ROCm."""
env = scrub_env(os.environ)
isolate_home(env, str(_managed_whisper_cpp_dir() / ".child_home"))
bin_dir = str(Path(binary).parent)
# A CUDA bundle needs the CUDA-from-PyTorch wheel dirs so libcudart/libcublas
# resolve at launch when they live only in site-packages/nvidia/*/lib. Placed
# after bin_dir so co-located libs still win; empty for other bundles.
cuda_runtime_dirs: list[str] = []
bundle_dir = Path(bin_dir)
has_cuda_module = any(
path.is_file()
for pattern in ("libggml-cuda.so*", "ggml-cuda*.dll")
for path in bundle_dir.glob(pattern)
)
if has_cuda_module:
try:
from utils.prebuilt.runtime_libs import python_runtime_dirs
cuda_runtime_dirs = python_runtime_dirs()
except Exception:
cuda_runtime_dirs = []
if sys.platform == "win32":
var, lead = "PATH", [bin_dir, *cuda_runtime_dirs]
elif sys.platform == "darwin":
var, lead = "DYLD_LIBRARY_PATH", [bin_dir]
else:
var, lead = "LD_LIBRARY_PATH", [bin_dir, *cuda_runtime_dirs]
wsl_rocm = _wsl_system_rocm_lib_dirs()
if wsl_rocm:
lead = [*wsl_rocm, bin_dir, *cuda_runtime_dirs]
env.setdefault("HSA_ENABLE_DXG_DETECTION", "1")
existing = [p for p in env.get(var, "").split(os.pathsep) if p]
env[var] = os.pathsep.join(_dedupe_existing_dirs([*lead, *existing]))
return env
# ---------------------------------------------------------------------------
# Model file download (single files; deliberately outside the Model Hub flow)
# ---------------------------------------------------------------------------
def _cached_model_path(model_id: str) -> Optional[str]:
"""Path of a fully downloaded GGML file in the shared HF cache, else None."""
from huggingface_hub import hf_hub_download
try:
return hf_hub_download(
repo_id = GGML_STT_REPOS[model_id],
filename = GGML_STT_MODELS[model_id],
local_files_only = True,
)
except Exception:
return None
class _GgmlDownloadState:
"""Tracks one background hf_hub_download of a curated GGML file."""
def __init__(self) -> None:
self._lock = threading.Lock()
self._thread: Optional[threading.Thread] = None
self._model_id: Optional[str] = None
self._error: Optional[str] = None
self._total_bytes: Optional[int] = None
self._etag: Optional[str] = None
def status(self) -> dict:
with self._lock:
downloading = self._thread is not None and self._thread.is_alive()
return {
"downloading": downloading,
"model": self._model_id if downloading else None,
"error": self._error,
"bytes_total": self._total_bytes if downloading else None,
"bytes_done": self._incomplete_bytes() if downloading else None,
}
def _incomplete_bytes(self) -> Optional[int]:
"""Best-effort progress: size of the in-flight blob in the HF cache.
hf_hub_download writes ``blobs/<etag>.incomplete``; prefer this file's
etag, else the largest in-flight blob.
"""
try:
from huggingface_hub.constants import HF_HUB_CACHE
# Caller may hold the non-reentrant self._lock; bare reads are safe.
model_id = self._model_id
if not model_id:
return None
repo_dir = (
Path(HF_HUB_CACHE)
/ f"models--{GGML_STT_REPOS[model_id].replace('/', '--')}"
/ "blobs"
)
if not repo_dir.is_dir():
return None
etag = self._etag
if etag:
target = repo_dir / f"{etag}.incomplete"
if target.is_file():
return target.stat().st_size
sizes = [p.stat().st_size for p in repo_dir.glob("*.incomplete") if p.is_file()]
return max(sizes) if sizes else None
except Exception:
return None
def start(
self,
model_id: str,
hf_token: Optional[str] = None,
) -> None:
model_id = resolve_ggml_model_id(model_id)
with self._lock:
if self._thread is not None and self._thread.is_alive():
if self._model_id == model_id:
return
raise SttModelIdError(
f"Another GGUF dictation model ('{self._model_id}') is still "
"downloading; wait for it to finish."
)
self._model_id = model_id
self._error = None
self._total_bytes = None
self._etag = None
thread = threading.Thread(target = self._run, args = (model_id, hf_token), daemon = True)
self._thread = thread
thread.start()
def _run(self, model_id: str, hf_token: Optional[str]) -> None:
repo_id = GGML_STT_REPOS[model_id]
filename = GGML_STT_MODELS[model_id]
try:
from huggingface_hub import (
get_hf_file_metadata,
hf_hub_download,
hf_hub_url,
)
try:
# One HEAD request for the total and etag.
meta = get_hf_file_metadata(hf_hub_url(repo_id, filename), token = hf_token or None)
with self._lock:
self._total_bytes = meta.size
self._etag = meta.etag
except Exception:
pass
hf_hub_download(
repo_id = repo_id,
filename = filename,
token = hf_token or None,
)
except Exception as exc:
logger.warning("GGUF STT download failed for %s: %s", model_id, exc)
with self._lock:
self._error = f"Download failed for '{model_id}'."
_download_state = _GgmlDownloadState()
def start_model_download(model: Optional[str], hf_token: Optional[str] = None) -> None:
_download_state.start(resolve_ggml_model_id(model), hf_token)
def download_status() -> dict:
return _download_state.status()
# ---------------------------------------------------------------------------
# WAV packaging
# ---------------------------------------------------------------------------
def _pcm_to_wav_bytes(decoded_audio) -> bytes:
"""Wrap decoded float32 mono 16 kHz PCM into an in-memory 16-bit WAV."""
import numpy as np
clipped = np.clip(decoded_audio, -1.0, 1.0)
pcm16 = (clipped * 32767.0).astype("<i2")
buf = io.BytesIO()
with wave.open(buf, "wb") as w:
w.setnchannels(1)
w.setsampwidth(2)
w.setframerate(_TARGET_SAMPLE_RATE)
w.writeframes(pcm16.tobytes())
return buf.getvalue()
# ---------------------------------------------------------------------------
# Sidecar
# ---------------------------------------------------------------------------
class GgmlSttSidecar:
"""Owns one whisper-server subprocess and proxies dictation to it."""
def __init__(self, keep_alive_seconds: float = STT_KEEP_ALIVE_SECONDS) -> None:
self._lock = threading.RLock()
self._process: Optional[subprocess.Popen] = None
self._port: Optional[int] = None
self._model_id: Optional[str] = None
self._idle_timer: Optional[threading.Timer] = None
self._idle_generation = 0
self._keep_alive_seconds = keep_alive_seconds
# Set while whisper-server starts so training admission can account for
# the accelerator memory it is about to bind. Read without the lock.
self._loading = False
# A still-starting whisper-server is cancellable so training can preempt
# it before it binds accelerator memory. Assigned inside self._lock but
# acted on without it: cancel_pending_load() runs while load() holds the
# lock, so the event is the source of truth and terminating the process
# is a best-effort fast path.
self._load_cancel_event: Optional[threading.Event] = None
self._starting_process: Optional[subprocess.Popen] = None
# Set before the updater waits for _lock, then kept set while it owns
# the lock and atomically replaces the managed install tree. New loads
# fail fast instead of starting a process from files being swapped.
self._update_in_progress = False
@property
def loaded_model(self) -> Optional[str]:
# Lock-free status read (like stt_sidecar.py): transcribe() holds
# self._lock for the whole inference call (up to
# _TRANSCRIBE_TIMEOUT_SECONDS), and status polls plus training admission
# must not block behind it. _process_alive() snapshots self._process
# before poll(), which subprocess guards with _waitpid_lock, so a
# concurrent unload is safe.
return self._model_id if self._process_alive() else None
@property
def device(self) -> Optional[str]:
return "whisper.cpp" if self._process_alive() else None
def is_loading(self) -> bool:
# True only while whisper-server is starting (seconds to bind its GPU
# backend); load() sets and clears the flag around that window.
return self._loading
@property
def keep_alive_seconds(self) -> float:
return self._keep_alive_seconds
def _process_alive(self) -> bool:
# Snapshot self._process once: a concurrent unload() nulls it under the
# lock, so lock-free readers would otherwise re-read None between the
# truthiness check and .poll().
process = self._process
return process is not None and process.poll() is None
# -- idle unload ------------------------------------------------------
def _cancel_idle_unload_locked(self) -> None:
self._idle_generation += 1
if self._idle_timer is not None:
self._idle_timer.cancel()
self._idle_timer = None
def _schedule_idle_unload_locked(self) -> None:
self._cancel_idle_unload_locked()
if not self._process_alive():
return
generation = self._idle_generation
timer = threading.Timer(self._keep_alive_seconds, self._idle_unload, args = (generation,))
timer.daemon = True
self._idle_timer = timer
timer.start()
def _idle_unload(self, generation: int) -> None:
with self._lock:
if generation != self._idle_generation:
return
logger.info("Unloading idle GGUF STT model %s", self._model_id)
self._release_locked()
# -- process lifecycle -------------------------------------------------
def _release_locked(self) -> None:
self._cancel_idle_unload_locked()
process = self._process
self._process = None
self._port = None
self._model_id = None
if process is not None and process.poll() is None:
process.terminate()
try:
process.wait(timeout = 10)
except subprocess.TimeoutExpired:
process.kill()
process.wait(timeout = 10)
if process is not None:
forget_pid(process.pid)
def unload(self) -> None:
with self._lock:
self._release_locked()
def _raise_if_update_in_progress(self) -> None:
if self._update_in_progress:
raise SttEngineUnavailableError(
"The local transcription runtime is being updated. Try dictation again shortly."
)
@contextmanager
def update_maintenance(self) -> Iterator[bool]:
"""Block new loads while the managed whisper.cpp tree is replaced.
The flag is published before waiting for an existing transcription to
release ``_lock``. Holding that lock across the yielded installer phase
prevents Windows from relocking the executable and prevents every host
from starting a process against a partially swapped tree. The yielded
value records whether a warm model had to be unloaded.
"""
self._update_in_progress = True
try:
with self._lock:
model_was_active = self._process_alive()
self._release_locked()
yield model_was_active
finally:
self._update_in_progress = False
def cancel_pending_load(self) -> bool:
# Preempt a starting whisper-server so training does not launch while it
# binds accelerator memory. load() holds self._lock for the whole startup,
# so act without the lock: signal abort and terminate the starting
# process. _wait_for_server observes the event and raises, then load()
# reaps the process and releases the lock.
if not self._loading:
return False
event = self._load_cancel_event
if event is None:
return False
event.set()
process = self._starting_process
if process is not None and process.poll() is None:
try:
process.terminate()
except Exception:
pass
return True
def wait_for_load_to_settle(self) -> None:
# load() holds self._lock across startup and cancel cleanup, so acquiring
# it blocks until a cancelled server is killed, reaped, and its
# accelerator memory released.
with self._lock:
pass
@staticmethod
def _reserve_free_port() -> tuple[socket.socket, int]:
"""Bind an ephemeral port and keep the socket held.
The caller closes the reservation immediately before spawning
whisper-server, shrinking the window in which another local process
could bind the port. SO_REUSEADDR lets the child rebind right after.
"""
s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
s.bind(("127.0.0.1", 0))
return s, s.getsockname()[1]
def _ensure_model_downloaded(self, model_id: str) -> str:
path = _cached_model_path(model_id)
if path is None:
raise SttModelNotDownloadedError(
f"STT model '{model_id}' (GGUF) is not downloaded. "
"Download it in Settings, then Voice, before loading it."
)
return path
def load(self, model: Optional[str] = None) -> None:
"""Start (or switch) whisper-server for the requested curated model."""
self._raise_if_update_in_progress()
model_id = resolve_ggml_model_id(model)
with self._lock:
self._raise_if_update_in_progress()
binary = ensure_engine_available()
if self._process_alive() and self._model_id == model_id:
self._schedule_idle_unload_locked()
return
model_path = self._ensure_model_downloaded(model_id)
self._release_locked()
reservation, port = self._reserve_free_port()
command = [binary, "-m", model_path, "--host", "127.0.0.1", "--port", str(port)]
marker = _whisper_install_marker(binary)
if _training_active():
# Keep whisper.cpp off the accelerator during training (like the
# Transformers sidecar's CPU choice) so a mid-training dictation
# cannot reclaim the VRAM training just freed.
command.append("--no-gpu")
elif marker is not None and marker.get("backend") == "cpu":
# A deliberate CPU install must stay CPU: the slim wiring links
# every llama ggml backend (including CUDA/ROCm), so without
# this flag a cpu-selected install would still grab the GPU.
command.append("--no-gpu")
logger.info(
"Starting whisper-server for STT model %s on 127.0.0.1:%s",
model_id,
port,
)
cancel_event = threading.Event()
self._load_cancel_event = cancel_event
self._loading = True
try:
# Release the reservation as late as possible: whisper-server
# binds the port moments after this close.
reservation.close()
process = subprocess.Popen(
command,
stdout = subprocess.DEVNULL,
stderr = subprocess.DEVNULL,
stdin = subprocess.DEVNULL,
# Co-located GPU libs on the loader path (WSL system HIP first),
# secrets scrubbed from the downloaded binary's env.
env = _whisper_server_child_env(binary),
# Die with Studio (Linux PDEATHSIG, Windows job) so a crash
# never orphans a server holding the model.
**child_popen_kwargs(),
)
self._starting_process = process
adopt_pid(process.pid) # terminate_all backstop for graceful exits
try:
self._wait_for_server(process, port, cancel_event)
except Exception:
if process.poll() is None:
process.kill()
process.wait(timeout = 10)
forget_pid(process.pid)
raise
self._process = process
self._port = port
self._model_id = model_id
self._schedule_idle_unload_locked()
finally:
reservation.close() # no-op when already released before spawn
self._loading = False
self._load_cancel_event = None
self._starting_process = None
@staticmethod
def _wait_for_server(
process: subprocess.Popen,
port: int,
cancel_event: Optional[threading.Event] = None,
) -> None:
deadline = time.monotonic() + _SERVER_START_TIMEOUT_SECONDS
while time.monotonic() < deadline:
if cancel_event is not None and cancel_event.is_set():
raise SttLoadCancelledError(
"GGUF STT model loading was cancelled so training could start."
)
if process.poll() is not None:
raise SttEngineUnavailableError(
"The local transcription runtime exited before becoming "
"ready; the model file may be corrupt or unsupported."
)
# Require a whisper-server-specific response twice, with the managed
# child alive around each probe. An arbitrary local process that won
# the bind race would otherwise be mistaken for the sidecar and
# receive the user's microphone audio.
if GgmlSttSidecar._probe_is_whisper_server(process, port) and (
GgmlSttSidecar._probe_is_whisper_server(process, port)
):
return
time.sleep(0.2)
raise SttEngineUnavailableError("The local transcription runtime did not start in time.")
@staticmethod
def _probe_is_whisper_server(process: subprocess.Popen, port: int) -> bool:
"""One readiness probe: our child is alive and the responder looks like
whisper.cpp's server (its index page and errors identify whisper)."""
if process.poll() is not None:
return False
try:
req = urllib.request.Request(f"http://127.0.0.1:{port}/", method = "GET")
with urllib.request.urlopen(req, timeout = 2) as response:
body = response.read(65536)
except Exception:
return False
if process.poll() is not None:
return False
return b"whisper" in body.lower()
# -- transcription ------------------------------------------------------
def transcribe(
self,
audio: bytes,
model: Optional[str] = None,
language: Optional[str] = None,
fast: bool = False,
) -> dict:
"""Transcribe encoded audio bytes via whisper-server.
Accepts any container PyAV can decode (same validation and caps as the
Transformers sidecar). Returns {text, language, duration, model}.
"""
self._raise_if_update_in_progress()
ensure_engine_available()
model_id = resolve_ggml_model_id(model)
lang = normalize_whisper_language(language)
known_languages = _known_whisper_languages()
if lang is not None and known_languages is not None and lang not in known_languages:
raise SttLanguageError(
f"Language '{language}' is not supported by STT model '{model_id}'."
)
# Reject a missing model before decoding so a long clip does not burn CPU
# only to 409 (matches the Transformers sidecar's preflight).
self._ensure_model_downloaded(model_id)
decoded_audio = _decode_audio_bounded(audio)
wav_bytes = _pcm_to_wav_bytes(decoded_audio)
with self._lock:
try:
self.load(model_id)
text = self._post_inference(wav_bytes, lang, fast)
finally:
self._schedule_idle_unload_locked()
duration = (len(decoded_audio) / _TARGET_SAMPLE_RATE) if len(decoded_audio) else None
return {
"text": text,
"language": lang,
"duration": duration,
"model": model_id,
}
def _post_inference(self, wav_bytes: bytes, lang: Optional[str], fast: bool) -> str:
boundary = uuid.uuid4().hex
fields = {
"temperature": "0.0",
"response_format": "json",
# Match the Transformers sidecar: 5-way beam search, greedy for fast.
"beam_size": "1" if fast else "5",
"language": lang or "auto",
}
parts: list[bytes] = []
for name, value in fields.items():
parts.append(
(
f"--{boundary}\r\nContent-Disposition: form-data; "
f'name="{name}"\r\n\r\n{value}\r\n'
).encode()
)
parts.append(
(
f"--{boundary}\r\nContent-Disposition: form-data; "
'name="file"; filename="dictation.wav"\r\n'
"Content-Type: audio/wav\r\n\r\n"
).encode()
+ wav_bytes
+ b"\r\n"
)
parts.append(f"--{boundary}--\r\n".encode())
body = b"".join(parts)
req = urllib.request.Request(
f"http://127.0.0.1:{self._port}/inference",
data = body,
headers = {"Content-Type": f"multipart/form-data; boundary={boundary}"},
)
try:
with urllib.request.urlopen(req, timeout = _TRANSCRIBE_TIMEOUT_SECONDS) as resp:
payload = json.load(resp)
except SttAudioDecodeError:
raise
except Exception as exc:
raise SttEngineUnavailableError(
"The local transcription runtime did not answer the request."
) from exc
text = payload.get("text")
if not isinstance(text, str):
raise SttAudioDecodeError("Could not decode the audio.")
# whisper.cpp joins segments with newlines; dictation wants one line.
return " ".join(part.strip() for part in text.splitlines() if part.strip()).strip()
_sidecar: Optional[GgmlSttSidecar] = None
def get_ggml_stt_sidecar() -> GgmlSttSidecar:
global _sidecar
if _sidecar is None:
_sidecar = GgmlSttSidecar()
return _sidecar

1142
studio/backend/core/inference/stt_sidecar.py Normal file
View file

File diff suppressed because it is too large Load diff

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

File diff suppressed because it is too large Load diff

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

@ -1,7 +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
"""Shared controller state for Studio local agentic tool loops.
"""Shared controller state for Unsloth local agentic tool loops.
This module is intentionally dependency-light: it owns only per-response
ledger state and value objects used by the GGUF and safetensors loops.
@ -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")

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

File diff suppressed because it is too large Load diff

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