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

Merge remote-tracking branch 'origin/diffusion-phase16-native-engine-routing' into diffusion-image-workflows

Browse source 
# Conflicts:
#	studio/backend/core/inference/diffusion.py
#	studio/backend/core/inference/diffusion_families.py
#	studio/backend/tests/test_sd_cpp_install.py
#	studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx
#	studio/frontend/src/features/images/api.ts
#	studio/frontend/src/features/images/images-page.tsx
#	studio/install_sd_cpp_prebuilt.py
This commit is contained in:
Daniel Han 2026-07-01 11:48:33 +00:00
commit 38ed3ce5b5
293 changed files with 23437 additions and 2827 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -209,7 +209,7 @@ jobs:
'peft>=0.18,<0.20' 'accelerate>=0.34,<2' \
ipython
# torchvision: unsloth_zoo.vision_utils imports it at module scope.
pip install --index-url https://download.pytorch.org/whl/cpu \
pip install --index-url https://download.pytorch.org/whl/cpu --extra-index-url https://pypi.org/simple \
'torch>=2.4,<2.11' 'torchvision<0.26'
# transformers + trl from the matrix combo.
pip install "$RESOLVED_TRANSFORMERS_SPEC"
@ -268,6 +268,10 @@ jobs:
tests/saving/test_save_shell_injection.py \
tests/saving/test_patch_saving_none_tokenizer.py \
tests/saving/test_fix_sentencepiece_gguf_robustness.py \
tests/saving/test_compressed_export_schemes.py \
tests/saving/test_export_api_surface.py \
tests/saving/test_export_dispatch.py \
tests/saving/test_imatrix_export.py \
tests/utils/test_attention_masks.py \
tests/utils/test_trunc_normal_patch.py \
tests/python/test_fast_language_model_text_only.py
@ -353,6 +357,10 @@ jobs:
tests/saving/test_save_shell_injection.py \
tests/saving/test_patch_saving_none_tokenizer.py \
tests/saving/test_fix_sentencepiece_gguf_robustness.py \
tests/saving/test_compressed_export_schemes.py \
tests/saving/test_export_api_surface.py \
tests/saving/test_export_dispatch.py \
tests/saving/test_imatrix_export.py \
tests/utils/test_attention_masks.py \
tests/utils/test_trunc_normal_patch.py \
tests/python/test_fast_language_model_text_only.py \
@ -2166,7 +2174,7 @@ jobs:
python -m pip install --upgrade pip
# Match the matrix job's torch path so unsloth_zoo's
# `import torch` resolves to the same CPU build.
pip install --index-url https://download.pytorch.org/whl/cpu \
pip install --index-url https://download.pytorch.org/whl/cpu --extra-index-url https://pypi.org/simple \
'torch>=2.4,<2.11' 'torchvision<0.26'
pip install \
'numpy<3' protobuf sentencepiece \

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

@ -163,7 +163,7 @@ jobs:
'pytest==9.0.3' \
'pytest-asyncio==1.3.0' \
'httpx==0.28.1'
pip install --index-url https://download.pytorch.org/whl/cpu \
pip install --index-url https://download.pytorch.org/whl/cpu --extra-index-url https://pypi.org/simple \
'torch==2.10.0'
# github.com occasionally 500s on the git fetch; retry the
# zoo install so a single upstream blip does not fail CI.
@ -231,99 +231,6 @@ jobs:
tests/studio/test_is_mlx_dispatch_gate.py \
tests/studio/test_mlx_training_worker_behaviors.py
# Studio prebuilt llama.cpp install + GGUF inference. Mirrors the
# path Studio's setup.sh takes on macOS since #5963: plan against
# the unslothai/llama.cpp fork's latest release, which ships the
# bin-macos-arm64 bundle plus the llama-prebuilt-manifest.json the
# default policy reads. After install, downloads a small published
# GGUF (unsloth/gemma-3-270m-it-GGUF, Q4_K_M) and validates
# llama-server /completion end to end. An install failure or a
# non-zero binary exit is an Unsloth/Studio bug.
- name: Studio prebuilt llama.cpp install + GGUF inference (Mac M1)
env:
# Withheld on PR: this step runs checked-out PR code; public GGUF still downloads.
HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }}
# install_llama_prebuilt.py hits the GitHub releases API to
# resolve the asset URL. Anonymous calls share the runner-IP
# rate-limit bucket and 403 quickly -- pass the workflow's
# automatic GITHUB_TOKEN to bump us to the 5000/hr authenticated
# bucket.
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
run: |
set -euo pipefail
INSTALL_DIR="$HOME/.unsloth-studio-prebuilt-test/llama.cpp"
rm -rf "$INSTALL_DIR"
# Mirror studio/setup.sh on macOS (the install.sh user path):
# it plans against the unslothai/llama.cpp fork's latest
# release with no policy or tag flags.
python studio/install_llama_prebuilt.py \
--install-dir "$INSTALL_DIR" \
--published-repo unslothai/llama.cpp
# Studio bundles only llama-server + llama-quantize from the
# prebuilt (not llama-cli) -- inference goes through
# llama-server's HTTP /completion endpoint. Validate both:
# llama-quantize --help proves the dynamic libs link, then
# spin up llama-server and POST a /completion request on a
# tiny published GGUF.
LLAMA_SERVER="$INSTALL_DIR/build/bin/llama-server"
LLAMA_QUANT="$INSTALL_DIR/build/bin/llama-quantize"
[ -x "$LLAMA_SERVER" ] || { echo "::error::llama-server missing at $LLAMA_SERVER"; find "$INSTALL_DIR/build" -type f | head -40; exit 1; }
[ -x "$LLAMA_QUANT" ] || { echo "::error::llama-quantize missing at $LLAMA_QUANT"; exit 1; }
echo "llama-server : $LLAMA_SERVER"
echo "llama-quantize: $LLAMA_QUANT"
"$LLAMA_QUANT" --help >/dev/null && echo " llama-quantize loads OK"
mkdir -p /tmp/ggufs
bash .github/scripts/hf-download-with-retry.sh \
'unsloth/gemma-3-270m-it-GGUF' \
'gemma-3-270m-it-Q4_K_M.gguf' \
/tmp/ggufs
PORT=18080
echo "=== starting llama-server on 127.0.0.1:$PORT ==="
"$LLAMA_SERVER" \
-m /tmp/ggufs/gemma-3-270m-it-Q4_K_M.gguf \
--host 127.0.0.1 \
--port "$PORT" \
-c 256 \
-n 16 \
--no-warmup \
> /tmp/llama-server.log 2>&1 &
SERVER_PID=$!
trap 'kill "$SERVER_PID" 2>/dev/null || true' EXIT
# Wait for /health to come up
for i in $(seq 1 30); do
if curl -sf "http://127.0.0.1:$PORT/health" >/dev/null 2>&1; then
echo " server up after ${i}s"
break
fi
sleep 1
done
if ! curl -sf "http://127.0.0.1:$PORT/health" >/dev/null 2>&1; then
echo "::error::llama-server never became healthy"
tail -40 /tmp/llama-server.log
exit 1
fi
PROMPT="Hello, my name is"
echo "=== POST /completion ==="
RESP=$(curl -sf -X POST "http://127.0.0.1:$PORT/completion" \
-H 'Content-Type: application/json' \
-d "{\"prompt\":\"$PROMPT\",\"n_predict\":16,\"temperature\":0,\"seed\":3407}")
echo "raw response (head): $(echo "$RESP" | head -c 600)"
CONTENT=$(echo "$RESP" | python -c "import json,sys; print(json.loads(sys.stdin.read()).get('content',''))")
echo "completion content: $CONTENT"
if [ -z "$CONTENT" ]; then
echo "::error::llama-server /completion returned empty content"
tail -40 /tmp/llama-server.log
exit 1
fi
echo "OK: Studio prebuilt llama.cpp on Mac M1 + GGUF /completion works"
# Real MLX training + inference smoke test. Trains
# unsloth/gemma-3-270m-it for 7 deterministic LoRA steps
# (batch_size=2, gradient_accumulation_steps=3) on a single
@ -338,6 +245,9 @@ jobs:
UNSLOTH_COMPILE_DISABLE: '1'
run: |
mkdir -p mlx_workdir
# Authenticate llama.cpp's release-API lookup (anonymous 403s on rate-limit);
# read-only GITHUB_TOKEN scoped here only, never to steps that run binaries.
GH_TOKEN="${{ secrets.GITHUB_TOKEN }}" GITHUB_TOKEN="${{ secrets.GITHUB_TOKEN }}" \
python tests/studio/run_real_mlx_smoke.py train \
--workdir "$PWD/mlx_workdir"
@ -406,3 +316,88 @@ jobs:
cat "$f" 2>/dev/null || echo "(missing)"
echo
done
# Validates the macOS prebuilt path Studio's setup.sh uses (#5963): install the
# unslothai/llama.cpp fork's latest release, download a small public GGUF, and
# check llama-server /completion end to end. Split and placed last so the
# untrusted binary runs only in the final smoke step, after every HF_TOKEN step,
# leaving no token-bearing step or shared workspace for a tampered prebuilt to
# corrupt. GH_TOKEN: releases API; HF_TOKEN (withheld on PR): probe + GGUF fetch.
- name: Studio prebuilt llama.cpp install + GGUF download (Mac M1)
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }}
run: |
set -euo pipefail
INSTALL_DIR="$HOME/.unsloth-studio-prebuilt-test/llama.cpp"
rm -rf "$INSTALL_DIR"
# Download only -- no llama-quantize / llama-server launch in this step.
python studio/install_llama_prebuilt.py \
--install-dir "$INSTALL_DIR" \
--published-repo unslothai/llama.cpp
mkdir -p /tmp/ggufs
bash .github/scripts/hf-download-with-retry.sh \
'unsloth/gemma-3-270m-it-GGUF' \
'gemma-3-270m-it-Q4_K_M.gguf' \
/tmp/ggufs
# Final step: runs the downloaded binaries with no secrets present, and clears
# the GitHub Actions command files so a tampered prebuilt cannot influence the job.
- name: Studio prebuilt llama.cpp GGUF inference smoke (Mac M1)
run: |
set -euo pipefail
unset GITHUB_ENV GITHUB_PATH GITHUB_OUTPUT GITHUB_STEP_SUMMARY
INSTALL_DIR="$HOME/.unsloth-studio-prebuilt-test/llama.cpp"
# Studio bundles only llama-server + llama-quantize (not llama-cli);
# inference goes through llama-server's HTTP /completion endpoint.
LLAMA_SERVER="$INSTALL_DIR/build/bin/llama-server"
LLAMA_QUANT="$INSTALL_DIR/build/bin/llama-quantize"
[ -x "$LLAMA_SERVER" ] || { echo "::error::llama-server missing at $LLAMA_SERVER"; find "$INSTALL_DIR/build" -type f | head -40; exit 1; }
[ -x "$LLAMA_QUANT" ] || { echo "::error::llama-quantize missing at $LLAMA_QUANT"; exit 1; }
echo "llama-server : $LLAMA_SERVER"
echo "llama-quantize: $LLAMA_QUANT"
"$LLAMA_QUANT" --help >/dev/null && echo " llama-quantize loads OK"
PORT=18080
echo "=== starting llama-server on 127.0.0.1:$PORT ==="
"$LLAMA_SERVER" \
-m /tmp/ggufs/gemma-3-270m-it-Q4_K_M.gguf \
--host 127.0.0.1 \
--port "$PORT" \
-c 256 \
-n 16 \
--no-warmup \
> /tmp/llama-server.log 2>&1 &
SERVER_PID=$!
trap 'kill "$SERVER_PID" 2>/dev/null || true' EXIT
# Wait for /health to come up
for i in $(seq 1 30); do
if curl -sf "http://127.0.0.1:$PORT/health" >/dev/null 2>&1; then
echo " server up after ${i}s"
break
fi
sleep 1
done
if ! curl -sf "http://127.0.0.1:$PORT/health" >/dev/null 2>&1; then
echo "::error::llama-server never became healthy"
tail -40 /tmp/llama-server.log
exit 1
fi
PROMPT="Hello, my name is"
echo "=== POST /completion ==="
RESP=$(curl -sf -X POST "http://127.0.0.1:$PORT/completion" \
-H 'Content-Type: application/json' \
-d "{\"prompt\":\"$PROMPT\",\"n_predict\":16,\"temperature\":0,\"seed\":3407}")
echo "raw response (head): $(echo "$RESP" | head -c 600)"
CONTENT=$(echo "$RESP" | python -c "import json,sys; print(json.loads(sys.stdin.read()).get('content',''))")
echo "completion content: $CONTENT"
if [ -z "$CONTENT" ]; then
echo "::error::llama-server /completion returned empty content"
tail -40 /tmp/llama-server.log
exit 1
fi
echo "OK: Studio prebuilt llama.cpp on Mac M1 + GGUF /completion works"

2
.github/workflows/notebooks-ci.yml vendored
View file

@ -263,7 +263,7 @@ jobs:
# unsloth_zoo.vision_utils imports PIL at module top, and the
# easiest way to get a torch-compatible PIL on a CPU runner is
# to let torchvision pull the right Pillow version.
pip install --index-url https://download.pytorch.org/whl/cpu \
pip install --index-url https://download.pytorch.org/whl/cpu --extra-index-url https://pypi.org/simple \
'torch>=2.8,<2.11' 'torchvision<0.26'
# Pin to the same versions update_all_notebooks.py installs in
# generated notebooks. Keep these in lockstep with PIN_TRL /

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

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

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

@ -76,7 +76,7 @@ jobs:
# Torch CPU + transformers are required by a chunk of the backend test
# suite (gpu_selection, kv_cache_estimation, utils). CPU-only torch
# keeps the install ~250 MB / ~1 min on a clean runner.
pip install --index-url https://download.pytorch.org/whl/cpu 'torch>=2.4,<2.11'
pip install --index-url https://download.pytorch.org/whl/cpu --extra-index-url https://pypi.org/simple 'torch>=2.4,<2.11'
pip install 'transformers>=4.51,<5.5'
- name: Backend tests
@ -137,7 +137,7 @@ jobs:
pyyaml jinja2 mammoth unpdf requests typer \
'numpy<3' pytest pytest-asyncio httpx
# torchvision: unsloth_zoo.vision_utils imports it at module scope.
pip install --index-url https://download.pytorch.org/whl/cpu \
pip install --index-url https://download.pytorch.org/whl/cpu --extra-index-url https://pypi.org/simple \
'torch>=2.4,<2.11' 'torchvision<0.26'
pip install 'transformers>=4.51,<5.5'
# bitsandbytes: hard import in unsloth/models/_utils.py. Recent

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

@ -185,13 +185,14 @@ jobs:
# Retry up to 3 times to absorb known macos-14 free-runner
# flakes: (1) Playwright Node 24 pipeTransport.js 'Unexpected
# end of JSON input' crash when the Chromium browser process
# dies mid-test, and (2) Chromium net::ERR_NO_BUFFER_SPACE
# when the runner's kernel briefly runs out of socket buffers.
# The retry FULLY resets Studio (kill, reset-password, reboot,
# wait /api/health, re-export bootstrap pw) before re-running
# the script. A real test failure (assertion / timeout) does
# NOT match either pattern so it bypasses retry and surfaces
# immediately.
# dies mid-test, (2) Chromium net::ERR_NO_BUFFER_SPACE when the
# runner's kernel briefly runs out of socket buffers, and (3) a
# goto 'interrupted by another navigation' when the SPA auth
# guard redirects mid-navigation. The retry FULLY resets Studio
# (kill, reset-password, reboot, wait /api/health, re-export
# bootstrap pw) before re-running the script. A real test failure
# (assertion / timeout) does NOT match any pattern so it bypasses
# retry and surfaces immediately.
run: |
mkdir -p logs/playwright
attempt=1
@ -204,8 +205,9 @@ jobs:
if [ "$rc" -eq 0 ]; then
break
fi
if { grep -q "Unexpected end of JSON input" logs/playwright_attempt_${attempt}.log \
|| grep -q "ERR_NO_BUFFER_SPACE" logs/playwright_attempt_${attempt}.log; } \
if { grep -q "Unexpected end of JSON input" logs/playwright_attempt_${attempt}.log \
|| grep -q "ERR_NO_BUFFER_SPACE" logs/playwright_attempt_${attempt}.log \
|| grep -q "interrupted by another navigation" logs/playwright_attempt_${attempt}.log; } \
&& [ "$attempt" -lt "$max_attempts" ]; then
echo "::warning::Playwright flake on attempt ${attempt}; resetting Studio and retrying..."
kill "${STUDIO_PID}" 2>/dev/null || true
@ -280,8 +282,8 @@ jobs:
STUDIO_UI_TURN_TIMEOUT_MS: '540000'
GGUF_REPO: ${{ env.GGUF_REPO }}
GGUF_VARIANT: ${{ env.GGUF_VARIANT }}
# Same flake-retry shape as "Drive the chat UI with Playwright"
# -- catches pipeTransport JSON crash and ERR_NO_BUFFER_SPACE.
# Same flake-retry shape as "Drive the chat UI with Playwright" -- catches
# pipeTransport JSON crash, ERR_NO_BUFFER_SPACE, and nav interrupts.
run: |
mkdir -p logs/playwright_extra
attempt=1
@ -294,8 +296,9 @@ jobs:
if [ "$rc" -eq 0 ]; then
break
fi
if { grep -q "Unexpected end of JSON input" logs/playwright_extra_attempt_${attempt}.log \
|| grep -q "ERR_NO_BUFFER_SPACE" logs/playwright_extra_attempt_${attempt}.log; } \
if { grep -q "Unexpected end of JSON input" logs/playwright_extra_attempt_${attempt}.log \
|| grep -q "ERR_NO_BUFFER_SPACE" logs/playwright_extra_attempt_${attempt}.log \
|| grep -q "interrupted by another navigation" logs/playwright_extra_attempt_${attempt}.log; } \
&& [ "$attempt" -lt "$max_attempts" ]; then
echo "::warning::Playwright flake on attempt ${attempt}; resetting Studio and retrying..."
kill "${STUDIO_EXTRA_PID}" 2>/dev/null || true

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

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

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

@ -1338,11 +1338,19 @@ jobs:
shell: pwsh
run: |
$ErrorActionPreference = 'Stop'
# A Program Files dir can hold a transient handle (Defender / MSBuild node)
# so Rename-Item intermittently fails with "Access is denied"; retry to ride it out.
function Rename-WithRetry($Path, $NewName) {
for ($i = 1; $i -le 6; $i++) {
try { Rename-Item -LiteralPath $Path -NewName $NewName -ErrorAction Stop; return }
catch { if ($i -eq 6) { throw }; Start-Sleep -Seconds 3 }
}
}
# Rename the Visual Studio install roots (incl. the Installer that holds
# vswhere.exe) so Find-VsBuildTools' vswhere + filesystem scan both miss.
foreach ($d in @("$env:ProgramFiles\Microsoft Visual Studio", "${env:ProgramFiles(x86)}\Microsoft Visual Studio")) {
if (Test-Path -LiteralPath $d) {
Rename-Item -LiteralPath $d -NewName ((Split-Path $d -Leaf) + '.vsoff')
Rename-WithRetry $d ((Split-Path $d -Leaf) + '.vsoff')
Write-Host "Hid VS: $d"
}
}
@ -1351,7 +1359,7 @@ jobs:
$hidden = @()
foreach ($c in (Get-Command cmake -All -ErrorAction SilentlyContinue)) {
if ($c.Source -and (Test-Path -LiteralPath $c.Source)) {
Rename-Item -LiteralPath $c.Source -NewName ((Split-Path $c.Source -Leaf) + '.off')
Rename-WithRetry $c.Source ((Split-Path $c.Source -Leaf) + '.off')
$hidden += $c.Source
Write-Host "Hid cmake: $($c.Source)"
}
@ -1376,7 +1384,7 @@ jobs:
- name: PyTorch CPU wheel installs and imports (no Visual Studio)
run: |
python -m pip install --upgrade pip
python -m pip install torch --index-url https://download.pytorch.org/whl/cpu
python -m pip install torch --index-url https://download.pytorch.org/whl/cpu --extra-index-url https://pypi.org/simple
python -c "import torch; print('torch', torch.__version__, 'cuda?', torch.cuda.is_available())"
- name: Install Studio (--local, --no-torch) with no build tools present
@ -1536,8 +1544,16 @@ jobs:
shell: pwsh
run: |
$ErrorActionPreference = 'Stop'
# Retry the rename: a Program Files dir can hold a transient handle that
# makes Rename-Item intermittently fail with "Access is denied".
function Rename-WithRetry($Path, $NewName) {
for ($i = 1; $i -le 6; $i++) {
try { Rename-Item -LiteralPath $Path -NewName $NewName -ErrorAction Stop; return }
catch { if ($i -eq 6) { throw }; Start-Sleep -Seconds 3 }
}
}
foreach ($d in @("$env:ProgramFiles\Microsoft Visual Studio", "${env:ProgramFiles(x86)}\Microsoft Visual Studio")) {
if (Test-Path -LiteralPath $d) { Rename-Item -LiteralPath $d -NewName ((Split-Path $d -Leaf) + '.vsoff'); Write-Host "Hid VS: $d" }
if (Test-Path -LiteralPath $d) { Rename-WithRetry $d ((Split-Path $d -Leaf) + '.vsoff'); Write-Host "Hid VS: $d" }
}
- name: Windows CUDA and ROCm prebuilts exist in unslothai/llama.cpp (what GPU users download, no VS)

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

@ -242,7 +242,7 @@ jobs:
run: |
python -m pip install --upgrade pip
# CPU torch (vllm/peft/st all depend on it).
pip install --index-url https://download.pytorch.org/whl/cpu \
pip install --index-url https://download.pytorch.org/whl/cpu --extra-index-url https://pypi.org/simple \
'torch>=2.4,<2.11' 'torchvision<0.26' 'torchcodec<0.10'
# torchcodec is a hard requirement on transformers 5.x:
# transformers/audio_utils.py:55 does

14
README.md
View file

@ -246,6 +246,20 @@ curl -fsSL https://unsloth.ai/install.sh | UNSLOTH_STUDIO_HOME=/abs/path sh
$env:UNSLOTH_STUDIO_HOME='C:\path'; irm https://unsloth.ai/install.ps1 | iex
```
On macOS, the installer defaults to the system certificate store (`UV_SYSTEM_CERTS=1`) so uv trusts the CAs in your Keychain, needed behind TLS-inspecting proxies (Cisco Umbrella, Zscaler, etc.). Opt out with:
```bash
curl -fsSL https://unsloth.ai/install.sh | UV_SYSTEM_CERTS=0 sh
```
Point the frontend build at a corporate npm mirror/proxy with `UNSLOTH_NPM_REGISTRY` (for the developer install behind a firewall that blocks `registry.npmjs.org`):
```bash
UNSLOTH_NPM_REGISTRY=https://artifactory.example.com/api/npm/npm/ ./install.sh --local
```
```powershell
$env:UNSLOTH_NPM_REGISTRY='https://artifactory.example.com/api/npm/npm/'; .\install.ps1 --local
```
It is threaded as `--registry` into the Studio frontend `npm`/`bun` installs; the supply-chain locks (7-day `min-release-age`, exact version pins) stay in force.
Cap Studio's native CPU thread pools on high-core hosts: `UNSLOTH_CPU_THREADS=8 unsloth studio -p 8888`.
#### Uninstall

15
build.sh
View file

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

38
install.sh
View file

@ -447,8 +447,12 @@ _on_install_exit() {
if [ "$_status" -ne 0 ]; then
_restore_studio_venv_replacement
fi
[ -n "${_UV_OVERRIDE_TMPDIR:-}" ] && rm -rf "$_UV_OVERRIDE_TMPDIR" 2>/dev/null || true
exit "$_status"
}
# Empty so an inherited value can never reach the trap's rm; only a temp dir
# this script creates below (Apple Silicon, spaced path) is ever removed.
_UV_OVERRIDE_TMPDIR=""
trap _on_install_exit EXIT
# ── Helper: download a URL to a file (supports curl and wget) ──
@ -1427,6 +1431,25 @@ fi
if [ "$OS" = "macos" ] && [ "$_ARCH" = "arm64" ]; then
_OVERRIDES_FILE="$(cd "$(dirname "$0" 2>/dev/null || echo ".")" && pwd)/studio/backend/requirements/single-env/overrides-darwin-arm64.txt"
if [ -f "$_OVERRIDES_FILE" ]; then
# uv splits UV_OVERRIDE on whitespace, so a repo path with whitespace
# truncates it and aborts every later uv call (issue #6503). Hand uv a copy.
case "$_OVERRIDES_FILE" in
*[[:space:]]*)
_UV_OVERRIDE_TMPDIR=$(mktemp -d 2>/dev/null) || _UV_OVERRIDE_TMPDIR=""
case "$_UV_OVERRIDE_TMPDIR" in
"") ;;
*[[:space:]]*) rm -rf "$_UV_OVERRIDE_TMPDIR" 2>/dev/null || true; _UV_OVERRIDE_TMPDIR="" ;;
*)
if cp "$_OVERRIDES_FILE" "$_UV_OVERRIDE_TMPDIR/overrides-darwin-arm64.txt" 2>/dev/null; then
_OVERRIDES_FILE="$_UV_OVERRIDE_TMPDIR/overrides-darwin-arm64.txt"
else
rm -rf "$_UV_OVERRIDE_TMPDIR" 2>/dev/null || true
_UV_OVERRIDE_TMPDIR=""
fi
;;
esac
;;
esac
export UV_OVERRIDE="$_OVERRIDES_FILE"
fi
fi
@ -1613,6 +1636,21 @@ export UV_HTTP_RETRIES
: "${UV_HTTP_TIMEOUT:=180}"
export UV_HTTP_TIMEOUT
# macOS: trust the system Keychain so uv uses SecureTransport instead of rustls.
# Required behind TLS-inspecting proxies (Cisco Umbrella, Zscaler, etc.) which
# present their own CA certificate. rustls (uv's default) ignores the Keychain
# and rejects intercepted connections with "invalid peer certificate: UnknownIssuer".
# Set both vars: UV_SYSTEM_CERTS is the modern one (uv >= 0.11), UV_NATIVE_TLS the
# legacy one understood by uv 0.8.16-0.10.x, which the installer keeps if already
# present (UV_MIN_VERSION) and which ignores UV_SYSTEM_CERTS. Mirror the choice onto
# both so it works on either uv. Opt out with UV_SYSTEM_CERTS=0.
if [ "$OS" = "macos" ]; then
: "${UV_SYSTEM_CERTS:=1}"
: "${UV_NATIVE_TLS:=$UV_SYSTEM_CERTS}"
fi
[ -n "${UV_SYSTEM_CERTS:-}" ] && export UV_SYSTEM_CERTS
[ -n "${UV_NATIVE_TLS:-}" ] && export UV_NATIVE_TLS
version_ge() {
# returns 0 if $1 >= $2
_a=$1

23
pyproject.toml
View file

@ -255,10 +255,6 @@ cu118onlytorch270 = [
"xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.30-cp310-cp310-manylinux_2_28_x86_64.whl ; python_version=='3.10' and ('linux' in sys_platform)",
"xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.30-cp311-cp311-manylinux_2_28_x86_64.whl ; python_version=='3.11' and ('linux' in sys_platform)",
"xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.30-cp312-cp312-manylinux_2_28_x86_64.whl ; python_version=='3.12' and ('linux' in sys_platform)",
"xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.30-cp39-cp39-win_amd64.whl ; python_version=='3.9' and (sys_platform == 'win32')",
"xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.30-cp310-cp310-win_amd64.whl ; python_version=='3.10' and (sys_platform == 'win32')",
"xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.30-cp311-cp311-win_amd64.whl ; python_version=='3.11' and (sys_platform == 'win32')",
"xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.30-cp312-cp312-win_amd64.whl ; python_version=='3.12' and (sys_platform == 'win32')",
]
cu126onlytorch270 = [
"xformers @ https://download.pytorch.org/whl/cu126/xformers-0.0.30-cp39-cp39-manylinux_2_28_x86_64.whl ; python_version=='3.9' and ('linux' in sys_platform)",
@ -282,7 +278,6 @@ cu128onlytorch270 = [
]
cu118onlytorch271 = [
"xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.31.post1-cp39-abi3-manylinux_2_28_x86_64.whl ; ('linux' in sys_platform)",
"xformers @ https://download.pytorch.org/whl/cu118/xformers-0.0.31.post1-cp39-abi3-win_amd64.whl ; (sys_platform == 'win32')",
]
cu126onlytorch271 = [
"xformers @ https://download.pytorch.org/whl/cu126/xformers-0.0.31.post1-cp39-abi3-manylinux_2_28_x86_64.whl ; ('linux' in sys_platform)",
@ -879,14 +874,12 @@ flashattentiontorch240abiFALSEcu12x = [
"flash-attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.4cxx11abiFALSE-cp310-cp310-linux_x86_64.whl ; ('linux' in sys_platform) and python_version == '3.10'",
"flash-attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.4cxx11abiFALSE-cp311-cp311-linux_x86_64.whl ; ('linux' in sys_platform) and python_version == '3.11'",
"flash-attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.4cxx11abiFALSE-cp312-cp312-linux_x86_64.whl ; ('linux' in sys_platform) and python_version == '3.12'",
"flash-attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.4cxx11abiFALSE-cp313-cp313-linux_x86_64.whl ; ('linux' in sys_platform) and python_version == '3.13'",
]
flashattentiontorch240abiTRUEcu12x = [
"flash-attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.4cxx11abiTRUE-cp39-cp39-linux_x86_64.whl ; ('linux' in sys_platform) and python_version == '3.9'",
"flash-attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.4cxx11abiTRUE-cp310-cp310-linux_x86_64.whl ; ('linux' in sys_platform) and python_version == '3.10'",
"flash-attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.4cxx11abiTRUE-cp311-cp311-linux_x86_64.whl ; ('linux' in sys_platform) and python_version == '3.11'",
"flash-attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.4cxx11abiTRUE-cp312-cp312-linux_x86_64.whl ; ('linux' in sys_platform) and python_version == '3.12'",
"flash-attn @ https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.4cxx11abiTRUE-cp313-cp313-linux_x86_64.whl ; ('linux' in sys_platform) and python_version == '3.13'",
]
intelgputorch260 = [
"unsloth_zoo[intelgpu]",
@ -1174,14 +1167,14 @@ intelgputorch2120 = [
"unsloth_zoo[intelgpu]",
"unsloth[huggingfacenotorch]",
"triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.1-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl#sha256=844d981cb1b3948085e8cfa62c74de9f100259f6131959aa70be49123b88ae81 ; platform_system == 'Linux' and python_version == '3.10' and platform_machine == 'x86_64'",
"triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.1-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl#sha256=a16b1d00e94ad87d62af3512e390348b8656419598004100c56028bf494f086b ; platform_system == 'Linux' and python_version == '3.11' and platform_machine == 'x86_64'",
"triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.1-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl#sha256=4e46e71e077cf483404a4c17ce40d71c5f0e13a81459139d4346ca427b1dd455 ; platform_system == 'Linux' and python_version == '3.12' and platform_machine == 'x86_64'",
"triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.1-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl#sha256=4fdaed1bafc51d3a2834656a3420a6686a74ea226508765a49bf15d58ff3a930 ; platform_system == 'Linux' and python_version == '3.13' and platform_machine == 'x86_64'",
"triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.1-cp310-cp310-win_amd64.whl#sha256=2778b46b22e9fa0916398db299a125027a1b2331c1173b3dd2b9e2cab6263a31 ; sys_platform == 'win32' and python_version == '3.10' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')",
"triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.1-cp311-cp311-win_amd64.whl#sha256=ad5b147d04ee0d40f3d4d32f85f5aa3a3beb6cd5799ca026d3d7f4afa3d9e24f ; sys_platform == 'win32' and python_version == '3.11' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')",
"triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.1-cp312-cp312-win_amd64.whl#sha256=d9482063af2a308543f23333e32edd738ea87cbb33ade68afda9ae0fd704ccd9 ; sys_platform == 'win32' and python_version == '3.12' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')",
"triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.1-cp313-cp313-win_amd64.whl#sha256=5d4d67f0deb1e851c01b293e602b8dcddad26ca2be61221cee3dc0e1aa0cdefd ; sys_platform == 'win32' and python_version == '3.13' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')",
"triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.1-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl#sha256=81ff0eb0c4fc8e19d2510b28c3e1d9382a3c7d6fdaf6a9f9631a93a030d841cf ; platform_system == 'Linux' and python_version == '3.10' and platform_machine == 'x86_64'",
"triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.1-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl#sha256=55574a68d275b85cd4d5cbf185084bae019ebf09c3f43b0bd2831b14935ec8e7 ; platform_system == 'Linux' and python_version == '3.11' and platform_machine == 'x86_64'",
"triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.1-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl#sha256=a31c058c5c2e78ebe490a2e69f2f50caec6b1307ac096e944f116fdc06819d9a ; platform_system == 'Linux' and python_version == '3.12' and platform_machine == 'x86_64'",
"triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.1-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl#sha256=e701a31efa0334775f357c98716f3821775aa944219f7888e13c2dfe2daabe2a ; platform_system == 'Linux' and python_version == '3.13' and platform_machine == 'x86_64'",
"triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.1-cp310-cp310-win_amd64.whl#sha256=0d7730651c3e52fbf3a430cc201455f0c6600dc72e681aec495f131ea44f341a ; sys_platform == 'win32' and python_version == '3.10' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')",
"triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.1-cp311-cp311-win_amd64.whl#sha256=8f4a63de73e3d632098f93c8f0bd77244958a47d7c5f728b8ff35f8a91fdb983 ; sys_platform == 'win32' and python_version == '3.11' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')",
"triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.1-cp312-cp312-win_amd64.whl#sha256=6589ece3adc2b1ab88d90ff1267afc25df5c7b868f0b633e732cac70df36cbde ; sys_platform == 'win32' and python_version == '3.12' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')",
"triton-xpu @ https://download.pytorch.org/whl/triton_xpu-3.7.1-cp313-cp313-win_amd64.whl#sha256=2fdf001a9b0575e8b1827127259bb9b13bf36e659882be74c2dfab46597d3e7a ; sys_platform == 'win32' and python_version == '3.13' and (platform_machine == 'AMD64' or platform_machine == 'x86_64')",
"torch @ https://download.pytorch.org/whl/xpu/torch-2.12.0%2Bxpu-cp310-cp310-linux_x86_64.whl#sha256=e8923cd1fe560472904b1461b745d2f1826bb9c1bc0808225d5f28a450e4d553 ; platform_system == 'Linux' and python_version == '3.10' and platform_machine == 'x86_64'",
"torch @ https://download.pytorch.org/whl/xpu/torch-2.12.0%2Bxpu-cp311-cp311-linux_x86_64.whl#sha256=f7c082b2fc9b61def594d30ea57762dc4a8bc7111a9a9593953ed948de242e28 ; platform_system == 'Linux' and python_version == '3.11' and platform_machine == 'x86_64'",

14
scripts/compare_engines.py
View file

@ -40,9 +40,13 @@ def bench_pytorch(repo, gguf, resolutions, steps, seed, iters):
backend = DiffusionBackend()
for speed in ("off", "default"):
backend.begin_load(repo, gguf_filename = gguf, speed_mode = speed)
deadline = time.time() + 1800 # 30 min: a stuck download/load must not hang forever
while backend.load_progress().get("phase") != "ready":
if backend.load_progress().get("phase") == "error":
raise RuntimeError(backend.load_progress())
prog = backend.load_progress()
if prog.get("phase") == "error":
raise RuntimeError(prog)
if time.time() > deadline:
raise TimeoutError(f"load timed out (last progress: {prog})")
time.sleep(0.5)
for res in resolutions:
@ -128,11 +132,13 @@ def main(argv = None) -> int:
)
p.add_argument(
"--vae",
default = "/mnt/disks/unslothai/ubuntu/workspace_81/sdcpp_assets/flux_vae/ae.safetensors",
default = None,
help = "VAE safetensors for sd.cpp (required when benchmarking the sd.cpp engine)",
)
p.add_argument(
"--llm",
default = "/mnt/disks/unslothai/ubuntu/workspace_81/sdcpp_assets/qwen3_te/Qwen3-4B-Instruct-2507-Q4_K_M.gguf",
default = None,
help = "text-encoder GGUF for sd.cpp (required when benchmarking the sd.cpp engine)",
)
p.add_argument("--resolutions", default = "512,1024")
p.add_argument("--steps", type = int, default = 8)

19
scripts/diffusion_bench.py
View file

@ -142,8 +142,10 @@ def _psnr(ref_png: Path, cand_png: Path) -> float:
import numpy as np
from PIL import Image
a = np.asarray(Image.open(ref_png).convert("RGB"), dtype = np.float64)
b = np.asarray(Image.open(cand_png).convert("RGB"), dtype = np.float64)
with Image.open(ref_png) as im_a:
a = np.asarray(im_a.convert("RGB"), dtype = np.float64)
with Image.open(cand_png) as im_b:
b = np.asarray(im_b.convert("RGB"), dtype = np.float64)
if a.shape != b.shape:
# Different geometry means the comparison is meaningless; report worst case.
return 0.0
@ -368,9 +370,14 @@ def _compare(args: argparse.Namespace) -> int:
print(" refusing noisy comparison (pass --force-compare to override).", flush = True)
return 2
# PSNR vs the stored reference image.
# PSNR vs the stored reference image. The baseline stores an absolute reference_png,
# which breaks if the baseline directory was copied/moved, so fall back to reference.png
# next to the baseline JSON. A still-missing reference is a failure below, not a silent
# pass -- otherwise the benchmark would report PASS having done no image comparison.
ref_png = Path(baseline.get("accuracy", {}).get("reference_png", ""))
psnr = _psnr(ref_png, args._image_out) if ref_png.exists() else float("nan")
if not ref_png.is_file():
ref_png = baseline_path.parent / "reference.png"
psnr = _psnr(ref_png, args._image_out) if ref_png.is_file() else float("nan")
base_gen = baseline.get("generate", {})
cur_gen = metrics["generate"]
@ -402,7 +409,9 @@ def _compare(args: argparse.Namespace) -> int:
)
if base_peak and cur_peak and vram_reg > args.max_vram_regression:
failures.append(f"peak VRAM +{vram_reg * 100:.1f}% > {args.max_vram_regression * 100:.0f}%")
if not math.isnan(psnr) and psnr < args.min_psnr:
if math.isnan(psnr):
failures.append("PSNR reference image missing; cannot verify output quality")
elif psnr < args.min_psnr:
failures.append(f"PSNR {psnr:.2f}dB < {args.min_psnr:.1f}dB (output changed)")
if failures:

17
scripts/diffusion_quality.py
View file

@ -186,6 +186,18 @@ def _wait_for_load(backend: Any, timeout_s: int = 3600) -> None:
def _hf_file_size_mib(repo: str, filename: str) -> Optional[int]:
# A local model dir / file: stat it directly. The Hub lookup below returns None for
# a local path, which would drop every candidate from _recommend (file_size_mib None).
try:
local = Path(repo).expanduser()
if local.is_dir():
f = local / filename
if f.is_file():
return int(f.stat().st_size // (1024 * 1024))
elif local.is_file():
return int(local.stat().st_size // (1024 * 1024))
except Exception:
pass
try:
from huggingface_hub import HfApi
info = HfApi().model_info(repo, files_metadata = True, token = os.environ.get("HF_TOKEN"))
@ -268,6 +280,11 @@ def _compare(
clip_sim.append(clip.image_similarity(img, ref))
def _mean(xs: list[float]) -> Optional[float]:
# Preserve +inf: an identical render (reference vs itself, or a lossless
# quant/offload) scores PSNR=inf, which is exactly the case this harness
# verifies; dropping it as non-finite would print "-" instead of "inf".
if xs and any(x == math.inf for x in xs):
return math.inf
finite = [x for x in xs if math.isfinite(x)]
return round(sum(finite) / len(finite), 4) if finite else None

9
scripts/fbcache_flux_probe.py
View file

@ -18,7 +18,7 @@ import numpy as np
BASE = "black-forest-labs/FLUX.1-dev"
PROMPT = "A cinematic photograph of a red fox in a snowy forest at dawn, highly detailed"
OUT = Path("/mnt/disks/unslothai/ubuntu/workspace_81/outputs/quant_research/fbcache_flux_images")
OUT = Path(__file__).resolve().parent.parent / "outputs" / "quant_research" / "fbcache_flux_images"
_LP = {"fn": None}
@ -101,8 +101,13 @@ def run(
from diffusers.hooks import apply_first_block_cache
apply_first_block_cache(pipe.transformer, FirstBlockCacheConfig(threshold = threshold))
if compile_:
# FBCache's per-step decision is a graph break, so a cached run must compile with
# fullgraph=False (mirroring the production path); fullgraph=True would fail the
# warmup compile and the row would silently fall back to an eager cached run,
# producing misleading speedup numbers.
fullgraph = threshold is None
try:
pipe.transformer.compile_repeated_blocks(fullgraph = True, dynamic = True)
pipe.transformer.compile_repeated_blocks(fullgraph = fullgraph, dynamic = True)
except Exception as exc: # noqa: BLE001
print(f" [{tag}] compile {type(exc).__name__}: {str(exc)[:80]}", flush = True)
try:

6
scripts/nvfp4_probe.py
View file

@ -16,7 +16,7 @@ import numpy as np
BASE = "Tongyi-MAI/Z-Image-Turbo"
PROMPT = "A cinematic photograph of a red fox in a snowy forest at dawn, highly detailed"
OUT = Path("/mnt/disks/unslothai/ubuntu/workspace_81/outputs/quant_research/nvfp4_images")
OUT = Path(__file__).resolve().parent.parent / "outputs" / "quant_research" / "nvfp4_images"
def _psnr(a, b):
@ -84,11 +84,15 @@ def main(argv = None) -> int:
p.add_argument("--seed", type = int, default = 42)
p.add_argument("--iters", type = int, default = 3)
p.add_argument("--min-feat", type = int, default = 512)
p.add_argument("--out-dir", default = None, help = "image output dir (default: repo outputs/)")
args = p.parse_args(argv)
steps, res, seed, mf = args.steps, args.res, args.seed, args.min_feat
import torch
import torch.nn as nn
global OUT
if args.out_dir:
OUT = Path(args.out_dir).expanduser()
OUT.mkdir(parents = True, exist_ok = True)
def filt(mod, fqn = ""):

19
scripts/perf_levers_probe.py
View file

@ -25,7 +25,7 @@ import numpy as np
BASE = "Tongyi-MAI/Z-Image-Turbo"
PROMPT = "A cinematic photograph of a red fox in a snowy forest at dawn, highly detailed"
OUT = Path("/mnt/disks/unslothai/ubuntu/workspace_81/outputs/quant_research/perf_levers_images")
OUT = Path(__file__).resolve().parent.parent / "outputs" / "quant_research" / "perf_levers_images"
_LP = {"fn": None}
@ -36,11 +36,14 @@ def _lpips(ref, arr):
import lpips
import torch
# Keep the metric model on CPU: caching it on CUDA leaves it resident across
# variants, and each run resets peak-memory stats, so its VRAM would be charged
# to (and reduce headroom for) every later variant's measurement.
if _LP["fn"] is None:
_LP["fn"] = lpips.LPIPS(net = "alex", verbose = False).cuda().eval()
_LP["fn"] = lpips.LPIPS(net = "alex", verbose = False).eval()
def t(x):
return (torch.from_numpy(x).float().permute(2, 0, 1).unsqueeze(0) / 127.5 - 1.0).cuda()
return torch.from_numpy(x).float().permute(2, 0, 1).unsqueeze(0) / 127.5 - 1.0
with torch.no_grad():
return float(_LP["fn"](t(ref), t(arr)).item())
@ -69,6 +72,12 @@ def _reset_inductor_flags():
ic.coordinate_descent_tuning = False
ic.coordinate_descent_check_all_directions = False
ic.epilogue_fusion = True
# Reset the int-mm fusion flag too, or it leaks from the inductor_flags variant into
# every later compiled row and the attention/fbcache measurements stop being isolated.
try:
ic.force_fuse_int_mm_with_mul = False
except Exception: # noqa: BLE001
pass
def _load():
@ -136,6 +145,8 @@ def run(
except Exception as exc: # noqa: BLE001
note = f"attn({attn})={type(exc).__name__}:{str(exc)[:60]}"
print(f" [{tag}] {note}", flush = True)
del pipe # free the resident pipe so a skipped variant doesn't leak VRAM
torch.cuda.empty_cache()
return None
if fbcache is not None:
try:
@ -143,6 +154,8 @@ def run(
apply_first_block_cache(pipe.transformer, FirstBlockCacheConfig(threshold = fbcache))
except Exception as exc: # noqa: BLE001
print(f" [{tag}] fbcache={type(exc).__name__}:{str(exc)[:60]}", flush = True)
del pipe
torch.cuda.empty_cache()
return None
try:
pipe.transformer.compile_repeated_blocks(fullgraph = True, dynamic = True)

7
scripts/perf_verify.py
View file

@ -151,7 +151,12 @@ def main(argv = None) -> int:
flush = True,
)
ok = (leak_psnr == float("inf")) and (def_t < off_t) and (_psnr(off_img, def_img) >= 30)
ok = (
(leak_psnr == float("inf"))
and (bal_psnr == float("inf")) # check 3: balanced must be bit-identical to off
and (def_t < off_t)
and (_psnr(off_img, def_img) >= 30)
)
print(f"\nPERF-VERIFY {'OK' if ok else 'CHECK'}", flush = True)
return 0 if ok else 1

2
scripts/prequant_probe.py
View file

@ -26,7 +26,7 @@ import numpy as np
BASE = "Tongyi-MAI/Z-Image-Turbo"
PROMPT = "A cinematic photograph of a red fox in a snowy forest at dawn, highly detailed"
ROOT = Path("/mnt/disks/unslothai/ubuntu/workspace_81/outputs/quant_research")
ROOT = Path(__file__).resolve().parent.parent / "outputs" / "quant_research"
CKPT = ROOT / "prequant_fp8" / "transformer_fp8_state.pt"
OUT = ROOT / "prequant_images"
MIN_FEAT = 512

14
scripts/quant_probe.py
View file

@ -27,7 +27,7 @@ REPO = "unsloth/Z-Image-Turbo-GGUF"
GGUF = "z-image-turbo-Q4_K_M.gguf"
BASE = "Tongyi-MAI/Z-Image-Turbo"
PROMPT = "A cinematic photograph of a red fox in a snowy forest at dawn, highly detailed"
OUT = Path("/mnt/disks/unslothai/ubuntu/workspace_81/outputs/quant_research/probe_images")
OUT = Path(__file__).resolve().parent.parent / "outputs" / "quant_research" / "probe_images"
def _psnr(a, b):
@ -39,17 +39,19 @@ _LPIPS = {"fn": None}
def _lpips(ref_arr, arr):
"""Perceptual LPIPS (alexnet) vs reference; lower is closer. None if unavailable."""
"""Perceptual LPIPS (alexnet) vs reference; lower is closer. None if unavailable.
Runs on CPU so the scorer never holds CUDA memory: each row resets peak VRAM, so a
resident GPU LPIPS module would inflate the reported load/gen VRAM and could even OOM."""
try:
import torch
import lpips
if _LPIPS["fn"] is None:
_LPIPS["fn"] = lpips.LPIPS(net = "alex", verbose = False).cuda().eval()
_LPIPS["fn"] = lpips.LPIPS(net = "alex", verbose = False).eval()
def t(x):
t = torch.from_numpy(x).float().permute(2, 0, 1).unsqueeze(0) / 127.5 - 1.0
return t.cuda()
return torch.from_numpy(x).float().permute(2, 0, 1).unsqueeze(0) / 127.5 - 1.0
with torch.no_grad():
return float(_LPIPS["fn"](t(ref_arr), t(arr)).item())
@ -111,7 +113,7 @@ def _quant_config(name):
return MXDynamicActivationMXWeightConfig(
activation_dtype = torch.float8_e4m3fn, weight_dtype = torch.float8_e4m3fn
)
except TypeError:
except (TypeError, AttributeError):
return MXDynamicActivationMXWeightConfig()
raise ValueError(name)

7
scripts/scan_packages.py
View file

@ -1208,9 +1208,10 @@ def check_js_file(content: str, filename: str, package: str) -> list[Finding]:
HIGH,
package,
filename,
f"Python wheel ships large ({len(content) // 1024} KB) JS bundle "
"(uncommon; manually review)",
"",
# Size stays in evidence, not the check label, so the baseline key
# does not drift when a wheel's bundle grows by a few KB.
"Python wheel ships large JS bundle (uncommon; manually review)",
f"{len(content) // 1024} KB JS bundle",
)
)
return findings

2
scripts/scan_packages_baseline.json
View file

@ -1181,7 +1181,7 @@
{
"package": "tensorboard",
"file": "tensorboard/plugins/projector/tf_projector_plugin/projector_binary.js",
"check": "Python wheel ships large (1918 KB) JS bundle (uncommon; manually review)",
"check": "Python wheel ships large JS bundle (uncommon; manually review)",
"severity": "HIGH",
"evidence": ""
},

9
scripts/verify_prequant_backend.py
View file

@ -19,17 +19,20 @@ from __future__ import annotations
import argparse
import logging
import os
import sys
import time
from pathlib import Path
import numpy as np
BACKEND = Path(__file__).resolve().parent.parent / "studio" / "backend"
_REPO = Path(__file__).resolve().parent.parent
_RESEARCH = _REPO / "outputs" / "quant_research"
BACKEND = _REPO / "studio" / "backend"
BASE = "Tongyi-MAI/Z-Image-Turbo"
CKPT = "/mnt/disks/unslothai/ubuntu/workspace_81/outputs/quant_research/prequant_fp8/transformer_fp8.pt"
CKPT = os.environ.get("PREQUANT_CKPT", str(_RESEARCH / "prequant_fp8" / "transformer_fp8.pt"))
PROMPT = "A cinematic photograph of a red fox in a snowy forest at dawn, highly detailed"
OUT = Path("/mnt/disks/unslothai/ubuntu/workspace_81/outputs/quant_research/prequant_verify_images")
OUT = Path(os.environ.get("PREQUANT_OUT_DIR", str(_RESEARCH / "prequant_verify_images")))
logging.basicConfig(level = logging.INFO, format = "%(message)s")
LOGGER = logging.getLogger("verify_prequant")

2
studio/Unsloth_Studio_Colab.ipynb
View file

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

7
studio/backend/assets/preview_page.html
View file

@ -283,6 +283,11 @@
</div>
<script>
const base = location.pathname.replace(/\/+$/, "");
// The capability token rides in ?k=; location.pathname drops it, so carry it
// onto the chat request explicitly. Not stored or logged.
const k = new URLSearchParams(location.search).get("k");
const chatUrl =
base + "/v1/chat/completions" + (k ? "?k=" + encodeURIComponent(k) : "");
const log = document.getElementById("log"),
thread = document.getElementById("thread"),
welcome = document.getElementById("welcome");
@ -328,7 +333,7 @@
out.innerHTML = '<span class="dots"><i></i><i></i><i></i></span>';
let acc = "";
try {
const r = await fetch(base + "/v1/chat/completions", {
const r = await fetch(chatUrl, {
method: "POST",
headers: { "Content-Type": "application/json" },
body: JSON.stringify({

145
studio/backend/auth/bootstrap_timeout.py Normal file
View file

@ -0,0 +1,145 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Auto-shutdown for an exposed first-run Studio whose admin password is unchanged.
On a fresh install the seeded bootstrap admin password stays a valid login
credential until first login changes it. When the web UI is put on the network
(``--secure`` / ``0.0.0.0``) and nobody completes that first-login change within
a deadline, tear Studio down so a fresh, unconfigured instance does not stay
publicly reachable indefinitely. If the password was changed, Studio keeps
running.
Scope: web UI launches only (never ``--api-only``, which authenticates by API
key rather than the admin password, and never Colab). Configurable via
``UNSLOTH_STUDIO_BOOTSTRAP_TIMEOUT`` (seconds; default 3600; ``0`` disables).
"""
import os
import sys
import threading
BOOTSTRAP_TIMEOUT_ENV_VAR = "UNSLOTH_STUDIO_BOOTSTRAP_TIMEOUT"
DEFAULT_BOOTSTRAP_TIMEOUT_SECONDS = 3600
def bootstrap_timeout_seconds(env = None) -> int:
"""Resolve the deadline in seconds. ``0`` (or invalid/negative) disables it.
A malformed value falls back to the default rather than disabling, so a typo
cannot silently remove the protection.
"""
env = os.environ if env is None else env
raw = env.get(BOOTSTRAP_TIMEOUT_ENV_VAR)
if raw is None or raw.strip() == "":
return DEFAULT_BOOTSTRAP_TIMEOUT_SECONDS
try:
value = int(raw)
except ValueError:
return DEFAULT_BOOTSTRAP_TIMEOUT_SECONDS
return value if value > 0 else 0
def _is_exposed_bind(host: str, secure: bool) -> bool:
"""True when this launch puts the web UI on the network (tunnel or non-loopback)."""
if secure:
return True
if host in ("0.0.0.0", "::"):
return True
try:
from utils.host_policy import is_external_host
except Exception:
return False
return bool(is_external_host(host))
def should_arm_bootstrap_timeout(
*,
host: str,
secure: bool,
api_only: bool,
frontend_served: bool,
is_colab: bool,
requires_change: bool,
timeout_seconds: int,
) -> bool:
"""Whether to arm the deadline: only for an exposed web UI whose seeded admin
password is still unchanged. Pure decision (no I/O) for cheap unit testing."""
if timeout_seconds <= 0:
return False
if api_only or not frontend_served or is_colab:
return False
if not requires_change:
return False
return _is_exposed_bind(host, secure)
def _format_duration(seconds: int) -> str:
"""Human-friendly duration for the shutdown message (seconds under a minute)."""
def _plural(n: int, unit: str) -> str:
return f"{n} {unit}{'' if n == 1 else 's'}"
if seconds < 60:
return _plural(seconds, "second")
minutes, rem = divmod(seconds, 60)
label = _plural(minutes, "minute")
if rem:
label += f" {_plural(rem, 'second')}"
return label
def enforce_bootstrap_password_deadline(
storage,
trigger_shutdown,
*,
timeout_seconds: int,
logger = None,
) -> bool:
"""Deadline handler: shut down iff the seeded admin password is still unchanged.
Returns True if it shut Studio down, False if it left it running (the
password was changed in time).
"""
try:
still_default = storage.requires_password_change(storage.DEFAULT_ADMIN_USERNAME)
except Exception:
return False
if not still_default:
return False # password changed in time -> leave Studio running
message = (
"\nUnsloth Studio was exposed on the network but its default admin "
f"password was not changed within {_format_duration(timeout_seconds)}. "
"Shutting down to avoid leaving an unsecured public instance running.\n"
"Next time, sign in and change the password on first login, or set "
f"{BOOTSTRAP_TIMEOUT_ENV_VAR}=0 to disable this timeout."
)
if logger is not None:
logger.warning(message)
print(message, file = sys.stderr, flush = True)
try:
trigger_shutdown()
except Exception as e: # shutdown is best-effort; never raise from the timer
if logger is not None:
logger.warning("Bootstrap-timeout shutdown failed: %s", e)
return True
def arm_bootstrap_timeout(
storage,
trigger_shutdown,
*,
timeout_seconds: int,
logger = None,
) -> "threading.Timer":
"""Start a daemon timer that enforces the deadline. Returns the Timer."""
timer = threading.Timer(
timeout_seconds,
enforce_bootstrap_password_deadline,
args = (storage, trigger_shutdown),
kwargs = {"timeout_seconds": timeout_seconds, "logger": logger},
)
timer.daemon = True
timer.start()
return timer

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

@ -110,6 +110,17 @@ def get_connection() -> sqlite3.Connection:
except OSError:
pass
conn.row_factory = sqlite3.Row
# WAL lets token reads run concurrently with refresh-token writes;
# busy_timeout bounds lock waits. Matches the other Studio SQLite stores.
# Set busy_timeout first: switching journal_mode needs a lock, so if a
# refresh-token write already holds one, journal_mode=WAL raises SQLITE_BUSY;
# with busy_timeout already in effect it waits instead of failing and leaving
# this connection on SQLite's default zero lock wait.
try:
conn.execute("PRAGMA busy_timeout=5000")
conn.execute("PRAGMA journal_mode=WAL")
except sqlite3.Error:
pass
conn.execute(
"""
CREATE TABLE IF NOT EXISTS auth_user (
@ -270,6 +281,63 @@ def compute_identity_proof(nonce: bytes, host: str, port: int) -> str:
return hmac.new(get_or_create_identity_secret(), msg, hashlib.sha256).hexdigest()
# Capability secret for public ``/p`` preview share links. HMAC(secret, ref)
# turns the deterministic preview ref into an unguessable bearer capability, so a
# guessed run/checkpoint name can't reach inference. Dedicated (not the per-user
# JWT secret) so rotating it revokes every shared link without touching logins.
_PREVIEW_LINK_SECRET_DB_KEY = "preview_link_secret"
_preview_link_secret_cache: Optional[bytes] = None
def get_or_create_preview_link_secret() -> bytes:
"""Return the preview-link signing secret (hex 32-byte row in app_secrets), creating it once."""
global _preview_link_secret_cache
if _preview_link_secret_cache is not None:
return _preview_link_secret_cache
conn = get_connection()
try:
row = conn.execute(
"SELECT value FROM app_secrets WHERE key = ?",
(_PREVIEW_LINK_SECRET_DB_KEY,),
).fetchone()
if row is None:
conn.execute(
"INSERT OR IGNORE INTO app_secrets (key, value) VALUES (?, ?)",
(_PREVIEW_LINK_SECRET_DB_KEY, secrets.token_hex(32)),
)
conn.commit()
row = conn.execute(
"SELECT value FROM app_secrets WHERE key = ?",
(_PREVIEW_LINK_SECRET_DB_KEY,),
).fetchone()
secret = bytes.fromhex(row["value"])
finally:
conn.close()
_preview_link_secret_cache = secret
return secret
def rotate_preview_link_secret() -> bytes:
"""Rotate the preview-link secret, immediately revoking every outstanding ``/p`` share link."""
global _preview_link_secret_cache
new_secret_hex = secrets.token_hex(32)
conn = get_connection()
try:
conn.execute(
"INSERT OR REPLACE INTO app_secrets (key, value) VALUES (?, ?)",
(_PREVIEW_LINK_SECRET_DB_KEY, new_secret_hex),
)
conn.commit()
finally:
conn.close()
secret = bytes.fromhex(new_secret_hex)
_preview_link_secret_cache = secret
return secret
_API_KEY_PBKDF2_ITERATIONS = 100_000
DESKTOP_SECRET_PREFIX = "desktop-"
_DESKTOP_SECRET_HASH_KEY = "desktop_secret_hash"

176
studio/backend/colab.py
View file

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

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

@ -28,6 +28,9 @@ from .constants import (
from .parse import apply_update, coerce_event, parse_log_message
from .types import Job
from .worker import run_job_process
from loggers import get_logger
logger = get_logger(__name__)
_CTX = mp.get_context("spawn")
@ -445,54 +448,86 @@ class JobManager:
events.append(coerce_event(q.get_nowait()))
except queue.Empty:
return events
except (EOFError, OSError, ValueError):
except Exception:
# Return what we have so the run still finalizes rather than wedging "active".
logger.exception(
"Data-recipe job pump: queue drain failed; finalizing with drained events"
)
return events
def _safe_handle_event(self, job: Job, event: dict) -> None:
"""Apply one event, swallowing any handler error so the pump can't die."""
try:
self._handle_event(job, event)
except Exception:
etype = event.get("type") if isinstance(event, dict) else type(event).__name__
logger.exception("Data-recipe job pump: failed to handle %s event; skipping", etype)
def _pump_loop(self) -> None:
"""Background thread: consumes worker events + updates job snapshot."""
"""Background thread: consume worker events and update the job snapshot.
Guarded so no single event can end the loop; it is the sole writer of the
snapshot the UI polls, so its death would freeze status/SSE.
"""
while True:
snap = self._snapshot()
if snap is None:
return
job, proc, mp_q = snap
event = self._read_queue_with_timeout(mp_q, timeout_sec = 0.25)
try:
event = self._read_queue_with_timeout(mp_q, timeout_sec = 0.25)
except Exception:
# If a read keeps raising after the worker died, finalize instead
# of spinning forever; only retry while the worker is still alive.
logger.exception("Data-recipe job pump: queue read failed; continuing")
if proc.is_alive():
time.sleep(0.1)
continue
event = None
if event is not None:
self._handle_event(job, event)
self._safe_handle_event(job, event)
continue
if proc.is_alive():
continue
for e in self._drain_queue(mp_q):
self._handle_event(job, e)
# Worker exited: drain + finalize, guarded so an error can't strand the run "active".
try:
for e in self._drain_queue(mp_q):
self._safe_handle_event(job, e)
retired_job: Job | None = None
with self._lock:
if self._job and self._job.status in {
"pending",
"active",
"cancelling",
}:
if self._job.status == "cancelling":
self._job.status = "cancelled"
else:
self._job.status = "error"
self._job.error = self._job.error or "process exited"
self._job.finished_at = time.time()
event_type = (
EVENT_JOB_CANCELLED if self._job.status == "cancelled" else EVENT_JOB_ERROR
)
self._emit(
{
"type": event_type,
"ts": time.time(),
"job_id": self._job.job_id,
}
)
retired_job = self._job
if retired_job is not None:
self._retire_workflow_key(retired_job)
retired_job: Job | None = None
with self._lock:
if self._job and self._job.status in {
"pending",
"active",
"cancelling",
}:
if self._job.status == "cancelling":
self._job.status = "cancelled"
else:
self._job.status = "error"
self._job.error = self._job.error or "process exited"
self._job.finished_at = time.time()
event_type = (
EVENT_JOB_CANCELLED
if self._job.status == "cancelled"
else EVENT_JOB_ERROR
)
self._emit(
{
"type": event_type,
"ts": time.time(),
"job_id": self._job.job_id,
}
)
retired_job = self._job
if retired_job is not None:
self._retire_workflow_key(retired_job)
except Exception:
logger.exception("Data-recipe job pump: finalization after worker exit failed")
return
def _handle_event(self, job: Job, event: dict) -> None:

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

@ -10,6 +10,7 @@ import tempfile
from loggers import get_logger
import os
import shutil
import contextlib
from pathlib import Path
from typing import Optional, Tuple, List
from unsloth import FastLanguageModel, FastVisionModel, _IS_MLX
@ -37,6 +38,65 @@ logger = get_logger(__name__)
_LLAMA_CPP_SCRIPTS_WARNING_EMITTED = False
def _supports_kwarg(fn, name):
"""True if `fn` accepts keyword `name` directly or via **kwargs."""
import inspect
try:
params = inspect.signature(fn).parameters
except (TypeError, ValueError):
return False
return name in params or any(p.kind == inspect.Parameter.VAR_KEYWORD for p in params.values())
def _compressed_export_supported():
"""True if the installed unsloth build can do FP8/NVFP4 compressed-tensors export."""
try:
import unsloth.save as _us
return hasattr(_us, "_normalize_compressed_method")
except Exception:
return False
def _hf_offline(timeout = 3):
"""True if export should avoid the Hub: honors the HF offline env vars, else does one
cheap TCP reachability probe so a network-down load uses local files / the HF cache
instead of hanging on connection timeouts. Proxy-aware (probes the proxy egress when
one is configured); disable the probe with UNSLOTH_OFFLINE_PROBE=0."""
_offline = {"1", "true", "yes", "on"}
if (
os.environ.get("HF_HUB_OFFLINE", "").strip().lower() in _offline
or os.environ.get("TRANSFORMERS_OFFLINE", "").strip().lower() in _offline
):
return True
if os.environ.get("UNSLOTH_OFFLINE_PROBE", "1").strip().lower() in {"0", "false", "no", "off"}:
return False # probe disabled -> assume online; loads still pass local_files_only on env
# Shared bounded, proxy-aware probe (also used by the export worker before version activation).
from utils.transformers_version import hf_endpoint_unreachable
if hf_endpoint_unreachable(timeout):
logger.warning("Hugging Face endpoint unreachable; loading checkpoint in offline mode")
return True
return False
# Reuse Unsloth's lock-guarded forced-offline context; no-op fallback if it moves.
try:
from unsloth.models.loader_utils import _force_hf_offline
except Exception:
import contextlib as _contextlib
@_contextlib.contextmanager
def _force_hf_offline():
yield
def _offline_window_if(local_files_only):
"""Forced-offline window when offline was detected, else a no-op context."""
return _force_hf_offline() if local_files_only else contextlib.nullcontext()
def _is_wsl():
"""Detect if running under Windows Subsystem for Linux."""
try:
@ -175,10 +235,19 @@ class ExportBackend:
model_id = base_model or checkpoint_path
# Token the type-detection probes too, else a gated multimodal base
# 404s here and falls through to the text loader.
self._audio_type = detect_audio_type(model_id, hf_token = token)
self.is_vision = not self._audio_type and is_vision_model(model_id, hf_token = token)
# Skip the Hub when offline so a no-internet export uses the local cache.
local_files_only = _hf_offline()
# Run the type-detection probes in the forced-offline window (else a gated
# base 404s); it covers is_vision_model's Hub reads + the transformers-5
# subprocess, and local_files_only makes detect_audio_type's requests.get skip.
with _offline_window_if(local_files_only):
self._audio_type = detect_audio_type(
model_id, hf_token = token, local_files_only = local_files_only
)
self.is_vision = not self._audio_type and is_vision_model(
model_id, hf_token = token, local_files_only = local_files_only
)
if self._audio_type == "csm":
from unsloth import FastModel
@ -193,6 +262,7 @@ class ExportBackend:
load_in_4bit = False,
trust_remote_code = trust_remote_code,
token = token,
local_files_only = local_files_only,
)
elif self._audio_type == "whisper":
@ -207,6 +277,7 @@ class ExportBackend:
auto_model = WhisperForConditionalGeneration,
trust_remote_code = trust_remote_code,
token = token,
local_files_only = local_files_only,
)
elif self._audio_type == "snac":
@ -218,6 +289,7 @@ class ExportBackend:
load_in_4bit = load_in_4bit,
trust_remote_code = trust_remote_code,
token = token,
local_files_only = local_files_only,
)
elif self._audio_type == "bicodec":
@ -230,6 +302,7 @@ class ExportBackend:
load_in_4bit = False,
trust_remote_code = trust_remote_code,
token = token,
local_files_only = local_files_only,
)
elif self._audio_type == "dac":
@ -241,6 +314,7 @@ class ExportBackend:
load_in_4bit = False,
trust_remote_code = trust_remote_code,
token = token,
local_files_only = local_files_only,
)
elif self.is_vision:
@ -252,6 +326,7 @@ class ExportBackend:
load_in_4bit = load_in_4bit,
trust_remote_code = trust_remote_code,
token = token,
local_files_only = local_files_only,
)
tokenizer = processor # vision: processor acts as tokenizer
@ -264,6 +339,7 @@ class ExportBackend:
load_in_4bit = load_in_4bit,
trust_remote_code = trust_remote_code,
token = token,
local_files_only = local_files_only,
)
if _IS_MLX:
@ -344,16 +420,33 @@ class ExportBackend:
)
output_path: Optional[str] = None
# compressed-tensors formats run save_pretrained_merged with an FP8/FP4 save_method and
# write to a sibling "<dir>-<suffix>" directory (for vLLM).
_COMPRESSED = {
"FP8 (compressed-tensors)": ("fp8", "fp8"),
"NVFP4 (compressed-tensors)": ("nvfp4", "nvfp4"),
}
is_compressed = format_type in _COMPRESSED
try:
if _IS_MLX:
if is_compressed:
return False, "Compressed-tensors export is not supported on macOS/MLX.", None
mlx_save_method = "merged_4bit" if format_type == "4-bit (FP4)" else "merged_16bit"
elif is_compressed:
if not _compressed_export_supported():
return (
False,
"Compressed-tensors (FP8/NVFP4) export requires an Unsloth build with "
"compressed-tensors support. Upgrade unsloth, or choose 16-bit.",
None,
)
save_method = _COMPRESSED[format_type][0]
elif format_type == "4-bit (FP4)":
save_method = "merged_4bit_forced"
elif self._audio_type == "whisper":
save_method = None
else:
if format_type == "4-bit (FP4)":
save_method = "merged_4bit_forced"
elif self._audio_type == "whisper":
save_method = None
else:
save_method = "merged_16bit"
save_method = "merged_16bit"
if save_directory:
save_directory = str(resolve_export_write_dir(save_directory))
@ -371,9 +464,15 @@ class ExportBackend:
save_directory, self.current_tokenizer, save_method = save_method
)
self._write_export_metadata(save_directory)
logger.info(f"Model saved successfully to {save_directory}")
output_path = str(Path(save_directory).resolve())
# Compressed export writes to the "<dir>-<suffix>" sibling; report that as output.
final_dir = (
f"{save_directory}-{_COMPRESSED[format_type][1]}"
if is_compressed
else save_directory
)
self._write_export_metadata(final_dir)
logger.info(f"Model saved successfully to {final_dir}")
output_path = str(Path(final_dir).resolve())
if push_to_hub:
if not repo_id or not hf_token:
@ -408,6 +507,32 @@ class ExportBackend:
token = hf_token,
private = private,
)
elif is_compressed and output_path and Path(output_path).is_dir():
# The compressed model was already built locally in output_path; upload it
# directly so we do not re-run the (expensive, OOM-prone) compression that
# push_to_hub_merged(save_method=fp8/nvfp4) would otherwise do a second time.
hf_api = HfApi(token = hf_token)
repo_id = PushToHubMixin._create_repo(
PushToHubMixin,
repo_id = repo_id,
private = private,
token = hf_token,
)
content = MODEL_CARD.format(
username = repo_id.split("/")[0],
base_model = getattr(self.current_model.config, "_name_or_path", "unknown"),
model_type = getattr(self.current_model.config, "model_type", "llm"),
method = format_type,
extra = "unsloth",
)
ModelCard(content).push_to_hub(
repo_id, token = hf_token, commit_message = "Unsloth Model Card"
)
hf_api.upload_folder(
folder_path = output_path,
repo_id = repo_id,
repo_type = "model",
)
else:
hub_save_method = save_method if save_method is not None else "merged_16bit"
self.current_model.push_to_hub_merged(
@ -565,6 +690,7 @@ class ExportBackend:
push_to_hub: bool = False,
repo_id: Optional[str] = None,
hf_token: Optional[str] = None,
imatrix_file = None,
) -> Tuple[bool, str, Optional[str]]:
"""
Export model in GGUF format.
@ -582,6 +708,19 @@ class ExportBackend:
if not self.current_model or not self.current_tokenizer:
return False, "No model loaded. Please select a checkpoint first.", None
# Only forward imatrix_file to an unsloth build that accepts it; otherwise even a plain
# no-imatrix export would fail with an unexpected-keyword error against an older unsloth.
if imatrix_file is not None and not _supports_kwarg(
self.current_model.save_pretrained_gguf, "imatrix_file"
):
return (
False,
"This Unsloth build does not support GGUF imatrix export. "
"Upgrade unsloth and unsloth_zoo, or disable the imatrix option.",
None,
)
imatrix_kw = {"imatrix_file": imatrix_file} if imatrix_file is not None else {}
output_path: Optional[str] = None
model_tmp_to_cleanup: Optional[str] = None
try:
@ -635,6 +774,7 @@ class ExportBackend:
_model_tmp,
self.current_tokenizer,
quantization_method = quant_method,
**imatrix_kw,
)
# Relocate the .gguf that convert_to_gguf wrote to cwd (repo root).
@ -701,6 +841,7 @@ class ExportBackend:
self.current_tokenizer,
quantization_method = quant_method,
token = hf_token,
**imatrix_kw,
)
logger.info(f"GGUF model pushed successfully to {repo_id}")

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

@ -499,6 +499,7 @@ class ExportOrchestrator:
push_to_hub: bool = False,
repo_id: Optional[str] = None,
hf_token: Optional[str] = None,
imatrix_file = None,
) -> Tuple[bool, str, Optional[str]]:
"""Export model in GGUF format."""
return self._run_export(
@ -509,6 +510,7 @@ class ExportOrchestrator:
"push_to_hub": push_to_hub,
"repo_id": repo_id,
"hf_token": hf_token,
"imatrix_file": imatrix_file,
},
)

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

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

72
studio/backend/core/inference/diffusion.py
View file

@ -82,6 +82,7 @@ from .diffusion_prequant import (
resolve_prequant_source,
)
from .diffusion_transformer_quant import (
DEFAULT_MIN_LINEAR_FEATURES,
dense_transformer_supported,
normalize_transformer_quant,
quantize_transformer,
@ -382,6 +383,22 @@ class DiffusionBackend:
base, rfilename, hf_token, cancel_event = self._cancel_event
)
@staticmethod
def _detect_family_for_pick(
repo_id: str, gguf_filename: Optional[str], family_override: Optional[str]
) -> Optional[DiffusionFamily]:
"""Detect the family from the repo id, falling back to the combined
path/filename for a direct local .gguf pick. The frontend splits such a
pick into (parent dir, basename), so the family keyword can live only in
the filename (e.g. /models/z-image-turbo-Q4_K_M.gguf) while the parent
directory carries none; scan it too when the directory alone is
undetectable. Only used as a fallback, so remote 'org/name' picks and
explicit overrides behave exactly as before."""
fam = detect_family(repo_id, family_override)
if fam is None and gguf_filename and not family_override:
fam = detect_family(f"{repo_id}/{gguf_filename}", family_override)
return fam
def validate_load_request(
self,
repo_id: str,
@ -397,7 +414,7 @@ class DiffusionBackend:
undetectable family, and ValueError/FileNotFoundError for a bad local path.
Touches no GPU, network, or state."""
kind = resolve_model_kind(gguf_filename, model_kind)
fam = detect_family(repo_id, family_override)
fam = self._detect_family_for_pick(repo_id, gguf_filename, family_override)
if fam is None:
raise ValueError(
f"'{repo_id}' is not a supported diffusion image model. Supported families: "
@ -419,7 +436,12 @@ class DiffusionBackend:
# missing one is an error here; a bare "org/name" id is a remote HF repo and
# is left for the background load to resolve.
local_root = Path(repo_id).expanduser()
path_shaped = repo_id.startswith(("/", "~", "./", "../")) or local_root.is_absolute()
# POSIX path-shaped, a "."/".." prefix (covers ./ ../ and their Windows .\ ..\
# forms), a Windows separator anywhere (never present in a bare "org/name" HF
# id), or an absolute path on this OS.
path_shaped = (
repo_id.startswith(("/", "\\", "~", ".")) or "\\" in repo_id or local_root.is_absolute()
)
if kind in ("gguf", "single_file"):
if not gguf_filename:
raise ValueError(f"a single-file checkpoint name is required for a '{kind}' load.")
@ -515,7 +537,9 @@ class DiffusionBackend:
# Resolve the base repo and estimate sizes on this thread (both network
# calls) so begin_load returns instantly; the bar shows raw bytes until
# the total lands. This is the only writer of _loading's fields here.
fam = detect_family(kwargs["repo_id"], kwargs.get("family_override"))
fam = self._detect_family_for_pick(
kwargs["repo_id"], kwargs.get("gguf_filename"), kwargs.get("family_override")
)
kind = resolve_model_kind(kwargs.get("gguf_filename"), kwargs.get("model_kind"))
if kind == "pipeline":
# The full pipeline IS the repo: from_pretrained pulls every component
@ -533,10 +557,12 @@ class DiffusionBackend:
kwargs.get("hf_token"),
kind = kind,
)
loading = self._loading
if loading is not None:
loading.base_repo = base
loading.expected_bytes = expected
with self._lock:
# Stamp progress only if this load is still current; a superseding
# load (or unload) has its own token and its own _LoadingState.
if self._load_token == token and self._loading is not None:
self._loading.base_repo = base
self._loading.expected_bytes = expected
# Download outside the lock so unload()/an eviction can preempt the
# multi-GB pull; load_pipeline below then assembles from the cache.
self._prefetch_files(
@ -558,9 +584,13 @@ class DiffusionBackend:
if self._load_token != token:
return
logger.error("diffusion.load_failed: %s", exc)
# Redact native paths: this error is surfaced verbatim via the
# load-progress poll, and Studio can run as a shared server.
from utils.native_path_leases import redact_native_paths
with self._lock:
if self._load_token == token and self._loading is not None:
self._loading.error = str(exc)
self._loading.error = redact_native_paths(str(exc))
def load_progress(self) -> dict[str, Any]:
"""Phase + downloaded/total bytes for the in-flight load (cache-scan based)."""
@ -613,7 +643,11 @@ class DiffusionBackend:
base_files.append(s.rfilename)
total += s.size or 0
return total, base_files
if gguf_filename:
# Skip the Hub size lookup for a LOCAL gguf path: model_info(repo_id) would
# raise on a filesystem path and (caught below) skip the base-repo lookup too,
# so the companion VAE/text-encoder files would never be prefetched and would
# instead download synchronously under the load lock.
if gguf_filename and not Path(repo_id).expanduser().exists():
info = api.model_info(repo_id, files_metadata = True, token = hf_token)
total += sum(s.size or 0 for s in info.siblings if s.rfilename == gguf_filename)
base_info = api.model_info(base_repo, files_metadata = True, token = hf_token)
@ -808,6 +842,18 @@ class DiffusionBackend:
# the quant noise floor), dense models stay bit-identical `off`. An
# explicit speed_mode (incl. "off") is honored verbatim.
effective_speed = resolve_speed_mode(speed_mode, is_gguf = kind == "gguf")
# A torchao-quantized dense transformer runs its matmuls through the
# regional torch.compile; UNcompiled (eager) it is ~30x slower and would
# lose to the GGUF fallback. A dense model otherwise resolves to `off`, so
# force at least `default` (regional compile) whenever the quant engaged,
# or the opt-in "fast" path silently commits an eager, pathologically slow
# pipeline.
if transformer_quant_engaged is not None and effective_speed == SPEED_OFF:
logger.info(
"diffusion.transformer_quant: forcing speed_mode=default "
"(quantized transformer must be compiled; eager is ~30x slower)"
)
effective_speed = SPEED_DEFAULT
# Opt-in speed optims run BEFORE placement (channels_last / compile
# must precede CPU offload). Snapshot the process-wide backend flags
# first so unload can restore them: TF32 / cudnn.benchmark are global,
@ -833,7 +879,10 @@ class DiffusionBackend:
pipe,
mode = transformer_cache,
threshold = transformer_cache_threshold,
quant_active = transformer_quant_engaged is not None,
# GGUF transformers are quantized too (the default Studio path), so the
# cache needs the higher quantized threshold to still trigger -- not just
# the dense-quant fast path.
quant_active = transformer_quant_engaged is not None or bool(gguf_filename),
logger = logger,
)
# Install the shared compile-safe eager patches (fused RMSNorm /
@ -1014,6 +1063,9 @@ class DiffusionBackend:
dtype = dtype,
hf_token = hf_token,
scheme = scheme,
# Reject a checkpoint built with a different Linear filter than the
# dense path uses, so the prequant and runtime-quant models match.
min_features = DEFAULT_MIN_LINEAR_FEATURES,
logger = logger,
)
if transformer is not None:

148
studio/backend/core/inference/diffusion_attention.py
View file

@ -55,7 +55,7 @@ def normalize_attention_backend(value: Optional[str]) -> Optional[str]:
Raises ValueError for an unsupported alias so a bad request is rejected cheaply."""
if value is None:
return ATTN_AUTO
normalized = str(value).strip().lower().replace("-", "_")
normalized = str(value).strip().lower()
if not normalized:
return ATTN_AUTO
if normalized not in ATTN_ALIASES:
@ -65,6 +65,46 @@ def normalize_attention_backend(value: Optional[str]) -> Optional[str]:
return normalized
# Backends diffusers validates only by *package* at set time (``_check_attention_backend_
# requirements`` checks the ``kernels`` install, not the GPU), but whose kernels need a
# specific CUDA arch at run time -- so an explicit request on the wrong card loads/sets fine
# and then crashes mid-generation. Gate them up front by a (min, max-exclusive) compute
# capability range. FlashAttention 3 is a Hopper-SM90 rewrite with no Blackwell kernel, so it
# needs an upper bound: an explicit flash3 on a B200 (SM100) must drop to native instead of
# setting fine then crashing at generation. FlashAttention 4 is Blackwell+ (no upper bound).
_ARCH_CAPABILITY: dict[str, tuple[tuple[int, int], Optional[tuple[int, int]]]] = {
"_flash_3_hub": ((9, 0), (10, 0)), # FlashAttention 3 -> Hopper (SM90) only
"flash_4_hub": ((10, 0), None), # FlashAttention 4 -> Blackwell (SM100)+
}
def _cuda_capability() -> Optional[tuple[int, int]]:
"""(major, minor) compute capability of the active CUDA device, or None if unknown."""
try:
import torch
if not torch.cuda.is_available():
return None
return tuple(torch.cuda.get_device_capability()) # type: ignore[return-value]
except Exception: # noqa: BLE001
return None
def _backend_arch_supported(backend: str) -> bool:
"""False only when ``backend`` needs a CUDA arch outside this device's supported range.
Unknown capability (no CUDA / detection failure) returns True so we never block on a
guess -- diffusers' own set-time check still guards the package, and a genuine run-time
failure falls back to native."""
bounds = _ARCH_CAPABILITY.get(backend)
if bounds is None:
return True
have = _cuda_capability()
if have is None:
return True
low, high = bounds
return have >= low and (high is None or have < high)
def _is_cuda_nvidia(target: Any) -> bool:
"""CUDA device on an NVIDIA (non-ROCm) build -- where cuDNN attention applies."""
if getattr(target, "device", None) != "cuda":
@ -87,13 +127,33 @@ def select_attention_backend(
alias = normalize_attention_backend(requested)
if alias != ATTN_AUTO:
backend = _ALIASES[alias]
return None if backend == "native" else backend
if backend == "native":
return None
# An arch-gated kernel (flash3/flash4) on a card that can't run it would set fine
# then crash mid-generation, so drop it to the native default up front.
if not _backend_arch_supported(backend):
return None
# cuDNN fused SDPA needs Ampere+ (SM80); diffusers accepts it on pre-SM80 cards
# (T4/V100) then fails at the first generation, so apply the same gate to an
# explicit cuDNN request as the auto path already does.
if backend == "_native_cudnn" and not _cudnn_attention_supported():
return None
return backend
# auto
if speed_active and _is_cuda_nvidia(target):
if speed_active and _is_cuda_nvidia(target) and _cudnn_attention_supported():
return "_native_cudnn"
return None
def _cudnn_attention_supported() -> bool:
"""cuDNN fused SDPA needs Ampere+ (SM80). On pre-SM80 NVIDIA cards (T4 SM75 /
V100 SM70) diffusers accepts ``_native_cudnn`` at set time but the kernel fails at
the first generation, so gate the auto-cuDNN upgrade on capability. Unknown
capability allows it (diffusers' set-time check + the run-time fallback still guard)."""
have = _cuda_capability()
return have is None or have >= (8, 0)
def apply_attention_backend(
pipe: Any,
backend: Optional[str],
@ -102,25 +162,85 @@ def apply_attention_backend(
) -> Optional[str]:
"""Set ``backend`` on ``pipe.transformer`` via the diffusers dispatcher.
Returns the backend actually engaged, or None when left at the default (either because
``backend`` was None or because the requested kernel was unavailable -> graceful
fallback to the diffusers default, never a load failure). Best-effort."""
if backend is None:
return None
Returns the backend actually engaged, or None when left at the native default (either
because ``backend`` was None or because the requested kernel was unavailable -> graceful
fallback, never a load failure).
diffusers keeps a *process-wide* active attention backend that ``set_attention_backend``
also updates, and a fresh transformer's processors follow it (their ``_attention_backend``
defaults to None). So a load that wants native must restore it explicitly: otherwise it
silently inherits a backend an earlier load pinned (e.g. cuDNN under a speed profile),
breaking the bit-identical/``off`` guarantee. Best-effort throughout."""
transformer = getattr(pipe, "transformer", None)
fn = getattr(transformer, "set_attention_backend", None)
if not callable(fn):
return None
if backend is not None:
try:
fn(backend)
# set_attention_backend also pins the backend in diffusers' process-wide
# registry. This transformer's own processors keep it locally (their
# _attention_backend is now explicit), so reset the global default back to
# native -- otherwise a later component whose processors are unconfigured
# (backend None) silently inherits this kernel.
_reset_global_backend_to_native(logger)
if logger is not None:
logger.info("diffusion.attention: backend=%s", backend)
return backend
except Exception as exc: # noqa: BLE001 — unavailable kernel -> restore native below
_warn(logger, backend, exc)
# No backend requested, or the requested one failed: pin the native default so a stale
# process-wide backend from a previous load can't leak into this one.
_restore_native_backend(fn, logger)
return None
def _active_attention_backend() -> Optional[str]:
"""The diffusers process-wide active attention backend name, or None if undeterminable."""
try:
fn(backend)
if logger is not None:
logger.info("diffusion.attention: backend=%s", backend)
return backend
except Exception as exc: # noqa: BLE001 — unavailable kernel -> diffusers default
_warn(logger, backend, exc)
from diffusers.models.attention_dispatch import _AttentionBackendRegistry
# get_active_backend() returns a (AttentionBackendName, fn) tuple (or None), so
# take element 0 and read its .value (e.g. "native"); reading .value off the
# tuple itself would yield a junk string that never compares equal to a name.
active = _AttentionBackendRegistry.get_active_backend()
if active is None:
return None
name = active[0] if isinstance(active, tuple) else active
return getattr(name, "value", str(name))
except Exception: # noqa: BLE001
return None
def _reset_global_backend_to_native(logger: Any) -> None:
"""Reset diffusers' process-wide active attention backend to native after a
successful per-transformer set, so a later component whose processors are
unconfigured (backend None) does not inherit this transformer's kernel. The
transformer's own processors keep the backend just set. Best-effort and silent:
if the diffusers internals move, the prior (leaking) behavior is unchanged."""
if _active_attention_backend() == ATTN_NATIVE:
return
try:
from diffusers.models.attention_dispatch import (
AttentionBackendName,
_AttentionBackendRegistry,
)
_AttentionBackendRegistry.set_active_backend(AttentionBackendName.NATIVE)
except Exception: # noqa: BLE001 — best-effort; leave the global as-is on any change
pass
def _restore_native_backend(set_backend_fn: Any, logger: Any) -> None:
"""Force the native default when the global active backend isn't already native."""
if _active_attention_backend() == ATTN_NATIVE:
return # already native -> avoid redundant work and an extra dispatcher warning
try:
set_backend_fn(ATTN_NATIVE)
except Exception as exc: # noqa: BLE001 — best-effort restore
_warn(logger, ATTN_NATIVE, exc)
def _warn(logger: Any, what: str, exc: Exception) -> None:
if logger is not None:
logger.warning("diffusion.attention: %s unavailable (%s); using default", what, exc)

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

@ -77,16 +77,22 @@ def apply_step_cache(
if threshold is not None
else (QUANT_FBCACHE_THRESHOLD if quant_active else DEFAULT_FBCACHE_THRESHOLD)
)
# Only engage via the transformer's native enable_cache (the diffusers CacheMixin path).
# That mixin is present exactly when the pipeline wraps the transformer call in a
# cache_context, which the First-Block-Cache hook requires at run time. The lower-level
# apply_first_block_cache hook would install on a non-CacheMixin transformer too (e.g.
# Z-Image), but its pipeline opens no cache_context, so the first generation would crash
# inside the hook -- so a model without enable_cache runs uncached per the best-effort
# contract instead of being reported as cached and then failing.
enable_cache = getattr(transformer, "enable_cache", None)
if not callable(enable_cache):
_warn(logger, mode, RuntimeError("transformer has no cache_context (not a CacheMixin)"))
return None
try:
from diffusers import FirstBlockCacheConfig
config = FirstBlockCacheConfig(threshold = thr)
enable_cache = getattr(transformer, "enable_cache", None)
if callable(enable_cache):
enable_cache(config)
else:
from diffusers.hooks import apply_first_block_cache
apply_first_block_cache(transformer, config)
enable_cache(config)
try:
transformer._unsloth_step_cache = f"{mode}@{thr}"
except Exception: # noqa: BLE001 — marker is best-effort

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

@ -14,6 +14,7 @@ diffusers classes and base repo needed to assemble the full pipeline.
from __future__ import annotations
import re
from dataclasses import dataclass, field
from pathlib import Path, PurePosixPath
from typing import Optional
@ -333,6 +334,6 @@ def resolve_local_gguf_child(repo_root: Path, gguf_filename: str) -> Path:
child = repo_root.joinpath(*rel.parts).resolve()
if child != repo_real and repo_real not in child.parents:
raise ValueError("gguf_filename must resolve to a file inside the repo.")
if not child.exists():
raise FileNotFoundError(f"'{gguf_filename}' not found under {repo_root}.")
if not child.is_file():
raise FileNotFoundError(f"'{gguf_filename}' is not a file under {repo_root}.")
return child

19
studio/backend/core/inference/diffusion_memory.py
View file

@ -164,7 +164,10 @@ def _cuda_memory(backend: str) -> tuple[Optional[int], Optional[int], str]:
free, total = torch.cuda.mem_get_info()
kind = "discrete_vram"
try:
props = torch.cuda.get_device_properties(0)
# Query the CURRENT device, not device 0: mem_get_info() above already
# reports the active device, so hardcoding 0 would inspect the wrong GPU
# (and misclassify discrete vs unified) when the active device isn't 0.
props = torch.cuda.get_device_properties(torch.cuda.current_device())
if bool(getattr(props, "integrated", False) or getattr(props, "is_integrated", False)):
kind = "unified_memory" # e.g. Jetson / integrated SoC
except Exception:
@ -454,12 +457,22 @@ def apply_memory_plan(
if plan.vae_slicing:
_enable_vae_saver(pipe, "enable_vae_slicing", "enable_slicing", logger)
def _fallback_to_model_offload() -> None:
# Group offload keeps the VAE resident, so the GROUP plan set vae_tiling=False.
# When group offload is unavailable and we drop to whole-module offload, the card
# is in the low-VRAM situation where the decode-time spike can OOM, so turn VAE
# tiling on now (if not already engaged) to cap it.
nonlocal tiling_engaged
pipe.enable_model_cpu_offload()
if not tiling_engaged:
tiling_engaged = _enable_vae_saver(pipe, "enable_vae_tiling", "enable_tiling", logger)
policy = plan.offload_policy
if policy == OFFLOAD_MODEL:
pipe.enable_model_cpu_offload()
elif policy == OFFLOAD_GROUP:
if not _apply_group_offload(pipe, device, logger):
pipe.enable_model_cpu_offload()
_fallback_to_model_offload()
policy = OFFLOAD_MODEL
elif policy == OFFLOAD_SEQUENTIAL:
try:
@ -471,7 +484,7 @@ def apply_memory_plan(
"falling back to whole-module offload",
exc,
)
pipe.enable_model_cpu_offload()
_fallback_to_model_offload()
policy = OFFLOAD_MODEL
else:
pipe.to(device)

33
studio/backend/core/inference/diffusion_precision.py
View file

@ -98,15 +98,46 @@ def quantize_text_encoders(
def _cast_fp8(encoder: Any, target: Any) -> None:
import re
import torch
from diffusers.hooks import apply_layerwise_casting
from diffusers.hooks.layerwise_casting import DEFAULT_SKIP_MODULES_PATTERN
# diffusers' layerwise casting stores each supported leaf module's weights in fp8 and
# upcasts them per forward. Two things on a transformers text encoder can push an fp8
# weight or activation into an op that can't handle it, and both crash only at
# generation (the load-time guard can't see them), so skip the offending modules:
skip = tuple(DEFAULT_SKIP_MODULES_PATTERN)
# (1) dtype-sensitive modules the encoder itself flags. T5 keeps "wo" in fp32: its
# gated feed-forward reads self.wo.weight.dtype and casts the activations to match
# BEFORE calling wo (transformers#20287), racing the forward-time upcast hook so
# F.linear sees an fp8 input against a bf16 weight. Names are literal substrings.
skip += tuple(re.escape(m) for m in (getattr(encoder, "_keep_in_fp32_modules", None) or ()))
# (2) an output projection tied to the input embedding. A CausalLM encoder (FLUX.2's
# Qwen3) ties lm_head.weight to embed_tokens.weight; lm_head is an nn.Linear so it
# gets cast to fp8 and, sharing one tensor, drags the embedding to fp8 with it. The
# embedding then emits fp8 activations that crash the first RMSNorm. Skip the tied
# projection so the shared tensor stays dense (lm_head is unused for prompt encoding).
get_out, get_in = getattr(encoder, "get_output_embeddings", None), getattr(encoder, "get_input_embeddings", None)
out_emb = get_out() if callable(get_out) else None
in_emb = get_in() if callable(get_in) else None
if out_emb is not None and in_emb is not None and out_emb.weight is in_emb.weight:
tied_name = next((n for n, m in encoder.named_modules() if m is out_emb), None)
if tied_name:
skip += (rf"^{re.escape(tied_name)}$",)
apply_layerwise_casting(
encoder,
storage_dtype = torch.float8_e4m3fn,
compute_dtype = target.dtype,
skip_modules_pattern = DEFAULT_SKIP_MODULES_PATTERN,
skip_modules_pattern = skip,
# Keep token-embedding tables (T5 "shared", Qwen "embed_tokens", etc.) full
# precision: the diffusers default pattern only skips vision pos/patch
# embeds, not nn.Embedding lookups, and fp8'ing those quantizes every prompt
# token straight to the coarse fp8 grid, hurting prompt fidelity.
skip_modules_classes = (torch.nn.Embedding,),
)

120
studio/backend/core/inference/diffusion_prequant.py
View file

@ -33,6 +33,58 @@ from typing import Any, Optional
# on-disk structure changes so an old/foreign artifact is rejected rather than mis-loaded.
PREQUANT_FORMAT = "unsloth_prequant_transformer_state_dict_v1"
# Loading a checkpoint ends in ``torch.load(weights_only=False)``, which executes arbitrary
# code embedded in the pickle. A hosted family *repo* checkpoint is first-party and trusted,
# but a ``source.kind == "path"`` can originate from the ``transformer_prequant_path`` field
# of a load request -- i.e. an authenticated API caller naming an arbitrary local file.
# Unpickling that is remote code execution, so a request-supplied path is unpickled ONLY when
# it resolves inside an operator-configured ALLOWLIST of directories. A bare on/off toggle is
# deliberately NOT accepted as a wildcard: enabling local checkpoints for one trusted
# directory must never also permit unpickling any other path a request happens to name. The
# trusted hosted-repo path is unaffected.
ALLOW_LOCAL_PREQUANT_PATH_ENV = "UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH"
_PREQUANT_TOGGLE_TOKENS = {"1", "true", "yes", "on", "0", "false", "no", "off"}
def _allowed_prequant_roots() -> list:
"""Operator-allowlisted directories whose pre-quant checkpoints may be unpickled.
Set ``UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH`` to one or more directories (separated by
``os.pathsep``). A bare truthy/falsey toggle is ignored on purpose -- it must name a
directory, so there is no "allow everything" mode."""
import os
raw = (os.environ.get(ALLOW_LOCAL_PREQUANT_PATH_ENV) or "").strip()
if not raw:
return []
roots = []
for part in raw.split(os.pathsep):
part = part.strip()
if not part or part.lower() in _PREQUANT_TOGGLE_TOKENS:
continue # a bare on/off value is not a directory -> never a wildcard allow
try:
roots.append(os.path.realpath(os.path.expanduser(part)))
except Exception: # noqa: BLE001 — a bad entry is simply not allowlisted
continue
return roots
def _local_prequant_path_allowed(path: str) -> bool:
"""True only when ``path`` resolves inside an operator-allowlisted directory; an
arbitrary request-supplied path is never unpickled. ``realpath`` first so a symlink
cannot point an allowlisted name at a file outside the allowed roots."""
import os
roots = _allowed_prequant_roots()
if not roots:
return False
try:
real = os.path.realpath(os.path.expanduser(path))
except Exception: # noqa: BLE001
return False
return any(real == r or real.startswith(r + os.sep) for r in roots)
@dataclass(frozen = True)
class PrequantSource:
@ -83,6 +135,7 @@ def load_prequantized_transformer(
dtype: Any,
hf_token: Optional[str] = None,
scheme: str,
min_features: Optional[int] = None,
logger: Any = None,
) -> Optional[Any]:
"""Load the pre-quantized transformer described by ``source`` onto ``device``.
@ -93,6 +146,21 @@ def load_prequantized_transformer(
ordinary unavailable artifact.
"""
try:
# weights_only=False (required below) executes pickle code, so a caller-supplied
# local path is unpickled ONLY when it resolves inside an operator-allowlisted
# directory. The hosted family repo is first-party and always allowed.
if source.kind == "path" and not _local_prequant_path_allowed(source.location):
_warn(
logger,
f"{scheme}:path",
RuntimeError(
"request-supplied local pre-quant path refused (unpickling an arbitrary "
f"file is unsafe); set {ALLOW_LOCAL_PREQUANT_PATH_ENV} to an allowlisted "
"directory containing trusted checkpoints to permit it",
),
)
return None
path = _resolve_checkpoint_path(source, hf_token)
if path is None:
return None
@ -100,11 +168,10 @@ def load_prequantized_transformer(
import torch
# torchao weight subclasses are not safetensors-serializable, so the checkpoint is
# a torch.save pickle. weights_only=False is required to rebuild those subclasses;
# only a configured family repo (first-party) or an explicit local path reaches
# here, which is the trust signal -- this never loads an arbitrary remote pickle.
# a torch.save pickle. weights_only=False is required to rebuild those subclasses.
# The local-path branch is gated above; the repo branch is a first-party artifact.
ckpt = torch.load(path, weights_only = False, map_location = "cpu")
if not _validate_checkpoint(ckpt, scheme, base, logger):
if not _validate_checkpoint(ckpt, scheme, base, logger, min_features = min_features):
return None
state_dict = ckpt["state_dict"]
@ -154,8 +221,19 @@ def _resolve_checkpoint_path(source: PrequantSource, hf_token: Optional[str]) ->
return None
def _validate_checkpoint(ckpt: Any, scheme: str, base: str, logger: Any) -> bool:
"""Reject a checkpoint that is the wrong format / scheme / base model."""
def _validate_checkpoint(
ckpt: Any,
scheme: str,
base: str,
logger: Any,
min_features: Optional[int] = None,
) -> bool:
"""Reject a checkpoint that is the wrong format / scheme / base model / filter.
``min_features`` (when given) is the runtime Linear-feature threshold: a checkpoint
built with a different ``--min-features`` quantises a different set of Linear layers,
so ``load_state_dict(assign=True)`` would silently install a model that does not match
what the dense path produces while status still reports the requested scheme. Reject it."""
if not isinstance(ckpt, dict) or ckpt.get("format") != PREQUANT_FORMAT:
_warn(logger, scheme, ValueError("unrecognised pre-quant checkpoint format"))
return False
@ -167,21 +245,41 @@ def _validate_checkpoint(ckpt: Any, scheme: str, base: str, logger: Any) -> bool
_warn(logger, scheme, ValueError(f"checkpoint scheme {meta.get('scheme')!r} != {scheme!r}"))
return False
ckpt_base = meta.get("base_model_id")
if ckpt_base and base and ckpt_base != base:
if ckpt_base and base and not _same_base_model(ckpt_base, base):
_warn(logger, scheme, ValueError(f"checkpoint base {ckpt_base!r} != {base!r}"))
return False
if min_features is not None:
ckpt_min = meta.get("min_features")
if ckpt_min is not None and int(ckpt_min) != int(min_features):
_warn(
logger,
scheme,
ValueError(f"checkpoint min_features {ckpt_min!r} != runtime {min_features!r}"),
)
return False
return True
def _same_base_model(a: str, b: str) -> bool:
"""Tolerant compare of two base-model ids: an exact match, or the same final
path/repo segment (so a local path or a fork id matches the canonical repo, e.g.
``/models/Z-Image-Turbo`` vs ``Tongyi-MAI/Z-Image-Turbo``)."""
def _tail(x: str) -> str:
return x.replace("\\", "/").rstrip("/").split("/")[-1].lower()
return a == b or _tail(a) == _tail(b)
def _has_meta_tensors(module: Any) -> bool:
"""True if any parameter or buffer is still on the meta device after loading."""
from itertools import chain
try:
for tensor in list(module.parameters()) + list(module.buffers()):
if getattr(tensor, "is_meta", False):
return True
return any(
getattr(t, "is_meta", False) for t in chain(module.parameters(), module.buffers())
)
except Exception: # noqa: BLE001
return False
return False
def _warn(logger: Any, what: str, exc: Exception) -> None:

82
studio/backend/core/inference/diffusion_speed.py
View file

@ -57,31 +57,51 @@ SPEED_MODES = (SPEED_OFF, SPEED_EAGER, SPEED_DEFAULT, SPEED_MAX)
def snapshot_backend_flags() -> Optional[dict]:
"""Capture the process-wide torch backend flags this layer may mutate, so the
caller can restore them on unload. None if torch is unavailable."""
caller can restore them on unload. None if torch is unavailable. Each flag is read
defensively so a build/platform missing one (e.g. no cuda.matmul on CPU/MPS) still
captures the rest -- otherwise a single missing attribute would skip the whole
snapshot and a real mutated flag would leak."""
try:
import torch
return {
"matmul_tf32": bool(torch.backends.cuda.matmul.allow_tf32),
"cudnn_tf32": bool(torch.backends.cudnn.allow_tf32),
"cudnn_benchmark": bool(torch.backends.cudnn.benchmark),
}
except Exception: # noqa: BLE001 — best-effort; no snapshot -> no restore
except Exception: # noqa: BLE001 — no torch -> nothing to snapshot/restore
return None
state: dict[str, bool] = {}
matmul = getattr(getattr(torch.backends, "cuda", None), "matmul", None)
if matmul is not None and hasattr(matmul, "allow_tf32"):
state["matmul_tf32"] = bool(matmul.allow_tf32)
cudnn = getattr(torch.backends, "cudnn", None)
if cudnn is not None:
if hasattr(cudnn, "allow_tf32"):
state["cudnn_tf32"] = bool(cudnn.allow_tf32)
if hasattr(cudnn, "benchmark"):
state["cudnn_benchmark"] = bool(cudnn.benchmark)
return state
def restore_backend_flags(state: Optional[dict]) -> None:
"""Restore the flags captured by ``snapshot_backend_flags``. No-op on None."""
"""Restore the flags captured by ``snapshot_backend_flags``. No-op on None. Each
flag is restored independently so one failure can't leave the others leaked."""
if not state:
return
try:
import torch
torch.backends.cuda.matmul.allow_tf32 = state["matmul_tf32"]
torch.backends.cudnn.allow_tf32 = state["cudnn_tf32"]
torch.backends.cudnn.benchmark = state["cudnn_benchmark"]
except Exception: # noqa: BLE001 — best-effort restore
except Exception: # noqa: BLE001 — no torch -> nothing to restore
return
def _set(obj: Any, attr: str, key: str) -> None:
if obj is not None and key in state and hasattr(obj, attr):
try:
setattr(obj, attr, state[key])
except Exception: # noqa: BLE001 — best-effort per-flag restore
pass
_set(
getattr(getattr(torch.backends, "cuda", None), "matmul", None), "allow_tf32", "matmul_tf32"
)
cudnn = getattr(torch.backends, "cudnn", None)
_set(cudnn, "allow_tf32", "cudnn_tf32")
_set(cudnn, "benchmark", "cudnn_benchmark")
def normalize_speed_mode(value: Optional[str]) -> str:
"""Lower/strip a requested speed mode (dashes ok); None / "" -> off."""
@ -158,6 +178,12 @@ def apply_speed_optims(
"compiled_dequant": False,
}
mode = normalize_speed_mode(speed_mode)
# TF32 is the one PROCESS-GLOBAL flag we flip (on max). Restore it whenever this
# load isn't max, so a later default/off diffusion load -- or chat inference in the
# same long-lived process -- doesn't silently inherit a prior max load's TF32 and
# lose the bit-identical default the regression harness checks.
if mode != SPEED_MAX:
_restore_tf32(logger)
if mode == SPEED_OFF:
return applied
@ -256,12 +282,22 @@ def _enable_cudnn_benchmark(logger: Any) -> bool:
except Exception as exc: # noqa: BLE001 — optimisation only
_warn(logger, "cudnn_benchmark", exc)
return False
# The TF32 flag values from before the first max load flipped them, so a later
# non-max load / unload can put the process back exactly as it found it (rather than
# forcing a hardcoded default that might clobber another component's choice).
_tf32_prev: Optional[tuple[bool, bool]] = None
def _enable_tf32(logger: Any) -> bool:
global _tf32_prev
try:
import torch
if _tf32_prev is None:
_tf32_prev = (
torch.backends.cuda.matmul.allow_tf32,
torch.backends.cudnn.allow_tf32,
)
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
return True
@ -270,6 +306,26 @@ def _enable_tf32(logger: Any) -> bool:
return False
def restore_tf32(logger: Any = None) -> None:
"""Put the process-global TF32 flags back to their pre-max-load values. No-op if
a max load never set them. Called on a non-max load and on unload."""
_restore_tf32(logger)
def _restore_tf32(logger: Any) -> None:
global _tf32_prev
if _tf32_prev is None:
return
try:
import torch
torch.backends.cuda.matmul.allow_tf32 = _tf32_prev[0]
torch.backends.cudnn.allow_tf32 = _tf32_prev[1]
except Exception as exc: # noqa: BLE001 — best-effort restore
_warn(logger, "tf32_restore", exc)
finally:
_tf32_prev = None
def _fuse_qkv(pipe: Any, logger: Any) -> bool:
for owner in (pipe, getattr(pipe, "transformer", None)):
fn = getattr(owner, "fuse_qkv_projections", None)

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

@ -106,13 +106,15 @@ _DATACENTER_GPU_TOKENS = frozenset(
{
"B200",
"B100",
"B300", # Blackwell Ultra data center
"GB200",
"GB300",
"GB10", # Blackwell data center
"H200",
"H100",
"H800",
"H20", # Hopper data center
"H20",
"GH200", # Grace-Hopper superchip (data center)
"A100",
"A800",
"A30",
@ -133,15 +135,22 @@ _DATACENTER_GPU_TOKENS = frozenset(
)
# Professional parts the rest of the backend treats as datacenter-class (see llama_cpp.py
# _DATACENTER_GPU_RE, which applies the same FP32-accum tuning to them). Matched as phrases
# because the marker spans tokens ("RTX PRO 6000", "RTX 6000 ADA"), so they must not be
# misread as consumer (which would put int8 ahead of fp8 and pick fast accumulate).
_PROFESSIONAL_GPU_MARKERS = ("RTX PRO 6000", "RTX 6000 ADA")
def _is_consumer_gpu(device: Any = None) -> bool:
"""Whether the active GPU is consumer / workstation class (GDDR), where fp8 FP32
accumulate is throughput-halved so fast (FP16) accumulate is a ~2x win. Data-center
HBM parts (recognised by name token) are not nerfed and return False, so they keep
the higher-precision default accumulate for free. Heuristic on the device name: a
GeForce / TITAN name is always consumer; a recognised data-center token is not;
anything else (workstation RTX, unknown) defaults to consumer -- the safe choice,
since fast accumulate is free on data-center and a win on consumer. Best-effort:
True on any probe failure."""
"""Whether the active GPU is consumer-class (GDDR), where fp8 FP32 accumulate is
throughput-halved so fast (FP16) accumulate is a ~2x win. Data-center HBM parts and
professional parts (recognised by name) are not nerfed and return False, so they keep
the higher-precision default accumulate and fp8 first. Heuristic on the device name: a
GeForce / TITAN name is always consumer; a recognised data-center token or professional
marker is not; anything else (unknown) defaults to consumer -- the safe choice, since
fast accumulate is free on data-center and a win on consumer. Best-effort: True on any
probe failure."""
try:
import re
@ -151,6 +160,8 @@ def _is_consumer_gpu(device: Any = None) -> bool:
return True
if "GEFORCE" in name or "TITAN" in name:
return True
if any(marker in name for marker in _PROFESSIONAL_GPU_MARKERS):
return False
tokens = set(re.split(r"[^A-Z0-9]+", name))
return not (tokens & _DATACENTER_GPU_TOKENS)
@ -302,7 +313,14 @@ def _make_quant_config(scheme: str, fast_accum: Optional[bool] = None) -> Any:
return Float8DynamicActivationFloat8WeightConfig()
if scheme == TQ_NVFP4:
from torchao.prototype.mx_formats import NVFP4DynamicActivationNVFP4WeightConfig
return NVFP4DynamicActivationNVFP4WeightConfig()
# Select the CUTLASS FP4 path, not the default Triton kernel: torchao defaults
# use_triton_kernel=True, which needs MSLK installed. On a Blackwell box with the
# CUTLASS FP4 extension but no MSLK, the default would make the smoke probe fail
# and silently fall back to GGUF instead of using the FP4 tensor cores.
try:
return NVFP4DynamicActivationNVFP4WeightConfig(use_triton_kernel = False)
except TypeError: # older torchao without the knob
return NVFP4DynamicActivationNVFP4WeightConfig()
if scheme == TQ_MXFP8:
import torch
from torchao.prototype.mx_formats import MXDynamicActivationMXWeightConfig
@ -310,7 +328,9 @@ def _make_quant_config(scheme: str, fast_accum: Optional[bool] = None) -> Any:
return MXDynamicActivationMXWeightConfig(
activation_dtype = torch.float8_e4m3fn, weight_dtype = torch.float8_e4m3fn
)
except TypeError:
except (TypeError, AttributeError):
# TypeError: older torchao without the explicit dtype knobs.
# AttributeError: a torch build without torch.float8_e4m3fn.
return MXDynamicActivationMXWeightConfig()
raise ValueError(f"unknown transformer quant scheme '{scheme}'")
@ -335,7 +355,7 @@ def make_filter_fn(min_features: int, exclude_name_tokens: tuple[str, ...] = ())
if in_features < min_features or out_features < min_features:
return False
if exclude_name_tokens:
name = fqn.lower()
name = fqn.lower() if fqn else ""
if any(tok in name for tok in exclude_name_tokens):
return False
return True

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

@ -771,11 +771,9 @@ class ExternalProviderClient:
self.base_url = self.base_url[: -len("/openai")]
self.api_key = api_key
self._timeout = httpx.Timeout(timeout, connect = 10.0)
# Disable read timeout on SSE streams: reasoning-heavy models pause
# tens of seconds between bytes while thinking, and httpx's read
# timeout is the per-byte gap, not wall clock. connect/write bounds
# still surface real network failures.
self._stream_timeout = httpx.Timeout(timeout, connect = 10.0, read = None)
# Generous per-byte read timeout: reasoning models pause tens of seconds
# between bytes, but a dead upstream must eventually error, not hang forever.
self._stream_timeout = httpx.Timeout(timeout, connect = 10.0, read = 300.0)
def _auth_headers(self) -> dict[str, str]:
"""Build authentication headers using the provider's registry config."""

6
studio/backend/core/inference/gpu_arbiter.py
View file

@ -36,7 +36,11 @@ def _evict_chat() -> None:
from routes.inference import get_llama_cpp_backend
llama = get_llama_cpp_backend()
if llama.is_loaded:
# is_active (process exists), not is_loaded (process exists AND healthy): a
# chat model still starting up holds/keeps allocating VRAM but isn't healthy
# yet, so gating on is_loaded would skip it and let the load race the
# diffusion pipeline. unload_model() sets _cancel_event and kills the process.
if llama.is_active:
llama.unload_model()
orchestrator = get_inference_backend()
if orchestrator.active_model_name:

16
studio/backend/core/inference/image_gallery.py
View file

@ -143,14 +143,24 @@ def list_images(limit: Optional[int] = None, offset: int = 0) -> list[dict[str,
except OSError:
return []
paths.sort(key = _mtime, reverse = True)
window = paths[offset:] if limit is None else paths[offset : offset + limit]
# Page over READABLE records, not raw files: filtering a foreign/corrupt PNG out of an
# already-sliced window would drop valid images that sort after it and make the route's
# has_more wrong. Read only as far as needed to fill the requested window.
# Known Phase-1 limit: this re-reads headers from the newest down to `offset+limit` on
# every page, so a deep infinite-scroll over a very large gallery (thousands of images,
# e.g. a long uncapped batch) is O(offset) header-opens per page. PIL opens are lazy
# (header only) and this runs off the event loop, so it's not a freeze; a later phase can
# switch to cursor-based paging (resume after the last-seen record) if it starts to bite.
want = None if limit is None else offset + limit
records = []
for path in window:
for path in paths:
meta = _read_meta(path)
if meta is None: # not one of ours (no recipe chunk) — skip
continue
records.append(_record(path.stem, meta))
return records
if want is not None and len(records) >= want:
break
return records[offset:] if limit is None else records[offset : offset + limit]
def delete(image_id: str) -> bool:

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

@ -1271,6 +1271,9 @@ class LlamaCppBackend:
self._cache_type_kv: Optional[str] = None
# Whether --split-mode tensor was applied on the active load.
self._tensor_parallel: bool = False
# Layer load kept multi-GPU only to honor a downgraded tensor request, so a
# later explicit tensor-off reloads instead of deduping to it (#6659).
self._layer_preserves_tensor_intent: bool = False
self._reasoning_default: bool = True
self._speculative_type: Optional[str] = None
# Canonical UI-facing mode the user requested
@ -1643,6 +1646,11 @@ class LlamaCppBackend:
"""Whether --split-mode tensor is active on the loaded server."""
return self._tensor_parallel
@property
def layer_preserves_tensor_intent(self) -> bool:
"""True when a downgraded tensor request kept this layer load multi-GPU."""
return self._layer_preserves_tensor_intent
@property
def speculative_type(self) -> Optional[str]:
return self._speculative_type
@ -2430,6 +2438,37 @@ class LlamaCppBackend:
# aborts a --split-mode tensor load, so it's dropped for the tensor attempt.
_TENSOR_PARALLEL_KV_TYPES = frozenset({"f16", "bf16", "f32"})
# (binary, mtime, model) that aborted on --split-mode tensor this process (#6415
# geometry limit, e.g. MQA n_head_kv=1). Model-keyed so one model's abort doesn't
# skip tensor for others; tensor is tried by default, recorded only on a real abort.
_tensor_split_abort_keys: set[tuple[str, int, str]] = set()
@classmethod
def _tensor_split_cache_key(
cls, binary: Optional[str], model: Optional[str]
) -> Optional[tuple[str, int, str]]:
"""(path, mtime_ns, model) key; ns mtime re-probes a same-second binary swap."""
if not binary or not model:
return None
try:
mtime = Path(binary).stat().st_mtime_ns
except OSError:
mtime = 0
return (binary, mtime, model)
@classmethod
def _tensor_split_aborts(cls, binary: Optional[str], model: Optional[str]) -> bool:
"""True if (binary, model) aborted on --split-mode tensor this session."""
key = cls._tensor_split_cache_key(binary, model)
return key is not None and key in cls._tensor_split_abort_keys
@classmethod
def _record_tensor_split_abort(cls, binary: Optional[str], model: Optional[str]) -> None:
"""Remember a (binary, model) that aborts on --split-mode tensor."""
key = cls._tensor_split_cache_key(binary, model)
if key is not None:
cls._tensor_split_abort_keys.add(key)
@staticmethod
def _windows_pip_nvidia_dll_dirs(prefix: str) -> list[str]:
"""Return DLL dirs from pip-installed CUDA wheels under
@ -2569,9 +2608,13 @@ class LlamaCppBackend:
usable_fraction: Optional[float] = None,
total_by_idx: Optional[dict[int, int]] = None,
per_device_overhead_bytes: int = 0,
min_gpus: int = 1,
) -> tuple[Optional[list[int]], bool]:
"""Pick GPU(s) for a model from estimated VRAM and free memory.
``min_gpus`` (default 1, capped at ``len(gpus)``) keeps a downgraded
tensor/multi-GPU request spread instead of collapsing to one card.
``model_size_bytes`` should include weights and estimated KV cache.
``usable_fraction`` (default ``_GPU_PIN_VRAM_FRACTION``) provides
headroom for compute buffers, CUDA context, and other runtime
@ -2590,9 +2633,11 @@ class LlamaCppBackend:
if not gpus:
return None, True
min_gpus = max(1, min(min_gpus, len(gpus)))
model_size_mib = model_size_bytes / (1024 * 1024)
if usable_fraction is None:
usable_fraction = LlamaCppBackend._GPU_PIN_VRAM_FRACTION
overhead_mib = per_device_overhead_bytes / (1024 * 1024)
# Per-GPU usable budget: free - (1-frac)*total when total is known, else
# the legacy free*frac (also covers a total-0 two-column probe).
@ -2606,19 +2651,26 @@ class LlamaCppBackend:
# card can have less usable room than a less-used small one.
ranked = sorted(gpus, key = lambda g: _usable(g[0], g[1]), reverse = True)
# Try 1 GPU at the usable-VRAM threshold.
if _usable(ranked[0][0], ranked[0][1]) >= model_size_mib:
# Cap a downgraded multi-GPU request to the usable count so it doesn't pull
# in a near-full card to hit min_gpus. No-op for the default min_gpus == 1.
usable_count = sum(1 for idx, free_mib in ranked if _usable(idx, free_mib) > overhead_mib)
min_gpus = max(1, min(min_gpus, usable_count or 1))
# Try 1 GPU at the usable-VRAM threshold (only when one device is allowed).
if min_gpus <= 1 and _usable(ranked[0][0], ranked[0][1]) >= model_size_mib:
return [ranked[0][0]], False
# Try N GPUs (accumulate usable memory from most-free). Each GPU past the
# first adds a fixed per-device overhead the pool must hold.
overhead_mib = per_device_overhead_bytes / (1024 * 1024)
# Try N GPUs (most-free first); each past the first adds per-device overhead.
# Require at least min_gpus devices before accepting a fit.
cumulative = 0.0
selected = []
for idx, free_mib in ranked:
selected.append(idx)
cumulative += _usable(idx, free_mib)
if cumulative >= model_size_mib + (len(selected) - 1) * overhead_mib:
if (
len(selected) >= min_gpus
and cumulative >= model_size_mib + (len(selected) - 1) * overhead_mib
):
return sorted(selected), False
# Too large even for all GPUs; let --fit handle it
@ -3147,9 +3199,10 @@ class LlamaCppBackend:
except (ValueError, OSError):
# Log file closed under us; tee silently.
pass
except (ValueError, OSError):
# Pipe closed -- process terminating.
pass
except Exception:
# Never let the drain thread die: a full stdout pipe can deadlock
# llama-server (Windows). Pipe-closed on exit is the common case.
logger.debug("llama-server stdout drain stopped", exc_info = True)
# GGUF KV type sizes for fast skipping
_GGUF_TYPE_SIZE = {
@ -3644,12 +3697,22 @@ class LlamaCppBackend:
hf_repo: str,
hf_variant: Optional[str] = None,
hf_token: Optional[str] = None,
force: bool = False,
allow_smaller_fallback: bool = True,
cancel_event: Optional[threading.Event] = None,
) -> str:
"""Download GGUF file(s) from HuggingFace. Returns local path.
Runs WITHOUT self._lock so unload_model() can set _cancel_event at
any time; checks it between each shard download.
``force`` re-fetches even when a (possibly stale) blob is cached.
``allow_smaller_fallback=False`` raises on low disk instead of silently
switching to a smaller quant. ``cancel_event`` overrides
``self._cancel_event`` so an update can use a private event without
touching the shared one; defaults to the shared event.
"""
cancel_event = cancel_event if cancel_event is not None else self._cancel_event
try:
import huggingface_hub # noqa: F401 -- presence check only
except ImportError:
@ -3715,21 +3778,22 @@ class LlamaCppBackend:
# cold whenever free disk is below the full weight footprint,
# even though nothing needs downloading.
already_cached_bytes = 0
for p in path_infos:
if not p.size:
continue
try:
cached_path = try_to_load_from_cache(hf_repo, p.path)
except Exception:
cached_path = None
if isinstance(cached_path, str) and os.path.exists(cached_path):
if not force:
for p in path_infos:
if not p.size:
continue
try:
on_disk = os.path.getsize(cached_path)
except OSError:
on_disk = 0
# Satisfied only when the full blob is present.
if on_disk >= p.size:
already_cached_bytes += p.size
cached_path = try_to_load_from_cache(hf_repo, p.path)
except Exception:
cached_path = None
if isinstance(cached_path, str) and os.path.exists(cached_path):
try:
on_disk = os.path.getsize(cached_path)
except OSError:
on_disk = 0
# Satisfied only when the full blob is present.
if on_disk >= p.size:
already_cached_bytes += p.size
total_download_bytes = max(0, total_bytes - already_cached_bytes)
@ -3752,6 +3816,13 @@ class LlamaCppBackend:
)
if total_download_bytes > free_bytes:
if not allow_smaller_fallback:
# Update path: never silently switch to a smaller quant;
# surface the disk shortfall for the requested variant.
raise RuntimeError(
f"Not enough disk space to download {gguf_filename}. "
f"Only {free_gb:.1f} GB free in {cache_dir}"
)
smaller = self._find_smallest_fitting_variant(
hf_repo,
free_bytes,
@ -3792,7 +3863,7 @@ class LlamaCppBackend:
)
logger.info(f"Resolving GGUF: {gguf_label}")
try:
if self._cancel_event.is_set():
if cancel_event.is_set():
raise RuntimeError("Cancelled")
dl_start = time.monotonic()
# Xet primary, HTTP fallback on stall; per-file so finished shards stay cached.
@ -3800,18 +3871,20 @@ class LlamaCppBackend:
hf_repo,
gguf_filename,
hf_token,
cancel_event = self._cancel_event,
cancel_event = cancel_event,
on_status = lambda m: logger.info(m),
force_download = force,
)
for shard in gguf_extra_shards:
if self._cancel_event.is_set():
if cancel_event.is_set():
raise RuntimeError("Cancelled")
logger.info(f"Resolving GGUF shard: {shard}")
hf_hub_download_with_xet_fallback(
hf_repo,
shard,
hf_token,
cancel_event = self._cancel_event,
cancel_event = cancel_event,
force_download = force,
)
except Exception as e:
if isinstance(e, RuntimeError) and "Cancelled" in str(e):
@ -3834,6 +3907,7 @@ class LlamaCppBackend:
hf_token: Optional[str],
pick: Callable[[list[str]], Optional[str]],
label: str,
cancel_event: Optional[threading.Event] = None,
) -> Optional[str]:
"""Resolve and fetch a companion GGUF (mmproj / MTP drafter) by name.
@ -3841,8 +3915,10 @@ class LlamaCppBackend:
(offline, same fallback as _download_gguf), then hf_hub_download.
Runs WITHOUT self._lock (like _download_gguf); honors _cancel_event so
an /unload between the main download and here skips the fetch.
``cancel_event`` overrides ``self._cancel_event`` (defaults to it).
"""
if self._cancel_event.is_set():
cancel_event = cancel_event if cancel_event is not None else self._cancel_event
if cancel_event.is_set():
return None
target: Optional[str] = None
@ -3851,7 +3927,7 @@ class LlamaCppBackend:
# Retry a transient listing blip; permanent repo/auth errors and offline
# mode are not retried (offline raises at once -> fall through to cache).
for attempt in range(3):
if self._cancel_event.is_set():
if cancel_event.is_set():
return None
try:
target = pick(list_repo_files(hf_repo, token = hf_token))
@ -3867,10 +3943,10 @@ class LlamaCppBackend:
logger.debug(f"Could not list repo files for {label}: {e}")
break
logger.debug(
f"Could not list repo files for {label} " f"(attempt {attempt + 1}/3): {e}"
f"Could not list repo files for {label} (attempt {attempt + 1}/3): {e}"
)
if attempt < 2:
self._cancel_event.wait(2**attempt)
cancel_event.wait(2**attempt)
if target is None:
try:
@ -3884,7 +3960,7 @@ class LlamaCppBackend:
except Exception as e:
logger.debug(f"Offline cache lookup for {label} failed: {e}")
if target is None or self._cancel_event.is_set():
if target is None or cancel_event.is_set():
return None
try:
@ -3894,7 +3970,7 @@ class LlamaCppBackend:
hf_repo,
target,
hf_token,
cancel_event = self._cancel_event,
cancel_event = cancel_event,
)
except Exception as e:
logger.warning(f"Could not download {label}: {e}")
@ -3905,11 +3981,13 @@ class LlamaCppBackend:
*,
hf_repo: str,
hf_token: Optional[str] = None,
cancel_event: Optional[threading.Event] = None,
) -> Optional[str]:
"""Download the mmproj (vision projection) file from a GGUF repo.
Prefers mmproj-F16.gguf, else any mmproj*.gguf. Returns the local
path, or None if none exists.
path, or None if none exists. ``cancel_event`` overrides
``self._cancel_event`` (defaults to it).
"""
def _pick_mmproj(candidates: list[str]) -> Optional[str]:
@ -3930,6 +4008,7 @@ class LlamaCppBackend:
hf_token = hf_token,
pick = _pick_mmproj,
label = "mmproj",
cancel_event = cancel_event,
)
def _download_mtp(
@ -4331,6 +4410,17 @@ class LlamaCppBackend:
)
)
@staticmethod
def _is_tensor_split_assert(output: str) -> bool:
"""True only for the #6415 split-axis warmup assert (GGML_BACKEND_SPLIT_AXIS_*),
not any ggml assert/abort, so an unrelated invariant isn't cached. stderr is
merged into output."""
text = (output or "").lower()
if "ggml_assert" not in text and "ggml_abort" not in text:
return False
# the split-axis enum token, unique to this assert (not the source file).
return "split_axis" in text
@staticmethod
def _is_signal_crash(returncode: Optional[int]) -> bool:
"""True only on a hard fault (SIGSEGV/SIGABRT/SIGILL/SIGFPE/SIGBUS or a
@ -4343,6 +4433,20 @@ class LlamaCppBackend:
return True
return -returncode in (4, 6, 7, 8, 11) # SIGILL SIGABRT SIGBUS SIGFPE SIGSEGV
@staticmethod
def _is_abort_exit(returncode: Optional[int]) -> bool:
"""Windows CRT abort() exit code (3) from GGML_ASSERT on MSVC -- not a POSIX
signal or 0xC0000000+ NTSTATUS."""
return returncode == 3
@classmethod
def _should_record_tensor_split_abort(cls, returncode: Optional[int], output: str) -> bool:
"""The #6415 split-axis abort: the marker plus a hard crash (POSIX signal or
Windows abort exit). Marker required so a generic crash isn't cached."""
return cls._is_tensor_split_assert(output) and (
cls._is_signal_crash(returncode) or cls._is_abort_exit(returncode)
)
@staticmethod
def _with_flash_attn_off(cmd: list[str]) -> Optional[list[str]]:
"""Return cmd with flash attention forced off, or None when its effective
@ -4487,6 +4591,8 @@ class LlamaCppBackend:
n_gpu_layers: Optional[int] = None, # caller compat, unused
n_parallel: int = 1,
extra_args: Optional[List[str]] = None,
# Route-level tensor->layer fallback retry: keep the layer split multi-GPU.
preserve_multi_gpu_on_layer: bool = False,
) -> bool:
"""Start llama-server with a GGUF model.
@ -4517,6 +4623,8 @@ class LlamaCppBackend:
"n_gpu_layers": n_gpu_layers,
"n_parallel": n_parallel,
"extra_args": list(extra_args) if extra_args is not None else None,
# Replayed by _respawn_if_dead so a downgraded model stays multi-GPU.
"preserve_multi_gpu_on_layer": preserve_multi_gpu_on_layer,
}
# Serialise the whole load so concurrent /load calls never leave two
# llama-server processes alive (#5401 / #5161). Doesn't block /unload.
@ -4540,6 +4648,7 @@ class LlamaCppBackend:
chat_template_override = chat_template_override,
extra_args = extra_args,
is_vision = is_vision,
preserve_multi_gpu_on_layer = preserve_multi_gpu_on_layer,
):
logger.info(
f"load_model: backend already in target state for "
@ -4625,6 +4734,9 @@ class LlamaCppBackend:
# Block-diffusion GGUFs (DiffusionGemma) cannot run on llama-server;
# serve them with the diffusion runner (same OpenAI-compat interface).
if self._is_diffusion:
# Not a tensor/layer GGUF: clear any preserved-fallback flag from a
# prior load (this path skips the command builder that clears it).
self._layer_preserves_tensor_intent = False
with self._lock:
if self._cancel_event.is_set():
logger.info("Load cancelled before diffusion server start")
@ -4779,6 +4891,9 @@ class LlamaCppBackend:
"image input will be disabled for this session"
)
model_size = None # set in the fit try; used by the APU RAM guard
# Layer-fallback min GPUs; raised below on a tensor downgrade. Bound
# before the try so the --fit-on except path still has it (no UnboundLocal).
_layer_min_gpus = 1
try:
gguf_size = self._get_gguf_size_bytes(model_path)
# Include GPU-loaded mmproj in the fit budget (#5825).
@ -5063,10 +5178,8 @@ class LlamaCppBackend:
_apple_budget_mib = self._apple_metal_memory_budget_bytes() // (1024 * 1024)
def _restore_after_tensor_downgrade():
# Tensor mode dropped a quantized KV and stripped the cache
# extras (it rejects quantized); layer split supports them, so
# restore the original type + extras (minus --split-mode) and
# clear the env flag so the layer launch re-emits them.
# Restore the quantized KV + extras tensor dropped (layer
# split supports them), minus --split-mode.
nonlocal cache_type_kv, _cache_type_from_env, extra_args
if _tensor_dropped_cache_type_kv is not None:
cache_type_kv = _tensor_dropped_cache_type_kv
@ -5077,13 +5190,22 @@ class LlamaCppBackend:
else extra_args
)
if tensor_parallel and effective_is_vision:
# The route fallback retry is tensor-off; keep it multi-GPU.
if preserve_multi_gpu_on_layer:
_layer_min_gpus = max(_layer_min_gpus, len(gpus))
if tensor_parallel and self._tensor_split_aborts(binary, model_identifier):
# Aborted on tensor for this model this session (#6415); skip
# tensor upfront, layer split serves it.
logger.info(
"Tensor parallelism skipped for vision model: "
"--split-mode tensor is incompatible with --mmproj "
"in the current llama.cpp build; using layer split."
"Tensor parallelism skipped: this llama.cpp build aborted "
"on --split-mode tensor for this model earlier this "
"session; using layer split across %d GPU(s).",
len(gpus),
)
tensor_parallel = False
# Keep the multi-GPU request (gated on it, not the cache).
_layer_min_gpus = max(_layer_min_gpus, len(gpus))
_restore_after_tensor_downgrade()
# Tensor mode replicates a compute buffer on every GPU, so drop
@ -5123,6 +5245,11 @@ class LlamaCppBackend:
len(gpus),
)
tensor_parallel = False
# GPUs below tensor's compute-buffer reserve can still do layer
# split, so keep multi-GPU (mirrors the budget/geometry drops);
# _select_gpus caps unusable cards.
if len(gpus) >= 2:
_layer_min_gpus = max(_layer_min_gpus, len(gpus))
# Layer split supports a quantized KV the tensor attempt
# dropped; restore the original cache type + extras (minus
# --split-mode) so the layer launch re-emits them.
@ -5159,8 +5286,12 @@ class LlamaCppBackend:
"per-device compute buffers; falling back to layer split."
)
tensor_parallel = False
# Restore the dropped quantized KV + original cache extras
# (minus --split-mode); layer split supports them.
# Weights needed >1 card, so keep multi-GPU across the
# usable tensor GPUs.
if len(tp_gpus) >= 2:
_layer_min_gpus = max(_layer_min_gpus, len(tp_gpus))
# Restore the dropped quantized KV + cache extras (minus
# --split-mode); layer split supports them.
_restore_after_tensor_downgrade()
if tensor_parallel and tp_gpus:
@ -5262,6 +5393,7 @@ class LlamaCppBackend:
usable_fraction = _pin_fraction,
total_by_idx = total_by_idx,
per_device_overhead_bytes = _pipeline_overhead_bytes,
min_gpus = _layer_min_gpus,
)
# No silent shrink: effective_ctx stays == requested_ctx.
else:
@ -5272,7 +5404,22 @@ class LlamaCppBackend:
ranked = sorted(
gpus, key = lambda g: _gpu_usable(g, pin_fraction), reverse = True
)
for n_gpus in range(1, len(ranked) + 1):
# Skips _select_gpus, so apply its cap: count only cards
# whose usable VRAM clears the per-device layer overhead.
_pipeline_overhead_mib = _pipeline_overhead_bytes / (1024 * 1024)
_auto_min_gpus = max(
1,
min(
_layer_min_gpus,
sum(
1
for g in ranked
if _gpu_usable(g, pin_fraction) > _pipeline_overhead_mib
)
or 1,
),
)
for n_gpus in range(_auto_min_gpus, len(ranked) + 1):
subset = ranked[:n_gpus]
pool_budget = _pool_budget_mib(subset, pin_fraction)
_ms = _subset_model_size(n_gpus)
@ -5302,7 +5449,7 @@ class LlamaCppBackend:
# at 131k may pin fine with a 4096 KV (#5106).
effective_ctx = min(4096, effective_ctx)
if effective_ctx > 0:
for n_gpus in range(1, len(ranked) + 1):
for n_gpus in range(_auto_min_gpus, len(ranked) + 1):
subset = ranked[:n_gpus]
kv = self._estimate_kv_cache_bytes(
effective_ctx,
@ -5338,6 +5485,7 @@ class LlamaCppBackend:
usable_fraction = _pin_fraction,
total_by_idx = total_by_idx,
per_device_overhead_bytes = _pipeline_overhead_bytes,
min_gpus = _layer_min_gpus,
)
if use_fit and not explicit_ctx:
# Weights don't fit on any subset; default UI to 4096
@ -5475,6 +5623,15 @@ class LlamaCppBackend:
"--no-context-shift",
]
# Report a clean public model id (matching GET /v1/models) rather
# than the raw -m path in llama-server's own /v1/models and the
# "model" field of its chat/completions responses.
from core.inference.model_ids import public_model_id
_alias = public_model_id(self._model_identifier or model_path)
if _alias:
cmd.extend(["--alias", _alias])
fully_gpu_offloaded = False
if use_fit:
cmd.extend(["--fit", "on"])
@ -5568,12 +5725,15 @@ class LlamaCppBackend:
]
)
self._tensor_parallel = True
self._layer_preserves_tensor_intent = False
logger.info(
"Tensor parallelism: --split-mode tensor, --tensor-split %s",
tp_tensor_split,
)
else:
self._tensor_parallel = False
# > 1 only when a tensor request was downgraded but kept multi-GPU.
self._layer_preserves_tensor_intent = _layer_min_gpus > 1
# Speculative decoding. See _build_speculative_flags for the
# mode resolution, benchmarks, and llama.cpp references.
@ -5857,7 +6017,17 @@ class LlamaCppBackend:
_startup_crashed = (
self._process.poll() is not None and self._process.returncode != 0
)
if _spawn_attempt == 0 and _fit_retry_allowed and _startup_crashed:
# A split-axis abort (#6415) is fit-independent: skip the
# --fit off retry and let the caller latch it.
_split_axis_crash = self._is_tensor_split_assert(
"\n".join(self._stdout_lines[-50:])
)
if (
_spawn_attempt == 0
and _fit_retry_allowed
and _startup_crashed
and not _split_axis_crash
):
logger.warning(
"llama-server crashed during startup (exit code %s) "
"with the default memory-fit step enabled; Studio "
@ -5903,6 +6073,21 @@ class LlamaCppBackend:
)
healthy = _spawn_and_wait(cmd)
# #6415 split-mode tensor warmup abort. Latch it on THIS first spawn:
# the flash-attn-off retry below can't run tensor (needs flash_attn),
# so its output drops the marker and recording later would miss it,
# looping every load. Record and raise to the route's layer fallback,
# skipping the futile flash-attn/MTP retries.
if not healthy and self._tensor_parallel and not self._cancel_event.is_set():
_ts_out = "\n".join(self._stdout_lines[-50:])
_ts_rc = self._process.poll() if self._process is not None else None
if self._should_record_tensor_split_abort(_ts_rc, _ts_out):
LlamaCppBackend._record_tensor_split_abort(binary, model_identifier)
self._kill_process()
raise RuntimeError(
"llama-server aborted on --split-mode tensor "
"(split-axis geometry); retrying with layer split."
)
# Flash-attention kernels hard-crash at startup on some ROCm/GPU
# builds (frequently inside the vision tower). Disabling FA keeps
# both vision and MTP, so retry that way before dropping either.
@ -6047,6 +6232,7 @@ class LlamaCppBackend:
# Read the crash code before _kill_process() clears _process.
_crash_rc = self._process.poll() if self._process is not None else None
self._kill_process()
# The #6415 split-axis abort is latched earlier (first spawn).
# Skip if a cancel/unload is pending (mirrors the MTP guard).
if (
launched_with_mmproj
@ -6478,6 +6664,7 @@ class LlamaCppBackend:
spec_draft_n_max: Optional[int] = None,
tensor_parallel: bool = False,
mtp_draft_path: Optional[str] = None,
preserve_multi_gpu_on_layer: bool = False,
) -> bool:
"""True iff the live server already satisfies these load kwargs.
@ -6520,6 +6707,17 @@ class LlamaCppBackend:
# server. An identical request would downgrade the same way.
if not _tensor_parallel_matches_loaded(extra_args, tensor_parallel, self._tensor_parallel):
return False
# Preserved tensor->layer fallback + an EXPLICIT tensor drop: reload so
# placement re-selects instead of keeping the all-GPU mask (mirrors the route,
# #6659). preserve_multi_gpu_on_layer carries the route's carry-forward decision
# (True for an implicit same-settings reload), so those still dedupe -- the HF
# auto-pick / local-dir flows skip the route guard and only reach here.
if (
self._layer_preserves_tensor_intent
and not _effective_tensor_parallel(extra_args, tensor_parallel)
and not preserve_multi_gpu_on_layer
):
return False
# Compare on the canonical requested mode. With --spec-type in
# extra_args the backend stores None; mirror that here.
@ -6631,6 +6829,7 @@ class LlamaCppBackend:
self._supports_tools = False
self._cache_type_kv = None
self._tensor_parallel = False
self._layer_preserves_tensor_intent = False
self._speculative_type = None
self._requested_spec_mode = None
self._spec_draft_n_max = None
@ -7100,7 +7299,13 @@ class LlamaCppBackend:
url = f"{self.base_url}/completion"
payload = {"prompt": "Hi", "n_predict": 4, "temperature": 0.0, "stream": False}
try:
resp = httpx.post(url, json = payload, timeout = timeout, headers = self._auth_headers)
resp = httpx.post(
url,
json = payload,
timeout = timeout,
headers = self._auth_headers,
trust_env = False,
)
except Exception as e:
logger.debug(f"MTP decode probe failed: {e}")
return False
@ -7252,7 +7457,9 @@ class LlamaCppBackend:
return False
try:
resp = httpx.get(url, timeout = 2.0)
# trust_env=False: skip ambient HTTP(S)_PROXY, which if it 503s
# for 127.0.0.1 loops the probe until timeout and hangs load.
resp = httpx.get(url, timeout = 2.0, trust_env = False)
if resp.status_code == 200:
return True
except (
@ -7299,7 +7506,7 @@ class LlamaCppBackend:
"""
url = f"{self.base_url}/props"
try:
resp = httpx.get(url, timeout = 5.0)
resp = httpx.get(url, timeout = 5.0, trust_env = False)
if resp.status_code != 200:
return None
settings = resp.json().get("default_generation_settings") or {}
@ -7379,7 +7586,9 @@ class LlamaCppBackend:
which differ only in how they parse the SSE body."""
stream_timeout = httpx.Timeout(connect = 10, read = 0.5, write = 10, pool = 10)
with httpx.Client(
timeout = stream_timeout, limits = httpx.Limits(max_keepalive_connections = 0)
timeout = stream_timeout,
limits = httpx.Limits(max_keepalive_connections = 0),
trust_env = False,
) as client:
first_token_deadline = time.monotonic() + _DEFAULT_FIRST_TOKEN_TIMEOUT_S
with self._stream_with_retry(
@ -8871,7 +9080,7 @@ class LlamaCppBackend:
system_text = _block_text(system)
try:
with httpx.Client(timeout = 10, headers = self._auth_headers) as client:
with httpx.Client(timeout = 10, headers = self._auth_headers, trust_env = False) as client:
def _tokenize(text: str) -> int:
r = client.post(
@ -8987,7 +9196,7 @@ class LlamaCppBackend:
"""Codec name on match, None on non-audio, raises on transport/JSON errors."""
if not self.is_loaded:
return None
with httpx.Client(timeout = 10, headers = self._auth_headers) as client:
with httpx.Client(timeout = 10, headers = self._auth_headers, trust_env = False) as client:
def _detok(tid: int) -> str:
# Non-200 means "marker not in vocab" -- keep probing.
@ -9102,7 +9311,9 @@ class LlamaCppBackend:
payload["n_probs"] = 1
with httpx.Client(
timeout = httpx.Timeout(300, connect = 10), headers = self._auth_headers
timeout = httpx.Timeout(300, connect = 10),
headers = self._auth_headers,
trust_env = False,
) as client:
resp = client.post(f"{self.base_url}/completion", json = payload)
if resp.status_code != 200:

6
studio/backend/core/inference/llama_http.py
View file

@ -22,11 +22,7 @@ _LIMITS = httpx.Limits(max_connections = 64, max_keepalive_connections = 32)
def _new_client() -> httpx.AsyncClient:
try:
return httpx.AsyncClient(limits = _LIMITS)
except Exception:
# Mirror external_provider: an unsupported env proxy scheme can raise.
return httpx.AsyncClient(limits = _LIMITS, trust_env = False)
return httpx.AsyncClient(limits = _LIMITS, trust_env = False)
# One client per running event loop: an httpx client binds its transport to the

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

@ -25,6 +25,11 @@ _DENYLIST_GROUPS: tuple[frozenset[str], ...] = (
# Model identity: Studio resolves it from LoadRequest; a second -m would
# load a different model than Studio thinks it loaded.
frozenset({"-m", "--model"}),
# Public model id: Studio sets a sanitized --alias so the OpenAI API never
# exposes the local .gguf path. A user-supplied alias is appended after
# Studio's and, with llama.cpp's last-wins parsing, would reintroduce the
# path leak this is meant to prevent.
frozenset({"-a", "--alias"}),
frozenset({"-mu", "--model-url"}),
frozenset({"-dr", "--docker-repo"}),
frozenset({"-hf", "-hfr", "--hf-repo"}),

71
studio/backend/core/inference/model_ids.py Normal file
View file

@ -0,0 +1,71 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Public model identifiers for the OpenAI-compatible API.
The exposed API must report a stable, clean model id rather than the absolute
on-disk path of a local GGUF. The internal identifier for a direct local load is
the absolute ``.gguf`` path, which leaks the host filesystem layout and is
awkward for clients to round-trip. ``public_model_id`` maps such an internal
identifier to a clean name while leaving Hugging Face repo ids (``org/model``)
and already-clean names untouched.
"""
from __future__ import annotations
import os
from typing import Optional
_GGUF_SUFFIX = ".gguf"
def _looks_like_path(identifier: str) -> bool:
"""True when *identifier* is a local filesystem path, not a HF repo id.
A repo id is ``org/model`` (a single forward slash, no leading separator, no
drive, no ``.gguf``). Anything ending in ``.gguf``, starting with a path
separator or a relative/home prefix (``./``, ``../``, ``~``), carrying a
Windows drive, or with three or more ``/`` segments is treated as a local
path.
"""
if identifier.lower().endswith(_GGUF_SUFFIX):
return True
if identifier.startswith(("/", "\\", "./", "../", ".\\", "..\\", "~")):
return True
if len(identifier) >= 2 and identifier[1] == ":": # Windows drive, e.g. C:\
return True
if identifier.count("/") >= 2 or "\\" in identifier:
return True
return False
def public_model_id(identifier: Optional[str]) -> Optional[str]:
"""Return a clean, path-free public id for *identifier*.
- Local GGUF path -> the file stem with ``.gguf`` stripped, e.g.
``/srv/models/Qwen3-30B-A3B-Q4_K_M.gguf`` -> ``Qwen3-30B-A3B-Q4_K_M``.
- HF repo id (``org/model``) and already-clean names -> returned unchanged.
- ``None`` / empty -> returned unchanged.
"""
if not identifier:
return identifier
if not _looks_like_path(identifier):
return identifier
name = os.path.basename(identifier.replace("\\", "/").rstrip("/"))
if name.lower().endswith(_GGUF_SUFFIX):
name = name[: -len(_GGUF_SUFFIX)]
return name or identifier
def model_id_matches(requested: Optional[str], internal: Optional[str]) -> bool:
"""Whether a client-supplied *requested* id refers to *internal*.
Accepts the clean public id (preferred) and, for backward compatibility, the
raw internal identifier (e.g. a legacy absolute path a client cached from an
older ``/v1/models`` response).
"""
if requested is None or internal is None:
return False
if requested == internal:
return True
return public_model_id(internal) == requested

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

@ -534,29 +534,34 @@ class InferenceOrchestrator:
except (EOFError, OSError, ValueError):
break
rid = resp.get("request_id")
rtype = resp.get("type", "")
# Sole consumer of the response queue; if it died every in-flight
# stream would hang, so never let routing kill the dispatcher.
try:
rid = resp.get("request_id")
rtype = resp.get("type", "")
# Status messages — log and skip
if rtype == "status":
logger.info("Subprocess status: %s", resp.get("message", ""))
continue
# Route to mailbox if a matching request_id exists
if rid:
with self._mailbox_lock:
mbox = self._mailboxes.get(rid)
if mbox is not None:
mbox.put(resp)
# Status messages: log and skip
if rtype == "status":
logger.info("Subprocess status: %s", resp.get("message", ""))
continue
# No matching mailbox (a _gen_lock reader or orphaned). Can't
# un-get from mp.Queue, so just log. (status was handled above.)
logger.debug(
"Dispatcher: no mailbox for request_id=%s type=%s, dropping",
rid,
rtype,
)
# Route to mailbox if a matching request_id exists
if rid:
with self._mailbox_lock:
mbox = self._mailboxes.get(rid)
if mbox is not None:
mbox.put(resp)
continue
# No matching mailbox; can't un-get from mp.Queue, so just log.
logger.debug(
"Dispatcher: no mailbox for request_id=%s type=%s, dropping",
rid,
rtype,
)
except Exception:
logger.exception("Inference dispatcher: failed to route a response; continuing")
continue
def _generate_dispatched(
self,

26
studio/backend/core/inference/sd_cpp_args.py
View file

@ -88,8 +88,11 @@ class SdCppGenParams:
prompt: str
negative_prompt: Optional[str] = None
width: int = 1024
height: int = 1024
# None = "unset": an image-conditioned run (img2img/inpaint/edit) then lets
# sd.cpp derive the size from the input image instead of forcing a resize; a
# plain txt2img run with unset dims falls back to 1024x1024 (see the builder).
width: Optional[int] = None
height: Optional[int] = None
steps: Optional[int] = None
cfg_scale: Optional[float] = None
guidance: Optional[float] = None
@ -236,7 +239,17 @@ def build_sd_cpp_command(
cmd += ["--lora-model-dir", params.lora_dir]
if params.lora_apply_mode:
cmd += ["--lora-apply-mode", params.lora_apply_mode]
cmd += ["--width", str(int(params.width)), "--height", str(int(params.height))]
# Emit explicit dims when given. For an image-conditioned run (img2img /
# inpaint / edit) that leaves them unset, omit the flags so sd.cpp derives the
# size from the input image (set_width_and_height_if_unset) rather than forcing
# a 1024x1024 resize/crop of the source. A plain txt2img run with unset dims
# keeps the prior 1024 default.
if params.width is not None or params.height is not None:
w = int(params.width) if params.width is not None else 1024
h = int(params.height) if params.height is not None else 1024
cmd += ["--width", str(w), "--height", str(h)]
elif not (params.init_img or params.ref_images):
cmd += ["--width", "1024", "--height", "1024"]
if params.steps is not None:
cmd += ["--steps", str(int(params.steps))]
if params.cfg_scale is not None:
@ -280,6 +293,11 @@ def build_sd_cpp_upscale_command(
raise ValueError("input_image is required for upscale")
if not params.upscale_model:
raise ValueError("upscale_model is required for upscale")
# A truthiness guard below would silently swallow repeats=0 and fall back to
# sd-cli's default of one pass, turning an explicit no-op into a real upscale.
# Reject it (and negatives) so the caller's intent isn't quietly changed.
if params.repeats < 1:
raise ValueError("repeats must be >= 1 for upscale")
cmd: list[str] = [
binary,
"--mode",
@ -289,7 +307,7 @@ def build_sd_cpp_upscale_command(
"--upscale-model",
params.upscale_model,
]
if params.repeats and params.repeats != 1:
if params.repeats != 1:
cmd += ["--upscale-repeats", str(int(params.repeats))]
if params.tile_size is not None:
cmd += ["--upscale-tile-size", str(int(params.tile_size))]

16
studio/backend/core/inference/sd_cpp_engine.py
View file

@ -125,7 +125,8 @@ def find_sd_cpp_binary() -> Optional[str]:
both engines look):
1. ``SD_CLI_PATH`` env -- a direct path to the binary.
2. ``UNSLOTH_SD_CPP_PATH`` env -- a stable-diffusion.cpp install dir.
3. ``~/.unsloth/stable-diffusion.cpp`` build layouts (the installer target).
3. The default install root build layouts (the installer target); honors
``UNSLOTH_STUDIO_HOME`` / ``STUDIO_HOME``, else ``~/.unsloth/stable-diffusion.cpp``.
4. ``./stable-diffusion.cpp`` in-tree build (developer checkout).
5. ``sd-cli`` (then legacy ``sd``) on PATH.
"""
@ -151,8 +152,17 @@ def find_sd_cpp_binary() -> Optional[str]:
if hit:
return hit
# 3. Default install root (sibling of ~/.unsloth/llama.cpp).
hit = _first_file(_layout_candidates(Path.home() / ".unsloth" / "stable-diffusion.cpp"))
# 3. Default install root. Honors UNSLOTH_STUDIO_HOME / STUDIO_HOME the same way
# the installer's default_install_dir does (base = the Studio home's parent), so
# a binary installed under a custom Studio root is discovered and side-by-side
# Studios stay isolated; falls back to the sibling of ~/.unsloth/llama.cpp.
studio_home = os.environ.get("UNSLOTH_STUDIO_HOME") or os.environ.get("STUDIO_HOME")
default_root = (
Path(studio_home).parent / "stable-diffusion.cpp"
if studio_home
else Path.home() / ".unsloth" / "stable-diffusion.cpp"
)
hit = _first_file(_layout_candidates(default_root))
if hit:
return hit

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

@ -1121,6 +1121,61 @@ def _autoinject_top_k() -> int:
return _AUTOINJECT_DEFAULT_TOP_K
def _thread_whole_doc_enabled(scope: dict) -> bool:
"""Whether a thread-attached file should be injected in full rather than
retrieved top-K. ``rag_scope.whole_doc=False`` disables it for this request."""
override = scope.get("whole_doc")
if override is False:
return False
try:
from core.rag import config as _rag_config
except Exception: # noqa: BLE001
return True
return _rag_config.THREAD_WHOLE_DOC
_IMAGE_PART_TOKEN_ESTIMATE = 1024
def _message_token_estimate(conversation: list[dict]) -> int:
"""Cheap prompt-size estimate for budget guards; exact tokenization happens later."""
total = 0
for msg in conversation:
content = msg.get("content")
if isinstance(content, str):
total += max(1, len(content) // 4)
elif isinstance(content, list):
for part in content:
if isinstance(part, dict):
if part.get("type") in ("image_url", "input_image"):
total += _IMAGE_PART_TOKEN_ESTIMATE
else:
total += max(1, len(str(part.get("text") or "")) // 4)
total += 4 # chat-template role / separator overhead estimate
return total
def _whole_doc_budget(scope: dict | None = None, conversation: list[dict] | None = None) -> int:
try:
from core.rag import config as _rag_config
except Exception: # noqa: BLE001
budget = 6000
else:
budget = _rag_config.WHOLE_DOC_MAX_TOKENS
if not scope:
return budget
context = _opt_int(scope.get("context_length") or scope.get("max_context_tokens"))
if context is None or context <= 0:
return budget
headroom = _opt_int(scope.get("response_headroom"))
if headroom is None:
headroom = max(1024, context // 4)
used = _message_token_estimate(conversation or [])
# Leave room for tool XML wrappers, citation metadata, and chat-template overhead.
available = context - headroom - used - 512
return min(budget, max(0, available))
def _last_user_text(conversation: list[dict]) -> str:
"""Plain text of the most recent user turn (text parts only)."""
for msg in reversed(conversation):
@ -1154,7 +1209,11 @@ def build_rag_autoinject(conversation: list[dict], rag_scope: dict | None) -> di
enabled = rag_scope.get("autoinject")
if enabled is None:
enabled = _autoinject_enabled()
if not enabled:
thread_id = rag_scope.get("thread_id")
whole_doc_requested = (
bool(thread_id) and not rag_scope.get("kb_id") and _thread_whole_doc_enabled(rag_scope)
)
if not enabled and not whole_doc_requested:
return None
query = _last_user_text(conversation)
if not query:
@ -1163,35 +1222,81 @@ def build_rag_autoinject(conversation: list[dict], rag_scope: dict | None) -> di
from storage import rag_db
if not rag_db.RAG_AVAILABLE:
return None
from core.rag.tool import search_for_autoinject
from core.rag.tool import render_sources, search_for_autoinject, whole_document_context
except Exception as exc: # noqa: BLE001
logger.warning("RAG auto-inject unavailable: %s", exc)
return None
text: str | None = None
sources: list[dict] = []
floor_override = rag_scope.get("autoinject_min_score")
floor = float(floor_override) if floor_override is not None else _autoinject_floor()
# Cap at the lean top_k, but honor a lower user setting.
lean_k = _autoinject_top_k()
sidebar_k = _opt_int(rag_scope.get("default_top_k"))
top_k = min(sidebar_k, lean_k) if sidebar_k is not None else lean_k
try:
found = search_for_autoinject(
query = query,
scope_kb_id = rag_scope.get("kb_id"),
scope_thread_id = rag_scope.get("thread_id"),
scope_project_id = rag_scope.get("project_id"),
top_k = top_k,
min_dense_score = floor,
**_scope_retrieval_kwargs(rag_scope),
)
except Exception as exc: # noqa: BLE001
logger.warning("RAG auto-inject retrieval failed: %s", exc)
return None
if not found:
logger.info("RAG auto-inject: no passage >= %.2f; skipping", floor)
# Whole-document mode: a thread-attached file under budget is injected in full so
# the model reads everything. A KB selection is exclusive, so whole-doc never
# preempts it; in a project chat the project sources are still retrieved top-K and
# appended under one citation numbering. Oversized files (or no thread doc) fall
# through to the combined top-K retrieval below.
if whole_doc_requested:
try:
budget = _whole_doc_budget(rag_scope, conversation)
whole = whole_document_context(
scope_thread_id = thread_id,
max_tokens = budget,
)
except Exception as exc: # noqa: BLE001
logger.warning("RAG whole-document context failed: %s", exc)
whole = None
if whole is not None:
text, sources = whole
project_id = rag_scope.get("project_id")
if project_id:
try:
proj = search_for_autoinject(
query = query,
scope_project_id = project_id,
top_k = top_k,
min_dense_score = floor,
**_scope_retrieval_kwargs(rag_scope),
)
except Exception as exc: # noqa: BLE001
logger.warning("RAG project retrieval (whole-doc companion) failed: %s", exc)
proj = None
if proj is not None:
merged = sources + proj[1]
merged_text = render_sources(merged)
if max(1, len(merged_text) // 4) <= budget:
sources = merged
text = merged_text
logger.info("RAG auto-inject: whole-document context (%d chunk(s))", len(sources))
if text is None and enabled:
try:
found = search_for_autoinject(
query = query,
scope_kb_id = rag_scope.get("kb_id"),
scope_thread_id = rag_scope.get("thread_id"),
scope_project_id = rag_scope.get("project_id"),
top_k = top_k,
min_dense_score = floor,
**_scope_retrieval_kwargs(rag_scope),
)
except Exception as exc: # noqa: BLE001
logger.warning("RAG auto-inject retrieval failed: %s", exc)
return None
if not found:
logger.info("RAG auto-inject: no passage >= %.2f; skipping", floor)
return None
text, sources = found
if text is None:
return None
text, sources = found
import json as _json
import uuid as _uuid
@ -1236,7 +1341,7 @@ def build_rag_autoinject(conversation: list[dict], rag_scope: dict | None) -> di
"content": text,
},
]
logger.info("RAG auto-inject: %d passage(s) >= %.2f for %r", len(sources), floor, query[:80])
logger.info("RAG auto-inject: %d passage(s) for %r", len(sources), query[:80])
return {"events": events, "messages": messages}

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

@ -1,9 +1,12 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Caption figures with the loaded vision model and splice the text into the page
so images are searchable via the normal FTS5 + dense path. No-op (never raises)
without a vision model or on failure; gated by ``config.CAPTION_IMAGES``."""
"""Vision-model helpers for ingestion: figure captioning and scanned-page OCR.
Both turn pixels into indexable text and are a no-op (never raise) without a loaded
vision model. They reuse the chat model's vision endpoint, so it must be served with
``--ubatch-size`` >= one image's tokens (some encoders, e.g. Gemma, attend
non-causally and abort otherwise); Studio's vision chat already requires this."""
from __future__ import annotations
@ -15,11 +18,54 @@ from . import config
logger = logging.getLogger(__name__)
_CAPTION_PROMPT = (
"Describe this figure or image from a document in one or two concise "
"sentences, for search indexing. State what it depicts (e.g. a diagram, "
"chart, table or photo) and its key content. Do not add commentary."
"Read this figure or image from a document for search indexing.\n"
"First, on a line 'TEXT:', transcribe every piece of visible text exactly as "
"written, in reading order: the title, axis labels and units, legend and series "
"names, EVERY box / node / arrow label, table headers and cells, equations, and "
"footnotes. List each distinct label even if it is small.\n"
"Then, on a line 'SUMMARY:', add one or two sentences on what it shows (chart "
"type and trend, diagram subject, table topic, or photo content).\n"
"Report only what is visible. Transcribe exactly; do not invent or guess any "
"text, label, or number."
)
_OCR_PROMPT = (
"Transcribe all text on this document page exactly as it appears, in reading "
"order, including any text inside figures, diagrams, charts, and tables (keep "
"table rows readable). Output only the transcribed text, with no commentary or "
"code fences. Preserve headings, lists, and line breaks. If the page has no "
"readable text, output nothing."
)
def _collapse_runaway(
text: str,
max_repeat: int = 3,
max_total: int = 8,
) -> str:
"""Cap runaway repetition: vision models sometimes loop a line many times. Keep
each distinct line to ``max_repeat`` in a row and ``max_total`` total, and collapse
blank-line floods, so a degenerate page cannot flood the index."""
out: list[str] = []
seen: dict[str, int] = {}
prev: str | None = None
run = 0
for line in text.splitlines():
key = line.strip()
if not key:
if prev == "": # collapse runs of blank lines to a single separator
continue
prev = ""
out.append("")
continue
run = run + 1 if key == prev else 1
prev = key
seen[key] = seen.get(key, 0) + 1
if run > max_repeat or seen[key] > max_total:
continue
out.append(line)
return "\n".join(out)
def vision_endpoint() -> tuple[str, str] | None:
"""``(base_url, model)`` for a loaded vision GGUF model, else None."""
@ -33,7 +79,28 @@ def vision_endpoint() -> tuple[str, str] | None:
return None
def _caption_one(base_url: str, model: str, image_bytes: bytes, timeout: float) -> str | None:
def _vision_auth_headers() -> dict | None:
"""Bearer header for the backend's API, or None. Vision calls share the chat
endpoint, so they need the same key under direct-stream (``--api-key``) mode."""
try:
from routes.inference import get_llama_cpp_backend
return get_llama_cpp_backend()._auth_headers or None
except Exception: # noqa: BLE001 - auth discovery must never break ingestion
return None
def _vision_complete(
base_url: str,
model: str,
image_bytes: bytes,
*,
prompt: str,
timeout: float,
max_tokens: int,
temperature: float = 0.0,
) -> str | None:
"""One image-in / text-out call to the loaded vision model's OpenAI-compatible
endpoint. Returns the stripped text or ``None`` on empty/failure (non-fatal)."""
import httpx
data_url = "data:image/png;base64," + base64.b64encode(image_bytes).decode("ascii")
@ -43,33 +110,62 @@ def _caption_one(base_url: str, model: str, image_bytes: bytes, timeout: float)
{
"role": "user",
"content": [
{"type": "text", "text": _CAPTION_PROMPT},
{"type": "text", "text": prompt},
{"type": "image_url", "image_url": {"url": data_url}},
],
}
],
"max_tokens": 200,
"temperature": 0.2,
"max_tokens": max_tokens,
# Deterministic by default: transcription must not randomly drop labels.
"temperature": temperature,
"stream": False,
# Off: thinking models would spend the budget reasoning, returning "".
"chat_template_kwargs": {"enable_thinking": False},
}
try:
r = httpx.post(f"{base_url}/v1/chat/completions", json = payload, timeout = timeout)
r = httpx.post(
f"{base_url}/v1/chat/completions",
json = payload,
timeout = timeout,
headers = _vision_auth_headers(),
)
r.raise_for_status()
text = r.json()["choices"][0]["message"]["content"]
return text.strip() or None
except Exception: # noqa: BLE001 - a failed caption is non-fatal
logger.debug("caption request failed", exc_info = True)
except Exception: # noqa: BLE001 - a failed vision call is non-fatal
logger.debug("vision request failed", exc_info = True)
return None
def _caption_one(base_url: str, model: str, image_bytes: bytes, timeout: float) -> str | None:
return _vision_complete(
base_url,
model,
image_bytes,
prompt = _CAPTION_PROMPT,
timeout = timeout,
max_tokens = config.CAPTION_MAX_TOKENS,
)
def _ocr_one(base_url: str, model: str, image_bytes: bytes, timeout: float) -> str | None:
return _vision_complete(
base_url,
model,
image_bytes,
prompt = _OCR_PROMPT,
timeout = timeout,
max_tokens = config.OCR_MAX_TOKENS,
)
def caption_images(
images: list, *, endpoint: tuple[str, str] | None = None
) -> dict[int, list[str]]:
"""Caption ``ParsedImage`` objects, keyed by 1-based page number; ``{}`` when
disabled, no vision model, or no images. Bounded by ``CAPTION_MAX_IMAGES``."""
if not config.CAPTION_IMAGES or not images:
"""Caption ``ParsedImage`` objects, keyed by 1-based page number; ``{}`` when there
are no images or no vision model. The caller (`ingestion._run`) owns the on/off
policy. Bounded by ``CAPTION_MAX_IMAGES``; each caption passes ``_collapse_runaway``."""
if not images:
return {}
ep = endpoint or vision_endpoint()
if ep is None:
@ -84,7 +180,50 @@ def caption_images(
caption = _caption_one(base_url, model, image_bytes, config.CAPTION_TIMEOUT_S)
if caption:
page = getattr(img, "page_number", None) or 0
out.setdefault(int(page), []).append(caption)
out.setdefault(int(page), []).append(_collapse_runaway(caption))
return out
def ocr_pages(
page_pngs: dict[int, bytes], *, endpoint: tuple[str, str] | None = None
) -> dict[int, str]:
"""OCR rendered page PNGs (keyed by 1-based page number) to text; ``{}`` when there
is no vision model or no pages. The caller (`ingestion._ocr_scanned_pages`) owns the
on/off policy. Bounded by ``OCR_MAX_PAGES``."""
if not page_pngs:
return {}
ep = endpoint or vision_endpoint()
if ep is None:
return {}
base_url, model = ep
out: dict[int, str] = {}
for page_num in sorted(page_pngs)[: config.OCR_MAX_PAGES]:
text = _ocr_one(base_url, model, page_pngs[page_num], config.OCR_TIMEOUT_S)
if text:
out[int(page_num)] = _collapse_runaway(text)
return out
def merge_page_captions(captions: dict[int, list[str]]) -> dict[int, list[str]]:
"""Merge a page's per-tile captions into one deduped block: drop lines repeated
across overlapping tiles (first kept, order preserved), then ``_collapse_runaway``,
so ``splice_captions`` adds a single figure block per page."""
out: dict[int, list[str]] = {}
for page, caps in captions.items():
seen: set[str] = set()
lines: list[str] = []
for cap in caps:
for line in (cap or "").splitlines():
stripped = line.strip()
key = stripped.lower()
if not stripped or key in seen:
continue
seen.add(key)
lines.append(stripped)
merged = _collapse_runaway("\n".join(lines))
if merged.strip():
out[page] = [merged]
return out

49
studio/backend/core/rag/config.py
View file

@ -17,13 +17,50 @@ TOP_K_DENSE = int(os.environ.get("RAG_TOP_K_DENSE", "30"))
TOP_K_HYBRID = int(os.environ.get("RAG_TOP_K_HYBRID", "10"))
RRF_K = int(os.environ.get("RAG_RRF_K", "60"))
UPLOAD_EXTS = {".pdf", ".txt", ".md", ".markdown", ".docx", ".html", ".htm"}
# Whole-document context: a thread-attached file under the token budget is injected
# in full (every chunk, in order) instead of top-K retrieval; above it, use retrieval.
THREAD_WHOLE_DOC = os.environ.get("RAG_THREAD_WHOLE_DOC", "1") == "1"
WHOLE_DOC_MAX_TOKENS = int(os.environ.get("RAG_WHOLE_DOC_MAX_TOKENS", "6000"))
# Figure captioning via the loaded vision model; off by default since each caption
# is a model call. MAX_IMAGES bounds per-doc cost.
CAPTION_IMAGES = os.environ.get("RAG_CAPTION_IMAGES", "0") == "1"
CAPTION_MAX_IMAGES = int(os.environ.get("RAG_CAPTION_MAX_IMAGES", "8"))
CAPTION_TIMEOUT_S = float(os.environ.get("RAG_CAPTION_TIMEOUT_S", "30"))
UPLOAD_EXTS = {".pdf", ".txt", ".md", ".markdown", ".docx", ".html", ".htm"}
# Reject uploads larger than this, so one pathological file can't drive unbounded parse
# + vision work at ingest. 0 disables the cap. Default 200 MB.
MAX_UPLOAD_BYTES = int(os.environ.get("RAG_MAX_UPLOAD_BYTES", str(200 * 1024 * 1024)))
# Extract PDF text as layout-aware Markdown (pymupdf4llm) instead of flat text, so
# tables, headings and lists survive into chunks and retrieval. Falls back to plain
# PyMuPDF text when off, when pymupdf4llm is missing, or when extraction fails.
PDF_MARKDOWN = os.environ.get("RAG_PDF_MARKDOWN", "1") == "1"
# Figure captioning via the loaded vision model: detected figures are transcribed +
# described so they become searchable. On by default, a no-op without a vision model;
# the chat's "Describe figures & charts" toggle overrides it per upload.
CAPTION_IMAGES = os.environ.get("RAG_CAPTION_IMAGES", "1") == "1"
# Total per-document tile budget (figure-bearing pages are tiled, see below).
CAPTION_MAX_IMAGES = int(os.environ.get("RAG_CAPTION_MAX_IMAGES", "24"))
CAPTION_TIMEOUT_S = float(os.environ.get("RAG_CAPTION_TIMEOUT_S", "60"))
# Larger than a one-line caption since captions transcribe every label. FIGURE_DPI is
# high enough to keep small box/axis labels legible when tiles are rendered.
CAPTION_MAX_TOKENS = int(os.environ.get("RAG_CAPTION_MAX_TOKENS", "768"))
FIGURE_DPI = int(os.environ.get("RAG_FIGURE_DPI", "200"))
# Figure pages are tiled into an overlapping ROWS x COLS grid of high-DPI tiles (plus
# an optional full page), so small labels and every sub-figure are covered without
# exact region detection. MAX_PAGES bounds figure pages; MAX_IMAGES bounds total tiles.
FIGURE_TILE_ROWS = int(os.environ.get("RAG_FIGURE_TILE_ROWS", "2"))
FIGURE_TILE_COLS = int(os.environ.get("RAG_FIGURE_TILE_COLS", "2"))
FIGURE_TILE_OVERLAP = float(os.environ.get("RAG_FIGURE_TILE_OVERLAP", "0.12"))
FIGURE_FULLPAGE = os.environ.get("RAG_FIGURE_FULLPAGE", "1") == "1"
CAPTION_MAX_PAGES = int(os.environ.get("RAG_CAPTION_MAX_PAGES", "4"))
# Scanned-PDF OCR: a page with little extractable text is rendered and transcribed by
# the vision model so it becomes searchable. Needs a vision model, else skipped (page
# stays empty). MIN_CHARS is the text length below which a page is treated as scanned.
OCR_SCANNED = os.environ.get("RAG_OCR_SCANNED", "1") == "1"
OCR_MIN_CHARS = int(os.environ.get("RAG_OCR_MIN_CHARS", "16"))
OCR_MAX_PAGES = int(os.environ.get("RAG_OCR_MAX_PAGES", "20"))
OCR_DPI = int(os.environ.get("RAG_OCR_DPI", "150"))
OCR_TIMEOUT_S = float(os.environ.get("RAG_OCR_TIMEOUT_S", "60"))
OCR_MAX_TOKENS = int(os.environ.get("RAG_OCR_MAX_TOKENS", "2048"))
# Embedder backend. "auto": sentence-transformers on a CUDA/ROCm GPU (torch fp16
# wins bulk indexing), else torch-free GGUF llama-server. Switching backends changes

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

@ -26,6 +26,11 @@ _jobs_lock = threading.Lock()
_EMBED_BATCH = 64 # bounds peak memory
# Poll with a timeout so the generator wakes periodically to detect a gone
# client or a terminal job whose worker died without the None sentinel.
_SSE_POLL_SECONDS = 1.0
_TERMINAL_JOB_STATUSES = {"completed", "failed"}
def _sha256_file(path: str) -> str:
h = hashlib.sha256()
@ -94,25 +99,108 @@ def _embed_all(texts: list[str], model_name: str | None):
return vectors
def _ocr_scanned_pages(
pages: list,
stored_path: str,
conn,
job_id: str,
ocr: bool | None = None,
) -> tuple[list, set[int]]:
"""Replace text on near-empty (scanned/image-only) PDF pages with vision-model OCR
so image PDFs become searchable. ``ocr`` overrides ``config.OCR_SCANNED`` per upload
(``None`` = config default); no-op without scanned pages or a vision model. OCR'd
pages have no text layer, so no preview highlight regions, but stay searchable.
Returns ``(pages, ocred)``: new ``Page`` objects for OCR'd pages (originals
otherwise) and the set of page numbers actually transcribed."""
if not (config.OCR_SCANNED if ocr is None else ocr):
return pages, set()
scanned = [
p.page_number
for p in pages
if p.page_number is not None and len((p.text or "").strip()) < config.OCR_MIN_CHARS
]
if not scanned or captioner.vision_endpoint() is None:
return pages, set()
if len(scanned) > config.OCR_MAX_PAGES:
logger.warning(
"OCR: %d scanned pages exceed OCR_MAX_PAGES=%d; pages past the cap stay "
"untranscribed (raise RAG_OCR_MAX_PAGES to cover them)",
len(scanned),
config.OCR_MAX_PAGES,
)
scanned = scanned[: config.OCR_MAX_PAGES]
_progress(conn, job_id, "ocr", 0.25)
page_pngs = parsers.render_pdf_pages(stored_path, scanned, dpi = config.OCR_DPI)
texts = captioner.ocr_pages(page_pngs)
if not texts:
return pages, set()
from .parsers import Page
out: list = []
ocred: set[int] = set()
for page in pages:
text = texts.get(page.page_number)
if text:
original = (page.text or "").strip()
merged = text if not original or original in text else f"{original}\n\n{text}"
out.append(Page(text = merged, page_number = page.page_number, char_count = len(merged)))
ocred.add(page.page_number)
else:
out.append(page)
return out, ocred
def _run(
job_id: str, document_id: str, scope: str, stored_path: str, model_name: str | None
job_id: str,
document_id: str,
scope: str,
stored_path: str,
model_name: str | None,
ocr: bool | None = None,
caption: bool | None = None,
) -> None:
conn = rag_db.get_connection()
try:
_progress(conn, job_id, "parsing", 0.1)
pages = parsers.parse(stored_path)
if config.CAPTION_IMAGES and stored_path.lower().endswith(".pdf"):
# Caption figures, splice into page text (no-op without a vision model).
is_pdf = stored_path.lower().endswith(".pdf")
ocred: set[int] = set()
if is_pdf:
pages, ocred = _ocr_scanned_pages(pages, stored_path, conn, job_id, ocr = ocr)
caption_on = config.CAPTION_IMAGES if caption is None else caption
# Skip all figure work (PDF rasterization included) without a vision model.
if caption_on and is_pdf and captioner.vision_endpoint() is not None:
# Tile figure pages, transcribe+describe each tile, then merge/dedup/splice
# into the page text so small labels and every sub-figure are captured.
try:
figures = parsers.render_pdf_figures(
stored_path, max_figures = config.CAPTION_MAX_IMAGES
fig_pages = parsers.pages_with_figures(
stored_path,
max_pages = config.CAPTION_MAX_PAGES,
# Skip only pages OCR actually transcribed (it covers them whole); a
# scanned figure page past the OCR cap or with empty OCR still tiles.
exclude_pages = ocred,
)
tiles = (
parsers.render_pdf_figure_tiles(
stored_path,
fig_pages,
dpi = config.FIGURE_DPI,
rows = config.FIGURE_TILE_ROWS,
cols = config.FIGURE_TILE_COLS,
overlap = config.FIGURE_TILE_OVERLAP,
fullpage = config.FIGURE_FULLPAGE,
max_tiles = config.CAPTION_MAX_IMAGES,
)
if fig_pages
else []
)
except Exception:
logger.warning("figure rendering failed for job %s", job_id, exc_info = True)
figures = []
if figures:
_progress(conn, job_id, "captioning", 0.2)
captions = captioner.caption_images(figures)
logger.warning("figure tiling failed for job %s", job_id, exc_info = True)
tiles = []
if tiles:
_progress(conn, job_id, "captioning", 0.28)
captions = captioner.merge_page_captions(captioner.caption_images(tiles))
pages = captioner.splice_captions(pages, captions)
_progress(conn, job_id, "chunking", 0.3)
@ -170,6 +258,8 @@ def start_ingestion(
*,
project_id: str | None = None,
model_name: str | None = None,
ocr: bool | None = None,
caption: bool | None = None,
) -> tuple[str, str]:
"""Create the document + job rows and spawn the worker, returning
``(document_id, job_id)``. A duplicate content hash in this scope returns the
@ -178,18 +268,34 @@ def start_ingestion(
if ext not in config.UPLOAD_EXTS:
raise ValueError(f"unsupported file type: {ext}")
# Reclaim queues for finished jobs so the registry stays bounded.
_reap_finished_jobs()
sha = _sha256_file(stored_path)
conn = rag_db.get_connection()
try:
existing = store.document_by_hash(conn, scope, sha)
if existing is not None:
job_id = _new_job(conn, existing, scope, status = "completed", progress = 1.0)
_remove_upload(stored_path)
with _jobs_lock:
_jobs[job_id] = queue.Queue()
_emit(job_id, {"type": "complete", "num_chunks": 0, "deduped": True})
_emit(job_id, None)
return existing, job_id
doc = store.get_document(conn, existing)
empty_completed = (
doc is not None and doc.get("status") == "completed" and not doc.get("num_chunks")
)
if empty_completed:
# A prior ingest of identical bytes yielded zero chunks (e.g. a scanned
# PDF uploaded before a vision model loaded). Re-ingest, don't dedupe.
store.delete_document(conn, existing)
_remove_upload(doc.get("stored_path"), keep_path = stored_path)
else:
job_id = _new_job(conn, existing, scope, status = "completed", progress = 1.0)
_remove_upload(stored_path)
with _jobs_lock:
_jobs[job_id] = queue.Queue()
_emit(
job_id,
{"type": "complete", "num_chunks": doc.get("num_chunks") or 0, "deduped": True},
)
_emit(job_id, None)
return existing, job_id
for failed in store.failed_documents_by_hash(conn, scope, sha):
store.delete_document(conn, failed["id"])
_remove_upload(failed.get("stored_path"), keep_path = stored_path)
@ -213,7 +319,7 @@ def start_ingestion(
_jobs[job_id] = queue.Queue()
threading.Thread(
target = _run,
args = (job_id, document_id, scope, stored_path, model_name),
args = (job_id, document_id, scope, stored_path, model_name, ocr, caption),
daemon = True,
).start()
return document_id, job_id
@ -248,26 +354,99 @@ def _new_job(
return job_id
def _reap_finished_jobs() -> None:
"""Drop per-job queues whose DB row already reached a terminal status.
Otherwise removed only by ``job_events`` after the ``None`` sentinel, so a
caller that polls ``/jobs/{id}`` instead of streaming would grow ``_jobs``
forever. Safe while streaming: ``job_events`` holds its queue reference.
"""
with _jobs_lock:
job_ids = list(_jobs.keys())
for jid in job_ids:
row = get_job_status(jid)
if row is not None and row.get("status") in _TERMINAL_JOB_STATUSES:
with _jobs_lock:
_jobs.pop(jid, None)
def job_events(job_id: str):
"""Yield job events for SSE; ends when the worker signals completion."""
"""Yield job events for SSE; ends when the worker signals completion.
Timed ``get`` so the generator can't block forever: it wakes to heartbeat,
to notice a disconnected client, and to stop on a terminal DB status (a hard
worker death that skipped the ``None`` sentinel). Drops the queue only on a
terminal exit, never on an early client disconnect.
It deliberately does *not* end on idle alone: a long silent stage (e.g.
embedding a large doc) is not a failure, and ending there would send
``[DONE]`` with the row still pending, which the client treats as completion.
The stream ends only on a terminal status, the ``None`` sentinel, or disconnect.
"""
with _jobs_lock:
q = _jobs.get(job_id)
if q is None:
return
while True:
event = q.get()
if event is None:
break
yield event
with _jobs_lock:
_jobs.pop(job_id, None)
terminal = False
try:
while True:
try:
event = q.get(timeout = _SSE_POLL_SECONDS)
except queue.Empty:
try:
row = get_job_status(job_id)
except Exception: # noqa: BLE001
# A transient status read (e.g. the DB momentarily locked) must
# not abort the stream: routes/rag.py would turn the raised
# exception into a terminal {type: error} frame and the UI would
# drop a document whose worker is still running. Heartbeat and
# retry on the next poll instead.
logger.warning(
"job_events status read failed for %s; continuing", job_id, exc_info = True
)
yield {"type": "heartbeat"}
continue
if row is None or row.get("status") in _TERMINAL_JOB_STATUSES:
# Worker finished (or row gone); stop and let the client reconcile via getJob.
terminal = True
break
yield {"type": "heartbeat"}
continue
if event is None:
terminal = True
break
yield event
finally:
# Drop the queue once nothing more will be emitted into it: either a
# terminal exit, or a disconnect after the job already finished (the UI
# stops on the terminal event, before [DONE], so terminal is still False
# here -- _run writes the terminal DB status before emitting it). Keep it
# only while the worker is still running, so an early disconnect can
# reconnect and resume its events.
if not terminal:
try:
row = get_job_status(job_id)
terminal = row is None or row.get("status") in _TERMINAL_JOB_STATUSES
except Exception: # noqa: BLE001
# Can't confirm terminality (transient DB error) -- keep the queue so
# a reconnect can resume rather than orphaning a live worker's events.
terminal = False
if terminal:
with _jobs_lock:
_jobs.pop(job_id, None)
def get_job_status(job_id: str) -> dict | None:
"""Read the persisted ingestion job row (status / stage / progress / error)."""
"""Read the persisted ingestion job row (status / stage / progress / error), plus
the document's ``num_chunks`` so a client polling to completion learns the chunk
count (the SSE ``complete`` frame carries it, but the poll/reconcile path does not)."""
conn = rag_db.get_connection()
try:
row = conn.execute("SELECT * FROM ingestion_jobs WHERE id=?", (job_id,)).fetchone()
row = conn.execute(
"SELECT j.*, d.num_chunks AS num_chunks FROM ingestion_jobs j "
"LEFT JOIN documents d ON d.id = j.document_id WHERE j.id=?",
(job_id,),
).fetchone()
return dict(row) if row else None
finally:
conn.close()

6
studio/backend/core/rag/locators.py
View file

@ -39,9 +39,11 @@ def _norm_token(token: str) -> str:
def _anchor_tokens(page_text: str, match: LocatorMatch) -> list[str]:
"""Normalized anchor tokens from the chunk's leading span. Drops first and last
token (boundaries often slice mid-word) when long enough."""
token (boundaries often slice mid-word) when long enough. Pipes are split out so
Markdown table cells (``|Q1|$1.2M|``) become individual words that match the PDF
word stream."""
segment = page_text[match.start : match.end]
raw = segment.split()
raw = segment.replace("|", " ").split()
if len(raw) >= MIN_ANCHOR_WORDS + 2:
raw = raw[1:-1]
tokens = [t for t in (_norm_token(w) for w in raw) if t]

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

@ -15,6 +15,8 @@ import os
from dataclasses import dataclass
from html.parser import HTMLParser
from . import config
logger = logging.getLogger(__name__)
@ -67,6 +69,28 @@ def _html(raw: str) -> list[Page]:
return [_page("\n".join(parser.out), 1)]
def _pdf_markdown(doc) -> list[str] | None:
"""Per-page layout-aware Markdown (tables, headings, lists) via pymupdf4llm; index
i maps to page i+1. Returns None when the lib is missing, extraction fails, or the
page count does not line up, so the caller falls back to plain PyMuPDF text."""
try:
import pymupdf4llm
except Exception:
return None
try:
chunks = pymupdf4llm.to_markdown(
doc,
page_chunks = True,
show_progress = False,
)
except Exception: # noqa: BLE001 - never let Markdown extraction break ingestion
logger.warning("pymupdf4llm extraction failed; using plain text", exc_info = True)
return None
if not isinstance(chunks, list) or len(chunks) != doc.page_count:
return None
return [str(c.get("text") or "") for c in chunks]
def _pdf(path: str, want_images: bool) -> tuple[list[Page], list[ParsedImage]]:
import fitz # PyMuPDF
@ -74,8 +98,11 @@ def _pdf(path: str, want_images: bool) -> tuple[list[Page], list[ParsedImage]]:
images: list[ParsedImage] = []
doc = fitz.open(path)
try:
md = _pdf_markdown(doc) if config.PDF_MARKDOWN else None
for i, page in enumerate(doc):
text = page.get_text("text") or ""
# Prefer layout-aware Markdown (keeps tables/headings legible for retrieval);
# fall back to plain text when Markdown is off, unavailable, or empty here.
text = (md[i] if md else "") or page.get_text("text") or ""
pages.append(_page(text, i + 1))
if want_images:
for img in page.get_images(full = True):
@ -118,63 +145,164 @@ def _merge_rects(boxes: list) -> list:
return merged
def render_pdf_figures(
path: str,
def _figure_boxes(
page,
*,
dpi: int = 130,
min_area_frac: float = 0.04,
min_side: float = 40.0,
max_figures: int = 8,
) -> list[ParsedImage]:
"""Detect figure regions and render each to a PNG for captioning.
) -> list:
"""Qualifying figure-region rectangles on a page: cluster vector drawings + raster
placements, merge overlaps, keep the page-spanning ones (area/side filtered)."""
boxes: list = []
try:
boxes.extend(info["bbox"] for info in page.get_image_info())
except Exception:
pass
try:
boxes.extend(page.cluster_drawings())
except Exception:
pass
if not boxes:
return []
page_area = page.rect.width * page.rect.height
keep: list = []
for box in _merge_rects(boxes):
if (
box.get_area() >= min_area_frac * page_area
and box.width >= min_side
and box.height >= min_side
):
keep.append(box)
return keep
Academic figures are vector, so raster extraction yields fragments; instead
cluster vector drawings + raster placements into boxes, keep the page-spanning
ones, and render them. Any failure yields [], never an exception.
"""
def pages_with_figures(
path: str,
*,
max_pages: int = 4,
min_area_frac: float = 0.04,
min_side: float = 40.0,
exclude_pages: set[int] | None = None,
) -> list[int]:
"""1-based page numbers with a qualifying figure region, capped at ``max_pages``;
drives figure tiling. ``exclude_pages`` (1-based) are skipped: those are the pages
OCR already transcribed whole, so tiling them would duplicate the vision work. Any
failure yields []."""
exclude = exclude_pages or set()
try:
import pymupdf
except Exception:
return []
out: list[ParsedImage] = []
try:
doc = pymupdf.open(path)
except Exception:
return []
pages: list[int] = []
try:
for i, page in enumerate(doc):
boxes: list = []
try:
boxes.extend(info["bbox"] for info in page.get_image_info())
except Exception:
pass
try:
boxes.extend(page.cluster_drawings())
except Exception:
pass
if not boxes:
if (i + 1) in exclude:
continue
page_area = page.rect.width * page.rect.height
for box in _merge_rects(boxes):
if (
box.get_area() >= min_area_frac * page_area
and box.width >= min_side
and box.height >= min_side
):
try:
pix = page.get_pixmap(dpi = dpi, clip = box)
out.append(
ParsedImage(
image_bytes = pix.tobytes("png"),
page_number = i + 1,
xref = 0,
)
if _figure_boxes(page, min_area_frac = min_area_frac, min_side = min_side):
pages.append(i + 1)
if len(pages) >= max_pages:
break
return pages
finally:
doc.close()
def render_pdf_figure_tiles(
path: str,
page_numbers,
*,
dpi: int = 200,
rows: int = 2,
cols: int = 2,
overlap: float = 0.12,
fullpage: bool = True,
max_tiles: int = 24,
) -> list[ParsedImage]:
"""Render figure-bearing pages as overlapping high-DPI tiles (plus an optional full
page), each a ``ParsedImage`` keyed by page number. Tiling keeps small labels legible
and covers every sub-figure without exact region detection. Any failure yields []."""
wanted = [int(n) for n in page_numbers]
if not wanted:
return []
rows, cols = max(1, int(rows)), max(1, int(cols)) # never divide by zero
try:
import pymupdf
except Exception:
return []
try:
doc = pymupdf.open(path)
except Exception:
return []
out: list[ParsedImage] = []
try:
for num in wanted:
if num < 1 or num > doc.page_count:
continue
page = doc[num - 1]
rect = page.rect
clips: list = [rect] if fullpage else []
cw, ch = rect.width / cols, rect.height / rows
ox, oy = cw * overlap, ch * overlap
for r in range(rows):
for c in range(cols):
clips.append(
pymupdf.Rect(
rect.x0 + c * cw - ox,
rect.y0 + r * ch - oy,
rect.x0 + (c + 1) * cw + ox,
rect.y0 + (r + 1) * ch + oy,
)
except Exception:
continue
if len(out) >= max_figures:
return out
& rect
)
for clip in clips:
try:
pix = page.get_pixmap(dpi = dpi, clip = clip)
out.append(ParsedImage(image_bytes = pix.tobytes("png"), page_number = num, xref = 0))
except Exception:
continue
if len(out) >= max_tiles:
return out
return out
finally:
doc.close()
def render_pdf_pages(
path: str,
page_numbers,
*,
dpi: int = 150,
) -> dict[int, bytes]:
"""Render whole PDF pages (given as 1-based numbers) to PNG bytes, keyed by
page number. Backs scanned-page OCR. Any failure yields ``{}`` (or skips that
page), never an exception.
"""
wanted = {int(n) for n in page_numbers}
if not wanted:
return {}
try:
import pymupdf
except Exception:
return {}
try:
doc = pymupdf.open(path)
except Exception:
return {}
out: dict[int, bytes] = {}
try:
for i, page in enumerate(doc):
num = i + 1
if num not in wanted:
continue
try:
pix = page.get_pixmap(dpi = dpi)
out[num] = pix.tobytes("png")
except Exception:
continue
return out
finally:
doc.close()

37
studio/backend/core/rag/store.py
View file

@ -292,3 +292,40 @@ def chunks_by_id(conn: sqlite3.Connection, ids) -> dict:
list(ids),
).fetchall()
return {r["id"]: r for r in rows}
def all_chunks_for_scope(conn: sqlite3.Connection, scope) -> list[dict]:
"""Every completed-document chunk for a scope, ordered document-then-index and
joined with the document filename. Backs whole-document context injection, so
it does no retrieval or embedding."""
scopes = _scopes(scope)
if not scopes:
return []
placeholders = ",".join("?" * len(scopes))
rows = conn.execute(
f"SELECT c.id, c.text, c.document_id, c.chunk_index, c.page_number, "
f"c.token_count, d.filename, d.created_at "
f"FROM chunks c JOIN documents d ON d.id=c.document_id "
f"WHERE c.scope IN ({placeholders}) AND d.status='completed' "
f"ORDER BY d.created_at, c.document_id, c.chunk_index",
list(scopes),
).fetchall()
return [dict(r) for r in rows]
def scope_token_estimate(conn: sqlite3.Connection, scope) -> int:
"""Upper-bound token total for a scope's completed chunks without hydrating text.
Mirrors ``all_chunks_for_scope`` + the ``tool._row_token_count`` fallback (stored
count, else length/4), so the whole-doc budget can be checked before loading text."""
scopes = _scopes(scope)
if not scopes:
return 0
placeholders = ",".join("?" * len(scopes))
row = conn.execute(
f"SELECT COALESCE(SUM(CASE WHEN c.token_count > 0 THEN c.token_count "
f"ELSE MAX(1, length(COALESCE(c.text, '')) / 4) END), 0) AS total "
f"FROM chunks c JOIN documents d ON d.id=c.document_id "
f"WHERE c.scope IN ({placeholders}) AND d.status='completed'",
list(scopes),
).fetchone()
return int(row["total"] or 0)

81
studio/backend/core/rag/tool.py
View file

@ -16,7 +16,13 @@ from xml.sax.saxutils import quoteattr
from storage import rag_db
from . import config, retrieval
from .store import kb_scope, project_scope, thread_scope
from .store import (
all_chunks_for_scope,
kb_scope,
project_scope,
scope_token_estimate,
thread_scope,
)
SEARCH_KNOWLEDGE_BASE_TOOL = {
"type": "function",
@ -90,6 +96,30 @@ def _format(rows, hits) -> tuple[str, list[dict]]:
return "\n\n".join(blocks), sources
def render_sources(sources: list[dict]) -> str:
"""Render a citation-source list to sequentially-numbered ``<chunk>`` blocks,
rewriting each source's ``citationId`` to match its 1-based position. Lets
independently-built source lists (a whole-document thread attachment plus
retrieved project passages) be merged under one citation numbering."""
blocks: list[str] = []
for i, s in enumerate(sources, 1):
s["citationId"] = i
src = quoteattr(s.get("filename") or "unknown")
page = s.get("page")
page_attr = f" page={quoteattr(str(page))}" if page else ""
blocks.append(f'<chunk id="{i}" source={src}{page_attr}>\n{s.get("text") or ""}\n</chunk>')
return "\n\n".join(blocks)
def _row_token_count(row) -> int:
"""Chunk token count for budgeting, falling back to a length estimate when the
stored count is missing or zero, so a malformed chunk cannot bypass the budget."""
tc = row["token_count"]
if tc:
return int(tc)
return max(1, len(row["text"] or "") // 4)
def search_knowledge_base_with_sources(
*,
query: str,
@ -186,6 +216,55 @@ def search_for_autoinject(
return (text, sources) if sources else None
def whole_document_context(
*, scope_thread_id: str | None = None, max_tokens: int
) -> tuple[str, list[dict]] | None:
"""Render EVERY chunk of the THREAD's attached documents (in order) as the same
``<chunk>`` blocks + citation source-map as retrieval, so the model reads the whole
file rather than top-K passages. Thread-attached files only: KB and project corpora
are search corpora, never whole-document, so this resolves the thread scope alone.
``None`` (caller falls back to retrieval) when there is no thread scope, no completed
chunks, or the total exceeds ``max_tokens``."""
if not scope_thread_id:
return None
# A non-positive budget means "never inject" (disable whole-doc via
# RAG_THREAD_WHOLE_DOC=0), not "inject the whole corpus unbounded".
if max_tokens <= 0:
return None
scope = thread_scope(scope_thread_id)
conn = rag_db.get_connection()
try:
# Cheap budget pre-check (SUM, no text hydration): reject an oversized attachment
# before loading the whole corpus; all_chunks_for_scope runs only once it fits.
if scope_token_estimate(conn, scope) > max_tokens:
return None
rows = all_chunks_for_scope(conn, scope)
finally:
conn.close()
if not rows:
return None
total = sum(_row_token_count(r) for r in rows)
if total > max_tokens:
return None
sources: list[dict] = [
{
"citationId": i,
"chunkId": r["id"],
"documentId": r["document_id"],
"filename": r["filename"] or "unknown",
"page": r["page_number"],
"text": r["text"] or "",
"score": None,
}
for i, r in enumerate(rows, 1)
]
rendered = render_sources(sources)
if max(1, len(rendered) // 4) > max_tokens:
return None
return rendered, sources
def search_knowledge_base(
*,
query: str,

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

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

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

@ -216,6 +216,9 @@ class TrainingBackend:
self._event_queue: Any = None
self._stop_queue: Any = None
self._pump_thread: Optional[threading.Thread] = None
# True while a pump thread should be running; cleared on intended exits.
# Left True after an abnormal death so _ensure_pump_alive spots a crash.
self._pump_running: bool = False
self._lock = threading.Lock()
# Progress state (updated by pump thread from subprocess events)
@ -289,10 +292,14 @@ class TrainingBackend:
logger.warning("Previous pump thread did not exit within 5s — refusing to start")
return False
self._pump_thread = None
# Clear a stale crash flag from a prior died pump so the watchdog can't
# treat this fresh setup as a recoverable death.
self._pump_running = False
# Build config dict for the subprocess
config = {
"model_name": kwargs["model_name"],
"project_name": kwargs.get("project_name"),
"training_type": kwargs.get("training_type", "LoRA/QLoRA"),
"hf_token": kwargs.get("hf_token", ""),
"load_in_4bit": kwargs.get("load_in_4bit", True),
@ -472,16 +479,21 @@ class TrainingBackend:
self._xet_fallback_used = False
self._needs_xet_respawn = False
# Assign subprocess handles after state reset.
self._event_queue = event_queue
self._stop_queue = stop_queue
self._proc = proc
# Eagerly create DB run row so it appears in history during model loading.
# Create the DB run row before the pump can consume events, so it appears
# in history during model loading and a fast terminal worker can't race the
# pump into a duplicate create/finalize. From here the pump only finalizes.
self._ensure_db_run_created()
self._pump_thread = threading.Thread(target = self._pump_loop, daemon = True)
self._pump_thread.start()
# Assign handles and start the pump together under the lock so a concurrent
# poll can't see a live _proc with no pump and spawn a duplicate.
new_pump = threading.Thread(target = self._pump_loop, daemon = True)
with self._lock:
self._pump_running = False
self._event_queue = event_queue
self._stop_queue = stop_queue
self._proc = proc
self._pump_thread = new_pump
new_pump.start()
return True
@ -606,6 +618,9 @@ class TrainingBackend:
except Exception:
logger.error("Failed to respawn training subprocess", exc_info = True)
with self._lock:
# No replacement pump will run; clear the flag so a later run can't
# inherit a stale _pump_running=True and spawn a duplicate.
self._pump_running = False
self._progress.is_training = False
self._progress.error = "Failed to recover stalled model download"
self._ensure_db_run_created()
@ -623,10 +638,44 @@ class TrainingBackend:
self._stop_queue = stop_queue
self._proc = new_proc
self._pump_thread = new_pump
new_pump.start()
# Start under the lock so _ensure_pump_alive can never observe the
# new pump as a not-yet-started (dead) thread and spawn a duplicate.
new_pump.start()
def _ensure_pump_alive(self) -> bool:
"""Restart the event pump if it crashed, even after the worker exited.
Defence in depth behind _pump_loop's guards. _pump_running stays True only
after an abnormal exit (the loop clears it on intended exits), so a True
flag plus a dead thread is an unambiguous crash. Restarts even after worker
exit so a fresh pump can drain the terminal events and finalize; otherwise
the run looks stuck "running" forever. Returns True if restarted.
"""
with self._lock:
if not self._pump_running:
return False
# A restarted pump needs the worker handle and queue to drain/finalize;
# their absence means nothing is left to recover.
if self._proc is None or self._event_queue is None:
return False
if self._pump_thread is not None and self._pump_thread.is_alive():
return False
logger.error(
"Training event pump thread died while the worker is still running; "
"restarting it so progress updates resume."
)
new_pump = threading.Thread(target = self._pump_loop, daemon = True)
self._pump_thread = new_pump
# Start under the lock so a concurrent _ensure_pump_alive can't see
# this thread as not-yet-started and spawn yet another pump.
new_pump.start()
return True
def is_training_active(self) -> bool:
"""Check if training is currently active."""
# Self-heal a crashed pump first: a dead pump must never leave the worker
# training invisibly behind a frozen UI. Cheap enough for per-second polls.
self._ensure_pump_alive()
with self._lock:
if self._proc is not None and self._proc.is_alive():
return True
@ -727,51 +776,87 @@ class TrainingBackend:
# Event pump (background thread)
# ------------------------------------------------------------------
def _safe_handle_event(self, event: dict) -> None:
"""Apply one event, swallowing any handler error.
The pump is the only writer of the progress state every status surface
reads, so a malformed event must never propagate and kill it.
"""
try:
self._handle_event(event)
except Exception:
etype = event.get("type") if isinstance(event, dict) else type(event).__name__
logger.exception("Training event pump: failed to handle %s event; skipping", etype)
def _pump_loop(self) -> None:
"""Background thread: consume events from subprocess → update state."""
"""Background thread: consume subprocess events and update state.
Sole writer of the in-memory progress state that /progress, /status,
/metrics and DB history read. If it exited while the worker still ran, the
run would burn GPU with events piling up while every surface froze. So no
single bad event or transient queue/DB error may end it; it returns only
through intended exits (worker gone, respawn handed off, finalized).
"""
self._pump_running = True
while True:
if self._proc is None or self._event_queue is None:
self._pump_running = False
return
event = self._read_queue(self._event_queue, timeout_sec = 0.25)
try:
event = self._read_queue(self._event_queue, timeout_sec = 0.25)
except Exception:
# If a read keeps raising after the worker died, fall through to
# finalize instead of spinning; only retry while the worker lives.
logger.exception("Training event pump: queue read failed; continuing")
if self._proc is not None and self._proc.is_alive():
time.sleep(0.1)
continue
event = None
if event is not None:
self._handle_event(event)
self._safe_handle_event(event)
continue
if self._proc.is_alive():
continue
# Process exited — drain remaining events.
for e in self._drain_queue(self._event_queue):
self._handle_event(e)
# Worker exited. Drain the backlog and finalize, guarded so a slow or
# failing DB write can't strand the thread; we return either way.
try:
for e in self._drain_queue(self._event_queue):
self._safe_handle_event(e)
# Model-load stall: respawn over HTTP instead of finalizing as failure.
# Runs on THIS exiting pump thread and starts a fresh pump (never joins
# the current thread); DB run-state is preserved.
if self._needs_xet_respawn:
self._needs_xet_respawn = False
self._respawn_worker_disable_xet()
return
# Model-load stall: respawn over HTTP instead of finalizing as failure.
# Starts a fresh pump on this thread (no self-join); it takes over
# _pump_running, so this exit leaves the flag set.
if self._needs_xet_respawn:
self._needs_xet_respawn = False
self._respawn_worker_disable_xet()
return
# Mark done if no explicit complete/error was received.
with self._lock:
if self._progress.is_training:
if self._should_stop:
self._progress.is_training = False
self._progress.status_message = "Training stopped."
else:
self._progress.is_training = False
self._progress.error = (
self._progress.error or "Training process exited unexpectedly"
)
# Mark done if no explicit complete/error was received.
with self._lock:
if self._progress.is_training:
if self._should_stop:
self._progress.is_training = False
self._progress.status_message = "Training stopped."
else:
self._progress.is_training = False
self._progress.error = (
self._progress.error or "Training process exited unexpectedly"
)
self._ensure_db_run_created()
self._finalize_run_in_db(
status = "stopped" if self._should_stop else "error",
error_message = None
if self._should_stop
else "Training process terminated unexpectedly",
)
self._ensure_db_run_created()
self._finalize_run_in_db(
status = "stopped" if self._should_stop else "error",
error_message = None
if self._should_stop
else "Training process terminated unexpectedly",
)
except Exception:
logger.exception("Training event pump: finalization after worker exit failed")
self._pump_running = False
return
def _handle_event(self, event: dict) -> None:
@ -1094,6 +1179,8 @@ class TrainingBackend:
except queue.Empty:
return None
except (EOFError, OSError, ValueError):
# A closed/broken queue reads as "no event"; any other error is left to
# _pump_loop's guarded block, which logs and backs off.
return None
@staticmethod
@ -1104,7 +1191,12 @@ class TrainingBackend:
events.append(q.get_nowait())
except queue.Empty:
return events
except (EOFError, OSError, ValueError):
except Exception:
# A drain error must not abort finalization: return what we have so
# the run finalizes rather than wedging "active" behind a dead worker.
logger.exception(
"Training event pump: queue drain failed; finalizing with drained events"
)
return events
# ------------------------------------------------------------------

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

@ -44,6 +44,7 @@ if sys.platform.startswith("linux") and "HSA_ENABLE_DXG_DETECTION" not in os.env
logger = get_logger(__name__)
from utils.hardware import apply_gpu_ids
from utils.training_runs import build_default_output_dir_name
from utils.wheel_utils import (
direct_wheel_url,
flash_attn_wheel_url,
@ -1787,11 +1788,14 @@ def _run_mlx_training(event_queue, stop_queue, config):
# ── 5. Build output dir ──
# Resolve to ~/.unsloth/studio/outputs/ so the export page finds it
from utils.paths import resolve_output_dir, ensure_dir, default_run_dir_name
from utils.paths import resolve_output_dir, ensure_dir
output_dir = config.get("output_dir", "")
if not output_dir:
output_dir = f"{default_run_dir_name(model_name)}_{int(time.time())}"
output_dir = build_default_output_dir_name(
model_name,
config.get("project_name"),
)
output_dir = str(resolve_output_dir(output_dir))
ensure_dir(Path(output_dir))
@ -3019,7 +3023,10 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) ->
resume_from_checkpoint
)
if not output_dir:
output_dir = f"{default_run_dir_name(model_name)}_{int(time.time())}"
output_dir = build_default_output_dir_name(
model_name,
config.get("project_name"),
)
output_dir = str(resolve_output_dir(output_dir))
ensure_dir(Path(output_dir))
@ -3500,7 +3507,10 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) ->
resume_from_checkpoint
)
if not output_dir:
output_dir = f"{default_run_dir_name(model_name)}_{int(time.time())}"
output_dir = build_default_output_dir_name(
model_name,
config.get("project_name"),
)
output_dir = str(resolve_output_dir(output_dir))
num_epochs = config.get("num_epochs", 2)

3
studio/backend/hub/schemas/inventory.py
View file

@ -27,6 +27,9 @@ class GgufVariantDetail(BaseModel):
downloaded: bool = Field(
False, description = "Whether this variant is already in the local HF cache"
)
update_available: bool = Field(
False, description = "Whether a newer main GGUF blob is available on Hugging Face"
)
partial: bool = Field(
False,
description = "Whether this variant has an in-progress (.incomplete) blob in cache",

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

@ -314,25 +314,50 @@ def register_worker(
worker_token = hf_token
def _watch() -> None:
finalize_worker_exit(
registry,
key,
proc,
hf_token = worker_token,
label = label,
log_prefix = log_prefix,
logger = logger,
repo_type = repo_type,
repo_id = repo_id,
transport = transport,
)
if registry.get_job(key).state in ("error", "cancelled"):
download_registry.purge_empty_marker_dir(
repo_type,
repo_id,
download_registry.variant_from_key(key),
try:
finalize_worker_exit(
registry,
key,
proc,
hf_token = worker_token,
label = label,
log_prefix = log_prefix,
logger = logger,
repo_type = repo_type,
repo_id = repo_id,
transport = transport,
)
hf_cache_scan.invalidate_hf_cache_scans()
except Exception:
# finalize_worker_exit is the only thing that clears running/cancelling;
# if it raises, force a terminal state so claim() isn't blocked until restart.
logger.exception("download watcher crashed for %s", key)
# finalize may have raised before reaping the worker; terminate the
# still-registered Popen first, else the terminal set_job clears the
# repo guard and a live worker would race a retry on the same repo.
try:
kill_and_reap_process(proc, label = label, logger = logger)
except Exception:
logger.exception("failed to reap worker after watcher crash for %s", key)
try:
registry.drop_process(key, proc)
except Exception:
logger.exception("failed to drop worker after watcher crash for %s", key)
try:
registry.set_job(key, "error", "download watcher crashed")
except Exception:
logger.exception("failed to mark %s errored after watcher crash", key)
finally:
try:
if registry.get_job(key).state in ("error", "cancelled"):
download_registry.purge_empty_marker_dir(
repo_type,
repo_id,
download_registry.variant_from_key(key),
)
except Exception:
logger.exception("post-finalize marker cleanup failed for %s", key)
finally:
hf_cache_scan.invalidate_hf_cache_scans()
threading.Thread(target = _watch, name = watch_name, daemon = True).start()
return True

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

@ -39,8 +39,10 @@ from hub.services.models.common import (
logger = get_logger(__name__)
_repo_size_cache: "OrderedDict[tuple[str, str], tuple[int, frozenset[str], float]]" = OrderedDict()
_repo_size_neg_cache: "OrderedDict[tuple[str, str], float]" = OrderedDict()
_repo_size_cache: "OrderedDict[tuple[str, str, str], tuple[int, frozenset[str], float]]" = (
OrderedDict()
)
_repo_size_neg_cache: "OrderedDict[tuple[str, str, str], float]" = OrderedDict()
_REPO_SIZE_CACHE_MAX = 256
_REPO_SIZE_POS_TTL = 60.0
_REPO_SIZE_NEG_TTL = 60.0
@ -52,7 +54,7 @@ def get_repo_snapshot_metadata_cached(
repo_id: str, hf_token: Optional[str] = None
) -> tuple[int, frozenset[str]]:
token_fp = hf_cache_scan.token_fingerprint(hf_token)
cache_key = (repo_id, token_fp)
cache_key = (repo_id, token_fp, "snapshot")
with _repo_size_cache_lock:
cached = _repo_size_cache.get(cache_key)
if cached is not None:
@ -119,6 +121,52 @@ def _repo_has_gguf_files(repo_info) -> bool:
return _repo_gguf_size_bytes(repo_info) > 0
def _cached_repo_file_name(file_obj) -> str:
file_path = getattr(file_obj, "file_path", None)
if file_path:
try:
path = Path(file_path)
parts = path.parts
snapshots_idx = max(i for i, part in enumerate(parts) if part == "snapshots")
if len(parts) > snapshots_idx + 2:
return Path(*parts[snapshots_idx + 2 :]).as_posix()
except Exception:
pass
return str(getattr(file_obj, "file_name", "")).replace("\\", "/")
def _repo_gguf_blob_map(repo_info, *, include_companions: bool = False) -> dict[str, set[str]]:
"""Map each cached GGUF file's repo-relative name to the SET of its local
blob hashes across all cached revisions.
HF names each local cache blob FILE by the file's etag (lfs.sha256 else
blob_id), so a local file's blob hash == ``Path(blob_path).name``. An updated
repo keeps BOTH the old and new revision snapshots until HF garbage-collects
them, so the same file resolves to several blobs; collecting them ALL (not
just the first one seen, since ``repo_info.revisions`` is a frozenset and
yields them in arbitrary order) lets the remote-vs-local diff treat the file
as current when the remote (``main``) blob is present in any cached revision.
Mirrors the ``cached_blob_ids`` membership test in routes/models.py.
By default this keeps the historical MAIN-GGUF-only behavior. GGUF update
checks opt into companions so a shared mmproj/MTP blob can be compared too.
"""
blob_map: dict[str, set[str]] = {}
for revision in repo_info.revisions:
for f in revision.files:
if include_companions:
if not _is_gguf_filename(f.file_name):
continue
elif not _is_main_gguf_filename(f.file_name):
continue
blob_path = getattr(f, "blob_path", None)
if not blob_path:
continue
name = _cached_repo_file_name(f)
blob_map.setdefault(name, set()).add(Path(blob_path).name)
return blob_map
def _prefer_cache_row(candidate: dict, existing: Optional[dict]) -> bool:
if existing is None:
return True

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

@ -6,6 +6,7 @@
from __future__ import annotations
import asyncio
import errno
from pathlib import Path
from typing import Optional
@ -15,7 +16,7 @@ from loggers import get_logger
from hub.utils import download_manifest
from hub.utils import download_registry
from hub.utils import inventory_scan as hf_cache_scan
from hub.utils.gguf import extract_quant_label
from hub.utils.gguf import extract_quant_label, extract_quant_token
from hub.utils.hf_cache_state import (
INCOMPLETE_SUFFIX,
purge_partial_repo,
@ -106,6 +107,76 @@ def _has_remaining_main_gguf(target_repo) -> bool:
)
def _remove_empty_variant_dirs(target_repos: list, variant: str) -> tuple[int, list[str]]:
"""Remove now-empty ``snapshots/<rev>/<quant>/`` folders for *variant* (the
quant label names the folder); only empty dirs go, so siblings are safe.
Returns (count removed, removal failures other than a concurrent refill)."""
variant_key = (extract_quant_token(variant) or variant).lower()
removed = 0
failures: list[str] = []
for target_repo in target_repos:
repo_path = getattr(target_repo, "repo_path", None)
if not repo_path:
continue
snapshots = Path(repo_path) / "snapshots"
if not snapshots.is_dir():
continue
try:
snap_dirs = [s for s in snapshots.iterdir() if s.is_dir() and not s.is_symlink()]
except OSError:
continue
for snap in snap_dirs:
try:
subs = list(snap.iterdir())
except OSError:
continue
for sub in subs:
try:
if sub.is_symlink() or not sub.is_dir():
continue
folder_quant = extract_quant_token(sub.name)
matches = (
folder_quant is not None and folder_quant.lower() == variant_key
) or sub.name.lower() == variant.lower()
if not matches or any(sub.iterdir()):
continue
except OSError:
continue
try:
sub.rmdir()
removed += 1
except OSError as e:
# A concurrent download refilling the dir (ENOTEMPTY) is not a
# failure; a read-only cache or locked dir is, so surface it.
if e.errno != errno.ENOTEMPTY:
failures.append(f"{sub.name}: {e}")
return removed, failures
def _remove_empty_snapshot_dirs(target_repos: list) -> tuple[int, list[str]]:
removed = 0
failures: list[str] = []
for target_repo in target_repos:
repo_path = getattr(target_repo, "repo_path", None)
if not repo_path:
continue
snapshots = Path(repo_path) / "snapshots"
if not snapshots.is_dir():
continue
try:
snap_dirs = [s for s in snapshots.iterdir() if s.is_dir() and not s.is_symlink()]
except OSError:
continue
for snap in snap_dirs:
try:
snap.rmdir()
removed += 1
except OSError as e:
if e.errno != errno.ENOTEMPTY:
failures.append(f"{snap.name}: {e}")
return removed, failures
def _delete_gguf_variant_from_repos(
repo_id: str,
variant: str,
@ -206,11 +277,26 @@ def _delete_gguf_variant_from_repos(
)
state_purged = download_manifest.purge_state("model", repo_id, variant)
# Reclaim the empty quant folder so it stops 404ing on delete.
removed_dirs, dir_failures = _remove_empty_variant_dirs(target_repos, variant)
removed_snap_dirs, snap_dir_failures = _remove_empty_snapshot_dirs(target_repos)
removed_dirs += removed_snap_dirs
dir_failures.extend(snap_dir_failures)
if dir_failures:
raise HTTPException(
status_code = 409,
detail = (
f"Couldn't fully delete {variant} for {repo_id}: "
f"{len(dir_failures)} folder(s) could not be removed "
"(read-only cache or in use). Try again."
),
)
if (
removed_snapshots == 0
and deleted_blobs == 0
and incomplete_result.deleted == 0
and not state_purged
and removed_dirs == 0
):
raise HTTPException(
status_code = 404,
@ -225,6 +311,181 @@ def _delete_gguf_variant_from_repos(
return {"status": "deleted", "repo_id": repo_id, "variant": variant}
def reclaim_replaced_gguf_variant(
repo_id: str,
variant: str,
keep_main_hashes: frozenset[str],
hf_token: Optional[str] = None,
) -> dict:
"""Prune stale main-GGUF files for a variant after a replacement verified.
This is intentionally narrower than user-driven delete: it removes only
same-variant main files whose local blob hash is not in *keep_main_hashes*,
then unlinks their blobs only if no remaining snapshot references them.
Shared companions and sibling variants are left intact.
"""
if not keep_main_hashes:
logger.info(
"Skipping stale GGUF reclaim for %s [%s]: current main hashes unresolved",
repo_id,
variant,
)
return {
"status": "skipped",
"repo_id": repo_id,
"variant": variant,
"reason": "unresolved_hashes",
}
if not _is_valid_repo_id(repo_id) or not _is_valid_gguf_variant(variant):
return {
"status": "skipped",
"repo_id": repo_id,
"variant": variant,
"reason": "invalid_target",
}
failures: list[str] = []
removed_snapshots = 0
deleted_blobs = 0
deleted_bytes = 0
variant_key = variant.lower()
try:
cache_scans = cache_inventory.all_hf_cache_scans()
except Exception as e:
logger.warning(
"Skipping stale GGUF reclaim for %s [%s]: cache scan failed: %s",
repo_id,
variant,
download_registry.scrub_secrets(str(e), hf_token = hf_token),
)
return {
"status": "skipped",
"repo_id": repo_id,
"variant": variant,
"reason": "scan_failed",
}
candidate_repos = [
repo_info
for hf_cache in cache_scans
for repo_info in hf_cache.repos
if str(getattr(repo_info, "repo_type", "")) == "model"
and str(getattr(repo_info, "repo_id", "")).lower() == repo_id.lower()
]
try:
matched_repo_ids = resolve_destructive_repo_ids(
repo_id,
[str(getattr(repo_info, "repo_id", "")) for repo_info in candidate_repos],
noun = "models",
)
except HTTPException as e:
detail = getattr(e, "detail", str(e))
logger.warning(
"Skipping stale GGUF reclaim for %s [%s]: %s",
repo_id,
variant,
download_registry.scrub_secrets(str(detail), hf_token = hf_token),
)
return {
"status": "skipped",
"repo_id": repo_id,
"variant": variant,
"reason": "ambiguous_repo",
}
target_repos = [
repo_info
for repo_info in candidate_repos
if str(getattr(repo_info, "repo_id", "")) in matched_repo_ids
]
for target_repo in target_repos:
repo_dir = Path(target_repo.repo_path) if getattr(target_repo, "repo_path", None) else None
stale_matches: list[tuple[Path, Optional[Path], str]] = []
matches = _repo_file_matches(
target_repo,
lambda name: _is_main_gguf_filename(name)
and extract_quant_label(name).lower() == variant_key,
)
for snap, blob, name in matches:
blob_hash = _blob_hash_from_path(blob) if blob is not None else None
if blob_hash is None or blob_hash in keep_main_hashes:
continue
stale_matches.append((snap, blob, name))
if not stale_matches:
continue
for snap, _blob, name in stale_matches:
try:
if _path_exists_or_symlink(snap):
snap.unlink()
removed_snapshots += 1
except OSError as e:
failures.append(f"{name}: {e}")
ref_counts = _snapshot_blob_reference_counts(repo_dir)
seen_blobs: set[Path] = set()
for _snap, blob, name in stale_matches:
if blob is None:
continue
try:
blob_key = blob.resolve()
except OSError:
blob_key = blob
if blob_key in seen_blobs:
continue
seen_blobs.add(blob_key)
if ref_counts.get(blob_key, 0) > 0:
continue
try:
if blob.exists():
deleted_bytes += blob.stat().st_size
blob.unlink()
deleted_blobs += 1
except OSError as e:
failures.append(f"{name}: {e}")
removed_dirs = 0
dir_failures: list[str] = []
if target_repos:
removed_dirs, dir_failures = _remove_empty_variant_dirs(target_repos, variant)
removed_snap_dirs, snap_dir_failures = _remove_empty_snapshot_dirs(target_repos)
removed_dirs += removed_snap_dirs
dir_failures.extend(snap_dir_failures)
failures.extend(dir_failures)
if failures:
logger.warning(
"Stale GGUF reclaim for %s [%s] left %d failure(s): %s",
repo_id,
variant,
len(failures),
"; ".join(failures[:3]),
)
if removed_snapshots or deleted_blobs or removed_dirs:
cache_inventory.invalidate_hf_cache_scans()
logger.info(
"Reclaimed stale GGUF %s [%s]: snapshots=%d blobs=%d dirs=%d freed=%.1f MB",
repo_id,
variant,
removed_snapshots,
deleted_blobs,
removed_dirs,
deleted_bytes / (1024 * 1024),
)
return {
"status": "reclaimed",
"repo_id": repo_id,
"variant": variant,
"removed_snapshots": removed_snapshots,
"deleted_blobs": deleted_blobs,
"removed_dirs": removed_dirs,
}
def _loaded_id_matches_repo(loaded_id: str, repo_id: str) -> bool:
"""True when *loaded_id* is *repo_id* or a file within it; ``/``-boundary aware so ``org/model`` doesn't match sibling ``org/model-v2``."""
rid = repo_id.lower()

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

@ -27,6 +27,7 @@ from hub.utils.gguf import (
extract_quant_label,
iter_hf_cache_snapshots,
is_big_endian_gguf_path,
list_empty_gguf_variant_dirs,
list_gguf_variants,
list_gguf_variants_from_hf_cache,
list_local_gguf_variants,
@ -290,6 +291,75 @@ def _partial_transport_for_variant(repo_id: str, variant: str) -> Optional[str]:
return hf_cache_scan.partial_transport_for("model", repo_id, variant)
def _local_main_gguf_blobs_by_quant(repo_id: str) -> dict[str, dict[str, set[str]]]:
"""Map quant -> repo-relative expected GGUF filename -> cached blob hashes.
Shared companions are copied into each main-quant bucket so update checks can
detect mmproj/MTP-only upstream changes without a separate remote call.
"""
result: dict[str, dict[str, set[str]]] = {}
companion_blobs: dict[str, set[str]] = {}
try:
from hub.services.models import cache_inventory
scans = cache_inventory.all_hf_cache_scans()
except Exception as e:
logger.warning("Failed to scan local GGUF blobs for %s: %s", repo_id, e)
return result
target_lower = repo_id.lower()
for hf_cache in scans:
for repo_info in hf_cache.repos:
if str(getattr(repo_info, "repo_type", "")) != "model":
continue
if str(getattr(repo_info, "repo_id", "")).lower() != target_lower:
continue
for path, hashes in cache_inventory._repo_gguf_blob_map(
repo_info,
include_companions = True,
).items():
normalized = str(path).replace("\\", "/")
if not hashes:
continue
if _is_mmproj_filename(normalized) or _is_mtp_drafter_path(normalized):
companion_blobs.setdefault(normalized, set()).update(
str(blob) for blob in hashes if blob
)
continue
quant = extract_quant_label(normalized).lower()
if is_big_endian_gguf_path(normalized, quant):
continue
bucket = result.setdefault(quant, {}).setdefault(normalized, set())
bucket.update(str(blob) for blob in hashes if blob)
if companion_blobs:
for local_blobs in result.values():
for path, hashes in companion_blobs.items():
local_blobs.setdefault(path, set()).update(hashes)
return result
def _variant_update_available_from_requirement(
local_blobs: dict[str, set[str]], requirement: Optional[_GgufVariantRequirement], variant: str
) -> bool:
if requirement is None or not local_blobs:
return False
local_by_posix = {path.replace("\\", "/"): blobs for path, blobs in local_blobs.items()}
for expected in requirement.expected_files:
path = str(expected.path).replace("\\", "/")
if not (
is_main_gguf_variant_path(path, variant)
or _is_mmproj_filename(path)
or _is_mtp_drafter_path(path)
):
continue
remote_blob = expected.sha256
if not remote_blob:
continue
local_set = local_by_posix.get(path)
if not local_set or remote_blob not in local_set:
return True
return False
def delete_variant_incomplete_blobs_result(
repo_id: str,
variant: str,
@ -334,6 +404,32 @@ def delete_variant_incomplete_blobs_result(
return VariantIncompleteDeleteResult(deleted = deleted, unresolved = False)
def _mark_empty_dir_cleanables(
repo_id: str, response: GgufVariantsResponse
) -> GgufVariantsResponse:
"""Surface empty leftover ``<quant>/`` folders (interrupted downloads) as
partial so the UI can delete them -- on local/offline paths too, not just a
remote listing. A listed quant is flipped to partial; an unlisted one is
appended as a zero-byte cleanable entry."""
try:
empty_labels = list_empty_gguf_variant_dirs(repo_id)
except Exception as e:
logger.warning(f"Failed to scan empty GGUF variant folders for {repo_id}: {e}")
return response
if not empty_labels:
return response
empty_by_key = {label.lower(): label for label in empty_labels}
variants = list(response.variants)
listed = {v.quant.lower() for v in variants}
for i, v in enumerate(variants):
if v.quant.lower() in empty_by_key and not v.downloaded and not v.partial:
variants[i] = v.model_copy(update = {"partial": True})
for key, label in sorted(empty_by_key.items()):
if key not in listed:
variants.append(GgufVariantDetail(filename = f"{label}.gguf", quant = label, partial = True))
return response.model_copy(update = {"variants": variants})
async def get_gguf_variants_response(
repo_id: str,
prefer_local_cache: bool = False,
@ -630,9 +726,12 @@ async def get_gguf_variants_response(
_partial_transport_for_variant(repo_id, variant.quant),
)
local_blobs_by_quant = _local_main_gguf_blobs_by_quant(repo_id)
def _variant_detail(v) -> GgufVariantDetail:
is_partial = v.quant in partial_quants
requirement = requirements_by_quant.get(v.quant.lower())
downloaded = _is_fully_downloaded(v) and not is_partial
return GgufVariantDetail(
filename = v.filename,
quant = v.quant,
@ -641,7 +740,13 @@ async def get_gguf_variants_response(
download_size_bytes = (
requirement.download_size_bytes if requirement is not None else v.size_bytes
),
downloaded = _is_fully_downloaded(v) and not is_partial,
downloaded = downloaded,
update_available = downloaded
and _variant_update_available_from_requirement(
local_blobs_by_quant.get(v.quant.lower(), {}),
requirement,
v.quant,
),
partial = is_partial,
partial_transport = (partial_quant_transports.get(v.quant) if is_partial else None),
)
@ -653,8 +758,28 @@ async def get_gguf_variants_response(
default_variant = default_variant,
)
def _compute_with_cleanables() -> GgufVariantsResponse:
skip = is_local_path(repo_id) or not _is_valid_repo_id(repo_id)
try:
response = _compute()
except Exception:
# Offline / metadata fetch failed with only an empty leftover
# <quant>/ folder cached: still surface it so the UI can delete it,
# otherwise re-raise the original error.
if skip:
raise
enriched = _mark_empty_dir_cleanables(
repo_id, GgufVariantsResponse(repo_id = repo_id, variants = [])
)
if enriched.variants:
return enriched
raise
if skip:
return response
return _mark_empty_dir_cleanables(repo_id, response)
try:
return await asyncio.to_thread(_compute)
return await asyncio.to_thread(_compute_with_cleanables)
except HTTPException:
raise
except Exception as e:

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

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

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

@ -1632,6 +1632,34 @@ def test_variant_partial_accepts_variant_filtered_legacy_hashes(monkeypatch, tmp
)
def test_variant_partial_accepts_completed_variant_in_non_latest_snapshot(monkeypatch, tmp_path):
"""A verified GGUF update can prune an older snapshot and make that old
directory the newest by mtime. The variant is still complete when another
snapshot satisfies its manifest."""
monkeypatch.setattr(state_dir, "cache_root", lambda: tmp_path / "state")
repo_dir = tmp_path / "cache" / "models--Org--Repo"
old_snapshot = repo_dir / "snapshots" / "old"
new_snapshot = repo_dir / "snapshots" / "new"
old_snapshot.mkdir(parents = True)
new_snapshot.mkdir(parents = True)
(old_snapshot / "model-Q8_0.gguf").write_bytes(b"sibling")
(new_snapshot / "model-Q4_K_M.gguf").write_bytes(b"new")
assert download_manifest.write_manifest(
"model",
"Org/Repo",
"Q4_K_M",
[download_manifest.ExpectedFile(path = "model-Q4_K_M.gguf", size = 3)],
"http",
)
assert not inventory_scan.is_variant_partial(
"Org/Repo",
"Q4_K_M",
snapshot_dir = old_snapshot,
repo_cache_dir = repo_dir,
)
def test_gguf_variants_partial_marker_overrides_size_only_downloaded(monkeypatch, tmp_path):
async def _run_inline(fn, *args, **kwargs):
return fn(*args, **kwargs)

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

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

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

@ -36,7 +36,10 @@ def sibling_sha256(sibling) -> Optional[str]:
value = lfs.get("sha256")
else:
value = getattr(lfs, "sha256", None)
return value if isinstance(value, str) and value else None
if isinstance(value, str) and value:
return value
blob_id = getattr(sibling, "blob_id", None)
return blob_id if isinstance(blob_id, str) and blob_id else None
def sibling_size(sibling) -> int:

46
studio/backend/hub/utils/inventory_scan.py
View file

@ -387,9 +387,55 @@ def _manifest_partial(
)
if resolved is None:
return True
if repo_type == "model" and variant is not None:
if download_manifest.verify_against_disk(manifest, resolved).ok:
return False
for candidate in _manifest_snapshot_dirs(repo_type, repo_id, repo_cache_dir):
if candidate == resolved:
continue
if download_manifest.verify_against_disk(manifest, candidate).ok:
return False
return True
return not download_manifest.verify_against_disk(manifest, resolved).ok
def _manifest_snapshot_dirs(
repo_type: RepoType,
repo_id: str,
repo_cache_dir: Optional[Path] = None,
) -> list[Path]:
repo_dirs = (
[repo_cache_dir]
if repo_cache_dir is not None
else list(iter_repo_cache_dirs(repo_type, repo_id))
)
snapshots: list[Path] = []
seen: set[str] = set()
for repo_dir in repo_dirs:
if repo_dir is None:
continue
snapshots_dir = repo_dir / "snapshots"
try:
if not snapshots_dir.is_dir():
continue
entries = list(snapshots_dir.iterdir())
except OSError:
continue
for entry in entries:
try:
if not entry.is_dir():
continue
resolved = entry.resolve()
except OSError:
continue
key = str(resolved)
if key in seen:
continue
seen.add(key)
snapshots.append(resolved)
return snapshots
def is_snapshot_partial(
repo_type: RepoType,
repo_id: str,

15
studio/backend/hub/workers/hf_download.py
View file

@ -653,6 +653,21 @@ def _download_gguf_variant(repo_id: str, variant: str, hf_token: str | None, mod
snapshot_path,
metadata_unavailable = metadata_unavailable,
)
if plan is not None:
try:
from hub.services.models.deletion import reclaim_replaced_gguf_variant
reclaim_replaced_gguf_variant(
repo_id,
variant,
plan.main_hashes,
hf_token,
)
except Exception as e:
print(
f"Verified GGUF update for {repo_id} [{variant}], but stale-cache "
f"reclaim failed ({type(e).__name__}: {e})",
file = sys.stderr,
)
def _download_dataset(repo_id: str, hf_token: str | None, mode: str) -> None:

84
studio/backend/main.py
View file

@ -441,9 +441,30 @@ def _start_llama_cpp_probes_if_enabled(app: FastAPI) -> None:
).start()
def _warm_rag_embedder() -> None:
"""Warm RAG embeddings without blocking backend readiness."""
try:
from storage import rag_db
if not rag_db.RAG_AVAILABLE:
return
from core.rag import embeddings
embeddings.warm()
except Exception:
pass
@asynccontextmanager
async def lifespan(app: FastAPI):
"""Startup: detect hardware, seed default admin if needed. Shutdown: clean up compiled cache."""
import time as _time
_lifespan_started = _time.perf_counter()
import structlog as _structlog
_lifespan_log = _structlog.get_logger(__name__)
clear_unsloth_compiled_cache()
# Remove stale .venv_overlay from old versions; switching now uses .venv_t5/.
@ -454,6 +475,11 @@ async def lifespan(app: FastAPI):
# Detect hardware first — sets the DEVICE global used everywhere.
detect_hardware()
_lifespan_log.info(
"lifespan hardware detection completed in %.1fms",
(_time.perf_counter() - _lifespan_started) * 1000,
)
# Apple Silicon with MLX missing => Train/Export are greyed out (chat-only).
# Reinstall mlx by name on a background thread (off the critical path) and
# re-detect, so a reinstall/update that dropped mlx self-heals. No-op
@ -465,7 +491,13 @@ async def lifespan(app: FastAPI):
import structlog as _structlog
_structlog.get_logger(__name__).debug("mlx autorepair skipped: %s", _mlx_exc)
# Reap download workers orphaned by a previous crash before new downloads start.
# Reap workers/runs orphaned by a previous crash before new work starts.
try:
from storage.studio_db import cleanup_orphaned_runs
cleanup_orphaned_runs()
except Exception as exc:
_lifespan_log.warning("cleanup_orphaned_runs failed at startup: %s", exc)
reap_hub_orphan_workers()
# llama.cpp probes: capability (MTP support) + freshness (release age).
@ -479,35 +511,23 @@ async def lifespan(app: FastAPI):
app.state.llama_cpp_freshness = None
_start_llama_cpp_probes_if_enabled(app)
from storage.studio_db import cleanup_orphaned_runs
try:
cleanup_orphaned_runs()
from storage.rag_db import reconcile_orphaned_ingestion_jobs
reconcile_orphaned_ingestion_jobs()
except Exception as exc:
import structlog
structlog.get_logger(__name__).warning("cleanup_orphaned_runs failed at startup: %s", exc)
_lifespan_log.warning("reconcile_orphaned_ingestion_jobs failed at startup: %s", exc)
_start_helper_precache_if_enabled()
threading.Thread(target = _warm_rag_embedder, daemon = True, name = "rag-embedder-warm").start()
# Warm the RAG embedder so the first upload skips the cold load. Non-fatal.
def _warm_rag_embedder():
try:
from storage import rag_db
if not rag_db.RAG_AVAILABLE:
return
from core.rag import embeddings
embeddings.warm()
except Exception:
pass
threading.Thread(target = _warm_rag_embedder, daemon = True).start()
# Initialize RSA key pair for API key encryption (external providers)
# Initialize RSA key pair for API key encryption (external providers).
from core.inference.key_exchange import init_key_pair
init_key_pair()
_lifespan_log.info(
"lifespan pre-auth setup completed in %.1fms",
(_time.perf_counter() - _lifespan_started) * 1000,
)
if storage.ensure_default_admin():
bootstrap_pw = storage.get_bootstrap_password()
@ -522,6 +542,11 @@ async def lifespan(app: FastAPI):
print("=" * 60 + "\n")
else:
app.state.bootstrap_password = storage.get_bootstrap_password()
_lifespan_log.info(
"lifespan startup completed in %.1fms",
(_time.perf_counter() - _lifespan_started) * 1000,
)
yield
from core.inference.llama_http import aclose as _close_llama_http
@ -909,6 +934,21 @@ install_api_error_handlers(app)
# ============ Health and System Endpoints ============
@app.get("/api/liveness")
async def liveness_check():
"""Cheap process liveness for desktop port validation."""
return {
"status": "alive",
"service": "Unsloth UI Backend",
"desktop_protocol_version": 1,
"desktop_manageability_version": 1,
"supports_desktop_auth": True,
"supports_desktop_backend_ownership": True,
"studio_root_id": _studio_root_id(),
**({"desktop_owner": owner} if (owner := _desktop_owner()) else {}),
}
@app.get("/api/health")
async def health_check(request: Request):
"""Liveness plus launcher capability bits; host fingerprint gated on a bearer.

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

@ -158,9 +158,15 @@ class ExportCommonOptions(BaseModel):
class ExportMergedModelRequest(ExportCommonOptions):
"""Request for exporting a merged PEFT model."""
format_type: Literal["16-bit (FP16)", "4-bit (FP4)"] = Field(
format_type: Literal[
"16-bit (FP16)",
description = "Export precision / format for the merged model",
"4-bit (FP4)",
"FP8 (compressed-tensors)",
"NVFP4 (compressed-tensors)",
] = Field(
"16-bit (FP16)",
description = "Export precision / format for the merged model. The compressed-tensors "
"options run llm-compressor for vLLM (FP8 is data-free; NVFP4 calibrates).",
)
@ -199,6 +205,15 @@ class ExportGGUFRequest(BaseModel):
None,
description = "Hugging Face token for GGUF upload",
)
imatrix: bool = Field(
False,
description = "Use an importance matrix (auto-downloads the upstream unsloth GGUF "
"imatrix). Required for the IQ low-bit quants such as iq2_xxs / iq4_xs.",
)
imatrix_path: Optional[str] = Field(
None,
description = "Path to a custom imatrix file; overrides the auto-download when set.",
)
class ExportLoRAAdapterRequest(ExportCommonOptions):

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

@ -106,8 +106,7 @@ class LoadRequest(BaseModel):
"Extra arguments forwarded verbatim to llama-server for GGUF models. "
"One token per list entry, e.g. ['--top-k', '20', '--seed', '42']. "
"Studio-managed flags (model identity, port, context length, GPU placement, "
"auth, --flash-attn, --no-context-shift, --jinja) are rejected. Ignored for "
"non-GGUF models."
"auth, UI/server mode) are rejected. Ignored for non-GGUF models."
),
)
@ -1756,12 +1755,18 @@ class DiffusionLoadRequest(BaseModel):
"scripts/build_prequant_checkpoint.py) for the requested transformer_quant "
"scheme. Loads the already-quantized weights with the dense bf16 never on the "
"GPU (~half the load VRAM and a smaller download). null uses the family's hosted "
"checkpoint if configured, else quantises the dense transformer at load time.",
"checkpoint if configured, else quantises the dense transformer at load time. "
"Loading a local path unpickles the file (arbitrary code execution), so it is "
"ignored unless the path resolves inside a directory the operator allowlisted "
"via UNSLOTH_ALLOW_LOCAL_PREQUANT_PATH (one or more directories, separated by "
"the OS path separator). A bare on/off value such as '1' is deliberately not "
"accepted -- it must name an allowed directory.",
)
attention_backend: Optional[
Literal[
"auto",
"native",
"sdpa",
"cudnn",
"flash",
"flash2",
@ -1776,7 +1781,7 @@ class DiffusionLoadRequest(BaseModel):
description = "Attention kernel via the diffusers dispatcher. auto picks the best "
"exact backend for the device (cuDNN fused attention on NVIDIA, ~1.18x and "
"near-lossless, when a speed profile is active; native SDPA elsewhere and when "
"speed=off). native forces default SDPA; cudnn/flash/flash3/flash4 are exact "
"speed=off). native (alias sdpa) forces default SDPA; cudnn/flash/flash3/flash4 are exact "
"(kernel/arch-gated); sage is INT8 attention (a small quality cost, consumer "
"friendly); xformers/aiter are memory-efficient (NVIDIA) / AMD ROCm. An "
"unavailable kernel falls back to the default.",
@ -1893,6 +1898,9 @@ class GalleryImage(BaseModel):
guidance: float = Field(..., description = "Guidance scale")
seed: int = Field(..., description = "Seed used")
batch_index: int = Field(0, description = "Position within its batch (0-based)")
batch_size: int = Field(
1, description = "Batch size used; with batch_index it lets restore replay this image"
)
model: Optional[str] = Field(None, description = "Model repo id that produced it")
created_at: float = Field(..., description = "Creation time (epoch seconds)")

4
studio/backend/models/models.py
View file

@ -136,9 +136,13 @@ class GgufVariantDetail(BaseModel):
filename: str = Field(..., description = "GGUF filename (e.g., 'gemma-3-4b-it-Q4_K_M.gguf')")
quant: str = Field(..., description = "Quantization label (e.g., 'Q4_K_M')")
size_bytes: int = Field(0, description = "File size in bytes")
download_size_bytes: int = Field(0, description = "Total bytes needed to download this variant")
downloaded: bool = Field(
False, description = "Whether this variant is already in the local HF cache"
)
update_available: bool = Field(
False, description = "Whether a newer version of this variant is available on HF"
)
class GgufVariantsResponse(BaseModel):

16
studio/backend/models/training.py
View file

@ -9,6 +9,8 @@ import re
from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
from typing import Any, Optional, List, Dict, Literal
from utils.training_runs import normalize_project_name
# ASCII integer, optional single sign. Rejects "++512" and Unicode digits
# ("５１２") that slip through str.isdigit() + int().
@ -97,6 +99,11 @@ class TrainingStartRequest(BaseModel):
model_name: str = Field(
..., description = "Model identifier (e.g., 'unsloth/llama-3-8b-bnb-4bit')"
)
project_name: Optional[str] = Field(
None,
max_length = 80,
description = "Optional user-defined project name appended to run folders and shown in history",
)
training_type: Literal["LoRA/QLoRA", "Full Finetuning", "Continued Pretraining"] = Field(
...,
description = "Training type: 'LoRA/QLoRA', 'Full Finetuning', or 'Continued Pretraining'",
@ -155,6 +162,11 @@ class TrainingStartRequest(BaseModel):
values.setdefault("train_split", values.pop("split"))
return values
@field_validator("project_name")
@classmethod
def _normalize_project_name(cls, value: Optional[str]) -> Optional[str]:
return normalize_project_name(value)
# NOTE: pydantic runs all `mode="after"` validators in definition order. A
# second one, `_check_steps_or_epochs`, is defined lower in this class; keep
# these cross-field checks order-independent so the two stay decoupled.
@ -588,6 +600,7 @@ class TrainingRunSummary(BaseModel):
id: str
status: Literal["running", "completed", "stopped", "error"]
model_name: str
project_name: Optional[str] = None
dataset_name: str
display_name: Optional[str] = None
started_at: str
@ -603,6 +616,9 @@ class TrainingRunSummary(BaseModel):
resumed_later: bool = False
has_preview_model: bool = False
preview_ref: Optional[str] = None
# HMAC capability token for the `/p/{preview_ref}` share link; None when not
# previewable. The frontend appends it as `?k=` so a guessed ref can't be used.
preview_sig: Optional[str] = None
class TrainingRunUpdateRequest(BaseModel):

5
studio/backend/requirements/no-torch-runtime.txt
View file

@ -73,4 +73,9 @@ pillow
# this file installs --no-deps; without them Studio runs with RAG disabled.
sqlite-vec==0.1.9
pymupdf==1.27.2.3
# 0.3.x keeps pymupdf-layout (which pulls onnxruntime) an optional extra; the
# lockstep 1.27.x line makes it a hard dep we do not need for to_markdown().
pymupdf4llm==0.3.4
python-docx==1.2.0
lxml==6.0.2

3
studio/backend/requirements/studio.txt
View file

@ -26,4 +26,7 @@ gguf
# extras-no-deps.txt; these add the lexical+dense store and document parsing.
sqlite-vec==0.1.9
pymupdf==1.27.2.3
# 0.3.x keeps pymupdf-layout (which pulls onnxruntime) an optional extra; the
# lockstep 1.27.x line makes it a hard dep we do not need for to_markdown().
pymupdf4llm==0.3.4
python-docx==1.2.0

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

@ -74,9 +74,81 @@ _LOGIN_WINDOW_SECONDS = 60.0
_LOGIN_MAX_FAILS = 5
_LOGIN_IP_MAX_FAILS = 30
_LOGIN_LOCKOUT_SECONDS = 60
# Bucket-dict cap. On overflow, prune stale entries; if still full the failure
# folds into the per-IP aggregate only.
# Bucket-dict cap. On overflow, reclaim expired buckets; a new IP that still can't
# fit falls back to a sharded overflow rather than evicting a hot bucket.
_LOGIN_MAX_BUCKETS = 4096
# Last full stale-sweep time; rate-limits the O(n) sweep under a burst of new IPs.
_LAST_IP_PRUNE = 0.0
# Sharded overflow for per-IP failures that can't get their own bucket while the
# dict is saturated. Each shard is a small fixed-capacity dict ``ip -> [count,
# window_start]``: a per-IP count (so a source is throttled, and cleared on
# success, by its own failures -- no cross-IP collateral) with hard-bounded
# memory and O(1) lookups. When a shard is full a new IP evicts the lowest-count
# entry (and starts clean, never inheriting its count) rather than growing without
# bound, so a high-cardinality spray can't blow memory/CPU the way a per-failure
# deque could; a persistent attacker keeps a high count and is never the one
# evicted.
_LOGIN_IP_OVERFLOW_SHARDS = 256
_LOGIN_IP_OVERFLOW_MAX = 64 # distinct IPs tracked per shard
_LOGIN_IP_OVERFLOW: list[dict] = [dict() for _ in range(_LOGIN_IP_OVERFLOW_SHARDS)]
def _overflow_shard(ip: str) -> dict:
return _LOGIN_IP_OVERFLOW[hash(ip) % _LOGIN_IP_OVERFLOW_SHARDS]
def _overflow_record(ip: str, now: float) -> int:
"""Record an overflow failure for ``ip`` and return its windowed count."""
shard = _overflow_shard(ip)
entry = shard.get(ip)
if entry is not None:
if now - entry[1] > _LOGIN_WINDOW_SECONDS:
entry[0], entry[1] = 1, now
else:
# Only "at or above the per-IP threshold" matters for blocking, so cap
# the count there. This also keeps the migration into a per-IP bucket
# bounded -- without the cap a saturated source could accrue an
# unbounded count, then materialize one deque entry per failure
# (``[start] * carried``) on the next attempt, allocating an arbitrarily
# large deque while holding the login lock.
entry[0] = min(entry[0] + 1, _LOGIN_IP_MAX_FAILS)
return entry[0]
if len(shard) >= _LOGIN_IP_OVERFLOW_MAX:
# Make room by dropping the lowest-count entry, but the new source starts
# clean -- never inherit the evicted IP's failures, or an unrelated source
# could be 429'd after one attempt. Worst case under a saturated shard is
# that a heavy hitter briefly resets, not that a bystander is blocked.
del shard[min(shard, key = lambda k: shard[k][0])]
shard[ip] = [1, now]
return 1
def _overflow_blocked(ip: str, now: float) -> int:
"""Seconds this IP is throttled by its own overflow count, or 0."""
shard = _overflow_shard(ip)
entry = shard.get(ip)
if entry is None:
return 0
if now - entry[1] > _LOGIN_WINDOW_SECONDS:
del shard[ip]
return 0
if entry[0] >= _LOGIN_IP_MAX_FAILS:
return max(1, int(_LOGIN_WINDOW_SECONDS - (now - entry[1])))
return 0
def _overflow_take(ip: str, now: float) -> tuple[int, float]:
"""Pop ip's overflow entry, returning its ``(count, window_start)`` so the
count can migrate into a fresh per-IP bucket. ``(0, now)`` if none/expired."""
entry = _overflow_shard(ip).pop(ip, None)
if entry is None or now - entry[1] > _LOGIN_WINDOW_SECONDS:
return 0, now
# Cap the carried count so the bucket migration never allocates more than the
# per-IP threshold worth of deque entries (defensive; _overflow_record already
# clamps, but keep the bound at the consumption site too).
return min(entry[0], _LOGIN_IP_MAX_FAILS), entry[1]
# Unrepresentable as a real username (leading NUL); folds unknown-user attempts
# into one slot so attacker cardinality can't blow the bucket dict.
_UNKNOWN_LOGIN_USER = "\x00unknown-user"
@ -169,13 +241,50 @@ def _prune_stale_buckets(now: float) -> None:
_LOGIN_BUCKETS.pop(key, None)
def _prune_stale_ip_buckets(now: float) -> None:
"""Drop empty / expired per-IP buckets to bound memory under spray.
The dict is otherwise reclaimed only on a successful login, so a failure-only
spray from many (or spoofed) IPs would grow it without bound.
"""
stale: list[str] = []
for bucket_ip, bucket in _LOGIN_IP_BUCKETS.items():
_prune_bucket(bucket, now)
if not bucket:
stale.append(bucket_ip)
for bucket_ip in stale:
_LOGIN_IP_BUCKETS.pop(bucket_ip, None)
def _record_login_failure(key: tuple[str, str]) -> int:
global _LAST_IP_PRUNE
now = time.monotonic()
ip, _username = key
with _LOGIN_BUCKETS_LOCK:
ip_bucket = _LOGIN_IP_BUCKETS.setdefault(ip, deque())
_prune_bucket(ip_bucket, now)
ip_bucket.append(now)
# Keep the dict bounded without disabling throttling and without letting a
# spray reset a hot bucket: for a new IP at the cap, reclaim expired buckets
# (rate-limited) to make room.
ip_bucket = _LOGIN_IP_BUCKETS.get(ip)
if ip_bucket is None and len(_LOGIN_IP_BUCKETS) >= _LOGIN_MAX_BUCKETS:
if now - _LAST_IP_PRUNE >= 1.0:
_prune_stale_ip_buckets(now)
_LAST_IP_PRUNE = now
if ip_bucket is None and len(_LOGIN_IP_BUCKETS) >= _LOGIN_MAX_BUCKETS:
# Still full -- every bucket is hot. Count this failure in the IP's
# bounded overflow shard instead of evicting a live one, so the spray
# stays throttled but can't push out (and reset) any IP's own counter.
ip_fails = _overflow_record(ip, now)
else:
if ip_bucket is None:
ip_bucket = _LOGIN_IP_BUCKETS[ip] = deque()
# Carry over any overflow failures this IP accrued while the dict
# was saturated, so straddling the overflow -> bucket transition
# can't double the effective per-IP limit.
carried, start = _overflow_take(ip, now)
ip_bucket.extend([start] * carried)
_prune_bucket(ip_bucket, now)
ip_bucket.append(now)
ip_fails = len(ip_bucket)
if key not in _LOGIN_BUCKETS and len(_LOGIN_BUCKETS) >= _LOGIN_MAX_BUCKETS:
_prune_stale_buckets(now)
@ -184,8 +293,8 @@ def _record_login_failure(key: tuple[str, str]) -> int:
_prune_bucket(account_bucket, now)
account_bucket.append(now)
return len(account_bucket)
# Bucket dict at cap; per-IP cap still applies via ip_bucket.
return len(ip_bucket)
# Both dicts at cap (sustained spray): fall back to the per-IP count.
return ip_fails
def _blocked_for(bucket: deque | None, now: float, max_fails: int) -> int:
@ -202,10 +311,16 @@ def _login_blocked(key: tuple[str, str]) -> int:
now = time.monotonic()
ip, _username = key
with _LOGIN_BUCKETS_LOCK:
return max(
_blocked_for(_LOGIN_BUCKETS.get(key), now, _LOGIN_MAX_FAILS),
# Honor the IP's overflow shard regardless of current dict capacity: a
# source counted there during saturation must stay throttled until those
# failures age out, even if a bucket later frees up -- otherwise a fresh
# bucket would reset it. Shards are empty outside saturation, so this is a
# no-op in the common case.
ip_blocked = max(
_blocked_for(_LOGIN_IP_BUCKETS.get(ip), now, _LOGIN_IP_MAX_FAILS),
_overflow_blocked(ip, now),
)
return max(_blocked_for(_LOGIN_BUCKETS.get(key), now, _LOGIN_MAX_FAILS), ip_blocked)
def _clear_login_bucket(key: tuple[str, str]) -> None:
@ -213,6 +328,10 @@ def _clear_login_bucket(key: tuple[str, str]) -> None:
with _LOGIN_BUCKETS_LOCK:
_LOGIN_BUCKETS.pop(key, None)
_LOGIN_IP_BUCKETS.pop(ip, None)
# A successful login resets the IP's throttle, including any overflow it
# accumulated during saturation (drop only this IP's entry, so a
# shard-mate's throttle is untouched).
_overflow_shard(ip).pop(ip, None)
# Sync def (not async): compute_identity_proof touches SQLite on the first call,

1
studio/backend/routes/chat_history.py
View file

@ -150,6 +150,7 @@ class ChatInferenceSettings(BaseModel):
maxSeqLength: Optional[float] = None
maxTokens: Optional[float] = None
systemPrompt: Optional[str] = None
systemVariables: Optional[str] = None
trustRemoteCode: Optional[bool] = None
fastMode: Optional[bool] = None

31
studio/backend/routes/data_recipe/seed.py
View file

@ -481,6 +481,37 @@ async def upload_unstructured_file(
error = "No extractable text found in file",
)
extracted_path.write_text(extracted_text, encoding = "utf-8")
except ImportError as e:
raw_path.unlink(missing_ok = True)
extracted_path.unlink(missing_ok = True)
missing = getattr(e, "name", None)
expected_missing = {".pdf": "pymupdf4llm", ".docx": "mammoth"}.get(ext)
if isinstance(e, ModuleNotFoundError) and missing == expected_missing:
logger.error(
"data_recipe.seed.text_extraction_dependency_missing",
error = str(e),
missing = missing,
exc_info = True,
)
return UnstructuredFileUploadResponse(
file_id = file_id,
filename = original_filename,
size_bytes = size_bytes,
status = "error",
error = f"Cannot read {ext} files: the '{missing}' package is not installed.",
)
logger.error(
"data_recipe.seed.text_extraction_failed",
error = str(e),
exc_info = True,
)
return UnstructuredFileUploadResponse(
file_id = file_id,
filename = original_filename,
size_bytes = size_bytes,
status = "error",
error = "Text extraction failed.",
)
except Exception as e:
raw_path.unlink(missing_ok = True)
extracted_path.unlink(missing_ok = True)

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

@ -343,6 +343,8 @@ async def export_gguf(
"""
try:
backend = get_export_backend()
# A custom path wins; otherwise the imatrix toggle requests the upstream auto-download.
imatrix_file = request.imatrix_path or (True if request.imatrix else None)
success, message, output_path = await asyncio.to_thread(
backend.export_gguf,
save_directory = request.save_directory,
@ -350,6 +352,7 @@ async def export_gguf(
push_to_hub = request.push_to_hub,
repo_id = request.repo_id,
hf_token = request.hf_token,
imatrix_file = imatrix_file,
)
if not success:

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

@ -683,7 +683,9 @@ try:
detect_reasoning_flags,
)
from core.inference.llama_server_args import (
_effective_tensor_parallel,
_tensor_parallel_matches_loaded,
parse_split_mode_override,
resolve_tensor_parallel,
strip_shadowing_flags,
validate_extra_args,
@ -718,7 +720,9 @@ except ImportError:
detect_reasoning_flags,
)
from core.inference.llama_server_args import (
_effective_tensor_parallel,
_tensor_parallel_matches_loaded,
parse_split_mode_override,
resolve_tensor_parallel,
strip_shadowing_flags,
validate_extra_args,
@ -1115,6 +1119,7 @@ from auth.authentication import get_current_subject
from state.tool_approvals import resolve_tool_decision
from core.inference.key_exchange import decrypt_api_key
from core.inference.model_ids import public_model_id
from core.inference.api_monitor import api_monitor
from core.inference.llama_http import nonstreaming_client
from core.inference.providers import get_base_url
@ -1242,15 +1247,13 @@ async def _authenticate_header_or_query(request: Request, token: Optional[str])
@studio_router.get("/artifact-preview-frame", include_in_schema = False)
async def artifact_preview_frame(
request: Request,
allow_network: bool = False,
token: Optional[str] = None,
):
"""Serve the opaque sandbox shell used for client-side HTML canvases."""
async def artifact_preview_frame(allow_network: bool = False):
"""Serve the opaque sandbox shell for client-side HTML canvases.
if allow_network:
await _authenticate_header_or_query(request, token)
No auth token by design: the URL is readable by the untrusted canvas via
location.href, and this static shell exposes no server resource (frame-ancestors
plus the sandbox already gate it), so the CSP is chosen from allow_network alone.
"""
csp = (
_ARTIFACT_PREVIEW_FRAME_NETWORK_CSP if allow_network else _ARTIFACT_PREVIEW_FRAME_STRICT_CSP
@ -2087,6 +2090,32 @@ def _should_strip_split_mode(request: LoadRequest, backend_extra: Optional[list[
)
def _carry_preserved_tensor_intent(
*, preserved: bool, same_model: bool, explicit_drop: bool
) -> bool:
"""Carry a preserved multi-GPU layer fallback forward only for a reload of the
SAME loaded model that doesn't explicitly drop tensor intent, so a fitting model
isn't collapsed to one GPU on a ctx-only change -- but an unrelated model switch
(without /unload) or an explicit tensor-off doesn't inherit it (#6659)."""
return preserved and same_model and not explicit_drop
def _is_explicit_tensor_drop(request: LoadRequest) -> bool:
"""True only when the request explicitly selects a non-tensor --split-mode (e.g.
layer/row/none), a deliberate departure from a preserved tensor->layer fallback.
A bare tensor_parallel field is NOT a drop: the Studio UI always sends it and echoes
the /load response's resolved value back, so after a fallback every reload carries
tensor_parallel=false even though the user never changed it -- treating that as a drop
would collapse the preserved multi-GPU placement on the next ctx/settings reload. An
empty clear is not a drop either (a fallback always stores --split-mode layer, never a
tensor split mode, so a clear never wipes tensor intent), nor is an unrelated extra
(--top-k) or inherit (None). tensor_parallel=true / --split-mode tensor re-engage
tensor. Shared by the already-loaded dedup and the load carry-forward (#6659)."""
override = parse_split_mode_override(request.llama_extra_args)
return override is not None and override.strip().lower() != "tensor"
def _request_matches_loaded_settings(
request: LoadRequest,
llama_backend: LlamaCppBackend,
@ -2125,6 +2154,13 @@ def _request_matches_loaded_settings(
effective_extra, request.tensor_parallel, llama_backend.tensor_parallel
):
return False
# Preserved tensor->layer fallback (both report tensor=off, so the check above
# matches): if the user now explicitly drops tensor intent, reload so placement
# re-selects instead of keeping the all-GPU mask (#6659). The effective check
# includes the env, so an env-only tensor (LLAMA_ARG_SPLIT_MODE=tensor) that
# can't actually be dropped falls through to the env-downgrade match, not a loop.
if llama_backend.layer_preserves_tensor_intent and _is_explicit_tensor_drop(request):
return False
# Spec decoding works on vision models too (MTP is mmproj-compatible,
# llama.cpp #22673; the old ``not is_vision`` gate is gone), so compare
# the real requested mode -- coercing vision to ``off`` here used to
@ -2827,6 +2863,48 @@ async def load_model(
hf_variant = config.gguf_variant,
)
# Tensor intent for this load: the request itself, or a preserved
# multi-GPU layer fallback carried across a reload of the SAME model that
# doesn't drop it (e.g. a ctx-only change), so a fitting model doesn't
# silently collapse to one GPU. Only an explicit non-tensor --split-mode
# override counts as the drop -- the tensor field echo / unrelated extras keep
# the preserved placement; the same-model guard stops a switch-without-unload
# inheriting the prior model's intent.
_explicit_tensor_drop = _is_explicit_tensor_drop(request)
# Compare the resolved config.identifier (what load_model stores), not the
# raw request id: from_identifier normalizes shorthands (adds unsloth/, fixes
# case), so a reload with the shorthand would otherwise miss the match and
# drop the carry-forward. #6659
_same_model_loaded = (
llama_backend.is_loaded
and (llama_backend.model_identifier or "").lower()
== (config.identifier or "").lower()
)
# model_identifier is variant-agnostic for HF repos and dir-level for a
# local multi-variant directory, so also require the loaded quant to match
# (path else variant, mirroring _already_in_target_state) -- otherwise a
# different variant inherits the prior one's preserved intent. #6659
if _same_model_loaded:
if config.gguf_file and llama_backend.gguf_path:
try:
_same_model_loaded = (
Path(llama_backend.gguf_path).resolve()
== Path(config.gguf_file).resolve()
)
except OSError:
_same_model_loaded = False
else:
_same_model_loaded = (llama_backend.hf_variant or "").lower() == (
config.gguf_variant or ""
).lower()
_tensor_intent_overall = _effective_tensor_parallel(
extra_llama_args, request.tensor_parallel
) or _carry_preserved_tensor_intent(
preserved = llama_backend.layer_preserves_tensor_intent,
same_model = _same_model_loaded,
explicit_drop = _explicit_tensor_drop,
)
# Run a single load attempt with the given tensor flag + extras.
async def _attempt_gguf_load(
tensor_parallel: bool, attempt_extra_args: Optional[list[str]]
@ -2840,6 +2918,12 @@ async def load_model(
**_source_load_kwargs,
**attempt_kwargs,
tensor_parallel = tensor_parallel,
# True on the layer fallback retry (tensor wanted overall but not on
# this attempt): keep multi-GPU. Mirrors the fallback's key.
preserve_multi_gpu_on_layer = bool(
_tensor_intent_overall
and not _effective_tensor_parallel(attempt_extra_args, tensor_parallel)
),
)
# Tensor parallelism is arch-gated in llama.cpp and crashes some loads
@ -3721,7 +3805,7 @@ async def generate_audio(
# Pick backend — both return (wav_bytes, sample_rate)
llama_backend = get_llama_cpp_backend()
if llama_backend.is_loaded and getattr(llama_backend, "_is_audio", False):
model_name = llama_backend.model_identifier
model_name = public_model_id(llama_backend.model_identifier)
gen = lambda: llama_backend.generate_audio_response(
text = text,
audio_type = llama_backend._audio_type,
@ -3739,7 +3823,7 @@ async def generate_audio(
model_info = backend.models.get(backend.active_model_name, {})
if not model_info.get("is_audio"):
raise HTTPException(status_code = 400, detail = "Active model is not an audio model.")
model_name = backend.active_model_name
model_name = public_model_id(backend.active_model_name)
gen = lambda: backend.generate_audio_response(
text = text,
temperature = payload.temperature,
@ -4504,6 +4588,14 @@ async def _proxy_to_external_provider(
except Exception as exc:
logger.error("external_provider.stream_error", error = str(exc))
api_monitor.fail(monitor_id, _friendly_error(exc))
# Surface the failure: a bare EOF (e.g. after a read timeout) is treated
# by the chat client as success, saving a partial answer with no error.
yield (
"data: "
+ json.dumps({"error": {"message": _friendly_error(exc), "type": "server_error"}})
+ "\n\n"
)
yield "data: [DONE]\n\n"
finally:
try:
await gen.aclose()
@ -4847,7 +4939,8 @@ async def openai_chat_completions(
return response
if using_gguf:
model_name = llama_backend.model_identifier or payload.model
# Echo a clean public id in the response, never the absolute .gguf path.
model_name = public_model_id(llama_backend.model_identifier) or payload.model
if getattr(llama_backend, "_is_audio", False):
if _wants_multiple_choices(payload):
_raise_unsupported_n("GGUF audio chat completions")
@ -4862,7 +4955,9 @@ async def openai_chat_completions(
status_code = 400,
detail = "No model loaded. Call POST /inference/load first.",
)
model_name = backend.active_model_name or payload.model
# Clean public id so the response never echoes a local path; the audio
# branch below receives this sanitized label too.
model_name = public_model_id(backend.active_model_name) or payload.model
if _wants_multiple_choices(payload):
_raise_unsupported_n("non-GGUF chat completions")
@ -6397,6 +6492,9 @@ async def serve_sandbox_file(
# OpenAI-Compatible Models Listing (/models → /v1/models)
# =====================================================================
# `owned_by` marker on every /v1/models entry (loaded and available alike).
_OWNED_BY = "unsloth-studio"
def _openai_model_objects() -> list[dict]:
"""The model objects GET /v1/models exposes (one per loaded local backend).
@ -6411,10 +6509,12 @@ def _openai_model_objects() -> list[dict]:
llama_backend = get_llama_cpp_backend()
if llama_backend.is_loaded:
entry = {
"id": llama_backend.model_identifier,
# Public id, never the absolute .gguf path (which leaks the host
# filesystem layout); see core.inference.model_ids.public_model_id.
"id": public_model_id(llama_backend.model_identifier),
"object": "model",
"created": _created,
"owned_by": "local",
"owned_by": _OWNED_BY,
}
_ctx = _positive_int_or_none(getattr(llama_backend, "context_length", None))
if _ctx is not None:
@ -6432,10 +6532,10 @@ def _openai_model_objects() -> list[dict]:
if backend.active_model_name:
model_info = backend.models.get(backend.active_model_name, {})
entry = {
"id": backend.active_model_name,
"id": public_model_id(backend.active_model_name),
"object": "model",
"created": _created,
"owned_by": "local",
"owned_by": _OWNED_BY,
}
_ctx = _positive_int_or_none(model_info.get("context_length"))
if _ctx is None:
@ -6453,15 +6553,86 @@ def _openai_model_objects() -> list[dict]:
return models
# Brief cache for the local-model filesystem scan so repeated /v1/models calls
# don't rescan the HF cache and models dirs on every request.
_CATALOG_CACHE: dict = {"at": 0.0, "models": []}
_CATALOG_TTL_S = 30.0
_CATALOG_LOCK = asyncio.Lock()
async def _cached_local_catalog() -> list:
"""Locally available models (models dir + HF caches + LM Studio + scan
folders), cached for a few seconds. Returns a list of LocalModelInfo.
The scan walks several directories and stats many files, so it runs in a
worker thread (asyncio.to_thread) -- calling it inline would block the event
loop and stall every concurrent request and in-flight inference stream. A
lock with a double-check collapses a burst of simultaneous /v1/models calls
into a single scan instead of one per request."""
# Validity is keyed on "at" (set only after a scan), not on list contents, so
# an empty/errored scan is still cached instead of rescanning on every poll.
now = time.monotonic()
if _CATALOG_CACHE["at"] and (now - _CATALOG_CACHE["at"]) <= _CATALOG_TTL_S:
return _CATALOG_CACHE["models"]
async with _CATALOG_LOCK:
now = time.monotonic()
if _CATALOG_CACHE["at"] and (now - _CATALOG_CACHE["at"]) <= _CATALOG_TTL_S:
return _CATALOG_CACHE["models"]
try:
from routes.models import collect_local_models
_CATALOG_CACHE["models"] = await asyncio.to_thread(
collect_local_models, Path("./models").resolve()
)
except Exception as exc:
logger.debug("model catalog scan failed: %s", exc)
_CATALOG_CACHE["models"] = []
# Stamp after the scan, not the pre-scan "now": a scan slower than the TTL
# would otherwise leave the cache already expired, so every waiter rescans.
_CATALOG_CACHE["at"] = time.monotonic()
return _CATALOG_CACHE["models"]
async def _openai_catalog_objects() -> list[dict]:
"""Every model the server knows about for ``GET /v1/models``: the loaded
model(s) plus locally available (downloaded/cached) models discovered by
scanning. Loaded entries keep their context fields and are marked
``loaded: true``. All ids are clean public ids (never absolute paths)."""
_created = int(time.time())
# Loaded models first (clean ids + context fields), marked loaded.
by_id: dict[str, dict] = {}
for entry in _openai_model_objects():
by_id[entry["id"]] = {**entry, "loaded": True}
# Locally available (downloaded/cached) models that are not already loaded.
for info in await _cached_local_catalog():
cid = getattr(info, "model_id", None) or public_model_id(getattr(info, "id", None))
if not cid or cid in by_id:
continue
obj = {
"id": cid,
"object": "model",
"created": _created,
"owned_by": _OWNED_BY,
"loaded": False,
}
display = getattr(info, "display_name", None)
if display:
obj["display_name"] = display
by_id[cid] = obj
return list(by_id.values())
@router.get("/models")
async def openai_list_models(current_subject: str = Depends(get_current_subject)):
"""
OpenAI-compatible model listing endpoint.
OpenAI-compatible model listing endpoint (``GET /v1/models``).
Returns the currently loaded model in the format expected by
OpenAI-compatible clients (``GET /v1/models``).
Lists every model available on this server -- the loaded model(s) plus
locally available (downloaded/cached) models -- not only what is resident in
memory. Each entry carries a clean public id and a ``loaded`` flag.
"""
return {"object": "list", "data": _openai_model_objects()}
return {"object": "list", "data": await _openai_catalog_objects()}
@router.get("/models/{model_id:path}")
@ -6469,13 +6640,37 @@ async def openai_retrieve_model(model_id: str, current_subject: str = Depends(ge
"""
OpenAI-compatible single-model retrieval endpoint (``GET /v1/models/{id}``).
Returns the bare model object when ``model_id`` matches a loaded local
model, or 404 model_not_found otherwise. Defined after the LIST route so
it does not shadow it; ``{model_id:path}`` keeps ids with slashes intact.
Returns the bare model object when ``model_id`` matches a known model
(loaded or locally available), or 404 model_not_found otherwise. Defined
after the LIST route so it does not shadow it; ``{model_id:path}`` keeps ids
with slashes intact.
"""
for model in _openai_model_objects():
from core.inference.model_ids import model_id_matches
# Loaded models resolve without a catalog scan (the common case); only build
# the full catalog -- which may hit the filesystem -- for unloaded ids.
for entry in _openai_model_objects():
if entry["id"] == model_id:
return {**entry, "loaded": True}
objects = await _openai_catalog_objects()
for model in objects:
if model["id"] == model_id:
return model
# Backward compatibility: a client may still send the legacy raw identifier
# (e.g. an absolute .gguf path cached from an older /v1/models). Resolve it to
# the clean object so it keeps working, without ever echoing the path back.
llama_backend = get_llama_cpp_backend()
backend = get_inference_backend()
for raw in (
llama_backend.model_identifier if llama_backend.is_loaded else None,
backend.active_model_name or None,
):
if raw and model_id_matches(model_id, raw):
clean = public_model_id(raw)
for model in objects:
if model["id"] == clean:
return model
raise HTTPException(
status_code = 404,
detail = openai_error_body(
@ -6544,7 +6739,10 @@ async def openai_completions(request: Request, current_subject: str = Depends(ge
# honor stream_options.include_usage per event, while keeping SSE
# framing and token bytes intact.
_include_usage = bool((body.get("stream_options") or {}).get("include_usage"))
client = httpx.AsyncClient(timeout = _llama_streaming_generation_timeout())
client = httpx.AsyncClient(
timeout = _llama_streaming_generation_timeout(),
trust_env = False,
)
resp = None
bytes_iter = None
disconnect_event = threading.Event()
@ -7412,6 +7610,15 @@ async def _responses_stream(
target_url = f"{llama_backend.base_url}/v1/chat/completions"
async def event_generator():
# Clean public id for every response envelope. Prefer the loaded model's
# id so the stream agrees with /v1/models, chat/completions and the
# non-streaming twin; fall back to a sanitized payload.model (a legacy
# raw .gguf path is stripped, never echoed back).
_clean_model = (
public_model_id(getattr(llama_backend, "model_identifier", None))
or public_model_id(payload.model)
or payload.model
)
full_text = ""
full_reasoning = ""
input_tokens = 0
@ -7573,7 +7780,7 @@ async def _responses_stream(
"object": "response",
"created_at": created_at,
"status": "failed",
"model": payload.model,
"model": _clean_model,
"output": _snapshot_output(),
"usage": {
"input_tokens": input_tokens,
@ -7597,7 +7804,7 @@ async def _responses_stream(
"object": "response",
"created_at": created_at,
"status": "in_progress",
"model": payload.model,
"model": _clean_model,
"output": [],
"usage": {"input_tokens": 0, "output_tokens": 0, "total_tokens": 0},
},
@ -7610,7 +7817,10 @@ async def _responses_stream(
# `async with`, explicit aclose of lines_iter BEFORE resp / client so
# the innermost httpcore byte stream is finalised in this task (not via
# the asyncgen GC in a sibling task).
client = httpx.AsyncClient(timeout = _llama_streaming_generation_timeout())
client = httpx.AsyncClient(
timeout = _llama_streaming_generation_timeout(),
trust_env = False,
)
resp = None
lines_iter = None
disconnect_watcher = None
@ -7637,7 +7847,7 @@ async def _responses_stream(
"object": "response",
"created_at": created_at,
"status": "failed",
"model": payload.model,
"model": _clean_model,
"output": [],
"error": {"code": 502, "message": _friendly_error(e)},
},
@ -7663,7 +7873,7 @@ async def _responses_stream(
"object": "response",
"created_at": created_at,
"status": "failed",
"model": payload.model,
"model": _clean_model,
"output": [],
"error": {
"code": resp.status_code,
@ -8012,7 +8222,7 @@ async def _responses_stream(
"object": "response",
"created_at": created_at,
"status": "completed",
"model": payload.model,
"model": _clean_model,
"output": _snapshot_output(),
"usage": {
"input_tokens": input_tokens,
@ -8284,7 +8494,13 @@ async def anthropic_messages(
),
)
model_name = getattr(llama_backend, "model_identifier", None) or payload.model
# Clean public id so /v1/messages never echoes the local .gguf path (and a
# legacy raw path sent as payload.model is sanitized rather than returned).
model_name = (
public_model_id(getattr(llama_backend, "model_identifier", None))
or public_model_id(payload.model)
or payload.model
)
message_id = f"msg_{uuid.uuid4().hex[:24]}"
# ── Translate Anthropic → OpenAI ──────────────────────────
@ -9114,6 +9330,7 @@ async def _anthropic_passthrough_stream(
client = httpx.AsyncClient(
timeout = _llama_streaming_generation_timeout(),
limits = httpx.Limits(max_keepalive_connections = 0),
trust_env = False,
)
resp = None
lines_iter = None
@ -9640,6 +9857,7 @@ async def _openai_passthrough_stream(
client = httpx.AsyncClient(
timeout = _llama_streaming_generation_timeout(),
limits = httpx.Limits(max_keepalive_connections = 0),
trust_env = False,
)
resp = None
_truncate_budget = (
@ -10053,6 +10271,29 @@ async def _openai_passthrough_non_streaming(
# ──────────────────────────────────────────────────────────────────────────
def _guard_diffusion_load_against_training() -> None:
"""Refuse loading an image model while a training run is active. Unlike chat,
a diffusion pipeline's VRAM can't be cheaply estimated before the load, so the
load is refused outright rather than fit-checked. No-op when training is
inactive or its state can't be read. Raises HTTP 409."""
from core.training import get_training_backend
try:
if not get_training_backend().is_training_active():
return
except Exception as e:
logger.warning("Could not check training state for image-load guard: %s", e)
return
raise HTTPException(
status_code = 409,
detail = (
"Can't load an image model while training is running: the diffusion "
"pipeline would compete with the training run for GPU memory. Training "
"was left untouched. Try again after training finishes."
),
)
@studio_router.post("/images/load", response_model = DiffusionStatusResponse)
async def load_diffusion_model(
request: DiffusionLoadRequest, current_subject: str = Depends(get_current_subject)
@ -10083,6 +10324,10 @@ async def load_diffusion_model(
family_override = request.family_override,
model_kind = kind,
)
# Refuse while training is running: a multi-GB diffusion pipeline would
# compete with the training subprocess for VRAM. The chat path does the
# same via _guard_chat_load_against_training; this is its image sibling.
_guard_diffusion_load_against_training()
# Pick the engine for this host (diffusers on GPU, native sd.cpp with no GPU),
# installing the sd-cli binary if needed -- all BEFORE evicting chat, so a
# native fallback never strands a half-loaded state. Non-GGUF kinds force diffusers.
@ -10196,6 +10441,9 @@ async def generate_diffusion_image(
# Position within the batch: shared timestamp, so the export
# filename needs this to stay unique.
"batch_index": index,
# The batch shares one seed, so reproducing image batch_index>0
# needs the original batch_size: persist it so restore can replay.
"batch_size": request.batch_size,
"model": result.get("repo_id"),
"created_at": created_at,
},

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

@ -12,6 +12,7 @@ import sys
import uuid
from pathlib import Path
from fastapi import APIRouter, Body, Depends, Header, HTTPException, Query
from pydantic import BaseModel
from typing import List, Optional
import structlog
from loggers import get_logger
@ -22,10 +23,27 @@ import re as _re
_VALID_REPO_ID = _re.compile(r"^[A-Za-z0-9._-]+/[A-Za-z0-9._-]+$")
class CachedModelRepo(BaseModel):
repo_id: str
size_bytes: int
last_modified: Optional[float] = None
class CachedModelsResponse(BaseModel):
cached: List[CachedModelRepo]
def _is_valid_repo_id(repo_id: str) -> bool:
return bool(_VALID_REPO_ID.fullmatch(repo_id))
def _normalize_hf_token(hf_token) -> Optional[str]:
if not isinstance(hf_token, str):
return None
token = hf_token.strip()
return token or None
def _safe_is_dir(path) -> bool:
"""``Path.is_dir()`` returning ``False`` instead of raising.
@ -74,6 +92,7 @@ if str(backend_path) not in sys.path:
sys.path.insert(0, str(backend_path))
from auth.authentication import get_current_subject
from hub.dependencies import get_hf_token
try:
from utils.models import (
@ -722,6 +741,94 @@ def _scan_ollama_dir(ollama_dir: Path, limit: Optional[int] = None) -> List[Loca
return found
def collect_local_models(models_root: Path) -> List[LocalModelInfo]:
"""Scan ``models_root``, the HF caches, LM Studio dirs, and user scan folders,
returning a deduplicated, hidden-filtered list of discovered local models.
Shared by ``GET /models/local`` (the model picker) and the OpenAI-compatible
catalog (``GET /v1/models``) so the UI and the API never drift. ``models_root``
must already be validated/trusted by the caller.
"""
from storage.studio_db import list_scan_folders
from utils.paths import (
hf_default_cache_dir,
legacy_hf_cache_dir,
lmstudio_model_dirs,
)
hf_cache_dir = _resolve_hf_cache_dir()
legacy_hf = legacy_hf_cache_dir()
hf_default = hf_default_cache_dir()
lm_dirs = lmstudio_model_dirs()
local_models = _scan_models_dir(models_root) + _scan_hf_cache(hf_cache_dir)
# Resolve once; an inaccessible aux cache must skip that scan, not 500.
hf_cache_real = _safe_resolve(hf_cache_dir)
legacy_real = _safe_resolve(legacy_hf)
default_real = _safe_resolve(hf_default)
# Scan legacy Unsloth HF cache for backward compatibility.
if _safe_is_dir(legacy_hf) and legacy_real != hf_cache_real:
local_models += _scan_hf_cache(legacy_hf)
# Scan HF system default cache (may differ under env overrides).
if _safe_is_dir(hf_default) and default_real != hf_cache_real and default_real != legacy_real:
local_models += _scan_hf_cache(hf_default)
# Scan LM Studio directories.
for lm_dir in lm_dirs:
local_models += _scan_lmstudio_dir(lm_dir)
# Scan user-added custom folders (per-folder cap).
_MAX_MODELS_PER_FOLDER = 200
try:
custom_folders = list_scan_folders()
except Exception as e:
logger.warning("Could not load custom scan folders: %s", e)
custom_folders = []
for folder in custom_folders:
folder_path = Path(folder["path"])
try:
# Filter Ollama .studio_links/ from generic scanners to
# avoid duplicates and leaking internal paths into the UI.
_generic = [
m
for m in (
_scan_models_dir(folder_path, limit = _MAX_MODELS_PER_FOLDER)
+ _scan_hf_cache(folder_path)
+ _scan_lmstudio_dir(folder_path)
)
if not any(p in (".studio_links", "ollama_links") for p in Path(m.path).parts)
]
custom_models = _generic
if len(custom_models) < _MAX_MODELS_PER_FOLDER:
custom_models += _scan_ollama_dir(
folder_path,
limit = _MAX_MODELS_PER_FOLDER - len(custom_models),
)
except OSError as e:
logger.warning("Skipping unreadable scan folder %s: %s", folder_path, e)
continue
local_models += [m.model_copy(update = {"source": "custom"}) for m in custom_models]
# Deduplicate, but always keep custom folder entries (keyed by
# (id, source)) so they show in the "Custom Folders" UI section
# even when the model is also in the HF cache.
deduped: dict[str, LocalModelInfo] = {}
for model in local_models:
key = f"{model.id}\x00custom" if model.source == "custom" else model.id
if key not in deduped:
deduped[key] = model
models = sorted(
deduped.values(),
key = lambda item: (item.updated_at or 0),
reverse = True,
)
return [m for m in models if not _is_hidden_model(m.id, m.path)]
@router.get("/local", response_model = LocalModelListResponse)
async def list_local_models(
models_dir: str = Query(
@ -770,78 +877,7 @@ async def list_local_models(
)
try:
local_models = _scan_models_dir(models_root) + _scan_hf_cache(hf_cache_dir)
# Resolve once; an inaccessible aux cache must skip that scan, not 500.
hf_cache_real = _safe_resolve(hf_cache_dir)
legacy_real = _safe_resolve(legacy_hf)
default_real = _safe_resolve(hf_default)
# Scan legacy Unsloth HF cache for backward compatibility.
if _safe_is_dir(legacy_hf) and legacy_real != hf_cache_real:
local_models += _scan_hf_cache(legacy_hf)
# Scan HF system default cache (may differ under env overrides).
if (
_safe_is_dir(hf_default)
and default_real != hf_cache_real
and default_real != legacy_real
):
local_models += _scan_hf_cache(hf_default)
# Scan LM Studio directories.
for lm_dir in lm_dirs:
local_models += _scan_lmstudio_dir(lm_dir)
# Scan user-added custom folders (per-folder cap).
from storage.studio_db import list_scan_folders
_MAX_MODELS_PER_FOLDER = 200
try:
custom_folders = list_scan_folders()
except Exception as e:
logger.warning("Could not load custom scan folders: %s", e)
custom_folders = []
for folder in custom_folders:
folder_path = Path(folder["path"])
try:
# Filter Ollama .studio_links/ from generic scanners to
# avoid duplicates and leaking internal paths into the UI.
_generic = [
m
for m in (
_scan_models_dir(folder_path, limit = _MAX_MODELS_PER_FOLDER)
+ _scan_hf_cache(folder_path)
+ _scan_lmstudio_dir(folder_path)
)
if not any(p in (".studio_links", "ollama_links") for p in Path(m.path).parts)
]
custom_models = _generic
if len(custom_models) < _MAX_MODELS_PER_FOLDER:
custom_models += _scan_ollama_dir(
folder_path,
limit = _MAX_MODELS_PER_FOLDER - len(custom_models),
)
except OSError as e:
logger.warning("Skipping unreadable scan folder %s: %s", folder_path, e)
continue
local_models += [m.model_copy(update = {"source": "custom"}) for m in custom_models]
# Deduplicate, but always keep custom folder entries (keyed by
# (id, source)) so they show in the "Custom Folders" UI section
# even when the model is also in the HF cache.
deduped: dict[str, LocalModelInfo] = {}
for model in local_models:
key = f"{model.id}\x00custom" if model.source == "custom" else model.id
if key not in deduped:
deduped[key] = model
models = sorted(
deduped.values(),
key = lambda item: (item.updated_at or 0),
reverse = True,
)
models = [m for m in models if not _is_hidden_model(m.id, m.path)]
models = collect_local_models(models_root)
# Tag each GGUF with its task so the Images picker can filter to diffusion.
models = [
m.model_copy(update = {"task": _local_model_task(m.path, m.model_format)}) for m in models
@ -2581,109 +2617,41 @@ async def get_gguf_variants(
..., description = "HuggingFace repo ID (e.g. 'unsloth/gemma-3-4b-it-GGUF')"
),
hf_token: Optional[str] = Query(None, description = "HuggingFace token for private repos"),
hf_token_header: Optional[str] = Depends(get_hf_token),
current_subject: str = Depends(get_current_subject),
):
"""List GGUF quantization variants for a HF repo or local directory.
Returns all variants with file sizes, vision support, and the
recommended default.
"""
"""List GGUF quantization variants for a HF repo or local directory."""
try:
from utils.models.model_config import is_local_path, list_local_gguf_variants
hf_token = _normalize_hf_token(hf_token_header) or _normalize_hf_token(hf_token)
from hub.services.models import gguf_variants as hub_gguf_variants
# Local directory path — scan filesystem.
if is_local_path(repo_id):
variants, has_vision = list_local_gguf_variants(repo_id)
filenames = [v.filename for v in variants]
best = _pick_best_gguf(filenames)
default_variant = _extract_quant_label(best) if best else None
return GgufVariantsResponse(
repo_id = repo_id,
variants = [
GgufVariantDetail(
filename = v.filename,
quant = v.quant,
size_bytes = v.size_bytes,
downloaded = True, # all local variants are downloaded
)
for v in variants
],
has_vision = has_vision,
default_variant = default_variant,
context_length = _read_native_context_length(repo_id, is_local = True),
)
# Remote HuggingFace repo — query HF API.
variants, has_vision = list_gguf_variants(repo_id, hf_token = hf_token)
filenames = [v.filename for v in variants]
best = _pick_best_gguf(filenames)
default_variant = _extract_quant_label(best) if best else None
# Per-snapshot so a split GGUF's shards must all sit in one snapshot;
# mmproj adapters are excluded so they can't inflate a quant's bytes.
cached_bytes_by_quant_per_snapshot: list[dict[str, int]] = []
try:
from huggingface_hub import constants as hf_constants
if not _is_valid_repo_id(repo_id):
raise ValueError(f"Invalid repo_id format: {repo_id}")
cache_dir = Path(hf_constants.HF_HUB_CACHE)
target = f"models--{repo_id.replace('/', '--')}".lower()
for entry in cache_dir.iterdir():
if entry.name.lower() == target:
snapshots = entry / "snapshots"
if snapshots.is_dir():
for snap in snapshots.iterdir():
by_quant: dict[str, int] = {}
for f in _iter_gguf_paths(snap):
if _is_mmproj_filename(f.name):
continue
try:
size = f.stat().st_size
except OSError:
continue # broken symlink / unreadable: skip
rel = f.relative_to(snap).as_posix()
q = _extract_quant_label(rel)
if _is_big_endian_gguf_path(rel, q):
continue
q = q.lower()
by_quant[q] = by_quant.get(q, 0) + size
if by_quant:
cached_bytes_by_quant_per_snapshot.append(by_quant)
break
except Exception:
pass
def _is_fully_downloaded(variant) -> bool:
if variant.size_bytes == 0:
return False
# Complete within one snapshot (tolerance for symlink size jitter).
quant = variant.quant.lower()
return any(
by_quant.get(quant, 0) >= variant.size_bytes * 0.99
for by_quant in cached_bytes_by_quant_per_snapshot
)
response = await hub_gguf_variants.get_gguf_variants_response(
repo_id,
hf_token = hf_token,
)
local = is_local_path(repo_id)
return GgufVariantsResponse(
repo_id = repo_id,
repo_id = response.repo_id,
variants = [
GgufVariantDetail(
filename = v.filename,
quant = v.quant,
size_bytes = v.size_bytes,
downloaded = _is_fully_downloaded(v),
download_size_bytes = int(
getattr(v, "download_size_bytes", v.size_bytes) or v.size_bytes
),
downloaded = bool(v.downloaded),
update_available = bool(getattr(v, "update_available", False)),
)
for v in variants
for v in response.variants
],
has_vision = has_vision,
default_variant = default_variant,
context_length = _read_native_context_length(repo_id, is_local = False),
has_vision = response.has_vision,
default_variant = response.default_variant,
context_length = _read_native_context_length(repo_id, is_local = local),
)
except HTTPException:
raise
except Exception as e:
logger.error(f"Error listing GGUF variants for '{repo_id}': {e}", exc_info = True)
raise HTTPException(
@ -3064,23 +3032,45 @@ def _repo_gguf_last_modified(repo_info) -> float:
# image GGUFs in its On Device list.
_DIFFUSION_GGUF_ARCHS = frozenset(
{
"flux",
"flux2",
"sd1",
"sd2",
"sd3",
"sdxl",
"stable_diffusion",
"lumina2",
"qwen_image",
# ONLY the families the diffusion backend can actually assemble (see
# diffusion_families._FAMILIES). Other on-device diffusion archs (SD1/2/3,
# SDXL, PixArt, Lumina2, AuraFlow, Wan, HunyuanVideo, ...) would pass this
# Images-picker filter and then fail validate_load with a 400, so they are
# deliberately excluded until the backend supports them.
"flux", # flux.1
"flux2", # flux.2-klein
"qwen_image", # qwen-image
"qwenimage",
"auraflow",
"pixart",
"hunyuan_video",
"wan",
"z_image", # z-image
"zimage",
}
)
# Known diffusion / image-video GGUF archs the backend can NOT assemble yet. These
# are the GGUF general.architecture values llama.cpp also has no architecture for,
# kept in sync with core.inference.llama_cpp.LlamaCppBackend._DIFFUSION_ARCHES
# (minus the loadable set above). Tagging them with a dedicated, non-loadable task
# keeps them OUT of the chat picker -- loading one as a chat model dies with
# "unknown model architecture" -- while also keeping them out of the Images picker
# (the task is not an IMAGE_GEN_TASK), where they would 400 in validate_load.
_UNSUPPORTED_DIFFUSION_GGUF_ARCHS = frozenset(
{
"sd1",
"sd3",
"sdxl",
"aura",
"hidream",
"cosmos",
"ltxv",
"hyvid",
"wan",
"lumina2",
}
)
# Task tag for the archs above; mirrored by the frontend NON_CHAT_TASKS gate.
_UNSUPPORTED_DIFFUSION_TASK = "image-diffusion-unsupported"
def _gguf_architecture(path: str) -> Optional[str]:
"""The GGUF ``general.architecture``, or None. Delegates to the shared,
@ -3094,12 +3084,21 @@ def _gguf_architecture(path: str) -> Optional[str]:
def _arch_to_task(arch: Optional[str]) -> Optional[str]:
if arch is None:
return None
return "text-to-image" if arch.lower() in _DIFFUSION_GGUF_ARCHS else "text-generation"
a = arch.lower()
if a in _DIFFUSION_GGUF_ARCHS:
return "text-to-image"
# A diffusion arch the backend can't assemble: hide it from chat (it would die
# in llama.cpp) without surfacing it in Images (it would 400 in validate_load).
if a in _UNSUPPORTED_DIFFUSION_GGUF_ARCHS:
return _UNSUPPORTED_DIFFUSION_TASK
return "text-generation"
def _repo_gguf_task(repo_info) -> Optional[str]:
"""HF pipeline task of a cached GGUF repo, from its architecture:
'text-to-image' for diffusion archs, else 'text-generation' (None if unreadable)."""
'text-to-image' for a loadable diffusion arch, the non-loadable diffusion tag
for a recognized-but-unsupported image arch, else 'text-generation' (None if
unreadable)."""
try:
for path in _iter_gguf_paths(Path(repo_info.repo_path)):
if _is_mmproj_filename(path.name):
@ -3210,10 +3209,14 @@ def _repo_is_diffusers(repo_info) -> bool:
return False
@router.get("/cached-models")
async def list_cached_models(current_subject: str = Depends(get_current_subject)):
@router.get("/cached-models", response_model = CachedModelsResponse)
async def list_cached_models(
current_subject: str = Depends(get_current_subject),
hf_token: Optional[str] = Depends(get_hf_token),
):
"""List non-GGUF model repos downloaded to HF cache, legacy Unsloth cache, and HF default cache."""
_WEIGHT_EXTENSIONS = (".safetensors", ".bin")
hf_token = _normalize_hf_token(hf_token)
try:
cache_scans = _all_hf_cache_scans()
@ -3234,20 +3237,16 @@ async def list_cached_models(current_subject: str = Depends(get_current_subject)
)
if total_size == 0:
continue
has_weights = any(
f.file_name.endswith(_WEIGHT_EXTENSIONS)
weight_files = [
f
for rev in repo_info.revisions
for f in rev.files
)
if not has_weights:
if f.file_name.endswith(_WEIGHT_EXTENSIONS)
]
if not weight_files:
continue
last_modified = max(
(
_blob_mtime(f)
for rev in repo_info.revisions
for f in rev.files
if f.file_name.endswith(_WEIGHT_EXTENSIONS)
),
(_blob_mtime(f) for f in weight_files),
default = 0.0,
)
key = repo_id.lower()
@ -3270,9 +3269,12 @@ async def list_cached_models(current_subject: str = Depends(get_current_subject)
repo_label = getattr(repo_info, "repo_id", "<unknown>")
logger.warning(f"Skipping cached model repo {repo_label}: {e}")
continue
# Newest download first; stable repo_id tie-break for equal/missing mtimes.
rows = list(seen_lower.values())
# Local-only list path: update checks are GGUF-only and happen lazily
# when a repo's variants are viewed.
cached = sorted(
seen_lower.values(),
rows,
key = lambda c: (-(c.get("last_modified") or 0.0), c["repo_id"].lower()),
)
return {"cached": cached}
@ -3326,6 +3328,24 @@ async def delete_cached_model(
except Exception:
pass
# Also refuse if the diffusion (Images) backend has this repo loaded; its
# delete guard is otherwise chat-only, so its GGUF could be removed from
# under a live pipeline. Repo-level match, like the chat guards above.
try:
from core.inference.diffusion import get_diffusion_backend
diffusion_status = get_diffusion_backend().status()
if diffusion_status.get("loaded") and diffusion_status.get("repo_id"):
loaded_id = str(diffusion_status["repo_id"]).lower()
if loaded_id == repo_id.lower() or loaded_id.startswith(repo_id.lower()):
raise HTTPException(
status_code = 400,
detail = "Unload the model before deleting",
)
except HTTPException:
raise
except Exception:
pass
try:
cache_scans = _all_hf_cache_scans()

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

@ -19,17 +19,68 @@ from auth.storage import DEFAULT_ADMIN_USERNAME
from models.inference import ChatCompletionRequest, LoadRequest
from routes.inference import load_model, openai_chat_completions
from state.tool_policy import tools_force_disabled
from utils.client_ip import client_ip
from utils.models.checkpoints import list_preview_targets, resolve_preview_checkpoint
from utils.preview_rate_limit import check_rate_limit
from utils.preview_sharing_settings import get_preview_sharing_enabled
from utils.preview_token import sign_preview_ref, verify_preview_ref
logger = get_logger(__name__)
router = APIRouter()
# Public (no key); resolve_preview_checkpoint pins `run` under outputs_root.
# One model loads at a time, so serialize load+generate across previews.
# A shared preview link is a public bearer capability; cap per-request generation
# so a single call can't tie up the (serialized) preview GPU indefinitely.
_PREVIEW_MAX_OUTPUT_TOKENS = 1024
# Capability-gated (signed ref required); resolve_preview_checkpoint pins `run`
# under outputs_root. One model loads at a time, so serialize load+generate.
_preview_lock = asyncio.Lock()
def _extract_token(request: Request) -> str | None:
"""Capability token from the ``?k=`` query (browser link + preview page) or an
``Authorization: Bearer`` header (OpenAI-compatible clients using it as api_key)."""
token = request.query_params.get("k")
if token:
return token
header = request.headers.get("authorization", "")
if header[:7].lower() == "bearer ":
return header[7:].strip() or None
return None
def _verify_or_404(run: str, checkpoint: str | None, request: Request) -> None:
"""Require a valid preview capability BEFORE any checkpoint resolve / model load.
Missing or invalid tokens get a generic 404 -- identical to a non-existent ref --
so the public surface never confirms whether a run/checkpoint exists. When an
admin has switched public sharing off, every public request 404s regardless of
token.
Verify the (cheap, no-I/O) capability first: an unauthenticated caller with a
bad/missing token is rejected without the kill-switch DB read, so spamming
``/p/...`` can't be used as an unbounded settings-DB sink, and the response is
identical whether or not sharing is enabled (no on/off oracle).
"""
ref = run if not checkpoint else f"{run}/{checkpoint}"
if not verify_preview_ref(ref, _extract_token(request)):
raise HTTPException(status_code = 404, detail = "Not found")
if not get_preview_sharing_enabled():
raise HTTPException(status_code = 404, detail = "Not found")
def _enforce_rate_limit(request: Request) -> None:
"""Throttle the GPU-backed preview chat per client IP (429 on exceed)."""
retry_after = check_rate_limit(client_ip(request))
if retry_after:
raise HTTPException(
status_code = 429,
detail = "Too many preview requests. Please slow down.",
headers = {"Retry-After": str(retry_after)},
)
def _resolve_or_4xx(run: str, checkpoint: str | None):
try:
return resolve_preview_checkpoint(run, checkpoint)
@ -49,6 +100,21 @@ def _sanitize_preview_payload(
# Normalize use_adapter (never trust the caller): pin True for LoRA, None for
# merged. _apply_adapter_state mutates the shared model without restoring, so an
# unpinned `false` would persist to later visitors who omit the field.
#
# Cap generation cost on this public, GPU-backed surface. Derive one effective
# limit (mirroring _effective_max_tokens: max_completion_tokens wins, else the
# legacy max_tokens) and pin BOTH fields to it, so a caller's lower limit is
# honored and neither field can exceed the ceiling.
requested = (
payload.max_completion_tokens
if payload.max_completion_tokens is not None
else payload.max_tokens
)
capped_max_tokens = (
min(requested, _PREVIEW_MAX_OUTPUT_TOKENS)
if requested is not None
else _PREVIEW_MAX_OUTPUT_TOKENS
)
return payload.model_copy(
update = {
"tools": None,
@ -67,6 +133,9 @@ def _sanitize_preview_payload(
"encrypted_api_key": None,
"provider_base_url": None,
"use_adapter": True if is_lora else None,
"max_tokens": capped_max_tokens,
"max_completion_tokens": capped_max_tokens,
"n": 1,
}
)
@ -105,15 +174,30 @@ async def _serve_chat(
@router.get("")
async def list_previews(request: Request, current_subject: str = Depends(get_current_subject)):
base = str(request.base_url)
sharing_on = get_preview_sharing_enabled()
previews = []
for target in list_preview_targets():
ref = quote(target["ref"], safe = "/")
previews.append({**target, "url": f"{base}p/{ref}/v1"})
return {"object": "list", "data": previews}
# Mint the capability for the authenticated owner: ``key`` for OpenAI
# clients (Bearer / api_key), ``share_url`` for the browser link. When
# public sharing is off, every public /p request 404s, so don't hand out
# dead credentials -- omit the capability and signal the disabled state.
token = sign_preview_ref(target["ref"]) if sharing_on else None
previews.append(
{
**target,
"url": f"{base}p/{ref}/v1",
"key": token,
"share_url": f"{base}p/{ref}?k={token}" if token else None,
}
)
return {"object": "list", "data": previews, "sharing_enabled": sharing_on}
@router.post("/{run}/v1/chat/completions")
async def preview_chat_latest(run: str, payload: ChatCompletionRequest, request: Request):
_verify_or_404(run, None, request)
_enforce_rate_limit(request)
return await _serve_chat(run, None, payload, request)
@ -121,6 +205,8 @@ async def preview_chat_latest(run: str, payload: ChatCompletionRequest, request:
async def preview_chat_checkpoint(
run: str, checkpoint: str, payload: ChatCompletionRequest, request: Request
):
_verify_or_404(run, checkpoint, request)
_enforce_rate_limit(request)
return await _serve_chat(run, checkpoint, payload, request)
@ -140,13 +226,17 @@ def _models_response(run: str, checkpoint: str | None):
}
# The models/page GET routes only stat the checkpoint dir (no GPU), so they are
# token-gated but not rate-limited; only the GPU-backed chat path is throttled.
@router.get("/{run}/v1/models")
async def preview_models_latest(run: str):
async def preview_models_latest(run: str, request: Request):
_verify_or_404(run, None, request)
return _models_response(run, None)
@router.get("/{run}/{checkpoint}/v1/models")
async def preview_models_checkpoint(run: str, checkpoint: str):
async def preview_models_checkpoint(run: str, checkpoint: str, request: Request):
_verify_or_404(run, checkpoint, request)
return _models_response(run, checkpoint)
@ -183,14 +273,24 @@ def _preview_page(run: str, checkpoint: str | None) -> HTMLResponse:
_resolve_or_4xx(run, checkpoint)
title = run if not checkpoint else f"{run}/{checkpoint}"
page = _PREVIEW_PAGE_HTML.replace("__TITLE__", html.escape(title))
return HTMLResponse(page, headers = {"Content-Security-Policy": _PREVIEW_PAGE_CSP})
# no-referrer: the capability token rides in the query string, so keep it out
# of the Referer header on any outbound navigation.
return HTMLResponse(
page,
headers = {
"Content-Security-Policy": _PREVIEW_PAGE_CSP,
"Referrer-Policy": "no-referrer",
},
)
@router.get("/{run}", response_class = HTMLResponse)
async def preview_page_latest(run: str):
async def preview_page_latest(run: str, request: Request):
_verify_or_404(run, None, request)
return _preview_page(run, None)
@router.get("/{run}/{checkpoint}", response_class = HTMLResponse)
async def preview_page_checkpoint(run: str, checkpoint: str):
async def preview_page_checkpoint(run: str, checkpoint: str, request: Request):
_verify_or_404(run, checkpoint, request)
return _preview_page(run, checkpoint)

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

@ -19,7 +19,7 @@ import secrets
import time
import uuid
from fastapi import APIRouter, Depends, File, HTTPException, Query, UploadFile
from fastapi import APIRouter, Depends, File, Form, HTTPException, Query, UploadFile
from fastapi.responses import FileResponse, StreamingResponse
from pydantic import BaseModel, Field
@ -62,13 +62,24 @@ def _save_upload(file: UploadFile) -> tuple[str, str]:
uploads = ensure_dir(rag_uploads_root())
stored_path = str(uploads / f"{uuid.uuid4().hex}{ext}")
size = 0
cap = config.MAX_UPLOAD_BYTES
too_big = False
with open(stored_path, "wb") as out:
while True:
block = file.file.read(1 << 20)
if not block:
break
size += len(block)
if cap and size > cap:
too_big = True
break
out.write(block)
if too_big:
os.remove(stored_path)
raise HTTPException(
status_code = 413,
detail = f"File exceeds the {cap // (1024 * 1024)} MB upload limit.",
)
if size == 0:
os.remove(stored_path)
raise HTTPException(status_code = 400, detail = "Uploaded file is empty.")
@ -207,6 +218,8 @@ def delete_knowledge_base(kb_id: str, subject: str = Depends(get_current_subject
async def upload_kb_document(
kb_id: str,
file: UploadFile = File(...),
ocr: bool | None = Form(None),
caption: bool | None = Form(None),
subject: str = Depends(get_current_subject),
) -> dict:
_require_rag()
@ -218,7 +231,7 @@ async def upload_kb_document(
conn.close()
stored_path, filename = _save_upload(file)
document_id, job_id = ingestion.start_ingestion(
store.kb_scope(kb_id), kb_id, None, filename, stored_path
store.kb_scope(kb_id), kb_id, None, filename, stored_path, ocr = ocr, caption = caption
)
return {"documentId": document_id, "jobId": job_id, "filename": filename}
@ -238,12 +251,20 @@ def list_kb_documents(kb_id: str, subject: str = Depends(get_current_subject)) -
async def upload_thread_document(
thread_id: str,
file: UploadFile = File(...),
ocr: bool | None = Form(None),
caption: bool | None = Form(None),
subject: str = Depends(get_current_subject),
) -> dict:
_require_rag()
stored_path, filename = _save_upload(file)
document_id, job_id = ingestion.start_ingestion(
store.thread_scope(thread_id), None, thread_id, filename, stored_path
store.thread_scope(thread_id),
None,
thread_id,
filename,
stored_path,
ocr = ocr,
caption = caption,
)
return {"documentId": document_id, "jobId": job_id, "filename": filename}
@ -263,6 +284,8 @@ def list_thread_documents(thread_id: str, subject: str = Depends(get_current_sub
async def upload_project_document(
project_id: str,
file: UploadFile = File(...),
ocr: bool | None = Form(None),
caption: bool | None = Form(None),
subject: str = Depends(get_current_subject),
) -> dict:
_require_rag()
@ -278,6 +301,8 @@ async def upload_project_document(
filename,
stored_path,
project_id = project_id,
ocr = ocr,
caption = caption,
)
return {"documentId": document_id, "jobId": job_id, "filename": filename}
@ -321,6 +346,7 @@ def job_status(job_id: str, subject: str = Depends(get_current_subject)) -> dict
"stage": row.get("stage"),
"progress": row.get("progress") or 0.0,
"error": row.get("error"),
"numChunks": row.get("num_chunks") or 0,
}

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

@ -8,6 +8,7 @@ from fastapi import APIRouter, Depends
from pydantic import BaseModel, ConfigDict, Field, field_validator
from auth.authentication import get_current_subject
from auth.storage import rotate_preview_link_secret
from loggers import get_logger
from utils.utils import safe_error_detail, log_and_http_error
from utils.personalization_settings import (
@ -31,6 +32,11 @@ from utils.helper_precache_settings import (
helper_model_disabled_by_env,
set_helper_precache_enabled,
)
from utils.preview_sharing_settings import (
DEFAULT_PREVIEW_SHARING_ENABLED,
get_preview_sharing_enabled,
set_preview_sharing_enabled,
)
router = APIRouter()
@ -122,6 +128,55 @@ def update_helper_precache(
return _helper_precache_response(enabled)
class PreviewLinkRotateResponse(BaseModel):
rotated: bool = True
@router.post("/preview-links/rotate", response_model = PreviewLinkRotateResponse)
def rotate_preview_links(
current_subject: str = Depends(get_current_subject),
) -> PreviewLinkRotateResponse:
"""Rotate the preview-link signing secret, revoking every previously shared `/p` link."""
rotate_preview_link_secret()
logger.info("settings.preview_links_rotated subject=%s", current_subject)
return PreviewLinkRotateResponse(rotated = True)
class PreviewSharingPayload(BaseModel):
enabled: bool
class PreviewSharingResponse(BaseModel):
enabled: bool
default_enabled: bool = DEFAULT_PREVIEW_SHARING_ENABLED
@router.get("/preview-sharing", response_model = PreviewSharingResponse)
def get_preview_sharing(
current_subject: str = Depends(get_current_subject),
) -> PreviewSharingResponse:
return PreviewSharingResponse(enabled = get_preview_sharing_enabled())
@router.put("/preview-sharing", response_model = PreviewSharingResponse)
def update_preview_sharing(
payload: PreviewSharingPayload, current_subject: str = Depends(get_current_subject)
) -> PreviewSharingResponse:
"""Enable/disable the public `/p` preview surface. When off, links 404 even with a token."""
try:
enabled = set_preview_sharing_enabled(payload.enabled)
except ValueError as exc:
raise log_and_http_error(
exc,
400,
safe_error_detail(exc, fallback = "Invalid preview sharing setting."),
event = "settings.update_preview_sharing_failed",
log = logger,
) from exc
logger.info("settings.preview_sharing_updated subject=%s enabled=%s", current_subject, enabled)
return PreviewSharingResponse(enabled = enabled)
def _is_bundled_avatar_url(value: str) -> bool:
parsed = urlsplit(value)
if parsed.scheme or parsed.netloc:

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

@ -68,6 +68,11 @@ class TrainingStopRequest(PydanticBaseModel):
router = APIRouter()
logger = get_logger(__name__)
# Consecutive 1s polls without a step update that count as a stall. Applied only
# once stepping: the pre-first-step phase (model load + tokenization) can take far
# longer, and timing out there made a healthy long-prep run look frozen.
_PROGRESS_STALL_TIMEOUT_POLLS = 1800 # ~30 min at 1 poll/sec
def _validate_local_dataset_paths(paths: list[str], label: str = "Local dataset") -> list[str]:
"""Resolve and validate a list of local dataset paths. Returns validated absolute paths."""
@ -250,6 +255,7 @@ async def start_training(
# Convert request to backend kwargs.
training_kwargs = {
"model_name": request.model_name,
"project_name": request.project_name,
"training_type": request.training_type,
"hf_token": request.hf_token or "",
"load_in_4bit": request.load_in_4bit,
@ -357,6 +363,27 @@ async def start_training(
except Exception as e:
logger.warning("Could not shut down export subprocess: %s", e)
try:
# A resident or in-flight diffusion (Images) pipeline also holds
# GPU memory the training run needs, and it can't be cheaply sized,
# so tear it down unconditionally like the export subprocess above
# (the chat block below fit-checks; diffusion can't). unload() is a
# no-op when nothing is loaded and also preempts an in-flight load;
# release the arbiter so it doesn't think the gone pipeline owns
# the GPU. Must precede the chat block, which early-returns.
from core.inference import gpu_arbiter
from core.inference.diffusion import get_diffusion_backend
diffusion = get_diffusion_backend()
if diffusion.is_loaded:
logger.info(
"Unloading diffusion (Images) model to free GPU memory for training"
)
diffusion.unload()
gpu_arbiter.release(gpu_arbiter.DIFFUSION)
except Exception as e:
logger.warning("Could not unload diffusion model for training: %s", e)
try:
from routes.training_vram import (
can_keep_chat_during_training,
@ -833,9 +860,20 @@ async def stream_training_progress(
# ── Live polling loop ────────────────────────────────────
last_step = resume_from_step if resume_from_step is not None else -1
no_update_count = 0
max_no_updates = 1800 # Timeout after 30 min (large models need compile time)
# The stall timeout applies only once the run is stepping (pre-step prep
# may legitimately emit no step for a long time). On reconnect to an
# already-stepping run, seed from the resume point / history, else a worker
# that hangs after step N never times out for a client that reconnects past it.
seen_live_step = (resume_from_step is not None and resume_from_step > 0) or bool(
backend.step_history
)
while backend.is_training_active():
# Client gone: end the generator without falling through to the final
# "complete" frame, which a buffered/proxy consumer could otherwise read
# as a finished run while training is still active.
if await request.is_disconnected():
return
try:
tp_inner = getattr(getattr(backend, "trainer", None), "training_progress", None)
live_step = (getattr(tp_inner, "step", 0) or 0) if tp_inner else 0
@ -871,6 +909,7 @@ async def stream_training_progress(
)
last_step = current_step
no_update_count = 0
seen_live_step = True
else:
no_update_count += 1
# Heartbeat every 10 seconds.
@ -913,8 +952,9 @@ async def stream_training_progress(
event_id = 0,
)
# Timeout check
if no_update_count > max_no_updates:
# Fires only once stepping: a long pre-first-step prep phase is not
# a stall, and ending the stream there made a healthy run look frozen.
if seen_live_step and no_update_count > _PROGRESS_STALL_TIMEOUT_POLLS:
logger.warning("Progress stream timeout - no updates received")
tp_timeout = getattr(
getattr(backend, "trainer", None), "training_progress", None

29
studio/backend/routes/training_history.py
View file

@ -6,6 +6,7 @@ Training history API routes — browse, view, and delete past training runs.
"""
import json
from typing import Optional
from fastapi import APIRouter, Depends, HTTPException, Query
from loggers import get_logger
@ -28,12 +29,30 @@ from storage.studio_db import (
update_run_display_name,
)
from utils.models.checkpoints import has_preview_model, preview_ref
from utils.preview_sharing_settings import get_preview_sharing_enabled
from utils.preview_token import sign_preview_ref
logger = get_logger(__name__)
router = APIRouter()
def _preview_fields(output_dir: Optional[str], sharing_on: bool) -> dict:
"""Previewability + the signed `/p` share ref for a run's output dir.
The signature is what makes the share link a capability: these routes are
authenticated, so only the run's owner ever receives it. When public sharing
is switched off, omit the signature so the UI hides the copy-link affordance
(and the link would 404 anyway). ``sharing_on`` is resolved once per request.
"""
ref = preview_ref(output_dir)
return {
"has_preview_model": has_preview_model(output_dir),
"preview_ref": ref,
"preview_sig": sign_preview_ref(ref) if (ref and sharing_on) else None,
}
@router.get("/runs", response_model = TrainingRunListResponse)
async def list_training_runs(
limit: int = Query(50, ge = 1, le = 200),
@ -42,14 +61,14 @@ async def list_training_runs(
):
"""List training runs, newest first."""
result = list_runs(limit = limit, offset = offset)
sharing_on = get_preview_sharing_enabled()
return TrainingRunListResponse(
runs = [
TrainingRunSummary(
**{
**r,
"can_resume": can_resume_run(r),
"has_preview_model": has_preview_model(r.get("output_dir")),
"preview_ref": preview_ref(r.get("output_dir")),
**_preview_fields(r.get("output_dir"), sharing_on),
}
)
for r in result["runs"]
@ -78,8 +97,7 @@ async def get_training_run_detail(run_id: str, current_subject: str = Depends(ge
**{
**{k: v for k, v in run.items() if k != "config_json"},
"can_resume": can_resume_run(run),
"has_preview_model": has_preview_model(run.get("output_dir")),
"preview_ref": preview_ref(run.get("output_dir")),
**_preview_fields(run.get("output_dir"), get_preview_sharing_enabled()),
}
),
config = config,
@ -111,8 +129,7 @@ async def update_training_run(
**{
**{k: v for k, v in refreshed.items() if k != "config_json"},
"can_resume": can_resume_run(refreshed),
"has_preview_model": has_preview_model(refreshed.get("output_dir")),
"preview_ref": preview_ref(refreshed.get("output_dir")),
**_preview_fields(refreshed.get("output_dir"), get_preview_sharing_enabled()),
}
)

227
studio/backend/run.py
View file

@ -253,12 +253,13 @@ def _verify_global_reachability(display_host: str, port: int) -> None:
local_url_c = "\033[38;5;108;1m" if use_color else "" # matches banner's URL color
reset = "\033[0m" if use_color else ""
url = f"http://{display_host}:{port}"
url = f"http://{_url_host(display_host)}:{port}"
# Private/loopback/link-local addresses aren't globally routable.
try:
addr = ipaddress.ip_address(display_host)
if addr.is_loopback or addr.is_private or addr.is_link_local:
_public_reachable = False
print(
f"{dim} Note: {display_host} is a private/LAN address -- "
f"reachable on this network only, not from the public internet."
@ -380,6 +381,20 @@ def _verify_global_reachability(display_host: str, port: int) -> None:
pass
def _display_host_for_bind(host: str) -> str:
return _resolve_external_ip() if host in ("0.0.0.0", "::") else host
def _loopback_bind_host_for(host: str) -> str:
return "::1" if host == "::" else "127.0.0.1"
def _url_host(host: str) -> str:
return (
f"[{host}]" if ":" in host and not (host.startswith("[") and host.endswith("]")) else host
)
def _tool_policy_notice(host: str, secure: bool, enable_tools: "Optional[bool]") -> str:
"""One-line tool-policy summary for the plain-server startup banner, so a
network-reachable launch is never silent about code execution."""
@ -416,7 +431,7 @@ def _emit_secure_startup_output(port: int, enable_tools: "Optional[bool]" = None
print("")
print("🦥 Unsloth Studio is running (secure)")
print("─" * 52)
_print_cloudflare_line()
_print_cloudflare_line(secure = True)
print(f" On this machine only: http://127.0.0.1:{port}/")
print("─" * 52)
_emit_tool_policy_notice("127.0.0.1", True, enable_tools)
@ -447,30 +462,108 @@ def _emit_startup_output(
_print_localhost_ipv6_mismatch_warning(localhost_mismatch_url, port)
elif wildcard_bind:
_verify_global_reachability(display_host, port)
_print_cloudflare_line()
_print_cloudflare_line(loopback_host = _loopback_bind_host_for(host))
_emit_tool_policy_notice(host, False, enable_tools)
print_studio_stop_hint()
def _print_cloudflare_line() -> None:
"""Print the Cloudflare quick-tunnel URL for 0.0.0.0 binds, if one is up.
Reads the module-level URL set by ``run_server``. Prints nothing when the
tunnel is disabled or failed -- failures are silently ignored. When the public
reachability probe just failed (``_public_reachable is False``) but the tunnel
is up, reword to point the user at the Cloudflare link as the way in.
"""
if not _cloudflare_url:
return
def _print_cloudflare_line(secure: bool = False, loopback_host: str = "127.0.0.1") -> None:
"""Print Cloudflare tunnel state for startup banners."""
from startup_banner import stdout_supports_color
accent = "\033[38;5;150;1m"
warn = "\033[38;5;215;1m"
reset = "\033[0m"
if _public_reachable is False:
line = f" Use the secure link access via Cloudflare instead: {_cloudflare_url}"
else:
line = f" Secure link access via Cloudflare: {_cloudflare_url}"
print(f"{accent}{line}{reset}" if stdout_supports_color() else line)
color = stdout_supports_color()
def _emit(text: str, style: str = "") -> None:
print(f"{style}{text}{reset}" if (color and style) else text)
if _cloudflare_url:
if _public_reachable is False:
_emit(f" Use the secure link access via Cloudflare instead: {_cloudflare_url}", accent)
else:
_emit(f" Secure link access via Cloudflare: {_cloudflare_url}", accent)
if not secure:
if _public_reachable is True:
_emit(
" Cloudflare tunnel: ON. This Cloudflare URL is PUBLIC, and the "
"raw port is also publicly reachable. --no-cloudflare disables "
f"only the Cloudflare URL; bind {loopback_host} or close firewall "
"access to keep Studio private.",
warn,
)
else:
_emit(
" Cloudflare tunnel: ON. This is a PUBLIC internet URL: anyone "
"who has it can reach this Studio. Relaunch with --no-cloudflare "
f"to disable the Cloudflare URL; bind {loopback_host} or close "
"firewall access to keep Studio private.",
warn,
)
return
if _cloudflare_requested:
if _public_reachable is True:
_emit(
" Cloudflare tunnel: requested but failed to start. The raw port is "
"still reachable from the public internet (see the reachability check "
"above): anyone who can reach it can access this Studio.",
warn,
)
elif _public_reachable is False:
_emit(
" Cloudflare tunnel: requested but failed to start. Studio is reachable "
"on your local network only (no public link).",
warn,
)
else:
_emit(
" Cloudflare tunnel: requested but failed to start. There is no "
"Cloudflare public link. Raw port reachability was not verified; "
f"bind {loopback_host} or close firewall access to keep Studio private.",
warn,
)
elif _cloudflare_flag:
if _public_reachable is True:
_emit(
" Cloudflare tunnel: OFF for this mode. The raw port is still "
"reachable from the public internet (see the reachability check above): "
"anyone who can reach it can access this Studio.",
warn,
)
elif _public_reachable is False:
_emit(
" Cloudflare tunnel: OFF for this mode. Studio is reachable on your "
"local network only (no public link)."
)
else:
_emit(
" Cloudflare tunnel: OFF for this mode. There is no Cloudflare public "
"link. Raw port reachability was not verified; "
f"bind {loopback_host} or close firewall access to keep Studio private.",
warn,
)
elif not _cloudflare_flag:
if _public_reachable is True:
_emit(
" Cloudflare tunnel: OFF (--no-cloudflare). The raw port is still "
"reachable from the public internet (see the reachability check above): "
"--no-cloudflare disables only the Cloudflare link, not the public bind.",
warn,
)
elif _public_reachable is False:
_emit(
" Cloudflare tunnel: OFF (--no-cloudflare). Studio is reachable on your "
"local network only. Omit --no-cloudflare to expose a public "
"Cloudflare HTTPS link."
)
else:
_emit(
" Cloudflare tunnel: OFF (--no-cloudflare). There is no Cloudflare "
"public link. Raw port reachability was not verified; "
f"bind {loopback_host} or close firewall access to keep Studio private.",
warn,
)
def _get_pid_on_port(port: int) -> "tuple[int, str] | None":
@ -697,7 +790,7 @@ _server_thread = None
# Shutdown event -- wakes the main loop on signal.
_shutdown_event = None
# trycloudflare.com URL for 0.0.0.0 binds (set by run_server, read by the banner);
# trycloudflare.com URL for wildcard binds (set by run_server, read by the banner);
# None when there is no tunnel (loopback, disabled, or a silently-ignored failure).
_cloudflare_url = None
@ -707,6 +800,9 @@ _cloudflare_url = None
# not decide (timeout, blocked, private address).
_public_reachable = None
_cloudflare_requested = False
_cloudflare_flag = True
_DEFAULT_FRONTEND_PATH = Path(__file__).resolve().parent.parent / "frontend" / "dist"
@ -880,12 +976,12 @@ def _cloudflare_tunnel_should_start(
) -> bool:
"""Whether to start the Cloudflare tunnel. --secure exposes only the tunnel
(loopback bind), so it tunnels even api-only (headless secure API serving);
otherwise tunnel only a 0.0.0.0 bind, never api-only (Tauri) or Colab."""
otherwise tunnel wildcard binds, never api-only (Tauri) or Colab."""
if is_colab or not cloudflare:
return False
if secure:
return True
return host == "0.0.0.0" and not api_only
return host in ("0.0.0.0", "::") and not api_only
def _apply_cli_tool_policy(enable_tools: "Optional[bool]") -> None:
@ -933,6 +1029,9 @@ def run_server(
"""
global _server, _server_thread, _shutdown_event
boot_started = time.perf_counter()
logger.info("run_server startup begin api_only=%s host=%s port=%s", api_only, host, port)
# Reap every child if the parent dies abnormally (terminal close, Task
# Manager kill, SIGKILL); must run before any child can spawn.
from utils.process_lifetime import initialize_parent_lifetime
@ -984,7 +1083,14 @@ def run_server(
from threading import Thread, Event
import uvicorn
import_started = time.perf_counter()
from main import app, setup_frontend, _IS_COLAB
logger.info(
"Imported FastAPI app in %.1fms",
(time.perf_counter() - import_started) * 1000,
)
from utils.paths import ensure_studio_directories
# Allow local stdio MCP servers on a loopback bind (the user's own machine),
@ -997,6 +1103,11 @@ def run_server(
# Create all standard directories on startup.
ensure_studio_directories()
logger.info(
"Ensured Studio directories in %.1fms",
(time.perf_counter() - boot_started) * 1000,
)
# Auto-find a free port if the requested one is in use.
if not _is_port_free(host, port):
original_port = port
@ -1057,9 +1168,14 @@ def run_server(
)
# Resolve once; shared by the log rewrite and banner.
display_host = _resolve_external_ip() if host == "0.0.0.0" else host
display_host = _display_host_for_bind(host)
_install_uvicorn_startup_log_rewrite(host, display_host)
logger.info(
"run_server pre-uvicorn setup completed in %.1fms",
(time.perf_counter() - boot_started) * 1000,
)
ready_event = Event()
startup_failed = Event()
startup_errors = []
@ -1068,6 +1184,10 @@ def run_server(
async def startup(self, *args, **kwargs):
await super().startup(*args, **kwargs)
if getattr(self, "started", False) and not self.should_exit:
logger.info(
"Uvicorn startup hook completed in %.1fms",
(time.perf_counter() - boot_started) * 1000,
)
ready_event.set()
# server_header=False suppresses uvicorn's "Server: uvicorn"; SecurityHeadersMiddleware sets its own.
@ -1093,13 +1213,10 @@ def run_server(
# backend, not whatever a proxy/tunnel exposed. For ephemeral binds (port==0)
# leave it unset so handlers fall back to the request scope / base_url.
app.state.server_port = port if port and port > 0 else None
# Direct (non-tunnel) base for the API panel; resolve 0.0.0.0 to the LAN IP.
# Direct (non-tunnel) base for the API panel; resolve wildcard binds to the LAN IP.
if port and port > 0:
_direct_host = _resolve_external_ip() if host in ("0.0.0.0", "::") else host
# Bracket IPv6 literals so the URL is valid (http://[2405:...]:port).
if ":" in _direct_host and not _direct_host.startswith("["):
_direct_host = f"[{_direct_host}]"
app.state.server_url = f"http://{_direct_host}:{port}"
_direct_host = _display_host_for_bind(host)
app.state.server_url = f"http://{_url_host(_direct_host)}:{port}"
else:
app.state.server_url = None
app.state.secure = secure
@ -1150,6 +1267,11 @@ def run_server(
_shutdown_event.set()
raise
logger.info(
"run_server uvicorn ready after %.1fms",
(time.perf_counter() - boot_started) * 1000,
)
_write_pid_file()
import atexit
@ -1163,11 +1285,12 @@ def run_server(
if api_only and emit_tauri_port:
print(f"TAURI_PORT={port}", flush = True)
# Free trycloudflare.com tunnel for 0.0.0.0 binds (the raw ip:port is often
# Free trycloudflare.com tunnel for wildcard binds (the raw ip:port is often
# unreachable). Started pre-banner and even when silent so the CLI banner can
# read app.state.cloudflare_url; torn down by _graceful_shutdown.
global _cloudflare_url
global _cloudflare_url, _cloudflare_requested, _cloudflare_flag
_cloudflare_url = None
_cloudflare_flag = cloudflare
app.state.cloudflare_url = None
_cloudflare_enabled = _cloudflare_tunnel_should_start(
cloudflare = cloudflare,
@ -1176,6 +1299,7 @@ def run_server(
api_only = api_only,
is_colab = _IS_COLAB,
)
_cloudflare_requested = _cloudflare_enabled
if _cloudflare_enabled:
try: # best-effort: any failure must not block startup
from cloudflare_tunnel import start_studio_tunnel, stop_studio_tunnel
@ -1199,6 +1323,43 @@ def run_server(
_graceful_shutdown(_server)
sys.exit(1)
# Time-box a freshly-exposed web UI: if nobody changes the seeded admin
# password within the deadline (default 1h), shut down rather than leave an
# unsecured public instance running. No-op for loopback, --api-only, Colab,
# an already-changed password, or UNSLOTH_STUDIO_BOOTSTRAP_TIMEOUT=0.
try:
from auth import storage as _auth_storage
from auth.bootstrap_timeout import (
arm_bootstrap_timeout,
bootstrap_timeout_seconds,
should_arm_bootstrap_timeout,
)
_bootstrap_timeout = bootstrap_timeout_seconds()
if should_arm_bootstrap_timeout(
host = host,
secure = secure,
api_only = api_only,
frontend_served = bool(frontend_path) and not api_only,
is_colab = _IS_COLAB,
requires_change = _auth_storage.requires_password_change(
_auth_storage.DEFAULT_ADMIN_USERNAME
),
timeout_seconds = _bootstrap_timeout,
):
arm_bootstrap_timeout(
_auth_storage,
_trigger_shutdown,
timeout_seconds = _bootstrap_timeout,
logger = logger,
)
logger.info(
"Studio will shut down in %ds unless the default admin password is changed.",
_bootstrap_timeout,
)
except Exception as e: # best-effort: never block startup on the timeout
logger.warning("Bootstrap timeout not armed: %s", e)
if not silent:
_emit_startup_output(host, port, display_host, secure = secure, enable_tools = enable_tools)
@ -1243,8 +1404,10 @@ def _build_arg_parser():
"--cloudflare",
action = argparse.BooleanOptionalAction,
default = True,
help = "Auto-create a free Cloudflare HTTPS tunnel when bound to 0.0.0.0 "
"(default on; --no-cloudflare to disable)",
help = "Auto-create a free Cloudflare HTTPS tunnel for non-api-only wildcard "
"binds (0.0.0.0 or ::), exposing Studio on a PUBLIC internet URL (default on). "
"Pass --no-cloudflare to disable that Cloudflare URL; it does not change a "
"public wildcard bind. --api-only keeps it off unless paired with --secure.",
)
parser.add_argument(
"--secure",

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

@ -119,6 +119,10 @@ def get_connection() -> sqlite3.Connection:
ensure_dir(db_path.parent)
conn = sqlite3.connect(str(db_path))
conn.row_factory = sqlite3.Row
# Wait for a lock instead of erroring immediately: a figure/scan-heavy ingest can
# hold its connection across many seconds of vision calls, and a concurrent ingest
# or autoinject read would otherwise hit "database is locked".
conn.execute("PRAGMA busy_timeout = 5000")
try:
conn.enable_load_extension(True)
sqlite_vec.load(conn)
@ -156,3 +160,71 @@ def vec_table_exists(conn: sqlite3.Connection) -> bool:
"SELECT 1 FROM sqlite_master WHERE type='table' AND name='chunks_vec'"
).fetchone()
return row is not None
def _delete_document_chunks(conn, document_id: str) -> None:
"""Delete a document's chunk rows (chunks/chunks_fts/chunks_vec), keeping the
documents row. Used when reconciling a half-ingested doc to failed: retrieval
filters by scope not status, so leftover chunks would stay citable."""
chunk_ids = [
r["id"]
for r in conn.execute(
"SELECT id FROM chunks WHERE document_id=?", (document_id,)
).fetchall()
]
if not chunk_ids:
return
has_vec = vec_table_exists(conn)
for chunk_id in chunk_ids:
conn.execute("DELETE FROM chunks_fts WHERE chunk_id=?", (chunk_id,))
if has_vec:
conn.execute("DELETE FROM chunks_vec WHERE chunk_id=?", (chunk_id,))
conn.execute("DELETE FROM chunks WHERE document_id=?", (document_id,))
def reconcile_orphaned_ingestion_jobs() -> int:
"""Fail ingestion jobs/documents left mid-flight by a crash so they stop
showing as stuck "processing" and become re-ingestible. Run at startup.
No-op without RAG. Returns the number of jobs reset.
"""
if not RAG_AVAILABLE:
return 0
conn = get_connection()
try:
rows = conn.execute(
"SELECT id, document_id FROM ingestion_jobs "
"WHERE status NOT IN ('completed', 'failed')"
).fetchall()
for row in rows:
doc = conn.execute(
"SELECT status FROM documents WHERE id=?", (row["document_id"],)
).fetchone()
if doc is not None and doc["status"] == "completed":
# Worker finished indexing before the crash but didn't retire the
# job row. Mark the job completed (not failed) and keep its chunks,
# so the UI's getJob fallback after restart doesn't flag a
# searchable document as a failed ingestion.
conn.execute(
"UPDATE ingestion_jobs SET status='completed', stage='done', "
"progress=1.0, error=NULL WHERE id=?",
(row["id"],),
)
continue
conn.execute(
"UPDATE ingestion_jobs SET status='failed', stage='error', "
"error='Server restarted during ingestion' WHERE id=?",
(row["id"],),
)
conn.execute(
"UPDATE documents SET status='failed' "
"WHERE id=? AND status NOT IN ('completed', 'failed')",
(row["document_id"],),
)
# A failed or still-in-flight doc must not leave citable chunks
# (retrieval filters by scope, not status); also drops any chunks of a
# doc already 'failed' before the crash.
_delete_document_chunks(conn, row["document_id"])
conn.commit()
return len(rows)
finally:
conn.close()

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

@ -23,6 +23,16 @@ from typing import Any, Iterable, Optional
from utils.paths import project_workspaces_root, studio_db_path, ensure_dir
from utils.training_runs import extract_project_name
def _extract_project_name_from_config_json(config_json: Optional[str]) -> Optional[str]:
if not config_json:
return None
try:
return extract_project_name(json.loads(config_json))
except (json.JSONDecodeError, TypeError):
return None
def _denied_path_prefixes() -> list[str]:
@ -680,6 +690,7 @@ def list_runs(limit: int = 50, offset: int = 0) -> dict:
runs = []
for row in rows:
run = dict(row)
run["project_name"] = _extract_project_name_from_config_json(run.get("config_json"))
sparkline = run.get("loss_sparkline")
if sparkline:
try:
@ -719,6 +730,7 @@ def get_run(id: str) -> Optional[dict]:
if row is None:
return None
run = dict(row)
run["project_name"] = _extract_project_name_from_config_json(run.get("config_json"))
sparkline = run.get("loss_sparkline")
if sparkline:
try:

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

@ -57,6 +57,15 @@ def test_media_type_and_status():
assert err.status_code == 503
def test_pooled_client_disables_proxy_env():
async def _scenario():
client = llama_http.nonstreaming_client()
assert client.trust_env is False
await llama_http.aclose()
asyncio.run(_scenario())
def test_pooled_client_reused_within_loop_and_recreated_after_close():
async def _scenario():
a = llama_http.nonstreaming_client()

185
studio/backend/tests/test_bootstrap_timeout.py Normal file
View file

@ -0,0 +1,185 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Coverage for the exposed-first-run auto-shutdown deadline.
Tests the env parsing, the pure arm/no-arm decision matrix, and the deadline
handler (shut down iff the seeded admin password is still unchanged). The
threading.Timer itself is not exercised; the handler is invoked directly.
"""
from types import SimpleNamespace
from auth.bootstrap_timeout import (
DEFAULT_BOOTSTRAP_TIMEOUT_SECONDS,
_format_duration,
bootstrap_timeout_seconds,
enforce_bootstrap_password_deadline,
should_arm_bootstrap_timeout,
)
# ── bootstrap_timeout_seconds ───────────────────────────────────────
def test_default_when_unset():
assert bootstrap_timeout_seconds(env = {}) == DEFAULT_BOOTSTRAP_TIMEOUT_SECONDS
def test_default_when_empty():
assert bootstrap_timeout_seconds(env = {"UNSLOTH_STUDIO_BOOTSTRAP_TIMEOUT": " "}) == (
DEFAULT_BOOTSTRAP_TIMEOUT_SECONDS
)
def test_explicit_value_parsed():
assert bootstrap_timeout_seconds(env = {"UNSLOTH_STUDIO_BOOTSTRAP_TIMEOUT": "1800"}) == 1800
def test_zero_disables():
assert bootstrap_timeout_seconds(env = {"UNSLOTH_STUDIO_BOOTSTRAP_TIMEOUT": "0"}) == 0
def test_negative_disables():
assert bootstrap_timeout_seconds(env = {"UNSLOTH_STUDIO_BOOTSTRAP_TIMEOUT": "-5"}) == 0
def test_invalid_falls_back_to_default():
# A typo must keep the protection, not silently disable it.
assert bootstrap_timeout_seconds(env = {"UNSLOTH_STUDIO_BOOTSTRAP_TIMEOUT": "abc"}) == (
DEFAULT_BOOTSTRAP_TIMEOUT_SECONDS
)
# ── should_arm_bootstrap_timeout matrix ─────────────────────────────
def _arm_kwargs(**overrides):
kwargs = dict(
host = "0.0.0.0",
secure = False,
api_only = False,
frontend_served = True,
is_colab = False,
requires_change = True,
timeout_seconds = 3600,
)
kwargs.update(overrides)
return kwargs
def test_arm_exposed_wildcard_web_ui():
assert should_arm_bootstrap_timeout(**_arm_kwargs()) is True
def test_arm_secure_loopback_bind():
# --secure forces a loopback bind but exposes a public tunnel.
assert should_arm_bootstrap_timeout(**_arm_kwargs(host = "127.0.0.1", secure = True)) is True
def test_no_arm_loopback_bind():
assert should_arm_bootstrap_timeout(**_arm_kwargs(host = "127.0.0.1", secure = False)) is False
def test_no_arm_api_only():
assert should_arm_bootstrap_timeout(**_arm_kwargs(api_only = True)) is False
def test_no_arm_no_frontend():
assert should_arm_bootstrap_timeout(**_arm_kwargs(frontend_served = False)) is False
def test_no_arm_colab():
assert should_arm_bootstrap_timeout(**_arm_kwargs(is_colab = True)) is False
def test_no_arm_password_already_changed():
assert should_arm_bootstrap_timeout(**_arm_kwargs(requires_change = False)) is False
def test_no_arm_timeout_disabled():
assert should_arm_bootstrap_timeout(**_arm_kwargs(timeout_seconds = 0)) is False
# ── enforce_bootstrap_password_deadline ─────────────────────────────
def _fake_storage(requires_change: bool):
return SimpleNamespace(
DEFAULT_ADMIN_USERNAME = "unsloth",
requires_password_change = lambda _username: requires_change,
)
def test_deadline_shuts_down_when_password_unchanged():
calls = []
result = enforce_bootstrap_password_deadline(
_fake_storage(requires_change = True),
lambda: calls.append("shutdown"),
timeout_seconds = 3600,
)
assert result is True
assert calls == ["shutdown"]
def test_deadline_keeps_running_when_password_changed():
calls = []
result = enforce_bootstrap_password_deadline(
_fake_storage(requires_change = False),
lambda: calls.append("shutdown"),
timeout_seconds = 3600,
)
assert result is False
assert calls == []
def test_deadline_swallows_shutdown_errors():
def _boom():
raise RuntimeError("shutdown failed")
# A failing shutdown must not propagate out of the timer thread.
result = enforce_bootstrap_password_deadline(
_fake_storage(requires_change = True),
_boom,
timeout_seconds = 3600,
)
assert result is True
# ── _format_duration ────────────────────────────────────────────────
def test_format_duration_sub_minute_uses_seconds():
assert _format_duration(30) == "30 seconds"
def test_format_duration_singular_second():
assert _format_duration(1) == "1 second"
def test_format_duration_exact_minutes():
assert _format_duration(60) == "1 minute"
assert _format_duration(3600) == "60 minutes"
def test_format_duration_minutes_and_seconds():
assert _format_duration(90) == "1 minute 30 seconds"
def test_shutdown_message_uses_formatted_duration():
# The deadline message must reflect the real timeout, not a rounded
# "minute(s)" placeholder. Capture the warning via a fake logger.
logged = []
class _Logger:
def warning(self, msg, *args):
logged.append(msg)
enforce_bootstrap_password_deadline(
_fake_storage(requires_change = True),
lambda: None,
timeout_seconds = 3600,
logger = _Logger(),
)
assert any("60 minutes" in m for m in logged)
assert not any("minute(s)" in m for m in logged)

127
studio/backend/tests/test_cached_gguf_routes.py
View file

@ -20,6 +20,7 @@ if "structlog" not in sys.modules:
)
import routes.models as models_route
from hub.services.models import gguf_variants as GV
def _repo(
@ -564,21 +565,32 @@ def test_gguf_variants_mmproj_does_not_mark_quant_downloaded(monkeypatch, tmp_pa
"""The per-quant 'downloaded' flag is driven by the real weight file in a
single snapshot; an mmproj vision adapter (matching a quant label) must
not make that quant appear downloaded."""
import huggingface_hub.constants as hf_constants
variants = [
SimpleNamespace(filename = "model-Q4_K_M.gguf", quant = "Q4_K_M", size_bytes = 10_000),
SimpleNamespace(filename = "model-F16.gguf", quant = "F16", size_bytes = 20_000),
SimpleNamespace(
filename = "model-Q4_K_M.gguf",
quant = "Q4_K_M",
display_label = None,
size_bytes = 10_000,
),
SimpleNamespace(
filename = "model-F16.gguf",
quant = "F16",
display_label = None,
size_bytes = 20_000,
),
]
monkeypatch.setattr(
models_route, "list_gguf_variants", lambda repo_id, hf_token = None: (variants, True)
GV,
"list_gguf_variants",
lambda repo_id, hf_token = None: (variants, True, []),
)
monkeypatch.setattr(hf_constants, "HF_HUB_CACHE", str(tmp_path))
monkeypatch.setattr(GV, "_local_main_gguf_blobs_by_quant", lambda _repo_id: {})
snap = tmp_path / "models--org--repo" / "snapshots" / "rev"
snap.mkdir(parents = True)
(snap / "model-Q4_K_M.gguf").write_bytes(b"x" * 10_000) # real weight, fully present
(snap / "mmproj-F16.gguf").write_bytes(b"y" * 20_000) # mmproj adapter, label "F16"
monkeypatch.setattr(GV, "iter_hf_cache_snapshots", lambda _repo_id: [snap])
result = asyncio.run(
models_route.get_gguf_variants(
@ -592,21 +604,32 @@ def test_gguf_variants_mmproj_does_not_mark_quant_downloaded(monkeypatch, tmp_pa
def test_gguf_variants_ignore_big_endian_siblings(monkeypatch, tmp_path):
import huggingface_hub.constants as hf_constants
siblings = [
SimpleNamespace(rfilename = "model-Q4_K_M-be.gguf", size = 100),
SimpleNamespace(rfilename = "model-Q4_K_M.gguf", size = 10),
]
monkeypatch.setattr(
"huggingface_hub.model_info",
lambda *_args, **_kwargs: SimpleNamespace(siblings = siblings),
GV,
"list_gguf_variants",
lambda repo_id, hf_token = None: (
[
SimpleNamespace(
filename = "model-Q4_K_M.gguf",
quant = "Q4_K_M",
display_label = None,
size_bytes = 10,
)
],
False,
siblings,
),
)
monkeypatch.setattr(hf_constants, "HF_HUB_CACHE", str(tmp_path))
monkeypatch.setattr(GV, "_local_main_gguf_blobs_by_quant", lambda _repo_id: {})
snap = tmp_path / "models--org--repo" / "snapshots" / "rev"
snap.mkdir(parents = True)
(snap / "model-Q4_K_M.gguf").write_bytes(b"x" * 10)
monkeypatch.setattr(GV, "iter_hf_cache_snapshots", lambda _repo_id: [snap])
result = asyncio.run(
models_route.get_gguf_variants(
@ -620,19 +643,25 @@ def test_gguf_variants_ignore_big_endian_siblings(monkeypatch, tmp_path):
def test_gguf_variants_cached_big_endian_does_not_satisfy_variant(monkeypatch, tmp_path):
import huggingface_hub.constants as hf_constants
variants = [
SimpleNamespace(filename = "model-Q4_K_M.gguf", quant = "Q4_K_M", size_bytes = 10),
SimpleNamespace(
filename = "model-Q4_K_M.gguf",
quant = "Q4_K_M",
display_label = None,
size_bytes = 10,
),
]
monkeypatch.setattr(
models_route, "list_gguf_variants", lambda repo_id, hf_token = None: (variants, False)
GV,
"list_gguf_variants",
lambda repo_id, hf_token = None: (variants, False, []),
)
monkeypatch.setattr(hf_constants, "HF_HUB_CACHE", str(tmp_path))
monkeypatch.setattr(GV, "_local_main_gguf_blobs_by_quant", lambda _repo_id: {})
snap = tmp_path / "models--org--repo" / "snapshots" / "rev"
snap.mkdir(parents = True)
(snap / "model-Q4_K_M-be.gguf").write_bytes(b"x" * 10)
monkeypatch.setattr(GV, "iter_hf_cache_snapshots", lambda _repo_id: [snap])
result = asyncio.run(
models_route.get_gguf_variants(
@ -706,3 +735,69 @@ def test_gguf_download_progress_counts_quant_subdir(monkeypatch, tmp_path):
assert result["downloaded_bytes"] == 20_000
assert result["progress"] == 1.0
def test_arch_to_task_hides_unsupported_diffusion_from_chat():
# Loadable diffusion archs -> the Images-picker task.
assert models_route._arch_to_task("flux") == "text-to-image"
assert models_route._arch_to_task("z_image") == "text-to-image"
assert models_route._arch_to_task("qwen_image") == "text-to-image"
# A real LLM arch stays a chat model; None passes through.
assert models_route._arch_to_task("llama") == "text-generation"
assert models_route._arch_to_task(None) is None
# Known-but-unsupported diffusion archs get a task that is NEITHER chat
# ("text-generation") NOR a loadable image task ("text-to-image"), so the chat
# picker hides them (they'd die in llama.cpp) and the Images picker leaves them
# out (they'd 400 in validate_load).
for arch in ("sdxl", "sd1", "sd3", "wan", "lumina2", "hidream", "cosmos"):
task = models_route._arch_to_task(arch)
assert task == models_route._UNSUPPORTED_DIFFUSION_TASK
assert task not in ("text-generation", "text-to-image")
# Drift guard: every diffusion arch llama.cpp rejects as a chat model must be
# classified here as some image task (loadable OR unsupported), never chat.
from core.inference.llama_cpp import LlamaCppBackend
classified = models_route._DIFFUSION_GGUF_ARCHS | models_route._UNSUPPORTED_DIFFUSION_GGUF_ARCHS
missing = {a for a in LlamaCppBackend._DIFFUSION_ARCHES if a.lower() not in classified}
assert not missing, f"diffusion archs would still show in chat: {missing}"
def test_delete_cached_refuses_diffusion_loaded_repo(monkeypatch):
# The cached-delete guard refuses deleting a repo the diffusion (Images)
# backend has loaded, mirroring the chat guard, so its GGUF can't be removed
# from under a live pipeline.
from fastapi import HTTPException
import core.inference.diffusion as diffusion_mod
import routes.inference as routes_inference
# Chat and orchestrator report nothing loaded; only diffusion holds the repo.
# delete_cached_model resolves get_inference_backend from the models module
# namespace, so patch it there (not on core.inference) to isolate that guard.
monkeypatch.setattr(
routes_inference,
"get_llama_cpp_backend",
lambda: SimpleNamespace(is_loaded = False, model_identifier = None),
)
monkeypatch.setattr(
models_route,
"get_inference_backend",
lambda: SimpleNamespace(active_model_name = None),
)
monkeypatch.setattr(
diffusion_mod,
"get_diffusion_backend",
lambda: SimpleNamespace(status = lambda: {"loaded": True, "repo_id": "org/Z-Image-GGUF"}),
)
try:
asyncio.run(
models_route.delete_cached_model(
repo_id = "org/Z-Image-GGUF",
variant = None,
current_subject = "u",
)
)
assert False, "expected HTTPException refusing the delete"
except HTTPException as e:
assert e.status_code == 400
assert "Unload the model before deleting" in e.detail

256
studio/backend/tests/test_checkpoints_scan.py Normal file
View file

@ -0,0 +1,256 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import json
import sqlite3
import sys
import types as _types
from pathlib import Path
_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
if _BACKEND_DIR not in sys.path:
sys.path.insert(0, _BACKEND_DIR)
_loggers_stub = _types.ModuleType("loggers")
_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name)
sys.modules.setdefault("loggers", _loggers_stub)
sys.modules.setdefault("structlog", _types.ModuleType("structlog"))
from utils.models import checkpoints as checkpoints_module
from utils.training_runs import build_default_output_dir_name
def _make_history_connection(db_path: Path) -> sqlite3.Connection:
conn = sqlite3.connect(str(db_path))
conn.row_factory = sqlite3.Row
return conn
def _setup_training_runs_table(db_path: Path) -> None:
conn = _make_history_connection(db_path)
try:
conn.execute(
"""
CREATE TABLE training_runs (
id TEXT PRIMARY KEY,
model_name TEXT NOT NULL,
config_json TEXT NOT NULL,
output_dir TEXT,
started_at TEXT NOT NULL
)
"""
)
conn.commit()
finally:
conn.close()
def _make_outputs_dir(tmp_path, monkeypatch) -> Path:
studio_home = tmp_path / "studio-home"
outputs_dir = studio_home / "outputs"
outputs_dir.mkdir(parents = True)
monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(studio_home))
return outputs_dir
def test_scan_checkpoints_uses_output_dir_history_for_base_model(tmp_path, monkeypatch):
outputs_dir = _make_outputs_dir(tmp_path, monkeypatch)
run_dir = outputs_dir / "custom-run"
run_dir.mkdir()
(run_dir / "config.json").write_text("{}")
db_path = tmp_path / "studio.db"
_setup_training_runs_table(db_path)
conn = _make_history_connection(db_path)
try:
conn.execute(
"""
INSERT INTO training_runs (id, model_name, config_json, output_dir, started_at)
VALUES (?, ?, ?, ?, ?)
""",
(
"run-1",
"unsloth/Llama-3.2-3B-Instruct",
"{}",
str(run_dir.resolve()),
"2026-04-09T00:00:00Z",
),
)
conn.commit()
finally:
conn.close()
monkeypatch.setattr(
checkpoints_module,
"get_connection",
lambda: _make_history_connection(db_path),
)
models = checkpoints_module.scan_checkpoints(outputs_dir = str(outputs_dir))
assert models[0][2]["base_model"] == "unsloth/Llama-3.2-3B-Instruct"
def test_scan_checkpoints_matches_project_suffixed_default_dir_against_history(
tmp_path, monkeypatch
):
outputs_dir = _make_outputs_dir(tmp_path, monkeypatch)
run_name = build_default_output_dir_name(
"unsloth/Llama-3.2-3B-Instruct",
"Customer Support",
timestamp = 1771227800,
)
run_dir = outputs_dir / run_name
run_dir.mkdir()
(run_dir / "config.json").write_text("{}")
db_path = tmp_path / "studio.db"
_setup_training_runs_table(db_path)
conn = _make_history_connection(db_path)
try:
conn.execute(
"""
INSERT INTO training_runs (id, model_name, config_json, output_dir, started_at)
VALUES (?, ?, ?, ?, ?)
""",
(
"run-2",
"unsloth/Llama-3.2-3B-Instruct",
json.dumps({"project_name": "Customer Support"}),
None,
"2026-04-09T00:00:00Z",
),
)
conn.commit()
finally:
conn.close()
monkeypatch.setattr(
checkpoints_module,
"get_connection",
lambda: _make_history_connection(db_path),
)
models = checkpoints_module.scan_checkpoints(outputs_dir = str(outputs_dir))
assert models[0][2]["base_model"] == "unsloth/Llama-3.2-3B-Instruct"
def test_scan_checkpoints_strips_project_suffix_without_history(tmp_path, monkeypatch):
outputs_dir = _make_outputs_dir(tmp_path, monkeypatch)
run_name = build_default_output_dir_name(
"unsloth/Llama-3.2-3B-Instruct",
"Customer Support",
timestamp = 1771227800,
)
run_dir = outputs_dir / run_name
run_dir.mkdir()
(run_dir / "config.json").write_text("{}")
db_path = tmp_path / "studio.db"
_setup_training_runs_table(db_path)
monkeypatch.setattr(
checkpoints_module,
"get_connection",
lambda: _make_history_connection(db_path),
)
models = checkpoints_module.scan_checkpoints(outputs_dir = str(outputs_dir))
assert models[0][2]["base_model"] == "unsloth/Llama-3.2-3B-Instruct"
def test_scan_checkpoints_preserves_project_marker_in_model_without_history(tmp_path, monkeypatch):
outputs_dir = _make_outputs_dir(tmp_path, monkeypatch)
run_name = build_default_output_dir_name(
"org/foo__project-bar",
timestamp = 1771227800,
)
run_dir = outputs_dir / run_name
run_dir.mkdir()
(run_dir / "config.json").write_text("{}")
db_path = tmp_path / "studio.db"
_setup_training_runs_table(db_path)
monkeypatch.setattr(
checkpoints_module,
"get_connection",
lambda: _make_history_connection(db_path),
)
models = checkpoints_module.scan_checkpoints(outputs_dir = str(outputs_dir))
assert models[0][2]["base_model"] == "org/foo__project-bar"
def test_scan_checkpoints_preserves_legacy_folder_name_fallback(tmp_path, monkeypatch):
outputs_dir = _make_outputs_dir(tmp_path, monkeypatch)
run_dir = outputs_dir / "unsloth_Llama-3.2-3B-Instruct_1771227800"
run_dir.mkdir()
(run_dir / "config.json").write_text("{}")
db_path = tmp_path / "studio.db"
_setup_training_runs_table(db_path)
monkeypatch.setattr(
checkpoints_module,
"get_connection",
lambda: _make_history_connection(db_path),
)
models = checkpoints_module.scan_checkpoints(outputs_dir = str(outputs_dir))
assert models[0][2]["base_model"] == "unsloth/Llama-3.2-3B-Instruct"
def test_scan_checkpoints_prefers_exact_history_match_over_newer_suffix(tmp_path, monkeypatch):
outputs_dir = _make_outputs_dir(tmp_path, monkeypatch)
run_dir = outputs_dir / "unsloth_Test_1771227800"
run_dir.mkdir()
(run_dir / "config.json").write_text("{}")
copied_dir = tmp_path / "copied" / run_dir.name
copied_dir.mkdir(parents = True)
db_path = tmp_path / "studio.db"
_setup_training_runs_table(db_path)
conn = _make_history_connection(db_path)
try:
conn.execute(
"""
INSERT INTO training_runs (id, model_name, config_json, output_dir, started_at)
VALUES (?, ?, ?, ?, ?)
""",
(
"run-exact",
"correct/base",
"{}",
str(run_dir.resolve()),
"2026-04-09T00:00:00Z",
),
)
conn.execute(
"""
INSERT INTO training_runs (id, model_name, config_json, output_dir, started_at)
VALUES (?, ?, ?, ?, ?)
""",
(
"run-suffix",
"wrong/base",
"{}",
str(copied_dir.resolve()),
"2026-04-10T00:00:00Z",
),
)
conn.commit()
finally:
conn.close()
monkeypatch.setattr(
checkpoints_module,
"get_connection",
lambda: _make_history_connection(db_path),
)
models = checkpoints_module.scan_checkpoints(outputs_dir = str(outputs_dir))
assert models[0][2]["base_model"] == "correct/base"

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