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

Merge remote-tracking branch 'origin/main' into agent/harden-auto-permission-network-gates

Browse source
This commit is contained in:
Michael Han 2026-07-21 17:32:35 -07:00
commit 95a84bf74f
287 changed files with 34036 additions and 5951 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -36,6 +36,23 @@ AGENT="${2:?usage: agent-guides-drive.sh <mode> <agent>}"
# Determinism (seed/temp) is applied at the server level by
# serve-unsloth-run.sh --extra; agents inherit it through the API.
TIMEOUT="${AGENT_INVOKE_TIMEOUT:-180}"
# opencode is the slow outlier. Unlike the print-mode agents (claude -p, codex
# exec) it runs a full turn AND a separate small_model call to name the session,
# so one connection reply takes ~8 min on a CPU-served 4B -- right at the shared
# 600s cap, so the cell flaked when a run drifted past a ~480s success. Give it
# headroom (still well under the 40-min job budget); the fast agents keep the
# tight cap that still catches a real headless-TTY hang.
case "$AGENT" in
opencode)
# Double it, but only for a bare-integer seconds value. A GNU timeout(1)
# duration suffix (s/m/h/d, including floats like 0.5s) is left unchanged so
# the arithmetic never sees a non-number; timeout(1) parses it directly.
case "$TIMEOUT" in
*[!0-9]*) ;;
*) TIMEOUT=$(( TIMEOUT * 2 )) ;;
esac
;;
esac
# Claude refuses --dangerously-skip-permissions outside a sandbox; the CI runner
# IS the sandbox, so declare it (mirrors unslothai/scripts launcher.sh). Harmless

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

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

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

@ -19,6 +19,7 @@ on:
- 'install.sh'
- 'pyproject.toml'
- 'tests/studio/**'
- '.github/scripts/run-studio-permission-browser.sh'
- '.github/workflows/studio-mac-ui-smoke.yml'
push:
branches: [main, pip]
@ -96,7 +97,7 @@ jobs:
- name: Assert llama.cpp loads on this macOS
run: bash .github/scripts/assert-llama-loads.sh
- name: Install Playwright + Chromium
- name: Install Playwright browsers
# No --with-deps on Mac: that flag installs Linux apt packages.
# GitHub-hosted macos-14 ships the system frameworks Chromium
# needs already.
@ -112,7 +113,7 @@ jobs:
# in-script retry recover from any residual flakes.
run: |
pip install 'playwright>=1.55,<1.58'
python -m playwright install chromium
python -m playwright install chromium webkit
- name: Patch Playwright pipeTransport.js to tolerate malformed JSON
# In Playwright 1.55-1.58, pipeTransport.js does
@ -244,6 +245,10 @@ jobs:
kill "${STUDIO_PID}" 2>/dev/null || true
sleep 2
- name: Cross-browser permission controls
run: |
bash .github/scripts/run-studio-permission-browser.sh 18895 webkit
- name: Reset auth + boot Unsloth for extra UI tests (port 18897)
run: |
unsloth studio reset-password
@ -343,5 +348,7 @@ jobs:
logs/studio_extra.log
logs/install.log
logs/playwright
logs/playwright-permissions-*
logs/playwright_extra
logs/studio-permissions-*.log
retention-days: 7

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

@ -27,6 +27,7 @@ on:
# The Playwright test files themselves -- a PR that ONLY edits
# the test must still trigger UI CI.
- 'tests/studio/**'
- '.github/scripts/run-studio-permission-browser.sh'
- '.github/workflows/studio-ui-smoke.yml'
push:
branches: [main, pip]
@ -107,13 +108,10 @@ jobs:
set -o pipefail
bash install.sh --local --no-torch 2>&1 | tee logs/install.log
- name: Install Playwright + Chromium
- name: Install Playwright browsers
run: |
pip install 'playwright>=1.45'
# --with-deps installs the OS-level runtime libs Chromium
# needs (libnss3, libxkbcommon, etc.). About 30 s on a
# warm runner.
python -m playwright install --with-deps chromium
python -m playwright install --with-deps chromium firefox webkit
- name: Reset auth + boot Unsloth
run: |
@ -182,6 +180,12 @@ jobs:
kill "${STUDIO_PID}" 2>/dev/null || true
sleep 2
- name: Cross-browser permission controls
run: |
bash .github/scripts/run-studio-permission-browser.sh 18893 firefox
bash .github/scripts/run-studio-permission-browser.sh 18893 webkit
bash .github/scripts/run-studio-permission-browser.sh 18893 chromium chrome
# The chat UI test ends by clicking the Shutdown menuitem, which
# leaves the server dead. The extra UI test (Compare / Recipes /
# Export / Unsloth / Settings) needs a fresh Unsloth, so we boot a
@ -233,6 +237,54 @@ jobs:
kill "${STUDIO_EXTRA_PID}" 2>/dev/null || true
sleep 2
# Model-picker per-model-config regression (PR #7207 re-land of #6647).
# Fourth Unsloth on its own port; loads the tiny GGUF and drives the
# picker's run-settings surface: Context Length persists across a reload,
# Reset clears the stored override (never pins it), and the infra models
# (RAG embedder + llama.cpp probe) stay hidden from the picker.
- name: Reset auth + boot Unsloth for model-config tests (port 18898)
run: |
unsloth studio reset-password
mkdir -p logs
UNSLOTH_API_ONLY=1 unsloth studio -H 127.0.0.1 -p 18898 \
> logs/studio_modelcfg.log 2>&1 &
echo "STUDIO_MODELCFG_PID=$!" >> "$GITHUB_ENV"
- name: Wait for /api/health on 18898
run: |
for i in $(seq 1 180); do
if curl -fs "http://127.0.0.1:18898/api/health" > /tmp/health4.json; then
jq -e '.status == "healthy"' /tmp/health4.json && break
fi
sleep 1
done
jq -e '.status == "healthy"' /tmp/health4.json
- name: Pass bootstrap pw for model-config test
run: |
NEW="CIModelCfg-$(python -c 'import secrets; print(secrets.token_urlsafe(16))')"
echo "::add-mask::$NEW"
echo "STUDIO_MODELCFG_NEW_PW=$NEW" >> "$GITHUB_ENV"
- name: Drive model-picker per-model-config with Playwright
env:
BASE_URL: http://127.0.0.1:18898
STUDIO_NEW_PW: ${{ env.STUDIO_MODELCFG_NEW_PW }}
PW_ART_DIR: logs/playwright_modelcfg
STUDIO_UI_STRICT: '1'
GGUF_REPO: ${{ env.GGUF_REPO }}
GGUF_VARIANT: ${{ env.GGUF_VARIANT }}
STUDIO_MODEL_HINT: gemma-3-270m
run: |
mkdir -p logs/playwright_modelcfg
python tests/studio/playwright_model_config.py
- name: Stop fourth Unsloth
if: always()
run: |
kill "${STUDIO_MODELCFG_PID}" 2>/dev/null || true
sleep 2
# IME + multilingual paste regression (issue #5318 / PR #5327).
# Third Unsloth on its own port so a hang here cannot poison the
# earlier UI tests. No GGUF -- the bug surface is the composer.
@ -293,10 +345,14 @@ jobs:
path: |
logs/studio.log
logs/studio_extra.log
logs/studio_modelcfg.log
logs/studio_ime.log
logs/install.log
logs/server-logs/
logs/playwright
logs/playwright-permissions-*
logs/playwright_extra
logs/playwright_modelcfg
logs/playwright_ime
logs/studio-permissions-*.log
retention-days: 7

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

@ -19,6 +19,7 @@ on:
- 'install.ps1'
- 'pyproject.toml'
- 'tests/studio/**'
- '.github/scripts/run-studio-permission-browser.sh'
- '.github/workflows/studio-windows-ui-smoke.yml'
push:
branches: [main, pip]
@ -345,6 +346,10 @@ jobs:
kill "${STUDIO_PID}" 2>/dev/null || true
sleep 2
- name: Edge permission controls
run: |
bash .github/scripts/run-studio-permission-browser.sh 18895 chromium msedge
- name: Reset auth + boot Unsloth for extra UI tests (port 18897)
run: |
unsloth studio reset-password
@ -402,5 +407,7 @@ jobs:
logs/studio_extra.log
logs/install.log
logs/playwright
logs/playwright-permissions-*
logs/playwright_extra
logs/studio-permissions-*.log
retention-days: 7

62
README.md
View file

@ -11,6 +11,7 @@ Unsloth Studio lets you run and train models locally.
<p align="center">
<a href="#-features">Features</a> •
<a href="#-unsloth-news">News</a> •
<a href="#-install">Quickstart</a> •
<a href="#-free-notebooks">Notebooks</a> •
<a href="https://unsloth.ai/docs">Documentation</a>
@ -47,15 +48,44 @@ Unsloth Studio (Beta) lets you run and train text, [audio](https://unsloth.ai/do
* [Auto set inference settings](https://unsloth.ai/docs/new/studio/chat#auto-parameter-tuning) and customize chat templates.
* We work directly with teams behind [gpt-oss](https://docs.unsloth.ai/new/gpt-oss-how-to-run-and-fine-tune#unsloth-fixes-for-gpt-oss), [Qwen3](https://www.reddit.com/r/LocalLLaMA/comments/1kaodxu/qwen3_unsloth_dynamic_ggufs_128k_context_bug_fixes/), [Llama 4](https://github.com/ggml-org/llama.cpp/pull/12889), [Mistral](https://huggingface.co/mistralai/Mistral-Medium-3.5-128B/discussions/18), [Gemma 1-3](https://news.ycombinator.com/item?id=39671146), and [Phi-4](https://unsloth.ai/blog/phi4), where we’ve fixed bugs that improve model accuracy.
* Chat with images, audio, PDFs, code, DOCX and more. [Connect API providers](https://unsloth.ai/docs/integrations/connections) (OpenAI, Anthropic) or servers (vLLM, Ollama).
* [**Compare any two models**](https://unsloth.ai/docs/new/studio/chat#model-arena) side by side with the same prompt.
* **OpenAI/Anthropic-compatible APIs**: Serve local models through `/v1/chat/completions`, `/v1/responses` and `/v1/messages`.
* **Connect local models to agents**: Use `unsloth start` with Claude Code, Codex, Hermes and more.
* **Web/PDF search** can read PDF papers, manuals and other PDF results.
* **GGUF hardware controls**: Choose GPUs/layers, offload MoE experts, use multi-GPU or Tensor Parallelism.
* The opt-in **MCP control endpoint** lets AI clients manage models, training, recipes and exports.
### Training
* Train and RL **500+ models** up to **2x faster** with up to **70% less VRAM**, with no accuracy loss.
* Custom Triton and mathematical **kernels**. See some collabs we did with [PyTorch](https://unsloth.ai/docs/get-started/reinforcement-learning-rl-guide/fp8-reinforcement-learning) and [Hugging Face](https://unsloth.ai/docs/new/faster-moe).
* Train and RL **500+ models** up to **2x faster** with **70% less VRAM**; MoE up to **12x faster**.
* Train and run RL on [AMD GPUs](https://unsloth.ai/docs/basics/amd) across Windows, WSL and Linux.
* **Data Recipes**: [Auto-create datasets](https://unsloth.ai/docs/new/studio/data-recipe) from **PDF, CSV, DOCX** etc. Edit data in a visual-node workflow.
* **[Reinforcement Learning](https://unsloth.ai/docs/get-started/reinforcement-learning-rl-guide)** (RL): The most efficient [RL](https://unsloth.ai/docs/get-started/reinforcement-learning-rl-guide) library, using **80% less VRAM** for GRPO, [FP8](https://unsloth.ai/docs/get-started/reinforcement-learning-rl-guide/fp8-reinforcement-learning) etc.
* Supports full fine-tuning, RL, pretraining, 4-bit, 16-bit and, FP8 training.
* **[Reinforcement Learning](https://unsloth.ai/docs/get-started/reinforcement-learning-rl-guide)** uses **80% less VRAM** for GRPO, FP8 and vision RL, with 7x longer contexts.
* [**Long-context training**](https://unsloth.ai/docs/new/3x-faster-training-packing): **3x faster**, 30% less VRAM and 500K+ context.
* Supports LoRA/QLoRA, full fine-tuning, RL, pretraining, 4-bit, 16-bit and FP8.
* Custom Triton and mathematical **kernels** built with PyTorch and Hugging Face.
* **Observability**: Monitor training live, track loss and GPU usage and customize graphs.
* [Multi-GPU](https://unsloth.ai/docs/basics/multi-gpu-training-with-unsloth) training is supported, with major improvements coming soon.
## 🚀 Unsloth Start
[Unsloth Start](https://unsloth.ai/docs/integrations/unsloth-start) connects [Claude Code](https://unsloth.ai/docs/basics/claude-code), [Codex](https://unsloth.ai/docs/basics/codex) and other agents to local models with one command.
Start Unsloth, load a model, open your project folder, then run:
```bash
unsloth start claude
```
Replace `claude` with any supported agent:
| Agent | Command |
| --- | --- |
| Claude Code | `unsloth start claude` |
| OpenAI Codex | `unsloth start codex` |
| Hermes Agent | `unsloth start hermes` |
| OpenClaw | `unsloth start openclaw` |
| OpenCode | `unsloth start opencode` |
| Pi Coding Agent | `unsloth start pi` |
## 📥 Install
Unsloth can be used in two ways: through **[Unsloth Studio](https://unsloth.ai/docs/new/studio/)**, the web UI, or through **Unsloth Core**, the code-based version. Each has different requirements.
@ -65,7 +95,8 @@ Unsloth Studio (Beta) works on **Windows, Linux, WSL** and **macOS**.
* **CPU:** Supported for Chat and Data Recipes currently
* **NVIDIA:** Training works on RTX 30/40/50, Blackwell, DGX Spark, Station and more
* **macOS:** Training, MLX and GGUF inference are ALL supported.
* **AMD:** Chat + Data works. Train with [Unsloth Core](#unsloth-core-code-based). Unsloth Studio support is out soon.
* **AMD:** Training, RL, chat and deployment work on Windows, WSL and Linux. [Read the AMD guide](https://unsloth.ai/docs/basics/amd).
* **Vulkan:** GGUF inference is supported on [compatible GPUs, including Intel GPUs](https://github.com/unslothai/unsloth/pull/5819).
* **Multi-GPU:** Available now, with a major upgrade on the way
#### macOS, Linux, WSL:
@ -122,7 +153,7 @@ You can use the same Docker image as Unsloth Studio.
#### AMD, Intel:
For RTX 50x, B200, 6000 GPUs: `uv pip install unsloth --torch-backend=auto`. Read our guides for: [Blackwell](https://unsloth.ai/docs/blog/fine-tuning-llms-with-blackwell-rtx-50-series-and-unsloth) and [DGX Spark](https://unsloth.ai/docs/blog/fine-tuning-llms-with-nvidia-dgx-spark-and-unsloth). <br>
To install Unsloth on **AMD** and **Intel** GPUs, follow our [AMD Guide](https://unsloth.ai/docs/get-started/install/amd) and [Intel Guide](https://unsloth.ai/docs/get-started/install/intel).
To install Unsloth on **AMD** and **Intel** GPUs, follow our [AMD Guide](https://unsloth.ai/docs/basics/amd) and [Intel Guide](https://unsloth.ai/docs/get-started/install/intel).
## 📒 Free Notebooks
@ -148,13 +179,20 @@ Read our [guide](https://unsloth.ai/docs/get-started/fine-tuning-llms-guide). Ad
- See detailed documentation for Unsloth [here](https://unsloth.ai/docs)
## 🦥 Unsloth News
- **Connections**: Connect any API provider (OpenAI, Anthropic) or server (vLLM, Ollama). [Guide](https://unsloth.ai/docs/integrations/connections)
- **MTP**: Run Qwen3.6 MTP in Unsloth. MTP settings are autoset specific to your hardware. [Guide](https://unsloth.ai/docs/models/qwen3.6#mtp-guide)
- **API inference endpoint**: Deploy and run local LLMs in Claude Code, Codex tools. [Guide](https://unsloth.ai/docs/basics/api)
- **Qwen3.6**: Qwen3.6-35B-A3B can now be trained and run in Unsloth Studio. [Blog](https://unsloth.ai/docs/models/qwen3.6)
- **Gemma 4**: Run and train Google’s new models directly in Unsloth. [Blog](https://unsloth.ai/docs/models/gemma-4)
- **AMD training**: Train, run RL, chat and deploy on AMD GPUs across Windows, WSL and Linux. [Guide](https://unsloth.ai/docs/basics/amd)
- **GGUF hardware controls**: Choose GPU/layer placement, offload MoE experts and use multi-GPU or Tensor Parallelism. [#6414](https://github.com/unslothai/unsloth/pull/6414)
- **Local models for any agent**: Use `unsloth start` with Claude Code, Codex, Hermes, OpenCode, OpenClaw, Pi and more through Unsloth's OpenAI- and Anthropic-compatible APIs. [Guide](https://unsloth.ai/docs/basics/api)
- **MCP control endpoint**: Let compatible clients manage models, training, recipes, checkpoints and exports. [#7191](https://github.com/unslothai/unsloth/pull/7191)
- **Local inference reliability**: Resume long chats faster, recover stalled downloads and reuse existing GGUF files. [#7204](https://github.com/unslothai/unsloth/pull/7204) • [#6858](https://github.com/unslothai/unsloth/pull/6858) • [#7209](https://github.com/unslothai/unsloth/pull/7209)
- **New models**: [Qwen-AgentWorld](https://huggingface.co/unsloth/Qwen-AgentWorld-35B-A3B-GGUF), [Ornith](https://huggingface.co/unsloth/models?search=ornith), [Kimi K2.7 Code](https://unsloth.ai/docs/models/kimi-k2.7-code) and [MiniMax M3](https://unsloth.ai/docs/models/minimax-m3)
- **GLM-5.2**: Run Z.ai's 744B-parameter, 1M-context open model locally with Unsloth Dynamic GGUFs. [Guide](https://unsloth.ai/docs/models/glm-5.2)
- **DeepSeek-V4**: Run DeepSeek-V4-Flash locally with corrected multi-turn and tool-calling behavior. [Guide](https://unsloth.ai/docs/models/deepseek-v4)
- **DiffusionGemma**: Run and fine-tune Google's diffusion language model with 1.8x faster inference in Unsloth Studio. [Guide](https://unsloth.ai/docs/models/diffusiongemma)
- **Qwen3.6**: Run and train Qwen3.6 with MTP for 1.4-2.2x faster inference and NVFP4 quants for supported GPUs. [Guide](https://unsloth.ai/docs/models/qwen3.6)
- **Gemma 4**: Run and train Gemma 4 text, image and audio models with QAT, MTP, GGUF and MLX support. [Guide](https://unsloth.ai/docs/models/gemma-4)
- **MCP servers**: Connect local models to files, apps, databases and external tools through Model Context Protocol. [Guide](https://unsloth.ai/docs/basics/mcp)
- **Connections**: Mix local models with API providers (OpenAI, Anthropic) or servers (vLLM, Ollama) in the same interface. [Guide](https://unsloth.ai/docs/integrations/connections)
- **Introducing Unsloth Studio**: our new web UI for running and training LLMs. [Blog](https://unsloth.ai/docs/new/studio)
- **Qwen3.5** - 0.8B, 2B, 4B, 9B, 27B, 35-A3B, 112B-A10B are now supported. [Guide + notebooks](https://unsloth.ai/docs/models/qwen3.5/fine-tune)
- Train **MoE LLMs 12x faster** with 35% less VRAM - DeepSeek, GLM, Qwen and gpt-oss. [Blog](https://unsloth.ai/docs/new/faster-moe)
- **Embedding models**: Unsloth now supports ~1.8-3.3x faster embedding fine-tuning. [Blog](https://unsloth.ai/docs/new/embedding-finetuning) • [Notebooks](https://unsloth.ai/docs/get-started/unsloth-notebooks#embedding-models)
- New **7x longer context RL** vs. all other setups, via our new batching algorithms. [Blog](https://unsloth.ai/docs/new/grpo-long-context)

174
install.ps1
View file

@ -53,7 +53,8 @@ function Install-UnslothStudio {
param([string]$TorchIndexUrl)
if ($SkipTorch) { return "none" }
if ([string]::IsNullOrWhiteSpace($TorchIndexUrl)) { return "none" }
$leaf = ($TorchIndexUrl.TrimEnd('/') -split '/')[-1].ToLowerInvariant()
# Drop query/fragment first so a token-authenticated pin classifies by family.
$leaf = (($TorchIndexUrl -split '[?#]', 2)[0].TrimEnd('/') -split '/')[-1].ToLowerInvariant()
if (@("cpu", "cu118", "cu124", "cu126", "cu128", "cu130") -contains $leaf) { return $leaf }
if ($leaf -match '^rocm[0-9]+\.[0-9]+$') { return $leaf }
return "auto"
@ -62,7 +63,8 @@ function Install-UnslothStudio {
function Get-TauriGpuBranch {
param([string]$TorchIndexFamily)
if ($SkipTorch) { return "no_torch" }
if ($TorchIndexFamily -like "cu*") { return "cuda" }
# Require a digit after "cu" so /current or /custom isn't branded CUDA (parity ^cu[0-9]).
if ($TorchIndexFamily -match '^cu[0-9]') { return "cuda" }
if ($TorchIndexFamily -like "rocm*") { return "rocm" }
if ($TorchIndexFamily -eq "cpu") { return "cpu" }
return "unknown"
@ -467,22 +469,35 @@ function Install-UnslothStudio {
}
}
# Redact index-URL credentials (userinfo + ?query= + #fragment) from captured installer
# output before printing on failure; uv/pip errors echo the failing --index-url verbatim.
# Mirrors the other installers. Verbose mode streams uncaptured, so it isn't redacted.
function Redact-InstallOutput {
param([string]$Text)
if (-not $Text) { return $Text }
$Text = $Text -replace '(https?://)[^/@\s`]+@', '$1<redacted>@'
$Text = $Text -replace '([?&][^=\s&`]+)=[^&#\s`]+', '$1=<redacted>'
# A #token=... fragment is as sensitive as a query; URL-anchored.
return $Text -replace '(https?://[^\s`#]+)#[^\s`]+', '$1#<redacted>'
}
# Run native commands quietly by default to match install.sh behavior.
# Full command output is shown only when --verbose / UNSLOTH_VERBOSE=1.
function Invoke-InstallCommand {
param(
[Parameter(Mandatory = $true)][ScriptBlock]$Command
)
# Installer-pinned index installs (torch) must beat an inherited uv mirror
# (#6898): when the command pins an index, clear every uv index env var so
# it wins, then restore in finally. Other installs keep the user's mirror.
# Installer-pinned index installs (torch) must beat an inherited uv mirror (#6898):
# for --default-index, clear the uv index env vars (restore in finally) and set
# UV_NO_CONFIG=1 so a uv.toml/pyproject index can't outrank the CLI pin (uv 0.10).
$savedUvIndex = $null
if ($Command.ToString() -match '--default-index') {
$savedUvIndex = @{}
foreach ($n in 'UV_DEFAULT_INDEX', 'UV_INDEX_URL', 'UV_INDEX', 'UV_EXTRA_INDEX_URL') {
foreach ($n in 'UV_DEFAULT_INDEX', 'UV_INDEX_URL', 'UV_INDEX', 'UV_EXTRA_INDEX_URL', 'UV_TORCH_BACKEND', 'UV_FIND_LINKS', 'UV_CONFIG_FILE', 'UV_NO_CONFIG') {
$savedUvIndex[$n] = [Environment]::GetEnvironmentVariable($n)
Remove-Item "Env:$n" -ErrorAction SilentlyContinue
}
$env:UV_NO_CONFIG = '1'
}
$prevEap = $ErrorActionPreference
$ErrorActionPreference = "Continue"
@ -493,17 +508,23 @@ function Install-UnslothStudio {
# Merge stderr into stdout so progress/warning output stays visible
# without flipping $? on successful native commands (PS 5.1 treats
# stderr records as errors that set $? = $false even on exit code 0).
& $Command 2>&1 | Out-Host
# Redact per record: uv echoes index URLs (credentials and all) in
# its errors, and verbose mode must not bypass the quiet path's
# redaction. ForEach-Object/Out-Host leave $LASTEXITCODE untouched.
& $Command 2>&1 | ForEach-Object { Redact-InstallOutput "$_" } | Out-Host
} else {
$output = & $Command 2>&1 | Out-String
if ($LASTEXITCODE -ne 0) {
Write-Host $output -ForegroundColor Red
Write-Host (Redact-InstallOutput $output) -ForegroundColor Red
}
}
return [int]$LASTEXITCODE
} finally {
$ErrorActionPreference = $prevEap
if ($savedUvIndex) { foreach ($n in $savedUvIndex.Keys) { if ($null -ne $savedUvIndex[$n]) { Set-Item "Env:$n" $savedUvIndex[$n] } } }
if ($savedUvIndex) {
Remove-Item "Env:UV_NO_CONFIG" -ErrorAction SilentlyContinue
foreach ($n in $savedUvIndex.Keys) { if ($null -ne $savedUvIndex[$n]) { Set-Item "Env:$n" $savedUvIndex[$n] } }
}
}
}
@ -1960,10 +1981,31 @@ exit 0
# On an AMD GPU (no NVIDIA), surface the optional WSL-ROCm driver hint.
if (-not $HasNvidiaSmi -and ($ROCmGfxArch -or $ROCmGpuLabel)) { Show-AmdWslDriverHint }
# Trim trailing slashes from the URL PATH only, preserving ?query / #fragment: a whole-URL
# TrimEnd corrupts a token ending in "/", a single strip leaves .../cu128// empty. Shared.
function Trim-IndexPathSlashes {
param([string]$Url)
$value = $Url.Trim()
$idx = $value.IndexOfAny([char[]]@('?', '#'))
if ($idx -lt 0) {
return $value.TrimEnd('/')
}
return $value.Substring(0, $idx).TrimEnd('/') + $value.Substring($idx)
}
# ── Choose the correct PyTorch index URL based on driver CUDA version ──
# Mirrors Get-PytorchCudaTag in setup.ps1.
function Get-TorchIndexUrl {
$baseUrl = if ($env:UNSLOTH_PYTORCH_MIRROR) { $env:UNSLOTH_PYTORCH_MIRROR.TrimEnd('/') } else { "https://download.pytorch.org/whl" }
# Explicit pin -- skip ALL GPU probing (headless / CI / cross-install).
# UNSLOTH_TORCH_INDEX_URL wins (full URL, verbatim); _FAMILY is the leaf appended
# to the mirror base. Matches install.sh / install_python_stack.py.
if (-not [string]::IsNullOrWhiteSpace($env:UNSLOTH_TORCH_INDEX_URL)) {
return (Trim-IndexPathSlashes $env:UNSLOTH_TORCH_INDEX_URL)
}
if (-not [string]::IsNullOrWhiteSpace($env:UNSLOTH_TORCH_INDEX_FAMILY)) {
return "$baseUrl/$($env:UNSLOTH_TORCH_INDEX_FAMILY.Trim().Trim('/'))"
}
if (-not $NvidiaSmiExe) { return "$baseUrl/cpu" }
try {
$output = Invoke-NvidiaSmiBounded $NvidiaSmiExe
@ -1984,6 +2026,25 @@ exit 0
return "$baseUrl/cu126"
}
# Strip userinfo AND query/fragment so an authenticated pin never leaks. Shared with
# _strip_index_url_credentials (install.sh / py / setup.ps1).
function Remove-IndexUrlCredentials {
param([string]$Url)
$sep = $Url.IndexOf('://')
if ($sep -lt 0) { return $Url }
$scheme = $Url.Substring(0, $sep)
$rest = $Url.Substring($sep + 3)
# Drop query / fragment (may hold auth tokens).
$q = $rest.IndexOfAny([char[]]('?', '#'))
if ($q -ge 0) { $rest = $rest.Substring(0, $q) }
$slash = $rest.IndexOf('/')
$authority = if ($slash -ge 0) { $rest.Substring(0, $slash) } else { $rest }
$at = $authority.LastIndexOf('@')
$host_ = if ($at -ge 0) { $authority.Substring($at + 1) } else { $authority }
if ($slash -ge 0) { return "${scheme}://${host_}$($rest.Substring($slash))" }
return "${scheme}://${host_}"
}
# ── Torch flavor helpers (to repair a stale CPU / wrong-CUDA wheel) ──
# torch.__version__ -> flavor tag (cuXXX / rocm / cpu); untagged wheel = cpu,
# matching setup.ps1's stale-venv parse.
@ -2002,11 +2063,13 @@ exit 0
param([string]$TorchIndexUrl, [string]$ROCmIndexUrl)
if (-not [string]::IsNullOrWhiteSpace($ROCmIndexUrl)) { return 'rocm' }
if ([string]::IsNullOrWhiteSpace($TorchIndexUrl)) { return $null }
$leaf = ($TorchIndexUrl.TrimEnd('/') -split '/')[-1].ToLowerInvariant()
# Drop query/fragment first so .../cu128?token=x classifies as cu128 (else it reinstalls every run).
$leaf = (($TorchIndexUrl -split '[?#]', 2)[0].TrimEnd('/') -split '/')[-1].ToLowerInvariant()
if ($leaf -match '^cu\d+$') { return $leaf }
if ($leaf -eq 'cpu') { return 'cpu' }
if ($leaf -match '^rocm') { return 'rocm' }
if ($leaf -match '^gfx') { return 'rocm' }
# gfx must be followed by a digit (an architecture leaf); gfx-private is custom.
if ($leaf -match '^gfx[0-9]') { return 'rocm' }
return $null
}
@ -2041,6 +2104,10 @@ exit 0
} catch { return $null }
}
# An explicit pin is authoritative: the AMD ROCm reroute below must not rewrite it
# (e.g. a deliberate cpu pin on an AMD host).
$TorchIndexPinned = (-not [string]::IsNullOrWhiteSpace($env:UNSLOTH_TORCH_INDEX_URL)) -or `
(-not [string]::IsNullOrWhiteSpace($env:UNSLOTH_TORCH_INDEX_FAMILY))
$TorchIndexUrl = Get-TorchIndexUrl
# ── GPU arch → newest compatible Windows ROCm wheel release ──
@ -2052,13 +2119,19 @@ exit 0
# Override with UNSLOTH_ROCM_WINDOWS_MIRROR for air-gapped / mirror installs.
$ROCmIndexUrl = $null
$ROCmTorchFloor = $null
if (($HasROCm -or $ROCmGfxArch) -and $TorchIndexUrl -like "*/cpu" -and -not $SkipTorch) {
$PinnedRocmVisionSpec = $null
$PinnedRocmAudioSpec = $null
if (-not $TorchIndexPinned -and ($HasROCm -or $ROCmGfxArch) -and $TorchIndexUrl -like "*/cpu" -and -not $SkipTorch) {
$amdIndexBase = if ($env:UNSLOTH_ROCM_WINDOWS_MIRROR) { $env:UNSLOTH_ROCM_WINDOWS_MIRROR.TrimEnd('/') } else { "https://repo.amd.com/rocm/whl" }
$archFamilyMap = @{
"gfx1201" = "gfx120X-all"; "gfx1200" = "gfx120X-all" # RDNA 4
"gfx1151" = "gfx1151"; "gfx1150" = "gfx1150" # RDNA 3.5 (Strix Halo/Point)
"gfx1103" = "gfx110X-all"; "gfx1102" = "gfx110X-all" # RDNA 3
"gfx1101" = "gfx110X-all"; "gfx1100" = "gfx110X-all"
"gfx1036" = "gfx103X-all"; "gfx1035" = "gfx103X-all" # RDNA 2 (RX 6000)
"gfx1034" = "gfx103X-all"; "gfx1033" = "gfx103X-all"
"gfx1032" = "gfx103X-all"; "gfx1031" = "gfx103X-all"
"gfx1030" = "gfx103X-all"
"gfx90a" = "gfx90a"; "gfx908" = "gfx908" # MI200/MI100
}
# gfx120X (RDNA 4) and gfx1151/gfx1150 (Strix) have a null-pointer bug in
@ -2102,6 +2175,32 @@ exit 0
}
}
# A gfx*/rocm pin skips the auto-reroute above, but the generic CPU/CUDA install below
# would use torch>=2.4,<2.11 and pull a known-bad wheel on the gfx115x/gfx120x/rocm>=7.2
# indexes (the _grouped_mm bug). Route a pinned ROCm index through the ROCm path.
if ($TorchIndexPinned -and -not $ROCmIndexUrl -and -not $SkipTorch) {
$_pinLeaf = (($TorchIndexUrl -split '[?#]', 2)[0].TrimEnd('/') -split '/')[-1].ToLower()
$_pinRocm211 = $false
# Anchor ($) so a suffixed custom leaf (rocm7.2-private) falls through to verbatim.
if ($_pinLeaf -match '^rocm(\d+)\.(\d+)$') {
# Only KNOWN-2.11 rocm (rocm7.2) gets the floor. Matches Test-RocmKnown211Version.
$_pinRocm211 = ([int]$Matches[1] -eq 7 -and [int]$Matches[2] -eq 2)
}
# Only the 2.11-allowlist gfx arches need the floor; others publish <2.11 and stay bare.
$_pinGfx211 = @('gfx120x-all', 'gfx1151', 'gfx1150') -contains $_pinLeaf
if ($_pinGfx211 -or $_pinRocm211) {
$ROCmIndexUrl = $TorchIndexUrl
$ROCmTorchFloor = "torch>=2.11.0,<2.12.0"
$PinnedRocmVisionSpec = "torchvision>=0.26.0,<0.27.0"
$PinnedRocmAudioSpec = "torchaudio>=2.11.0,<2.12.0"
substep "pinned ROCm index ($_pinLeaf) -- enforcing $ROCmTorchFloor" "Cyan"
} elseif ($_pinLeaf -match '^gfx[0-9]' -or $_pinLeaf -match '^rocm[0-9]+(\.[0-9]+)?$') {
# Other gfx / older rocm (<=7.1) ship torch <2.11; route via the ROCm path with
# bare specs. Only EXACT rocm<digits>/gfx* are families; a suffixed leaf is verbatim.
$ROCmIndexUrl = $TorchIndexUrl
}
}
if ($ROCmIndexUrl) {
$TorchIndexFamily = "rocm"
} else {
@ -2164,14 +2263,14 @@ exit 0
}
if ($_Migrated) {
# Migrated env: force-reinstall unsloth+unsloth-zoo to ensure clean state
# in the new venv location, while preserving existing torch/CUDA
# Migrated env: force-reinstall unsloth+unsloth-zoo for a clean state, preserving
# existing torch/CUDA unless the flavor repair below re-lands it.
Write-TauriLog "STEP" "Installing unsloth"
substep "upgrading unsloth in migrated environment..."
if ($SkipTorch) {
# No-torch: install unsloth + unsloth-zoo with --no-deps, then
# runtime deps (typer, safetensors, transformers, etc.) with --no-deps.
$baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated no-torch)" { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.7.3" "unsloth-zoo>=2026.7.3" }
$baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated no-torch)" { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.7.4" "unsloth-zoo>=2026.7.4" }
if ($baseInstallExit -eq 0) {
# Resolve pydantic WITH deps so pip pins pydantic-core
# to the matching version (no-torch-runtime.txt below
@ -2185,7 +2284,7 @@ exit 0
}
}
} else {
$baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated)" { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.7.3" "unsloth-zoo>=2026.7.3" }
$baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (migrated)" { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.7.4" "unsloth-zoo>=2026.7.4" }
}
if ($baseInstallExit -ne 0) {
Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red
@ -2210,22 +2309,24 @@ exit 0
substep "skipping PyTorch (--no-torch flag set)." "Yellow"
} elseif ($ROCmIndexUrl) {
Write-TauriLog "STEP" "Installing PyTorch (AMD ROCm Windows)"
substep "installing PyTorch from $ROCmIndexUrl..."
substep "installing PyTorch from $(Remove-IndexUrlCredentials $ROCmIndexUrl)..."
$torchSpec = if ($ROCmTorchFloor) { $ROCmTorchFloor } else { "torch" }
# Pin the companions to match $torchSpec; bare names can resolve an
# ABI-incompatible torchvision/torchaudio on AMD's per-arch index.
$visionSpec = if ($ROCmGfxArch -and $torchvisionFloorMap.ContainsKey($ROCmGfxArch)) { $torchvisionFloorMap[$ROCmGfxArch] } else { "torchvision" }
$audioSpec = if ($ROCmGfxArch -and $torchaudioFloorMap.ContainsKey($ROCmGfxArch)) { $torchaudioFloorMap[$ROCmGfxArch] } else { "torchaudio" }
$visionSpec = if ($PinnedRocmVisionSpec) { $PinnedRocmVisionSpec } elseif ($ROCmGfxArch -and $torchvisionFloorMap -and $torchvisionFloorMap.ContainsKey($ROCmGfxArch)) { $torchvisionFloorMap[$ROCmGfxArch] } else { "torchvision" }
$audioSpec = if ($PinnedRocmAudioSpec) { $PinnedRocmAudioSpec } elseif ($ROCmGfxArch -and $torchaudioFloorMap -and $torchaudioFloorMap.ContainsKey($ROCmGfxArch)) { $torchaudioFloorMap[$ROCmGfxArch] } else { "torchaudio" }
$torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch (AMD ROCm)" { uv pip install --python $VenvPython --force-reinstall --default-index $ROCmIndexUrl $torchSpec $visionSpec $audioSpec }
if ($torchInstallExit -ne 0) {
# Transient AMD-index failure: fall back to a CPU base so the install
# still completes; Unsloth setup retries ROCm afterwards.
# Transient AMD-index failure: fall back to a CPU base (Unsloth setup retries
# ROCm). Use an explicit CPU index -- for a pinned ROCm index $TorchIndexUrl IS
# the ROCm mirror, so reusing it would just retry it.
$CpuFallbackIndexUrl = if ($env:UNSLOTH_PYTORCH_MIRROR) { "$($env:UNSLOTH_PYTORCH_MIRROR.TrimEnd('/'))/cpu" } else { "https://download.pytorch.org/whl/cpu" }
substep "ROCm PyTorch install failed (exit $torchInstallExit); using a CPU base, Unsloth setup retries ROCm." "Yellow"
# --force-reinstall: a failed ROCm install can leave an unpinned ROCm
# torch (e.g. 2.10.0+rocm on gfx110X/gfx90a) that still satisfies the CPU
# torch>= range, so without it uv would keep the ROCm build and only swap
# the companions -- a mismatched venv the flavor-repair block won't fix.
$torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch (CPU fallback)" { uv pip install --python $VenvPython --force-reinstall "torch>=2.4,<2.11.0" torchvision torchaudio --default-index $TorchIndexUrl }
$torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch (CPU fallback)" { uv pip install --python $VenvPython --force-reinstall "torch>=2.4,<2.11.0" "torchvision>=0.19,<0.26.0" "torchaudio>=2.4,<2.11.0" --default-index $CpuFallbackIndexUrl }
if ($torchInstallExit -ne 0) {
Write-Host "[ERROR] Failed to install PyTorch (ROCm and CPU base both failed, exit code $torchInstallExit)" -ForegroundColor Red
return (Exit-InstallFailure "Failed to install PyTorch (exit code $torchInstallExit)" $torchInstallExit)
@ -2238,8 +2339,14 @@ exit 0
}
} else {
Write-TauriLog "STEP" "Installing PyTorch"
substep "installing PyTorch ($TorchIndexUrl)..."
$torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch" { uv pip install --python $VenvPython "torch>=2.4,<2.11.0" torchvision torchaudio --default-index $TorchIndexUrl }
substep "installing PyTorch ($(Remove-IndexUrlCredentials $TorchIndexUrl))..."
# Bound the companions to the capped torch on EVERY index, cu<digits>
# families included: torchaudio 2.11 dropped its exact torch pin from
# the wheel metadata, so a bare companion next to torch<2.11 can
# resolve a mismatched 2.11.0 build. Mirrors install.sh.
$_pinVisionSpec = "torchvision>=0.19,<0.26.0"
$_pinAudioSpec = "torchaudio>=2.4,<2.11.0"
$torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch" { uv pip install --python $VenvPython "torch>=2.4,<2.11.0" $_pinVisionSpec $_pinAudioSpec --default-index $TorchIndexUrl }
if ($torchInstallExit -ne 0) {
Write-Host "[ERROR] Failed to install PyTorch (exit code $torchInstallExit)" -ForegroundColor Red
return (Exit-InstallFailure "Failed to install PyTorch (exit code $torchInstallExit)" $torchInstallExit)
@ -2251,7 +2358,7 @@ exit 0
if ($SkipTorch) {
# No-torch: install unsloth + unsloth-zoo with --no-deps, then
# runtime deps (typer, safetensors, transformers, etc.) with --no-deps.
$baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (no-torch)" { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.7.3" "unsloth-zoo>=2026.7.3" }
$baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (no-torch)" { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.7.4" "unsloth-zoo>=2026.7.4" }
if ($baseInstallExit -eq 0) {
# Same pydantic-with-deps trick as the migrated branch.
$baseInstallExit = Invoke-InstallCommandRetry -Label "install pydantic" { uv pip install --python $VenvPython pydantic }
@ -2263,7 +2370,7 @@ exit 0
}
}
} elseif ($StudioLocalInstall) {
$baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (local)" { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.7.3" "unsloth-zoo>=2026.7.3" }
$baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (local)" { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.7.4" "unsloth-zoo>=2026.7.4" }
} else {
$baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth" { uv pip install --python $VenvPython --upgrade-package unsloth -- "$PackageName" }
}
@ -2291,7 +2398,7 @@ exit 0
Write-TauriLog "STEP" "Installing unsloth"
substep "installing unsloth (this may take a few minutes)..."
if ($StudioLocalInstall) {
$baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (auto torch backend)" { uv pip install --python $VenvPython "unsloth-zoo>=2026.7.3" "unsloth>=2026.7.3" --torch-backend=auto }
$baseInstallExit = Invoke-InstallCommandRetry -Label "install unsloth (auto torch backend)" { uv pip install --python $VenvPython "unsloth-zoo>=2026.7.4" "unsloth>=2026.7.4" --torch-backend=auto }
if ($baseInstallExit -ne 0) {
Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red
return (Exit-InstallFailure "Failed to install unsloth (exit code $baseInstallExit)" $baseInstallExit)
@ -2317,6 +2424,13 @@ exit 0
}
}
$installedPackageVersion = (& $VenvPython -c "from importlib.metadata import version; import sys; print(version(sys.argv[1]))" $PackageName 2>$null | Out-String).Trim()
if ($LASTEXITCODE -eq 0 -and $installedPackageVersion) {
step $PackageName "$installedPackageVersion installed"
} else {
substep "[WARN] installed $PackageName version could not be determined" "Yellow"
}
# ── Enforce the installed torch flavor matches the detected GPU build ──
# PEP 440 ignores the +cpu/+cuXXX/+rocm local label in a version range, so uv
# keeps a stale torch==X+cpu against a CUDA index and setup.ps1 then loops on
@ -2335,8 +2449,8 @@ exit 0
$rocmSpec = if ($ROCmTorchFloor) { $ROCmTorchFloor } else { "torch" }
# Pin companions like the fresh ROCm path (bare names can pull an
# ABI-incompatible torchvision/torchaudio from the per-arch index).
$visionSpec = if ($ROCmGfxArch -and $torchvisionFloorMap.ContainsKey($ROCmGfxArch)) { $torchvisionFloorMap[$ROCmGfxArch] } else { "torchvision" }
$audioSpec = if ($ROCmGfxArch -and $torchaudioFloorMap.ContainsKey($ROCmGfxArch)) { $torchaudioFloorMap[$ROCmGfxArch] } else { "torchaudio" }
$visionSpec = if ($PinnedRocmVisionSpec) { $PinnedRocmVisionSpec } elseif ($ROCmGfxArch -and $torchvisionFloorMap -and $torchvisionFloorMap.ContainsKey($ROCmGfxArch)) { $torchvisionFloorMap[$ROCmGfxArch] } else { "torchvision" }
$audioSpec = if ($PinnedRocmAudioSpec) { $PinnedRocmAudioSpec } elseif ($ROCmGfxArch -and $torchaudioFloorMap -and $torchaudioFloorMap.ContainsKey($ROCmGfxArch)) { $torchaudioFloorMap[$ROCmGfxArch] } else { "torchaudio" }
substep "PyTorch flavor mismatch (installed $installedTorchTag, need ROCm) -- reinstalling correct build..." "Yellow"
$torchFixExit = Invoke-InstallCommand { uv pip install --python $VenvPython --force-reinstall --default-index $ROCmIndexUrl $rocmSpec $visionSpec $audioSpec }
if ($torchFixExit -ne 0) {
@ -2347,7 +2461,7 @@ exit 0
} elseif ($expectedTorchTag -ne 'rocm') {
# CUDA: stale +cpu (or wrong cuXXX) against a CUDA index -> reinstall triplet.
substep "PyTorch flavor mismatch (installed $installedTorchTag, need $expectedTorchTag) -- reinstalling correct build..." "Yellow"
$torchFixExit = Invoke-InstallCommand { uv pip install --python $VenvPython "torch>=2.4,<2.11.0" torchvision torchaudio --default-index $TorchIndexUrl --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio }
$torchFixExit = Invoke-InstallCommand { uv pip install --python $VenvPython "torch>=2.4,<2.11.0" "torchvision>=0.19,<0.26.0" "torchaudio>=2.4,<2.11.0" --default-index $TorchIndexUrl --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio }
if ($torchFixExit -ne 0) {
Write-Host "[ERROR] Failed to reinstall PyTorch with the correct CUDA build (exit code $torchFixExit)" -ForegroundColor Red
return (Exit-InstallFailure "Failed to reinstall PyTorch ($expectedTorchTag) (exit code $torchFixExit)" $torchFixExit)

586
install.sh
View file

@ -159,18 +159,58 @@ run_maybe_quiet() {
fi
}
# Trim trailing slashes from the URL PATH only, preserving ?query / #fragment: a whole-URL
# strip corrupts a token ending in "/", a single strip leaves .../cu128// empty. Shared.
_trim_index_path_slashes() {
_tips_v="$1"
case "$_tips_v" in
*[?#]*)
_tips_head="${_tips_v%%[?#]*}"
_tips_tail="${_tips_v#"$_tips_head"}"
;;
*)
_tips_head="$_tips_v"
_tips_tail=""
;;
esac
while [ -n "$_tips_head" ] && [ "${_tips_head%/}" != "$_tips_head" ]; do
_tips_head="${_tips_head%/}"
done
printf '%s%s' "$_tips_head" "$_tips_tail"
}
# Redact index-URL credentials (userinfo + ?query= + #fragment) from captured installer
# output before printing on failure; uv/pip errors echo the failing --index-url verbatim.
# Mirrors the other installers. Verbose mode streams uncaptured, so it isn't redacted.
_redact_install_output() {
sed -E \
-e 's#(https?://)[^/@[:space:]`]+@#\1<redacted>@#g' \
-e 's#([?&][^=[:space:]&`]+)=[^&#[:space:]`]+#\1=<redacted>#g' \
-e 's|(https?://[^[:space:]`#]+)#[^[:space:]`]+|\1#<redacted>|g' \
"$@"
}
run_install_cmd() {
_label="$1"
shift
# Installer-pinned index installs (torch) must beat an inherited uv mirror
# (#6898): when we pass --default-index, neutralize every uv index env var so
# the pinned index wins. Other installs keep the user's mirror.
# Installer-pinned index installs (torch) must beat an inherited uv mirror (#6898):
# for --default-index, neutralize the uv index/backend/config vars (UV_TORCH_BACKEND
# redirects torch; UV_NO_CONFIG=1 + dropping UV_CONFIG_FILE stops a uv.toml/pyproject
# index outranking the CLI pin, uv 0.10).
case " $* " in
*" --default-index "*) set -- env -u UV_DEFAULT_INDEX -u UV_INDEX_URL -u UV_INDEX -u UV_EXTRA_INDEX_URL "$@" ;;
*" --default-index "*) set -- env -u UV_DEFAULT_INDEX -u UV_INDEX_URL -u UV_INDEX -u UV_EXTRA_INDEX_URL -u UV_TORCH_BACKEND -u UV_FIND_LINKS -u UV_CONFIG_FILE UV_NO_CONFIG=1 "$@" ;;
esac
if _is_verbose; then
"$@" && return 0
_rc=$?
# Stream through the redactor: uv echoes index URLs (credentials and
# all) in its errors, and verbose mode previously bypassed the
# redaction the quiet path applies. The rc file preserves the
# command's exit code across the pipe without relying on pipefail
# (this script runs under plain sh).
_rcf=$(mktemp)
{ "$@" 2>&1; printf '%s' "$?" > "$_rcf"; } | _redact_install_output
_rc=$(cat "$_rcf" 2>/dev/null || echo 1)
rm -f "$_rcf"
[ "${_rc:-1}" -eq 0 ] 2>/dev/null && return 0
step "error" "$_label failed (exit code $_rc)" "$C_ERR" >&2
return "$_rc"
fi
@ -178,7 +218,7 @@ run_install_cmd() {
"$@" >"$_log" 2>&1 && { rm -f "$_log"; return 0; }
_rc=$?
step "error" "$_label failed (exit code $_rc)" "$C_ERR" >&2
cat "$_log" >&2
_redact_install_output "$_log" >&2
rm -f "$_log"
return $_rc
}
@ -257,7 +297,7 @@ _install_bnb_rocm() {
fi
_bnb_rc=$?
if _is_verbose; then
cat "$_bnb_log" >&2
_redact_install_output "$_bnb_log" >&2
fi
rm -f "$_bnb_log"
step "warning" "$_label (pre-release) failed (exit code $_bnb_rc)" "$C_WARN" >&2
@ -310,6 +350,11 @@ _tauri_torch_index_family() {
return
fi
_diag_url="${1:-}"
# Strip query/fragment AND a trailing slash before classifying (like _torch_index_url_leaf):
# a token isn't echoed into [TAURI:DIAG], and .../cu128/?token=x still classifies as cu128.
_diag_url="${_diag_url%%\?*}"
_diag_url="${_diag_url%%#*}"
_diag_url="${_diag_url%/}"
case "$_diag_url" in
*/cu118) echo "cu118" ;;
*/cu124) echo "cu124" ;;
@ -343,7 +388,8 @@ _tauri_gpu_branch() {
return
fi
case "$_diag_family" in
cu*) echo "cuda" ;;
# Require a digit after cu so /current or /custom isn't branded CUDA (parity ^cu[0-9]).
cu[0-9]*) echo "cuda" ;;
rocm*)
if [ "$_diag_radeon" = true ]; then
echo "rocm_radeon"
@ -472,11 +518,13 @@ _on_install_exit() {
_restore_studio_venv_replacement
fi
[ -n "${_UV_OVERRIDE_TMPDIR:-}" ] && rm -rf "$_UV_OVERRIDE_TMPDIR" 2>/dev/null || true
[ -n "${_UNSLOTH_TORCH_OVERRIDES:-}" ] && rm -f "$_UNSLOTH_TORCH_OVERRIDES" 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.
# Empty so an inherited value never reaches the trap's rm; only temp paths this
# script creates below (spaced-path dir, torch-trio overrides) are removed.
_UV_OVERRIDE_TMPDIR=""
_UNSLOTH_TORCH_OVERRIDES=""
trap _on_install_exit EXIT
# ── Helper: download a URL to a file (supports curl and wget) ──
@ -1573,6 +1621,12 @@ _has_usable_nvidia_gpu() {
# the STUDIO_HOME mkdir/venv so the origin distro is untouched.
_maybe_reroute_strixhalo_to_2404() {
[ "${OS:-}" = "wsl" ] || return 0
# An explicit index pin skips every GPU-driven reroute (same contract as
# the later Radeon/Strix guard): the pin is honored in THIS distro rather
# than probing the GPU and switching distributions. Whitespace-only
# overrides do not gate (parity with get_torch_index_url).
_rr_pin=$(printf '%s' "${UNSLOTH_TORCH_INDEX_URL:-}${UNSLOTH_TORCH_INDEX_FAMILY:-}" | tr -d '[:space:]')
[ -n "$_rr_pin" ] && return 0
[ "${SKIP_TORCH:-false}" = "false" ] || return 0
[ "${UNSLOTH_SKIP_ROCM_WSL_SETUP:-0}" = "1" ] && return 0
[ "${UNSLOTH_WSL_REROUTED:-0}" = "1" ] && return 0
@ -1634,6 +1688,10 @@ _maybe_reroute_strixhalo_to_2404() {
# Forward explicit ROCm-bootstrap consent (e.g. Tauri) so the child auto-enables the
# GPU instead of falling back to the desktop-app prompt path.
[ "${UNSLOTH_ROCM_WSL_AUTO:-0}" = "1" ] && _rr_exports="$_rr_exports; export UNSLOTH_ROCM_WSL_AUTO=1"
# Forward a pinned torch index into the rerouted distro; dropping it would
# silently revert the child install to auto-detection.
[ -n "${UNSLOTH_TORCH_INDEX_URL:-}" ] && _rr_exports="$_rr_exports; export UNSLOTH_TORCH_INDEX_URL=$(_rr_q "$UNSLOTH_TORCH_INDEX_URL")"
[ -n "${UNSLOTH_TORCH_INDEX_FAMILY:-}" ] && _rr_exports="$_rr_exports; export UNSLOTH_TORCH_INDEX_FAMILY=$(_rr_q "$UNSLOTH_TORCH_INDEX_FAMILY")"
[ "$_SKIP_AUTOSTART" = true ] && _rr_exports="$_rr_exports; export UNSLOTH_SKIP_AUTOSTART=1"
_rr_args=""
[ "$PACKAGE_NAME" != "unsloth" ] && _rr_args="$_rr_args --package $(_rr_q "$PACKAGE_NAME")"
@ -1821,6 +1879,8 @@ tauri_log "STEP" "Creating virtual environment"
mkdir -p "$STUDIO_HOME"
_MIGRATED=false
# Empty so an inherited value can never masquerade as a probed torch version.
_PREV_TORCH_VER=""
if [ -x "$VENV_DIR/bin/python" ]; then
# why: matching guard to the .venv branch below -- in env-mode
@ -1838,6 +1898,12 @@ if [ -x "$VENV_DIR/bin/python" ]; then
echo " Move it aside or choose an empty UNSLOTH_STUDIO_HOME." >&2
exit 1
fi
# Record the existing venv's torch BEFORE the replacement moves it aside: a re-run
# rebuilds the venv for clean state, but must keep the torch release the user
# already has (see _previous_torch_pin below). Last line only: sitecustomize or
# import-hook noise on stdout must not corrupt the version.
_PREV_TORCH_VER=$("$VENV_DIR/bin/python" -c \
"import torch; print(torch.__version__)" 2>/dev/null | tail -n 1 || true)
# New layout already exists — replace only after preserving rollback copy.
substep "preserving existing environment for rollback..."
_start_studio_venv_replacement "$VENV_DIR"
@ -1991,6 +2057,15 @@ if [ "$SKIP_TORCH" = false ] && [ "$OS" = "macos" ] && [ "$_ARCH" = "arm64" ]; t
TORCH_CONSTRAINT="torch>=2.6,<2.11.0"
fi
fi
# Companion (torchvision/torchaudio) constraints, bounded to torch's window.
# torchaudio 2.11 dropped its exact torch pin, so a bare companion next to a
# <2.11-capped torch resolves torchaudio 2.11 (verified: cpu leaf installed
# torch 2.10.0+cpu with torchaudio 2.11.0+cpu). torchvision still exact-pins
# torch and self-corrects, but is bounded for symmetry. Widened alongside the
# cu* torch window below; the torch-2.11 AMD paths (rocm7.2 / per-gfx / Strix)
# pin their own trio.
TORCHVISION_CONSTRAINT="torchvision>=0.19,<0.26.0"
TORCHAUDIO_CONSTRAINT="torchaudio>=2.4,<2.11.0"
# ── Resolve repo root (for --local installs) ──
_REPO_ROOT="$(cd "$(dirname "$0" 2>/dev/null || echo ".")" && pwd)"
@ -2059,6 +2134,24 @@ _has_amd_rocm_gpu() {
get_torch_index_url() {
_base="${UNSLOTH_PYTORCH_MIRROR:-https://download.pytorch.org/whl}"
_base="${_base%/}"
# Explicit override -- skip ALL GPU probing (headless / container / CI / cross-install).
# UNSLOTH_TORCH_INDEX_URL wins (full URL, verbatim); _FAMILY is the leaf (cpu, cu128, ...)
# appended to the mirror base. Trim whitespace so a whitespace-only value is unset.
_url="${UNSLOTH_TORCH_INDEX_URL:-}"
_url="${_url#"${_url%%[![:space:]]*}"}"; _url="${_url%"${_url##*[![:space:]]}"}"
if [ -n "$_url" ]; then
# Trim trailing PATH slashes (a multi-slash path 404s on strict pip proxies) while
# preserving a ?query/#fragment token (a whole-URL strip would eat a "/"-ending token).
_url=$(_trim_index_path_slashes "$_url")
echo "$_url"; return
fi
_family="${UNSLOTH_TORCH_INDEX_FAMILY:-}"
_family="${_family#"${_family%%[![:space:]]*}"}"; _family="${_family%"${_family##*[![:space:]]}"}"
if [ -n "$_family" ]; then
while [ "${_family#/}" != "$_family" ]; do _family="${_family#/}"; done
while [ "${_family%/}" != "$_family" ]; do _family="${_family%/}"; done
echo "$_base/$_family"; return
fi
# macOS: always CPU (no CUDA support)
case "$(uname -s)" in Darwin) echo "$_base/cpu"; return ;; esac
# Try nvidia-smi -- require the binary to actually list a usable GPU.
@ -2187,16 +2280,155 @@ _torch_flavor_tag() {
esac
}
# Final path segment of a wheel index URL ($1), lowercased, query/fragment stripped first
# so a token-authenticated pin (.../cu128?token=x) classifies as cu128 (else it reinstalls
# every update). Classification only. Shared with the py / ps1 leaf extractors.
_torch_index_url_leaf() {
_tl_u="${1%%\?*}"
_tl_u="${_tl_u%%#*}"
# Strip ALL trailing slashes, not one: .../rocm7.2// must yield rocm7.2, not an empty leaf.
while [ -n "$_tl_u" ] && [ "${_tl_u%/}" != "$_tl_u" ]; do
_tl_u="${_tl_u%/}"
done
printf '%s' "${_tl_u##*/}" | tr '[:upper:]' '[:lower:]'
}
# True (exit 0) when a lowercased leaf is an EXACT pip ROCm family: rocm<digits>[.<digits>]
# or a gfx ARCHITECTURE leaf (gfx followed by a digit: gfx90a, gfx1151, gfx120x-all). A leaf
# that merely starts with rocm/gfx (rocm7.2-private, gfx-private) is a custom verbatim pin.
# Matches the py / ps1 sides.
_is_pip_rocm_family_leaf() {
case "$1" in
gfx[0-9]*) return 0 ;;
rocm[0-9]*)
# Exact rocm<digits>[.<digits>]: both major and minor must be non-empty all-digits
# (rocm7., rocm7.2.1, rocm7.2-private are all custom pins, not a family).
_rocm_rest="${1#rocm}"
case "$_rocm_rest" in
*.*.*) return 1 ;;
*.*)
_rocm_minor="${_rocm_rest#*.}"
case "${_rocm_rest%%.*}" in "" | *[!0-9]*) return 1 ;; esac
case "$_rocm_minor" in "" | *[!0-9]*) return 1 ;; esac
;;
*[!0-9]*) return 1 ;;
esac
return 0
;;
*) return 1 ;;
esac
}
# Whether release base $1 (X.Y[.Z...]) falls inside constraint window $2
# ("torch>=A.B[.C],<D.E.F"). Compares at major.minor granularity, which is exact
# for the windows this script uses (ceilings are always X.Y.0); a non-.0 ceiling
# would only make this conservative (excludes the whole ceiling minor). Anything
# unparseable answers "no" so the caller fails toward the supported range.
_torch_release_in_window() {
_trw_con="$2"
case "$_trw_con" in
"torch>="*",<"*) ;;
*) echo "no"; return ;;
esac
_trw_floor="${_trw_con#torch>=}"; _trw_floor="${_trw_floor%%,*}"
_trw_ceil="${_trw_con##*,<}"
_v_maj="${1%%.*}"; _v_rest="${1#*.}"; _v_min="${_v_rest%%.*}"
_f_maj="${_trw_floor%%.*}"; _f_rest="${_trw_floor#*.}"; _f_min="${_f_rest%%.*}"
_c_maj="${_trw_ceil%%.*}"; _c_rest="${_trw_ceil#*.}"; _c_min="${_c_rest%%.*}"
for _trw_n in "$_v_maj" "$_v_min" "$_f_maj" "$_f_min" "$_c_maj" "$_c_min"; do
case "$_trw_n" in ''|*[!0-9]*) echo "no"; return ;; esac
done
if [ "$_v_maj" -gt "$_f_maj" ] || { [ "$_v_maj" -eq "$_f_maj" ] && [ "$_v_min" -ge "$_f_min" ]; }; then
if [ "$_v_maj" -lt "$_c_maj" ] || { [ "$_v_maj" -eq "$_c_maj" ] && [ "$_v_min" -lt "$_c_min" ]; }; then
echo "yes"
return
fi
fi
echo "no"
}
# Keep the previous venv's torch on a re-run: echo "torch==X.Y.Z" when the probed
# version ($1) is inside the active constraint window ($2), else "". The RELEASE is kept
# regardless of flavor tag; the pin installs from the freshly chosen index, so flavor
# follows the machine (cpu <-> cuda, cu126 -> cu130, PyPI bare -> +cu130) while the
# release follows the user. Gating on flavor was wrong: a PyPI torch reports a BARE
# version (on Linux the PyPI wheel IS CUDA), misclassified "cpu", so a healthy 2.10 on a
# cu130 host was moved to 2.11. Per-leaf floors still win (rocm7.2 / gfx >=2.11 for the
# Strix _grouped_mm fix, out-of-window manual installs) and are never pinned; the caller's
# _PREV_FALLBACK_CONSTRAINT installs the newest supported release when the index lacks the
# exact one. Opt out with UNSLOTH_TORCH_UPGRADE=1.
_previous_torch_pin() {
_ptp_ver="$1"
_ptp_con="$2"
[ -n "$_ptp_ver" ] || { echo ""; return; }
[ "${UNSLOTH_TORCH_UPGRADE:-0}" = "1" ] && { echo ""; return; }
_ptp_base="${_ptp_ver%%+*}"
# Base must be a plain numeric release (X.Y[.Z]); probe noise and
# nightly/dev/source builds (2.11.0.dev20250704, 2.9.0a0) must never
# become a pin -- no stable index carries them, so pinning would only
# print "keeping it" and then burn a doomed resolve before falling back.
case "$_ptp_base" in
*[!0-9.]* | *..* | .* | *.) echo ""; return ;;
[0-9]*.[0-9]*) ;;
*) echo ""; return ;;
esac
[ "$(_torch_release_in_window "$_ptp_base" "$_ptp_con")" = "yes" ] || { echo ""; return; }
echo "torch==$_ptp_base"
}
# Install torch from TORCH_INDEX_URL honoring a kept-release pin: with _PREV_TORCH_PIN
# set, TORCH_CONSTRAINT is the exact previous release; fall back to the supported range
# if the index lacks it (pruned mirror) rather than failing. Used by every --default-index
# path (NVIDIA cu*, AMD rocm/gfx fallbacks, cpu/mac, ROCm repairs) so preservation is
# uniform. Extra args (e.g. --force-reinstall) are passed through to uv.
_install_torch_default_index() {
if [ -n "$_PREV_TORCH_PIN" ]; then
# Pair the companions with the kept torch minor: torchaudio no longer
# exact-pins torch in its metadata, so leaving it unconstrained resolves
# a newer mismatched build (a kept torch 2.9.0 pulled torchaudio 2.11.0).
_itdi_base="${_PREV_TORCH_PIN#torch==}"
_itdi_minor="${_itdi_base#*.}"
_itdi_minor="${_itdi_minor%%.*}"
_itdi_tv="torchvision"
_itdi_ta="torchaudio"
case "$_itdi_base" in
2.*)
_itdi_tv="torchvision==0.$((_itdi_minor + 15)).*"
_itdi_ta="torchaudio==2.${_itdi_minor}.*"
;;
esac
if ! run_install_cmd_retry "install PyTorch (kept release)" uv pip install --python "$_VENV_PY" "$TORCH_CONSTRAINT" "$_itdi_tv" "$_itdi_ta" \
--default-index "$TORCH_INDEX_URL" "$@"; then
substep "[WARN] $_PREV_TORCH_PIN is not installable from $(_strip_index_url_credentials "$TORCH_INDEX_URL") -- installing the newest supported release instead" "$C_WARN"
TORCH_CONSTRAINT="$_PREV_FALLBACK_CONSTRAINT"
_PREV_TORCH_PIN=""
run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" "$TORCH_CONSTRAINT" "$TORCHVISION_CONSTRAINT" "$TORCHAUDIO_CONSTRAINT" \
--default-index "$TORCH_INDEX_URL" "$@"
fi
else
run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" "$TORCH_CONSTRAINT" "$TORCHVISION_CONSTRAINT" "$TORCHAUDIO_CONSTRAINT" \
--default-index "$TORCH_INDEX_URL" "$@"
fi
}
# Expected tag from the index leaf ($1): cuXXX / cpu / rocm (rocmX.Y and gfx* ->
# rocm). Empty on an unknown leaf (odd mirror) so the repair safely no-ops.
_expected_torch_flavor_tag() {
_u="${1%/}"
_leaf="${_u##*/}"
_leaf=$(_torch_index_url_leaf "$1")
case "$_leaf" in
cu[0-9]*) echo "$_leaf" ;;
cpu) echo "cpu" ;;
rocm*|gfx*) echo "rocm" ;;
*) echo "" ;;
cu[0-9]*)
# Exact cu + digits only; a cu*-suffixed leaf (cu128-private) -> "" (custom),
# else a correct +cu128 wheel is force-reinstalled every run.
case "${_leaf#cu}" in
*[!0-9]*) echo "" ;;
*) echo "$_leaf" ;;
esac
;;
cpu) echo "cpu" ;;
# Exact rocm/gfx families only; a custom rocm*-suffixed leaf -> "" (custom).
*)
if _is_pip_rocm_family_leaf "$_leaf"; then echo "rocm"; else echo ""; fi
;;
esac
}
@ -2206,14 +2438,42 @@ _expected_torch_flavor_tag() {
# fresh-install paths above already use -- so a stale wheel is auto-repairable.
# Unknown/odd-mirror leaves -> no, so we warn rather than risk a wrong reinstall.
_torch_index_repairable() {
_u="${1%/}"
_leaf="${_u##*/}"
_leaf=$(_torch_index_url_leaf "$1")
case "$_leaf" in
cu[0-9]*|rocm[0-9]*|gfx*) echo "yes" ;;
*) echo "no" ;;
cu[0-9]*) echo "yes" ;;
# Only EXACT rocm/gfx families resolve via --default-index; a suffixed leaf is verbatim.
*)
if _is_pip_rocm_family_leaf "$_leaf"; then echo "yes"; else echo "no"; fi
;;
esac
}
# Remove credentials from a wheel index URL ($1) so an authenticated pin never leaks:
# drops userinfo AND query/fragment; scheme/host/path stay exact. Shared with py / ps1.
_strip_index_url_credentials() {
_sic_url="$1"
case "$_sic_url" in
*://*) ;;
*) printf '%s' "$_sic_url"; return ;;
esac
_sic_scheme="${_sic_url%%://*}"
_sic_rest="${_sic_url#*://}"
# Drop query / fragment (may hold auth tokens).
_sic_rest="${_sic_rest%%\?*}"
_sic_rest="${_sic_rest%%#*}"
_sic_auth="${_sic_rest%%/*}"
# Drop user:pass@ userinfo if present.
case "$_sic_auth" in
*@*) _sic_host="${_sic_auth##*@}" ;;
*) _sic_host="$_sic_auth" ;;
esac
if [ "$_sic_auth" = "$_sic_rest" ]; then
printf '%s://%s' "$_sic_scheme" "$_sic_host"
else
printf '%s://%s/%s' "$_sic_scheme" "$_sic_host" "${_sic_rest#*/}"
fi
}
get_radeon_wheel_url() {
# Only meaningful on Linux. Picks a repo.radeon.com base URL whose listing
# contains torch wheels. Tries paths like rocm-rel-7.2.1/, rocm-rel-7.2/,
@ -2459,7 +2719,19 @@ _maybe_bootstrap_rocm_wsl() {
[ -n "$_rw_tmp" ] && rm -f "$_rw_tmp"
return 0
}
_maybe_bootstrap_rocm_wsl || true
# When the caller pins the wheel index (UNSLOTH_TORCH_INDEX_URL / _FAMILY), honour it
# everywhere: skip the WSL ROCm bootstrap and the Radeon/Strix reroute below (which would
# re-probe the GPU and overwrite the pin). Trim whitespace first (parity with
# get_torch_index_url): a whitespace-only override is unset there, so must not flip this true.
_torch_index_pinned=false
_ti_url_trim="${UNSLOTH_TORCH_INDEX_URL:-}"
_ti_url_trim="${_ti_url_trim#"${_ti_url_trim%%[![:space:]]*}"}"; _ti_url_trim="${_ti_url_trim%"${_ti_url_trim##*[![:space:]]}"}"
_ti_family_trim="${UNSLOTH_TORCH_INDEX_FAMILY:-}"
_ti_family_trim="${_ti_family_trim#"${_ti_family_trim%%[![:space:]]*}"}"; _ti_family_trim="${_ti_family_trim%"${_ti_family_trim##*[![:space:]]}"}"
if [ -n "$_ti_url_trim" ] || [ -n "$_ti_family_trim" ]; then
_torch_index_pinned=true
fi
[ "$_torch_index_pinned" = true ] || _maybe_bootstrap_rocm_wsl || true
TORCH_INDEX_URL=$(get_torch_index_url)
@ -2470,24 +2742,74 @@ TORCH_INDEX_URL=$(get_torch_index_url)
# whose base path happens to contain "rocm" or "gfx" must not mislabel a
# cu*/cpu index as ROCm (radeon repo URLs end in rocm-rel-X.Y/, Strix
# overrides in gfxNNNN/, so the trailing slash is stripped first).
_torch_index_leaf="${TORCH_INDEX_URL%/}"
# Lowercase the leaf so every gfx*/rocm*/cu* arm matches regardless of case (canonical AMD
# RDNA4 leaf is gfx120X-all). CUDA is branded only on a real cu[0-9]* leaf, so a mirror
# leaf (/current) does NOT commit a CUDA backend; an unknown leaf leaves the var unset so
# the stack probes the GPU. Query/fragment dropped first, then ALL trailing slashes (in
# lockstep with the shared _torch_index_url_leaf extractor).
_torch_index_leaf="${TORCH_INDEX_URL%%\?*}"
_torch_index_leaf="${_torch_index_leaf%%#*}"
# Strip ALL trailing slashes, not one: .../cu128// must yield cu128, not an empty leaf.
while [ -n "$_torch_index_leaf" ] && [ "${_torch_index_leaf%/}" != "$_torch_index_leaf" ]; do
_torch_index_leaf="${_torch_index_leaf%/}"
done
_torch_index_leaf="${_torch_index_leaf##*/}"
_torch_index_leaf=$(printf '%s' "$_torch_index_leaf" | tr '[:upper:]' '[:lower:]')
case "$_torch_index_leaf" in
rocm*|gfx*) export UNSLOTH_TORCH_BACKEND="rocm" ;;
cpu) export UNSLOTH_TORCH_BACKEND="cpu" ;;
*) export UNSLOTH_TORCH_BACKEND="cuda" ;;
cu[0-9]*) export UNSLOTH_TORCH_BACKEND="cuda" ;;
# Unknown leaf (odd mirror, /current): unset so a stale inherited value can't leak and
# the stack probes the GPU.
*) unset UNSLOTH_TORCH_BACKEND ;;
esac
# rocm7.2 ships torch 2.11.0 -- adjust the constraint to allow it.
# All other ROCm tags and CUDA stay within <2.11.0.
case "$TORCH_INDEX_URL" in
*/rocm7.2) TORCH_CONSTRAINT="torch>=2.11.0,<2.12.0" ;;
# Whether TORCH_INDEX_URL names an actual pip ROCm family (rocm<digit>* / gfx*), gating the
# ROCm-only side effects below (AMD bitsandbytes, ROCm-torch repair). Digit-gated so a leaf
# merely STARTING with "rocm" isn't force-repaired from the wrong path.
if _is_pip_rocm_family_leaf "$_torch_index_leaf"; then
_torch_index_is_rocm_family=true
else
_torch_index_is_rocm_family=false
fi
# rocm7.2 and the per-gfx indexes with the _grouped_mm <2.11 bug (gfx120X-all, gfx1151,
# gfx1150) ship torch 2.11.0 -- raise the floor (also covers a pinned override that skipped
# the Strix reroute). Pin the companions too: the per-gfx index publishes them independently
# and a bare name can resolve a 2.12 ABI-mismatched wheel. Match on the FINAL leaf so a
# custom mirror with a gfx/rocm7.2 path segment but a cu*/cpu family isn't forced.
case "$_torch_index_leaf" in
rocm7.2|gfx120x-all|gfx1151|gfx1150)
TORCH_CONSTRAINT="torch>=2.11.0,<2.12.0"
TORCHVISION_CONSTRAINT="torchvision>=0.26.0,<0.27.0"
TORCHAUDIO_CONSTRAINT="torchaudio>=2.11.0,<2.12.0"
;;
# CUDA cu12x/cu13x indexes ship torch 2.11.x: widen the ceiling to <2.12.0 (matches
# _CUDA_TORCH_PKG_SPEC) and widen the companions with it so the trio stays paired.
cu[0-9]*)
TORCH_CONSTRAINT="torch>=2.4,<2.12.0"
TORCHVISION_CONSTRAINT="torchvision>=0.19,<0.27.0"
TORCHAUDIO_CONSTRAINT="torchaudio>=2.4,<2.12.0"
;;
esac
# A pinned custom/unknown-leaf index (/simple, /current, /cu128-private) has no curated
# companion set, so bound torchvision/torchaudio to the same <2.11 range the Python path pins
# (else a mirror with newer companions resolves a 2.12 ABI-mismatched wheel). Known families
# keep their curated companions above (_expected_torch_flavor_tag returns "" only for custom).
if [ "$_torch_index_pinned" = true ] && \
[ -z "$(_expected_torch_flavor_tag "$TORCH_INDEX_URL")" ]; then
TORCHVISION_CONSTRAINT="torchvision>=0.19,<0.26.0"
TORCHAUDIO_CONSTRAINT="torchaudio>=2.4,<2.11.0"
fi
# Auto-detect GPU for AMD ROCm based
# get_torch_index_url must have chosen */rocm*
# (gfx in rocminfo or amd-smi list). Then require rocminfo "Marketing Name:.*Radeon".
# Skipped when the index is pinned: an explicit override must not be rerouted to the
# Radeon/Strix repos by GPU probing.
_amd_gpu_radeon=false
if [ "$_torch_index_pinned" = false ]; then
case "$TORCH_INDEX_URL" in
*/rocm*)
if _has_amd_rocm_gpu && command -v rocminfo >/dev/null 2>&1 && \
@ -2564,10 +2886,31 @@ case "$TORCH_INDEX_URL" in
done
TORCH_INDEX_URL="${_amd_strix_base}/${_strix_gfx}/"
TORCH_CONSTRAINT="torch>=2.11.0,<2.12.0"
# Pin companions to 2.11 (per-gfx index publishes them independently).
TORCHVISION_CONSTRAINT="torchvision>=0.26.0,<0.27.0"
TORCHAUDIO_CONSTRAINT="torchaudio>=2.11.0,<2.12.0"
_amd_gpu_radeon=false
fi
;;
esac
fi # _torch_index_pinned guard (Radeon + Strix reroute)
# Re-run over an existing install: keep the previous venv's torch RELEASE; the fresh
# index above supplies the right flavor for this machine. Evaluated HERE, after every
# index/constraint decision including the Strix reroute, so the window checked is the
# final one and a raised floor (rocm7.2 / Strix gfx) rejects an older release.
# _PREV_FALLBACK_CONSTRAINT keeps the range so the install can fall back when the exact
# release is not on the chosen index (mirrors may prune old wheels). Skipped for --no-torch.
_PREV_TORCH_PIN=""
_PREV_FALLBACK_CONSTRAINT="$TORCH_CONSTRAINT"
if [ "$SKIP_TORCH" = false ]; then
_prev_pin=$(_previous_torch_pin "$_PREV_TORCH_VER" "$TORCH_CONSTRAINT")
if [ -n "$_prev_pin" ]; then
_PREV_TORCH_PIN="$_prev_pin"
TORCH_CONSTRAINT="$_prev_pin"
substep "existing install has torch $_PREV_TORCH_VER -- keeping it (set UNSLOTH_TORCH_UPGRADE=1 to get the newest release)"
fi
fi
_TAURI_TORCH_INDEX_FAMILY=$(_tauri_torch_index_family "$TORCH_INDEX_URL")
if [ "$_amd_gpu_radeon" = true ] && [ "$SKIP_TORCH" = false ]; then
_TAURI_TORCH_INDEX_FAMILY="radeon"
@ -2697,7 +3040,7 @@ case "$TORCH_INDEX_URL" in
if [ "$_amd_gpu_radeon" = true ]; then
substep "wheels: repo.radeon.com (Radeon)"
else
substep "wheels: $TORCH_INDEX_URL"
substep "wheels: $(_strip_index_url_credentials "$TORCH_INDEX_URL")"
fi
;;
esac
@ -2705,9 +3048,46 @@ esac
# ── Install unsloth directly into the venv (no activation needed) ──
tauri_log "STEP" "Installing PyTorch"
_VENV_PY="$VENV_DIR/bin/python"
# A released unsloth wheel can pin an older torch (unsloth 2026.7.2 declares
# torch<2.11.0); a with-deps PyPI resolve then downgrades the whole trio,
# swapping the pinned +cuXXX/+rocm build for PyPI's default. The flavor guard
# below misses this (PyPI's torch 2.10 default is itself cu128-flavored), so
# freeze the trio via uv --overrides (overrides replace dependency requirements
# during resolution) while unsloth's other deps resolve normally. Sets
# _UNSLOTH_TORCH_OVERRIDES from the trio in the venv; every with-deps unsloth
# install (migrated and fresh) must call this before resolving and rm it after.
_build_unsloth_torch_overrides() {
_UNSLOTH_TORCH_OVERRIDES=""
[ "$SKIP_TORCH" = false ] || return 0
_torch_trio_pins=$("$_VENV_PY" -c "
from importlib.metadata import version, PackageNotFoundError
for _p in ('torch', 'torchvision', 'torchaudio'):
try:
print(_p + '==' + version(_p))
except PackageNotFoundError:
pass
" 2>/dev/null) || _torch_trio_pins=""
case "$_torch_trio_pins" in
torch==*)
_UNSLOTH_TORCH_OVERRIDES=$(mktemp)
printf '%s\n' "$_torch_trio_pins" > "$_UNSLOTH_TORCH_OVERRIDES"
# The CLI --overrides flag replaces any UV_OVERRIDE env file (same
# uv setting; macOS arm64 exports one here), so fold its pins in.
# awk, not cat: it drops inherited torch-trio lines (uv intersects
# duplicate overrides, so a conflicting pin would make resolution
# unsatisfiable) and newline-terminates the last line so an
# unterminated file cannot join two requirements into one.
for _ov_file in ${UV_OVERRIDE:-}; do
[ -f "$_ov_file" ] && awk '!/^[[:space:]]*torch(vision|audio)?([[:space:]<>=!~;@[]|$)/' "$_ov_file" >> "$_UNSLOTH_TORCH_OVERRIDES"
done
;;
esac
}
if [ "$_MIGRATED" = true ]; then
# Migrated env: force-reinstall unsloth+unsloth-zoo to ensure clean state
# in the new venv location, while preserving existing torch/CUDA
# Migrated env: force-reinstall unsloth+unsloth-zoo for a clean state, preserving
# existing torch/CUDA unless the ROCm repair below fires.
substep "upgrading unsloth in migrated environment..."
if [ "$SKIP_TORCH" = true ]; then
# No-torch: install unsloth + unsloth-zoo with --no-deps (current
@ -2716,7 +3096,7 @@ if [ "$_MIGRATED" = true ]; then
# to prevent transitive torch resolution.
run_install_cmd_retry "install unsloth (migrated no-torch)" uv pip install --python "$_VENV_PY" --no-deps \
--reinstall-package unsloth --reinstall-package unsloth-zoo \
"unsloth>=2026.7.3" "unsloth-zoo>=2026.7.3"
"unsloth>=2026.7.4" "unsloth-zoo>=2026.7.4"
# Resolve pydantic WITH deps so pip pins pydantic-core to the
# matching version (no-torch-runtime.txt below is --no-deps).
# All transitive deps are torch-free.
@ -2729,9 +3109,13 @@ if [ "$_MIGRATED" = true ]; then
else
# Pin mlx-lm away from 0.31.3 here too: a curl-piped migration has no
# overrides file, so UV_OVERRIDE is unset and this positional is the only cover.
_build_unsloth_torch_overrides
run_install_cmd_retry "install unsloth (migrated)" uv pip install --python "$_VENV_PY" \
${_UNSLOTH_TORCH_OVERRIDES:+--overrides "$_UNSLOTH_TORCH_OVERRIDES"} \
--reinstall-package unsloth --reinstall-package unsloth-zoo \
"unsloth>=2026.7.3" "unsloth-zoo>=2026.7.3" ${_MLX_LM_EXCLUDE_ARG:-}
"unsloth>=2026.7.4" "unsloth-zoo>=2026.7.4" ${_MLX_LM_EXCLUDE_ARG:-}
[ -n "$_UNSLOTH_TORCH_OVERRIDES" ] && rm -f "$_UNSLOTH_TORCH_OVERRIDES"
_UNSLOTH_TORCH_OVERRIDES=""
fi
if [ "$STUDIO_LOCAL_INSTALL" = true ]; then
substep "overlaying local repo (editable)..."
@ -2744,21 +3128,14 @@ if [ "$_MIGRATED" = true ]; then
# AMD ROCm: install bitsandbytes even in migrated environments so
# existing ROCm installs gain the AMD bitsandbytes build without a
# fresh reinstall.
if [ "$SKIP_TORCH" = false ]; then
case "$TORCH_INDEX_URL" in
*/rocm*|*/gfx*)
_install_bnb_rocm "install bitsandbytes (AMD)" "$_VENV_PY"
# Repair ROCm torch if overwritten during migrated install
_has_hip=$("$_VENV_PY" -c "import torch; print(getattr(torch.version,'hip','') or '')" 2>/dev/null || true)
if [ -z "$_has_hip" ]; then
substep "repairing ROCm torch (overwritten by dependency resolution)..."
run_install_cmd_retry "repair ROCm torch" uv pip install --python "$_VENV_PY" \
"$TORCH_CONSTRAINT" torchvision torchaudio \
--default-index "$TORCH_INDEX_URL" \
--force-reinstall
fi
;;
esac
if [ "$SKIP_TORCH" = false ] && [ "$_torch_index_is_rocm_family" = true ]; then
_install_bnb_rocm "install bitsandbytes (AMD)" "$_VENV_PY"
# Repair ROCm torch if overwritten during migrated install
_has_hip=$("$_VENV_PY" -c "import torch; print(getattr(torch.version,'hip','') or '')" 2>/dev/null || true)
if [ -z "$_has_hip" ]; then
substep "repairing ROCm torch (overwritten by dependency resolution)..."
_install_torch_default_index --force-reinstall
fi
fi
elif [ -n "$TORCH_INDEX_URL" ]; then
# Fresh: Step 1 - install torch from explicit index (skip when --no-torch or Intel Mac)
@ -2820,7 +3197,42 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
_ta_ver=$(_extract_version "$_ta_whl" "torchaudio")
_radeon_versions_match=false
if [ -n "$_torch_ver" ] && [ -n "$_tv_ver" ] && [ -n "$_ta_ver" ]; then
# Kept release (_PREV_TORCH_PIN) wins here too: pick its exact
# patch (else the newest patch of its minor) plus the paired
# vision/audio wheels. Any gap falls back to the newest-trio
# search below, mirroring _install_torch_default_index, so a
# rerun never drifts to another release nor below the kept one.
if [ -n "$_PREV_TORCH_PIN" ]; then
_prev_kept_base="${_PREV_TORCH_PIN#torch==}"
_prev_kept_minor="${_prev_kept_base#*.}"
_prev_kept_minor="${_prev_kept_minor%%.*}"
case "$_prev_kept_minor" in
''|*[!0-9]*) ;;
*)
_kept_torch=$(_pick_radeon_wheel "torch" "${_prev_kept_base}" 2>/dev/null) || _kept_torch=""
[ -z "$_kept_torch" ] && { _kept_torch=$(_pick_radeon_wheel "torch" "2.${_prev_kept_minor}." 2>/dev/null) || _kept_torch=""; }
_kept_tv=$(_pick_radeon_wheel "torchvision" "0.$((_prev_kept_minor + 15))." 2>/dev/null) || _kept_tv=""
_kept_ta=$(_pick_radeon_wheel "torchaudio" "2.${_prev_kept_minor}." 2>/dev/null) || _kept_ta=""
if [ -n "$_kept_torch" ] && [ -n "$_kept_tv" ] && [ -n "$_kept_ta" ]; then
_torch_whl=$_kept_torch
_tv_whl=$_kept_tv
_ta_whl=$_kept_ta
_tri_whl=""
_radeon_versions_match=true
# Say so when the listing pruned the exact patch
# and a same-series build is installed instead.
case "$(printf '%s' "${_kept_torch##*/}" | sed 's/%2[Bb]/+/g')" in
"torch-${_prev_kept_base}"[+-]*) ;;
*) substep "kept release ${_prev_kept_base} is not in the Radeon listing -- installing the closest 2.${_prev_kept_minor} series build instead" ;;
esac
else
substep "[WARN] Radeon repo lacks a complete wheel set for kept $_PREV_TORCH_PIN -- installing the newest compatible set instead" "$C_WARN"
fi
;;
esac
fi
if [ "$_radeon_versions_match" != true ] && \
[ -n "$_torch_ver" ] && [ -n "$_tv_ver" ] && [ -n "$_ta_ver" ]; then
_torch_minor=${_torch_ver#*.}
_ta_minor=${_ta_ver#*.}
_tv_minor=${_tv_ver#*.}
@ -2877,10 +3289,8 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
if [ -z "$_torch_whl" ] || [ -z "$_tv_whl" ] || [ -z "$_ta_whl" ] || \
[ "$_radeon_versions_match" != true ]; then
substep "[WARN] Radeon repo lacks a compatible wheel set for this Python; falling back to ROCm index ($TORCH_INDEX_URL)" "$C_WARN"
run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" \
"$TORCH_CONSTRAINT" torchvision torchaudio \
--default-index "$TORCH_INDEX_URL"
substep "[WARN] Radeon repo lacks a compatible wheel set for this Python; falling back to ROCm index ($(_strip_index_url_credentials "$TORCH_INDEX_URL"))" "$C_WARN"
_install_torch_default_index
else
substep "installing PyTorch from Radeon repo (${_RADEON_BASE_URL})..."
# Pass explicit wheel URLs so the matched trio is
@ -2900,42 +3310,34 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
fi
fi
else
substep "[WARN] Radeon repo unavailable; falling back to ROCm index ($TORCH_INDEX_URL)" "$C_WARN"
run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" \
"$TORCH_CONSTRAINT" torchvision torchaudio \
--default-index "$TORCH_INDEX_URL"
substep "[WARN] Radeon repo unavailable; falling back to ROCm index ($(_strip_index_url_credentials "$TORCH_INDEX_URL"))" "$C_WARN"
_install_torch_default_index
fi
else
substep "[WARN] Radeon GPU detected but could not detect full ROCm version; falling back to ROCm index" "$C_WARN"
run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" \
"$TORCH_CONSTRAINT" torchvision torchaudio \
--default-index "$TORCH_INDEX_URL"
_install_torch_default_index
fi
else
substep "installing PyTorch ($TORCH_INDEX_URL)..."
run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" "$TORCH_CONSTRAINT" torchvision torchaudio \
--default-index "$TORCH_INDEX_URL"
substep "installing PyTorch ($(_strip_index_url_credentials "$TORCH_INDEX_URL"))..."
_install_torch_default_index
fi
# AMD ROCm: install bitsandbytes (once, after torch, for all ROCm paths).
# Gate on SKIP_TORCH=false so a user running with --no-torch on a ROCm
# host stays in GGUF-only mode rather than pulling in bitsandbytes,
# which is only useful once torch is present for training.
if [ "$SKIP_TORCH" = false ]; then
case "$TORCH_INDEX_URL" in
*/rocm*|*/gfx*)
_install_bnb_rocm "install bitsandbytes (AMD)" "$_VENV_PY"
;;
esac
if [ "$SKIP_TORCH" = false ] && [ "$_torch_index_is_rocm_family" = true ]; then
_install_bnb_rocm "install bitsandbytes (AMD)" "$_VENV_PY"
fi
# Fresh: Step 2 - install unsloth, preserving pre-installed torch
# Fresh: Step 2 - install unsloth, preserving the torch Step 1 installed
tauri_log "STEP" "Installing Unsloth"
substep "installing unsloth (this may take a few minutes)..."
_build_unsloth_torch_overrides
if [ "$SKIP_TORCH" = true ]; then
# No-torch: install unsloth + unsloth-zoo with --no-deps, then
# runtime deps (typer, safetensors, transformers, etc.) with --no-deps.
run_install_cmd_retry "install unsloth (no-torch)" uv pip install --python "$_VENV_PY" --no-deps \
--upgrade-package unsloth --upgrade-package unsloth-zoo \
"unsloth>=2026.7.3" "unsloth-zoo>=2026.7.3"
"unsloth>=2026.7.4" "unsloth-zoo>=2026.7.4"
# Same pydantic-with-deps trick as the migrated branch.
run_install_cmd_retry "install pydantic (with deps for compatible core)" \
uv pip install --python "$_VENV_PY" pydantic
@ -2953,7 +3355,8 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
fi
elif [ "$STUDIO_LOCAL_INSTALL" = true ]; then
run_install_cmd_retry "install unsloth (local)" uv pip install --python "$_VENV_PY" \
--upgrade-package unsloth "unsloth>=2026.7.3" "unsloth-zoo>=2026.7.3"
${_UNSLOTH_TORCH_OVERRIDES:+--overrides "$_UNSLOTH_TORCH_OVERRIDES"} \
--upgrade-package unsloth "unsloth>=2026.7.4" "unsloth-zoo>=2026.7.4"
substep "overlaying local repo (editable)..."
run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps
substep "overlaying unsloth-zoo from git main..."
@ -2962,30 +3365,26 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
"unsloth-zoo @ git+https://github.com/unslothai/unsloth-zoo"
else
run_install_cmd_retry "install unsloth" uv pip install --python "$_VENV_PY" \
${_UNSLOTH_TORCH_OVERRIDES:+--overrides "$_UNSLOTH_TORCH_OVERRIDES"} \
--upgrade-package unsloth -- "$PACKAGE_NAME" ${_MLX_LM_EXCLUDE_ARG:-}
fi
[ -n "$_UNSLOTH_TORCH_OVERRIDES" ] && rm -f "$_UNSLOTH_TORCH_OVERRIDES"
_UNSLOTH_TORCH_OVERRIDES=""
# AMD ROCm: repair torch if the unsloth/unsloth-zoo install pulled in
# CUDA torch from PyPI, overwriting the ROCm wheels installed in Step 1.
if [ "$SKIP_TORCH" = false ]; then
case "$TORCH_INDEX_URL" in
*/rocm*|*/gfx*)
_has_hip=$("$_VENV_PY" -c "import torch; print(getattr(torch.version,'hip','') or '')" 2>/dev/null || true)
if [ -z "$_has_hip" ]; then
substep "repairing ROCm torch (overwritten by dependency resolution)..."
run_install_cmd_retry "repair ROCm torch" uv pip install --python "$_VENV_PY" \
"$TORCH_CONSTRAINT" torchvision torchaudio \
--default-index "$TORCH_INDEX_URL" \
--force-reinstall
fi
;;
esac
if [ "$SKIP_TORCH" = false ] && [ "$_torch_index_is_rocm_family" = true ]; then
_has_hip=$("$_VENV_PY" -c "import torch; print(getattr(torch.version,'hip','') or '')" 2>/dev/null || true)
if [ -z "$_has_hip" ]; then
substep "repairing ROCm torch (overwritten by dependency resolution)..."
_install_torch_default_index --force-reinstall
fi
fi
else
# Fallback: GPU detection failed to produce a URL -- let uv resolve torch
tauri_log "STEP" "Installing Unsloth"
substep "installing unsloth (this may take a few minutes)..."
if [ "$STUDIO_LOCAL_INSTALL" = true ]; then
run_install_cmd_retry "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" "unsloth-zoo>=2026.7.3" "unsloth>=2026.7.3" --torch-backend=auto
run_install_cmd_retry "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" "unsloth-zoo>=2026.7.4" "unsloth>=2026.7.4" --torch-backend=auto
substep "overlaying local repo (editable)..."
run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps
substep "overlaying unsloth-zoo from git main..."
@ -2997,6 +3396,15 @@ else
fi
fi
_installed_package_version=$("$_VENV_PY" -c \
'from importlib.metadata import version; import sys; print(version(sys.argv[1]))' \
"$PACKAGE_NAME" 2>/dev/null || true)
if [ -n "$_installed_package_version" ]; then
step "$PACKAGE_NAME" "$_installed_package_version installed"
else
substep "[WARN] installed $PACKAGE_NAME version could not be determined" "$C_WARN"
fi
# ── Enforce the installed torch flavor matches the detected GPU build ──
# PEP 440 ignores the +cpu/+cuXXX/+rocm local label in a version range, so uv
# keeps a stale torch==X+cpu against a GPU index and the venv silently trains on
@ -3014,9 +3422,7 @@ if [ "$SKIP_TORCH" = false ] && [ -n "${TORCH_INDEX_URL:-}" ]; then
if [ -n "$_installed_torch_tag" ] && [ "$_installed_torch_tag" != "$_expected_torch_tag" ] \
&& [ "$(_torch_index_repairable "$TORCH_INDEX_URL")" = "yes" ]; then
substep "PyTorch flavor mismatch (installed $_installed_torch_tag, need $_expected_torch_tag) -- reinstalling correct build..."
run_install_cmd "reinstall PyTorch ($_expected_torch_tag)" uv pip install --python "$_VENV_PY" \
"$TORCH_CONSTRAINT" torchvision torchaudio \
--default-index "$TORCH_INDEX_URL" \
_install_torch_default_index \
--reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio
_installed_torch_ver=$("$_VENV_PY" -c "import torch; print(torch.__version__)" 2>/dev/null || true)
_installed_torch_tag=""
@ -3027,7 +3433,7 @@ if [ "$SKIP_TORCH" = false ] && [ -n "${TORCH_INDEX_URL:-}" ]; then
substep "[WARN] PyTorch is CPU-only but a $_expected_torch_tag GPU build was expected for this machine." "$C_WARN"
substep "[WARN] Training and GPU inference will run on CPU until this is fixed." "$C_WARN"
substep "[WARN] Re-run this installer, or reinstall the GPU build manually:" "$C_WARN"
substep "[WARN] uv pip install --python \"$_VENV_PY\" \"$TORCH_CONSTRAINT\" torchvision torchaudio --default-index $TORCH_INDEX_URL --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio" "$C_WARN"
substep "[WARN] uv pip install --python \"$_VENV_PY\" \"$TORCH_CONSTRAINT\" \"$TORCHVISION_CONSTRAINT\" \"$TORCHAUDIO_CONSTRAINT\" --default-index $(_strip_index_url_credentials "$TORCH_INDEX_URL") --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio" "$C_WARN"
fi
fi
fi

6
pyproject.toml
View file

@ -74,7 +74,7 @@ triton = [
]
huggingfacenotorch = [
"unsloth_zoo>=2026.7.3",
"unsloth_zoo>=2026.7.4",
"wheel>=0.42.0",
"packaging",
"numpy",
@ -95,7 +95,7 @@ huggingfacenotorch = [
]
huggingface = [
"unsloth[huggingfacenotorch]",
"unsloth_zoo>=2026.7.3",
"unsloth_zoo>=2026.7.4",
"torchvision",
"unsloth[triton]",
]
@ -580,7 +580,7 @@ colab-ampere-torch220 = [
"flash-attn>=2.6.3 ; ('linux' in sys_platform)",
]
colab-new = [
"unsloth_zoo>=2026.7.3",
"unsloth_zoo>=2026.7.4",
"packaging",
"tyro",
"transformers>=4.51.3,!=4.52.0,!=4.52.1,!=4.52.2,!=4.52.3,!=4.53.0,!=4.54.0,!=4.55.0,!=4.55.1,!=4.57.0,!=4.57.4,!=4.57.5,!=5.0.0,!=5.1.0,<=5.5.0",

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

File diff suppressed because it is too large Load diff

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

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

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

@ -70,6 +70,8 @@ _DENYLIST_GROUPS: tuple[frozenset[str], ...] = (
# llama-server's own built-in tools flag would silently stack on top of
# Unsloth's --enable-tools / --disable-tools policy resolver.
frozenset({"--tools"}),
# Slot-state dir: Studio owns it for KV persistence across idle unload.
frozenset({"--slot-save-path"}),
)
_DENYLIST: frozenset[str] = frozenset().union(*_DENYLIST_GROUPS)
@ -186,12 +188,25 @@ _SPLIT_MODE_FLAGS: frozenset[str] = frozenset({"-sm", "--split-mode"})
_TENSOR_SPLIT_FLAGS: frozenset[str] = frozenset({"-ts", "--tensor-split"})
_SPLIT_SHADOWING_FLAGS: frozenset[str] = _SPLIT_MODE_FLAGS | _TENSOR_SPLIT_FLAGS
# GPU-offload flags. Stripped only when the GPU Memory mode owns offload
# (manual emits --fit / --gpu-layers / --n-cpu-moe); in auto, a user's
# inherited -ngl is respected (the offload_overridden path), so this group is
# opt-in, not default. Layer flags are shared with llama_cpp's override
# detection; the MoE flags are strip-only (manual's --n-cpu-moe slider owns them).
_LAYER_OFFLOAD_FLAGS: frozenset[str] = frozenset(
{"-ngl", "--gpu-layers", "--n-gpu-layers", "-fit", "--fit"}
)
_MOE_OFFLOAD_FLAGS: frozenset[str] = frozenset({"-ncmoe", "--n-cpu-moe", "-cmoe", "--cpu-moe"})
_OFFLOAD_SHADOWING_FLAGS: frozenset[str] = _LAYER_OFFLOAD_FLAGS | _MOE_OFFLOAD_FLAGS
_SHADOWING_FLAGS: frozenset[str] = (
_CONTEXT_FLAGS | _CACHE_FLAGS | _SPEC_FLAGS | _TEMPLATE_FLAGS | _SPLIT_SHADOWING_FLAGS
)
# Shadowing flags that take no value -- strip the flag only, not the next token.
_BOOLEAN_SHADOWING_FLAGS: frozenset[str] = frozenset({"--spec-default", "--jinja", "--no-jinja"})
_BOOLEAN_SHADOWING_FLAGS: frozenset[str] = frozenset(
{"--spec-default", "--jinja", "--no-jinja", "-cmoe", "--cpu-moe"}
)
def parse_ctx_override(args: Optional[Iterable[str]]) -> Optional[int]:
@ -424,6 +439,8 @@ def strip_shadowing_flags(
strip_spec: bool = True,
strip_template: bool = True,
strip_split_mode: bool = True,
strip_tensor_split: bool = False,
strip_offload: bool = False,
) -> list[str]:
"""Strip flags that shadow first-class Unsloth settings.
@ -432,6 +449,12 @@ def strip_shadowing_flags(
(same for cache / spec / template / split-mode). Each ``strip_*``
toggle controls one group; the route only strips groups whose
first-class field the caller actually supplied.
``strip_split_mode`` removes both ``--split-mode`` and the coupled
``--tensor-split`` (the Tensor Parallelism toggle owns the whole split).
``strip_tensor_split`` removes ``--tensor-split`` *alone*, so manual mode can
replace an inherited per-GPU ratio while leaving the user's ``--split-mode``
row/none/layer choice intact.
"""
shadowing: set[str] = set()
if strip_context:
@ -444,6 +467,10 @@ def strip_shadowing_flags(
shadowing |= _TEMPLATE_FLAGS
if strip_split_mode:
shadowing |= _SPLIT_SHADOWING_FLAGS
if strip_tensor_split:
shadowing |= _TENSOR_SPLIT_FLAGS
if strip_offload:
shadowing |= _OFFLOAD_SHADOWING_FLAGS
tokens = [str(a) for a in (args or [])]
out: list[str] = []

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

@ -54,9 +54,8 @@ class GenStreamError(str):
"""A stream chunk carrying a real backend/generation error, not model text.
Subclasses str so existing display/logging consumers are unaffected, while
callers that must abort a distributed run on error (raise_on_streamed_error)
can distinguish a real error from model output whose visible text starts with
"Error:" by checking isinstance(chunk, GenStreamError).
callers can distinguish a real error from model output whose visible text
starts with "Error:" by checking isinstance(chunk, GenStreamError).
"""
__slots__ = ("public",)

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

@ -158,6 +158,16 @@ def list_documents(conn: sqlite3.Connection, scope: str) -> list[dict]:
return [dict(r) for r in rows]
def list_all_documents(conn: sqlite3.Connection) -> list[dict]:
"""Every uploaded document across all scopes (KBs, threads, projects)."""
rows = conn.execute(
"SELECT id, scope, kb_id, thread_id, project_id, filename, sha256, status, error, "
"num_chunks, stored_path, created_at "
"FROM documents ORDER BY created_at DESC"
).fetchall()
return [dict(r) for r in rows]
def get_document(conn: sqlite3.Connection, document_id: str) -> dict | None:
row = conn.execute("SELECT * FROM documents WHERE id=?", (document_id,)).fetchone()
return dict(row) if row else None

173
studio/backend/core/training/resume.py
View file

@ -4,6 +4,8 @@
"""Helpers for validating resumable training outputs."""
import json
import pickletools
import zipfile
from pathlib import Path
from typing import Optional
@ -33,21 +35,158 @@ def _checkpoint_step(path: Path) -> int:
return -1
def get_resume_checkpoint_path(path_value: str) -> Optional[str]:
_MODEL_FILES = (
"adapter_model.safetensors",
"adapter_model.bin",
"model.safetensors",
"pytorch_model.bin",
)
_MODEL_INDEXES = ("model.safetensors.index.json", "pytorch_model.bin.index.json")
def _valid_state_file(path: Path, require_tensor: bool = True) -> bool:
try:
if not path.is_file() or path.stat().st_size == 0:
return False
if path.suffix == ".safetensors":
try:
from safetensors import SafetensorError, safe_open
except ImportError:
return False
try:
with safe_open(str(path), framework = "np") as state:
return bool(state.keys())
except SafetensorError:
return False
if path.suffix in {".bin", ".pt"}:
with zipfile.ZipFile(path) as state:
infos = state.infolist()
names = [info.filename for info in infos]
data_name = next(
(name for name in names if name == "data.pkl" or name.endswith("/data.pkl")),
None,
)
if data_name is None:
return False
data_prefix = data_name.removesuffix("data.pkl") + "data/"
operations = list(pickletools.genops(state.read(data_name)))
if not operations or operations[-1][0].name != "STOP":
return False
if not require_tensor:
return True
# Require a non-empty tensor record; a zero-byte one fails torch.load.
return any(
info.filename.startswith(data_prefix)
and not info.is_dir()
and info.file_size > 0
for info in infos
)
# Unrecognized state-file formats are not usable resume state.
return False
except (OSError, ValueError, zipfile.BadZipFile):
return False
def _checkpoint_state(path: Path) -> Optional[int]:
try:
state = json.loads((path / "trainer_state.json").read_text(encoding = "utf-8"))
step = state.get("global_step") if isinstance(state, dict) else None
except (OSError, UnicodeDecodeError, json.JSONDecodeError):
return None
if isinstance(step, bool) or not isinstance(step, int) or step < 0:
return None
directory_step = _checkpoint_step(path)
return step if directory_step < 0 or step == directory_step else None
_INDEX_SHARD_SUFFIX = {
"model.safetensors.index.json": ".safetensors",
"pytorch_model.bin.index.json": ".bin",
}
def _valid_indexed_shard(checkpoint: Path, shard: object, expected_suffix: str) -> bool:
# Shard must be a relative, in-format path contained in the checkpoint dir.
if not isinstance(shard, str) or not shard:
return False
if Path(shard).is_absolute() or Path(shard).suffix != expected_suffix:
return False
try:
root = checkpoint.resolve(strict = True)
candidate = (checkpoint / shard).resolve(strict = True)
candidate.relative_to(root)
except (OSError, ValueError):
return False
return _valid_state_file(candidate)
def _has_model_state(path: Path) -> bool:
if any(_valid_state_file(path / name) for name in _MODEL_FILES):
return True
for name in _MODEL_INDEXES:
try:
index = json.loads((path / name).read_text(encoding = "utf-8"))
shards = set(index["weight_map"].values())
except (
AttributeError,
OSError,
KeyError,
TypeError,
UnicodeDecodeError,
json.JSONDecodeError,
):
continue
expected_suffix = _INDEX_SHARD_SUFFIX[name]
if shards and all(_valid_indexed_shard(path, shard, expected_suffix) for shard in shards):
return True
return False
def is_resume_checkpoint_valid(
path: Path,
expected_step: Optional[int] = None,
backend: Optional[str] = None,
) -> bool:
step = _checkpoint_state(path) if path.is_dir() else None
step_valid = step is not None and (expected_step is None or step == expected_step)
if backend == "mlx":
valid_bundle = _valid_state_file(path / "adapters.safetensors") and _valid_state_file(
path / "optimizer_state.safetensors"
)
else:
valid_bundle = (
_has_model_state(path)
# optimizer/scheduler state can be validly tensor-free (e.g. SGD without
# momentum); _has_model_state still requires real model tensors.
and _valid_state_file(path / "optimizer.pt", require_tensor = False)
and _valid_state_file(path / "scheduler.pt", require_tensor = False)
)
if backend is None and not valid_bundle:
valid_bundle = _valid_state_file(path / "adapters.safetensors") and _valid_state_file(
path / "optimizer_state.safetensors"
)
return step_valid and valid_bundle
def get_resume_checkpoint_path(
path_value: str, expected_step: Optional[int] = None
) -> Optional[str]:
path = resolve_output_dir(path_value)
if not _is_under_outputs(path) or not path.is_dir():
return None
if (path / "trainer_state.json").is_file():
if is_resume_checkpoint_valid(path, expected_step):
return str(path)
checkpoints = [
child
for child in path.glob("checkpoint-*")
if child.is_dir() and (child / "trainer_state.json").is_file()
]
if not checkpoints:
return None
return str(max(checkpoints, key = _checkpoint_step))
checkpoints = sorted(path.glob("checkpoint-*"), key = _checkpoint_step, reverse = True)
return next(
(
str(checkpoint)
for checkpoint in checkpoints
if _checkpoint_step(checkpoint) >= 0
and is_resume_checkpoint_valid(checkpoint, expected_step)
),
None,
)
def normalize_resume_output_dir(path_value: str) -> str:
@ -78,9 +217,17 @@ def _uses_s3_dataset(run: dict) -> bool:
def can_resume_run(run: dict) -> bool:
if run.get("resumed_later"):
return False
# Set when a stop-and-save failed to write a current-step checkpoint.
if run.get("resume_blocked"):
return False
if _uses_s3_dataset(run):
return False
status = run.get("status")
if status == "error":
# A save-time crash can report final_step == total_steps with no artifacts; checkpoint state alone decides resumability.
return has_resume_state(run.get("output_dir"))
final_step = run.get("final_step")
total_steps = run.get("total_steps")
has_remaining_steps = (
@ -89,8 +236,4 @@ def can_resume_run(run: dict) -> bool:
or total_steps <= 0
or final_step < total_steps
)
return (
run.get("status") == "stopped"
and has_remaining_steps
and has_resume_state(run.get("output_dir"))
)
return status == "stopped" and has_remaining_steps and has_resume_state(run.get("output_dir"))

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

@ -3425,15 +3425,19 @@ class UnslothTrainer:
logger.info(
f"CPT: using UnslothTrainer with embedding_learning_rate={embedding_lr}\n"
)
cpt_args = _UnslothTrainingArguments(
embedding_learning_rate = embedding_lr,
**config_args,
)
if config_args.get("packing", False):
cpt_args.packing_strategy = "wrapped"
logger.info("CPT packing strategy: wrapped\n")
trainer_kwargs = {
"model": self.model,
"tokenizer": sft_tokenizer,
"train_dataset": dataset["dataset"],
"data_collator": data_collator,
"args": _UnslothTrainingArguments(
embedding_learning_rate = embedding_lr,
**config_args,
),
"args": cpt_args,
}
if eval_dataset is not None:
trainer_kwargs["eval_dataset"] = eval_dataset

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

@ -761,6 +761,7 @@ class TrainingBackend:
# Left True after an abnormal death so _ensure_pump_alive spots a crash.
self._pump_running: bool = False
self._lock = threading.Lock()
self._run_intent_lock = threading.RLock()
# Stop watchdog: after a stop is requested, escalates to force_terminate()
# if the worker does not exit on its own within a bounded time. The watched
@ -773,6 +774,7 @@ class TrainingBackend:
self._progress = TrainingProgress()
self._should_stop = False
self._cancel_requested = False # True only for stop(save=False)
self._cancel_cleanup_output_dir: Optional[str] = None
# Throttled training-status logging to the server log (not one line/step).
self._last_progress_log_ts: float = 0.0
@ -792,6 +794,8 @@ class TrainingBackend:
# Job metadata
self.current_job_id: Optional[str] = None
self._output_dir: Optional[str] = None
self._resume_source_run_id: Optional[str] = None
self._terminal_finalize_payload: Optional[dict] = None
# DB persistence
self._metric_buffer: list[dict] = []
@ -819,6 +823,7 @@ class TrainingBackend:
job_id: str,
*,
before_spawn = None,
resume_source_run_id: Optional[str] = None,
**kwargs,
) -> bool:
"""Spawn a subprocess to run the full training pipeline.
@ -956,6 +961,7 @@ class TrainingBackend:
self.current_job_id = job_id
self._should_stop = False
self._cancel_requested = False
self._cancel_cleanup_output_dir = None
self._complete_seen.clear()
self._progress = TrainingProgress(
is_training = True, status_message = "Initializing training..."
@ -972,7 +978,10 @@ class TrainingBackend:
self.eval_loss_history.clear()
self.eval_step_history.clear()
self.eval_enabled = False
self._output_dir = None
self._output_dir = config.get("output_dir") if resume_source_run_id else None
self._progress.output_dir = self._output_dir
self._resume_source_run_id = resume_source_run_id
self._terminal_finalize_payload = None
self._metric_buffer.clear()
self._run_finalized = False
self._db_run_created = False
@ -990,6 +999,17 @@ class TrainingBackend:
# 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()
if resume_source_run_id and not self._db_run_created:
if proc.is_alive():
proc.terminate()
proc.join(timeout = 5.0)
if proc.is_alive():
proc.kill()
proc.join(timeout = 2.0)
self._progress.is_training = False
self._progress.error = "Resume checkpoint is no longer available."
self._spawn_in_progress = False
return False
# 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.
@ -1011,28 +1031,75 @@ class TrainingBackend:
def stop_training(self, save: bool = True) -> bool:
"""Send stop signal to the training subprocess."""
self._should_stop = True
if not save:
self._cancel_requested = True
with self._lock:
if self._stop_queue is not None:
try:
self._stop_queue.put({"type": "stop", "save": save})
except (OSError, ValueError):
pass
# Update progress immediately for responsive UI.
self._progress.status_message = (
"Stopping training and saving checkpoint..." if save else "Cancelling training..."
)
# Guarantee the run finalizes even if the worker wedges after saving.
self._start_stop_watchdog(cancel = not save)
with self._run_intent_lock:
with self._lock:
run_id = self.current_job_id
if not save and run_id:
persist_error: Optional[Exception] = None
for attempt in range(_DB_FINALIZE_RETRIES):
try:
from storage.studio_db import mark_run_cancel_requested
self._ensure_db_run_created()
with self._lock:
terminal_payload = self._terminal_finalize_payload
if (
terminal_payload
and terminal_payload.get("expected_job_id") == run_id
):
return False
if not mark_run_cancel_requested(run_id):
if self._db_run_created:
return False
raise RuntimeError(
"Training run disappeared before cancellation persisted"
)
if self.current_job_id != run_id:
return False
self._should_stop = self._cancel_requested = True
self._cancel_cleanup_output_dir = self._output_dir
self._output_dir = self._progress.output_dir = None
persist_error = None
break
except Exception as exc:
persist_error = exc
if attempt + 1 < _DB_FINALIZE_RETRIES:
time.sleep(_DB_FINALIZE_RETRY_S)
if persist_error is not None:
raise RuntimeError("Failed to persist Stop-without-Save") from persist_error
with self._lock:
if self.current_job_id != run_id:
return False
if save or not run_id:
self._should_stop = True
if not save and not run_id:
self._cancel_requested = True
self._cancel_cleanup_output_dir = self._output_dir
self._output_dir = self._progress.output_dir = None
if self._stop_queue is not None:
try:
self._stop_queue.put({"type": "stop", "save": save})
except (OSError, ValueError):
pass
self._progress.status_message = (
"Stopping training and saving checkpoint..."
if save
else "Cancelling training..."
)
self._start_stop_watchdog(cancel = not save, expected_job_id = run_id)
return True
def _start_stop_watchdog(self, cancel: bool) -> None:
def _start_stop_watchdog(
self,
cancel: bool,
expected_job_id: Optional[str] = None,
) -> None:
"""Start a daemon that force-terminates the worker if a requested stop does not
exit on its own. No-op if no worker is alive or a live watchdog already watches
this proc (a stale watchdog on an old proc never blocks a new run's watcher)."""
with self._lock:
if expected_job_id is not None and self.current_job_id != expected_job_id:
return
proc = self._proc
if proc is None or not proc.is_alive():
return
@ -1113,8 +1180,9 @@ class TrainingBackend:
watched_job_id: Optional[str] = None,
) -> None:
"""Finalize parent state after a force-terminate so the UI leaves "Stopping..."
even if the worker is wedged in driver teardown; preserves output_dir so a saved
checkpoint is kept. No-ops if a new run already replaced the watched worker, so a
even if the worker is wedged in driver teardown; preserves output_dir on a save so
the checkpoint is kept, and clears it on a cancel (Stop without saving must not
offer resume/export). No-ops if a new run already replaced the watched worker, so a
stale watchdog never marks a fresh run stopped or drops its handle.
Supersession is checked on both the watched proc and job id: start_training sets
@ -1134,7 +1202,18 @@ class TrainingBackend:
return # a new run is already starting up; leave its state alone
run_id = self.current_job_id # == watched_job_id
self._progress.is_training = False
self._progress.status_message = "Training stopped."
terminal_payload = self._terminal_finalize_kwargs()
status = terminal_payload["status"]
error_message = terminal_payload.get("error_message")
output_dir = terminal_payload["output_dir"]
clear_output_dir = terminal_payload["clear_output_dir"]
resume_blocked = bool(terminal_payload.get("resume_blocked"))
with self._lock:
if self.current_job_id != run_id:
return
self._progress.status_message = error_message or "Training stopped."
if error_message:
self._progress.error = error_message
# Create the row if a start-time create failed (no-op otherwise; skips when the pump
# is mid-create, in which case its create-then-finalize records the run instead).
self._ensure_db_run_created()
@ -1148,7 +1227,8 @@ class TrainingBackend:
batch: list = []
final_step = final_loss = duration = None
loss_history: list = []
output_dir = self._output_dir
if clear_output_dir:
self._output_dir = self._progress.output_dir = None
if claim:
self._run_finalized = True # claim this run's finalize
batch = list(self._metric_buffer)
@ -1161,7 +1241,17 @@ class TrainingBackend:
loss_history = list(self.loss_history)
if claim:
self._finish_stopped_run(
run_id, output_dir, batch, final_step, final_loss, duration, loss_history
run_id,
output_dir,
batch,
final_step,
final_loss,
duration,
loss_history,
status = status,
error_message = error_message,
clear_output_dir = clear_output_dir,
resume_blocked = resume_blocked,
)
with self._lock:
if target_proc is None or self._proc is target_proc:
@ -1176,6 +1266,10 @@ class TrainingBackend:
final_loss: Optional[float],
duration: Optional[float],
loss_history: list,
status: str = "stopped",
error_message: Optional[str] = None,
clear_output_dir: bool = False,
resume_blocked: bool = False,
) -> None:
"""Record a force-stopped run finished by its captured id, from state snapshotted
under the lock. insert_metrics_batch upserts and finish_run is an idempotent UPDATE,
@ -1194,14 +1288,16 @@ class TrainingBackend:
sparkline = downsample(loss_history, 50)
finish_run(
id = run_id,
status = "stopped",
status = status,
ended_at = datetime.now(timezone.utc).isoformat(),
final_step = final_step,
final_loss = final_loss,
duration_seconds = duration,
loss_sparkline = _json.dumps(sparkline),
output_dir = output_dir,
error_message = None,
error_message = error_message,
clear_output_dir = clear_output_dir,
resume_blocked = resume_blocked,
)
return
except Exception:
@ -1231,7 +1327,7 @@ class TrainingBackend:
logger.info("Force-terminating training subprocess (pid=%s)", proc.pid)
proc.terminate()
cancelled = self._cancel_requested
output_dir = self._output_dir
output_dir = self._cancel_cleanup_output_dir or self._output_dir
if proc is not None:
proc.join(timeout = 5.0)
@ -1595,17 +1691,60 @@ class TrainingBackend:
)
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",
)
terminal_payload = self._terminal_finalize_kwargs()
with self._lock:
if terminal_payload["clear_output_dir"]:
self._output_dir = self._progress.output_dir = None
if terminal_payload.get("error_message"):
self._progress.error = terminal_payload["error_message"]
self._progress.status_message = terminal_payload["error_message"]
self._finalize_run_in_db(**terminal_payload)
except Exception:
logger.exception("Training event pump: finalization after worker exit failed")
self._pump_running = False
return
def _has_current_resume_checkpoint(self, output_dir, step) -> bool:
# A valid checkpoint at the current step means the stop-and-save landed on
# disk even if the worker died before confirming it.
if not output_dir or not isinstance(step, int) or step <= 0:
return False
from core.training.resume import get_resume_checkpoint_path
return get_resume_checkpoint_path(output_dir, expected_step = step) is not None
def _terminal_finalize_kwargs(self) -> dict:
with self._lock:
job_id = self.current_job_id
payload = self._terminal_finalize_payload
if payload and payload.get("expected_job_id") == job_id:
return dict(payload)
cancel, stopped = self._cancel_requested, self._should_stop
output_dir = None if cancel else self._output_dir
step = self._progress.step
existing_error = self._progress.error
status, error, blocked = (
("stopped", None, cancel)
if stopped
else (
"error",
existing_error or "Training process terminated unexpectedly",
False,
)
)
# Block only when no valid current-step checkpoint actually landed.
if stopped and not cancel and not self._has_current_resume_checkpoint(output_dir, step):
status = "error"
error = "Stop and Save ended before a valid current-step checkpoint was written."
blocked = True
return {
"status": status,
"error_message": error,
"output_dir": output_dir,
"clear_output_dir": cancel,
"resume_blocked": blocked,
"expected_job_id": job_id,
}
def _handle_event(self, event: dict) -> None:
"""Apply a subprocess event to local state.
@ -1764,6 +1903,15 @@ class TrainingBackend:
elif etype == "eval_configured":
self.eval_enabled = True
elif etype == "output_dir":
event_output_dir = event.get("output_dir")
if self._cancel_requested:
self._cancel_cleanup_output_dir = event_output_dir
self._output_dir = self._progress.output_dir = None
else:
self._output_dir = event_output_dir
db_action = "persist_output_dir"
elif etype == "status":
self._progress.status_message = event.get("message", "")
self._progress.is_training = True
@ -1778,7 +1926,12 @@ class TrainingBackend:
self._complete_seen.set()
self._progress.is_training = False
self._progress.is_completed = not stopped
self._output_dir = event.get("output_dir")
event_output_dir = event.get("output_dir")
if self._cancel_requested:
self._cancel_cleanup_output_dir = event_output_dir
self._output_dir = None
else:
self._output_dir = event_output_dir
self._progress.output_dir = self._output_dir
self._progress.status_message = msg
if not self._db_run_created and self.current_job_id and self._db_config:
@ -1788,11 +1941,16 @@ class TrainingBackend:
db_action_kwargs = {
"status": "stopped" if stopped else "completed",
"output_dir": self._output_dir,
"clear_output_dir": self._cancel_requested,
"expected_job_id": self.current_job_id,
}
self._terminal_finalize_payload = dict(db_action_kwargs)
elif etype == "error":
self._progress.is_training = False
self._progress.error = event.get("error", "Unknown error")
if self._cancel_requested:
self._output_dir = self._progress.output_dir = None
logger.error("Training error: %s", event.get("error"))
stack = event.get("stack", "")
if stack:
@ -1801,29 +1959,36 @@ class TrainingBackend:
db_action = "create_and_finalize"
else:
db_action = "finalize"
stop_save_failed = (
self._should_stop
and not self._cancel_requested
and not self._has_current_resume_checkpoint(
self._output_dir, self._progress.step
)
)
db_action_kwargs = {
"status": "stopped" if self._should_stop else "error",
"status": "stopped"
if self._should_stop
and not stop_save_failed
and not event.get("keep_error_status")
else "error",
"error_message": event.get("error", "Unknown error"),
"output_dir": self._output_dir,
"clear_output_dir": self._cancel_requested,
"resume_blocked": stop_save_failed or bool(event.get("resume_blocked")),
"expected_job_id": self.current_job_id,
}
self._terminal_finalize_payload = dict(db_action_kwargs)
# --- DB I/O outside the lock ---
if db_action == "create_run":
try:
from storage.studio_db import create_run
create_run(
id = db_action_kwargs["job_id"],
model_name = db_action_kwargs["model_name"],
dataset_name = db_action_kwargs["dataset_name"],
config_json = db_action_kwargs["config_json"],
started_at = db_action_kwargs["started_at"],
total_steps = db_action_kwargs["total_steps"],
)
self._db_run_created = True
self._ensure_db_run_created()
if self._db_run_created:
if db_action_kwargs["total_steps"]:
self._db_total_steps_set = True
except Exception:
logger.warning("Failed to create DB run record", exc_info = True)
self._persist_output_dir()
elif db_action == "persist_output_dir":
self._persist_output_dir()
elif db_action == "create_and_finalize":
self._ensure_db_run_created()
self._finalize_run_in_db(**db_action_kwargs)
@ -1842,6 +2007,22 @@ class TrainingBackend:
if etype == "progress":
self._log_training_progress()
def _persist_output_dir(self) -> None:
with self._lock:
if (
not self._output_dir
or not self.current_job_id
or not self._db_run_created
or self._cancel_requested
):
return
run_id, output_dir = self.current_job_id, self._output_dir
try:
from storage.studio_db import update_run_output_dir
update_run_output_dir(run_id, output_dir)
except Exception:
logger.warning("Failed to persist output_dir", exc_info = True)
def _log_training_progress(self) -> None:
"""One throttled training-status line to the server log (the per-step stream
still goes to the UI via SSE): first step, then at most every 30s, plus the
@ -1875,6 +2056,7 @@ class TrainingBackend:
caller create at a time, and ``_db_run_created`` is published only after
``create_run`` commits, so a concurrent finalize never runs ``finish_run`` against a
not-yet-inserted row (a zero-row UPDATE that would leave the run stuck as running)."""
self._run_intent_lock.acquire()
with self._lock:
if (
self._db_run_created
@ -1882,6 +2064,7 @@ class TrainingBackend:
or not self.current_job_id
or not self._db_config
):
self._run_intent_lock.release()
return
self._db_create_in_progress = True # only one caller creates
job_id = self.current_job_id
@ -1898,6 +2081,12 @@ class TrainingBackend:
or _s3_dataset_name(db_config.get("s3_dataset"))
or "unknown"
)
with self._lock:
if self.current_job_id != job_id:
return
output_dir = self._output_dir
cancel_requested = self._cancel_requested
resumed_from_run_id = self._resume_source_run_id
create_run(
id = job_id,
model_name = db_config["model_name"],
@ -1905,6 +2094,9 @@ class TrainingBackend:
config_json = _json.dumps(db_config),
started_at = started_at,
total_steps = total_steps,
output_dir = output_dir,
cancel_requested = cancel_requested,
resumed_from_run_id = resumed_from_run_id,
)
created = True
except Exception:
@ -1919,12 +2111,15 @@ class TrainingBackend:
if created:
self._db_run_created = True # publish only after the insert commits
self._db_create_in_progress = False
self._run_intent_lock.release()
def _finalize_run_in_db(
self,
status: str,
error_message: Optional[str] = None,
output_dir: Optional[str] = None,
clear_output_dir: bool = False,
resume_blocked: bool = False,
expected_job_id: Optional[str] = None,
) -> None:
"""Flush remaining metrics and mark a run finished in the DB. Claims the finalize
@ -1947,26 +2142,33 @@ class TrainingBackend:
duration = self._progress.elapsed_seconds
loss_history = list(self.loss_history)
self._flush_metrics_to_db(run_id = run_id)
try:
from storage.studio_db import finish_run
from utils.downsample import downsample
for attempt in range(_DB_FINALIZE_RETRIES):
try:
from storage.studio_db import finish_run
from utils.downsample import downsample
sparkline = downsample(loss_history, 50)
finish_run(
id = run_id,
status = status,
ended_at = datetime.now(timezone.utc).isoformat(),
final_step = final_step,
final_loss = final_loss,
duration_seconds = duration,
loss_sparkline = _json.dumps(sparkline),
output_dir = output_dir,
error_message = error_message,
)
except Exception:
with self._lock:
self._run_finalized = False # unclaim so a later flush can retry
logger.warning("Failed to finalize run in DB (status=%s)", status, exc_info = True)
finish_run(
id = run_id,
status = status,
ended_at = datetime.now(timezone.utc).isoformat(),
final_step = final_step,
final_loss = final_loss,
duration_seconds = duration,
loss_sparkline = _json.dumps(downsample(loss_history, 50)),
output_dir = output_dir,
error_message = error_message,
clear_output_dir = clear_output_dir,
resume_blocked = resume_blocked,
)
return
except Exception:
if attempt + 1 < _DB_FINALIZE_RETRIES:
time.sleep(_DB_FINALIZE_RETRY_S)
continue
with self._lock:
if self.current_job_id == run_id:
self._run_finalized = False
logger.warning("Failed to finalize run in DB (status=%s)", status, exc_info = True)
def _flush_metrics_to_db(self, run_id: Optional[str] = None) -> None:
"""Flush buffered metrics to the DB and update live progress. The target run id,

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

@ -1840,8 +1840,15 @@ def _run_mlx_training(event_queue, stop_queue, config):
# Resolve to ~/.unsloth/studio/outputs/ so the export page finds it
from utils.paths import ensure_dir
output_dir = _resolve_mlx_output_dir(config, model_name)
# Resume must land in the original run dir even when config lacks output_dir.
resume_dir = config.get("output_dir", "") or _output_dir_from_resume_checkpoint(
resume_from_checkpoint
)
output_dir = _resolve_mlx_output_dir(
{**config, "output_dir": resume_dir} if resume_dir else config, model_name
)
ensure_dir(Path(output_dir))
_emit_output_dir(event_queue, output_dir)
# ── 6. Create trainer ──
eval_steps_val = config.get("eval_steps", 0) or 0
@ -2067,6 +2074,17 @@ def _run_mlx_training(event_queue, stop_queue, config):
trainer.add_eval_callback(_on_eval)
_opt_ref = [None]
_orig_build_optimizer = getattr(trainer, "_build_optimizer", None)
if callable(_orig_build_optimizer):
def _capture_optimizer(total_steps):
_opt_ref[0] = _orig_build_optimizer(total_steps)
return _opt_ref[0]
trainer._build_optimizer = _capture_optimizer
# ── 11. Run training ──
gc.collect()
mx.synchronize()
@ -2082,31 +2100,58 @@ def _run_mlx_training(event_queue, stop_queue, config):
trainer.save_model = _save_model
# ── 12. Save and finalize ──
if trainer.stop_requested:
if not _stop_save[0]:
# Cancel (save=False): skip saving.
_send("complete", output_dir = None, status_message = "Training cancelled")
def _finish_tracking() -> None:
# Runs on every save/finalize exit so TB/W&B never leak on early return.
if tb_writer is not None:
try:
tb_writer.close()
except Exception:
pass
if wandb_run is not None:
try:
wandb_run.finish()
except Exception:
pass
def _stop_checkpoint_ok() -> bool:
if _write_mlx_stop_checkpoint(trainer, _opt_ref[0], output_dir):
return True
_send(
"error",
error = (
"Failed to save a resumable checkpoint after stop. "
"Model files were saved, but this run cannot be resumed."
),
# A user stop finalizes as 'stopped'; keep this failure's error status so history explains it.
keep_error_status = True,
# Older checkpoints are stale; resuming would roll back past this stop.
resume_blocked = True,
)
return False
try:
if trainer.stop_requested:
if not _stop_save[0]:
# Cancel (save=False): skip saving.
_send("complete", output_dir = None, status_message = "Training cancelled")
else:
_send("status", status_message = "Saving stopped model...")
mx.synchronize()
trainer.save_model(output_dir)
# Stop-and-save promises a resumable checkpoint, not just model files.
if not _stop_checkpoint_ok():
return
_send("complete", output_dir = output_dir, status_message = "Training stopped")
else:
_send("status", status_message = "Saving stopped model...")
_send("status", status_message = "Saving model...")
mx.synchronize()
trainer.save_model(output_dir)
_send("complete", output_dir = output_dir, status_message = "Training stopped")
else:
_send("status", status_message = "Saving model...")
mx.synchronize()
trainer.save_model(output_dir)
_send("complete", output_dir = output_dir, status_message = "Training completed")
if tb_writer is not None:
try:
tb_writer.close()
except Exception:
pass
if wandb_run is not None:
try:
wandb_run.finish()
except Exception:
pass
# A save-stop can race the natural final save; it made the same promise.
if trainer.stop_requested and _stop_save[0] and not _stop_checkpoint_ok():
return
_send("complete", output_dir = output_dir, status_message = "Training completed")
finally:
_finish_tracking()
def _is_current_process_apple_silicon() -> bool:
@ -3177,6 +3222,7 @@ def run_training_process(*, event_queue: Any, stop_queue: Any, config: dict) ->
)
output_dir = str(resolve_output_dir(output_dir))
ensure_dir(Path(output_dir))
_emit_output_dir(event_queue, output_dir)
tensorboard_dir = config.get("tensorboard_dir")
if config.get("enable_tensorboard", False):
@ -3296,6 +3342,61 @@ def _send_status(event_queue: Any, message: str) -> None:
)
def _emit_output_dir(event_queue: Any, output_dir: str) -> None:
try:
event_queue.put({"type": "output_dir", "output_dir": output_dir, "ts": time.time()})
except Exception:
pass
def _mlx_has_checkpoint_at_step(output_dir, step: int) -> bool:
if step <= 0:
return False
from core.training.resume import is_resume_checkpoint_valid
return is_resume_checkpoint_valid(
Path(output_dir) / f"checkpoint-{step}", expected_step = step, backend = "mlx"
)
def _write_mlx_stop_checkpoint(trainer, optimizer, output_dir) -> bool:
"""Write a full resume checkpoint for a stopped MLX run.
Returns True when a checkpoint for the current training step exists.
"""
step = int(getattr(trainer, "_global_step", 0) or 0)
# A periodic save or a resumed run may already cover the current step.
if _mlx_has_checkpoint_at_step(output_dir, step):
return True
if step <= 0 or optimizer is None:
return False
ckpt_dir = Path(output_dir) / f"checkpoint-{step}"
if ckpt_dir.is_symlink():
# Refuse a symlinked dir: it could redirect writes outside output_dir.
logger.error("Refusing to write MLX stop checkpoint through symlink: %s", ckpt_dir)
return False
try:
ckpt_dir.mkdir(parents = True, exist_ok = True)
from unsloth_zoo.mlx.utils import (
save_optimizer_state,
save_trainable_adapters,
save_trainer_state,
)
save_trainable_adapters(trainer.model, str(ckpt_dir))
save_optimizer_state(optimizer, str(ckpt_dir))
save_trainer_state(
{
"global_step": step,
"train_loss_history": list(getattr(trainer, "_train_loss_history", [])),
},
str(ckpt_dir),
)
logger.info("Saved stop checkpoint to %s", ckpt_dir)
except Exception:
logger.exception("Failed to write stop checkpoint under %s", output_dir)
return _mlx_has_checkpoint_at_step(output_dir, step)
def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) -> None:
"""Self-contained embedding model training pipeline.
@ -3660,6 +3761,7 @@ def _run_embedding_training(event_queue: Any, stop_queue: Any, config: dict) ->
config.get("project_name"),
)
output_dir = str(resolve_output_dir(output_dir))
_emit_output_dir(event_queue, output_dir)
num_epochs = config.get("num_epochs", 2)
batch_size = config.get("batch_size", 256)

2
studio/backend/hub/routes/__init__.py
View file

@ -5,8 +5,10 @@
from hub.routes.inventory import router as inventory_router
from hub.routes.datasets import router as datasets_router
from hub.routes.token import router as token_router
__all__ = [
"inventory_router",
"datasets_router",
"token_router",
]

11
studio/backend/hub/routes/inventory.py
View file

@ -28,6 +28,7 @@ from hub.schemas.inventory import (
CachedModelsResponse,
DeleteCachedModelResponse,
GgufVariantsResponse,
HiddenModelsResponse,
LocalModelListResponse,
ModelsFolderResponse,
RecommendedFoldersResponse,
@ -214,6 +215,16 @@ async def list_cached_models(
return await cache_inventory.list_cached_models_response(hf_token)
@router.get("/hidden-models", response_model = HiddenModelsResponse)
async def list_hidden_models(current_subject: str = Depends(get_current_subject)):
import asyncio
from routes.models import hidden_model_matchers
needles, exact_ids, exact_paths = await asyncio.to_thread(hidden_model_matchers)
return HiddenModelsResponse(needles = needles, exact_ids = exact_ids, exact_paths = exact_paths)
@router.delete(
"/delete-cached",
response_model = DeleteCachedModelResponse,

44
studio/backend/hub/routes/token.py Normal file
View file

@ -0,0 +1,44 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Hugging Face token validation endpoint."""
from __future__ import annotations
import asyncio
from typing import Literal, Optional
from fastapi import APIRouter, Depends, Request
from pydantic import BaseModel
from auth.authentication import get_current_subject
from hub.dependencies import get_hf_token
from utils.client_ip import client_ip
from utils.hf_token_validation import validate_hf_token
router = APIRouter()
class HfTokenValidationResponse(BaseModel):
status: Literal["missing", "valid", "invalid", "rate_limited", "unavailable"]
retry_after_seconds: Optional[int] = None
@router.post("/token/validate", response_model = HfTokenValidationResponse)
async def validate_token(
request: Request,
hf_token: Optional[str] = Depends(get_hf_token),
current_subject: str = Depends(get_current_subject),
):
if not hf_token:
return HfTokenValidationResponse(status = "missing")
result = await asyncio.to_thread(
validate_hf_token,
hf_token,
rate_key = f"{current_subject}:{client_ip(request)}",
)
return HfTokenValidationResponse(
status = result.status,
retry_after_seconds = result.retry_after_seconds,
)

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

@ -160,6 +160,7 @@ class CachedRepoBase(BaseModel):
repo_id: str
size_bytes: int = 0
cache_path: Optional[str] = None
last_modified: Optional[float] = None
partial: bool = False
partial_transport: Optional[str] = None
inventory_id: Optional[str] = None
@ -189,6 +190,12 @@ class CachedModelsResponse(BaseModel):
cached: List[CachedModelRepo] = Field(default_factory = list)
class HiddenModelsResponse(BaseModel):
needles: List[str] = Field(default_factory = list)
exact_ids: List[str] = Field(default_factory = list)
exact_paths: List[str] = Field(default_factory = list)
class AddScanFolderRequest(BaseModel):
"""Request body for adding a custom scan folder."""

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

@ -31,6 +31,7 @@ from hub.services.models.common import (
_is_checkpoint_weight_name,
_is_gguf_filename,
_is_main_gguf_filename,
_is_mmproj_filename,
_is_transformers_safetensors_weight_name,
_local_inventory_id,
_prefer_complete_larger,
@ -132,6 +133,39 @@ def _repo_has_gguf_files(repo_info) -> bool:
return _repo_gguf_size_bytes(repo_info) > 0
def _blob_mtime(file_obj) -> float:
ts = getattr(file_obj, "blob_last_modified", None)
if isinstance(ts, (int, float)) and ts > 0:
return float(ts)
blob_path = getattr(file_obj, "blob_path", None)
if blob_path:
try:
return float(Path(blob_path).stat().st_mtime)
except OSError:
pass
return 0.0
def _repo_gguf_last_modified(repo_info) -> float:
latest = 0.0
for revision in repo_info.revisions:
for f in revision.files:
if _is_main_gguf_filename(f.file_name):
latest = max(latest, _blob_mtime(f))
return latest
def _repo_has_mmproj(repo_info) -> bool:
# An mmproj file only makes a repo vision-capable when it is an actual GGUF
# projector; a non-GGUF sidecar (e.g. mmproj_config.json) does not, and the
# runtime's projector detection is GGUF-only.
return any(
_is_gguf_filename(f.file_name) and _is_mmproj_filename(f.file_name)
for revision in repo_info.revisions
for f in revision.files
)
def _cached_repo_file_name(file_obj) -> str:
file_path = getattr(file_obj, "file_path", None)
if file_path:
@ -291,6 +325,7 @@ def _scan_cached_gguf() -> list[dict]:
continue
key = repo_id.lower()
existing = seen_lower.get(key)
last_modified = _repo_gguf_last_modified(repo_info)
row = {
"repo_id": repo_id,
"size_bytes": max(total_size, variant_state_size),
@ -300,6 +335,9 @@ def _scan_cached_gguf() -> list[dict]:
# per-variant detail lives on GgufVariantDetail.
"partial_transport": None,
}
last_modified = max(last_modified, (existing or {}).get("last_modified", 0.0))
if last_modified > 0:
row["last_modified"] = last_modified
row.update(
_cache_inventory_fields(
repo_id,
@ -308,11 +346,20 @@ def _scan_cached_gguf() -> list[dict]:
requires_variant = True,
)
)
if _repo_has_mmproj(repo_info):
row["capabilities"]["supports_vision"] = True
# Visible infra variants remain management-only.
if is_hidden_infra:
row["capabilities"]["can_chat"] = False
if _prefer_cache_row(row, existing):
if existing and existing["capabilities"].get("supports_vision"):
row["capabilities"]["supports_vision"] = True
seen_lower[key] = row
else:
if last_modified > existing.get("last_modified", 0.0):
existing["last_modified"] = last_modified
if row["capabilities"].get("supports_vision"):
existing["capabilities"]["supports_vision"] = True
except Exception as e:
repo_label = getattr(repo_info, "repo_id", "<unknown>")
logger.warning(f"Skipping cached GGUF repo {repo_label}: {e}")
@ -340,13 +387,14 @@ class _CachedNonGgufPayload(NamedTuple):
size_bytes: int
has_runnable_weights: bool
model_format: ModelFormat
last_modified: float
def _repo_non_gguf_model_payload(repo_info) -> _CachedNonGgufPayload:
all_weight_blobs: dict[str, int] = {}
adapter_blobs: dict[str, int] = {}
safetensors_blobs: dict[str, int] = {}
checkpoint_blobs: dict[str, int] = {}
all_weight_blobs: dict[str, tuple[int, float]] = {}
adapter_blobs: dict[str, tuple[int, float]] = {}
safetensors_blobs: dict[str, tuple[int, float]] = {}
checkpoint_blobs: dict[str, tuple[int, float]] = {}
has_config = False
has_adapter_config = False
has_adapter_weights = False
@ -354,12 +402,15 @@ def _repo_non_gguf_model_payload(repo_info) -> _CachedNonGgufPayload:
has_transformers_safetensors = False
has_checkpoint = False
def _record_blob(target: dict[str, int], file_obj, rev_id: str, file_name: str) -> None:
def _record_blob(
target: dict[str, tuple[int, float]], file_obj, rev_id: str, file_name: str
) -> None:
blob_path = getattr(file_obj, "blob_path", None)
size = int(file_obj.size_on_disk or 0)
key = str(blob_path) if blob_path else f"{rev_id}:{file_name}"
target[key] = size
all_weight_blobs[key] = size
value = (size, _blob_mtime(file_obj))
target[key] = value
all_weight_blobs[key] = value
for revision in repo_info.revisions:
rev_id = getattr(revision, "commit_hash", None) or str(id(revision))
@ -403,18 +454,19 @@ def _repo_non_gguf_model_payload(repo_info) -> _CachedNonGgufPayload:
or "unknown"
)
if model_format == "adapter":
size_bytes = sum(adapter_blobs.values())
selected_blobs = adapter_blobs
elif model_format == "safetensors":
size_bytes = sum(safetensors_blobs.values())
selected_blobs = safetensors_blobs
elif model_format == "checkpoint":
size_bytes = sum(checkpoint_blobs.values())
selected_blobs = checkpoint_blobs
else:
size_bytes = sum(all_weight_blobs.values())
selected_blobs = all_weight_blobs
return _CachedNonGgufPayload(
size_bytes = size_bytes,
size_bytes = sum(size for size, _mtime in selected_blobs.values()),
has_runnable_weights = model_format != "unknown",
model_format = model_format,
last_modified = max((mtime for _size, mtime in selected_blobs.values()), default = 0.0),
)
@ -544,6 +596,12 @@ def _scan_cached_models() -> list[dict]:
),
**_cached_model_local_metadata(repo_path),
}
last_modified = max(
payload.last_modified,
(existing or {}).get("last_modified", 0.0),
)
if last_modified > 0:
row["last_modified"] = last_modified
row.update(
_cache_inventory_fields(
repo_id,
@ -553,6 +611,8 @@ def _scan_cached_models() -> list[dict]:
)
if _prefer_cache_row(row, existing):
seen_lower[key] = row
elif last_modified > existing.get("last_modified", 0.0):
existing["last_modified"] = last_modified
except Exception as e:
repo_label = getattr(repo_info, "repo_id", "<unknown>")
logger.warning(f"Skipping cached model repo {repo_label}: {e}")

66
studio/backend/main.py
View file

@ -289,6 +289,7 @@ from fastapi import Depends, FastAPI, HTTPException, Query, Request
from fastapi.middleware.cors import CORSMiddleware
from fastapi.staticfiles import StaticFiles
from fastapi.responses import FileResponse, HTMLResponse, Response
from starlette.middleware.gzip import GZipMiddleware
from pathlib import Path
from datetime import datetime
@ -312,7 +313,9 @@ from routes.preview import router as preview_router
from hub.routes import (
inventory_router as hub_inventory_router,
datasets_router as hub_datasets_router,
token_router as hub_token_router,
)
from picker.routes import templates_router as picker_templates_router
from hub.schemas.downloads import TransportCapabilities
from hub.utils.download_registry import (
get_download_transport_capabilities,
@ -547,8 +550,9 @@ async def lifespan(app: FastAPI):
threading.Thread(target = _warm_rag_embedder, daemon = True, name = "rag-embedder-warm").start()
# Idle auto-unload loop (no-op unless the OpenAI auto-unload TTL is set).
from core.inference.llama_keepwarm import idle_unload_loop
from core.inference.llama_keepwarm import idle_unload_loop, sweep_slot_save_dir
sweep_slot_save_dir()
app.state.idle_unload_task = asyncio.create_task(idle_unload_loop())
# Initialize RSA key pair for API key encryption (external providers).
@ -761,6 +765,7 @@ _BODY_PROTECTED_PREFIXES = (
"/v1/completions",
"/p/",
"/api/inference",
"/api/picker",
"/api/data-recipe",
"/api/datasets",
"/api/hub",
@ -992,6 +997,8 @@ app.include_router(rag_router, prefix = "/api/rag", tags = ["rag"])
app.include_router(training_history_router, prefix = "/api/train", tags = ["training-history"])
app.include_router(hub_inventory_router, prefix = "/api/hub", tags = ["hub"])
app.include_router(hub_datasets_router, prefix = "/api/hub/datasets", tags = ["hub"])
app.include_router(picker_templates_router, prefix = "/api/picker", tags = ["picker"])
app.include_router(hub_token_router, prefix = "/api/hub", tags = ["hub"])
# Re-wrap client-error responses on the /v1/* surface into OpenAI/Anthropic
# error envelopes; non-/v1 paths keep FastAPI's default {"detail": ...} shape.
@ -1148,17 +1155,35 @@ def _get_cached_system_gpu_info(logger) -> dict[str, Any]:
util = util_devices.get(idx, {})
total_vram = util.get("vram_total_gb") or dev.get("memory_total_gb") or 0
used_vram = util.get("vram_used_gb") or 0
# Keep None (usage unknown, e.g. Windows ROCm perf counter) so the UI
# shows unknown, not a fabricated 0 used / full free.
used_vram = util.get("vram_used_gb")
enriched_dev = dict(dev)
enriched_dev["vram_used_gb"] = used_vram
enriched_dev["vram_free_gb"] = round(total_vram - used_vram, 2) if total_vram else 0
enriched_dev["vram_free_gb"] = (
round(total_vram - used_vram, 2) if total_vram and used_vram is not None else None
)
enriched_dev["vram_utilization_pct"] = util.get("vram_utilization_pct")
enriched_devices.append(enriched_dev)
# Whether GGUF loads accept an explicit gpu_ids pick: /load and
# /validate 400 picks on XPU hosts (no visibility mask speaks torch-xpu
# ordinals) and on Vulkan-only builds (--device pins ggml's own
# ordinals), so the picker must not offer them.
try:
from core.inference.llama_cpp import LlamaCppBackend
from utils.hardware import DeviceType, get_device
gpu_ids_supported = (
get_device() != DeviceType.XPU and not LlamaCppBackend._is_vulkan_backend()
)
except Exception as e:
logger.debug(f"Could not resolve gpu_ids support: {e}")
gpu_ids_supported = True
gpu_info = {
"available": visibility_info.get("available", False),
"devices": enriched_devices,
"gguf_gpu_ids_supported": gpu_ids_supported,
}
_system_gpu_cache = (time.monotonic(), gpu_info)
return gpu_info
@ -1488,6 +1513,34 @@ def _should_inject_bootstrap(request: Request) -> bool:
return _is_local_bootstrap_request(request)
_IMMUTABLE_ASSET_CACHE_CONTROL = "public, max-age=31536000, immutable"
class ImmutableStaticFiles(StaticFiles):
"""Serve Vite's content-hashed assets without browser revalidation."""
def file_response(
self,
full_path,
stat_result,
scope,
status_code = 200,
):
response = super().file_response(full_path, stat_result, scope, status_code)
response.headers["Cache-Control"] = _IMMUTABLE_ASSET_CACHE_CONTROL
return response
class _AssetGZipMiddleware(GZipMiddleware):
"""Serve range requests uncompressed; gzip + 206 mislabels Content-Range."""
async def __call__(self, scope, receive, send):
if scope["type"] == "http" and any(key == b"range" for key, _ in scope["headers"]):
await self.app(scope, receive, send)
return
await super().__call__(scope, receive, send)
def setup_frontend(app: FastAPI, build_path: Path):
"""Mount frontend static files (optional)"""
if not build_path.exists():
@ -1495,7 +1548,12 @@ def setup_frontend(app: FastAPI, build_path: Path):
assets_dir = build_path / "assets"
if assets_dir.exists():
app.mount("/assets", StaticFiles(directory = assets_dir), name = "assets")
assets_app = _AssetGZipMiddleware(
ImmutableStaticFiles(directory = assets_dir),
minimum_size = 1024,
compresslevel = 6,
)
app.mount("/assets", assets_app, name = "assets")
def _build_index_response(request: Request) -> Response:
content = (build_path / "index.html").read_bytes()

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

@ -18,6 +18,8 @@ from pydantic import (
model_validator,
)
from picker.schemas import MAX_CHAT_TEMPLATE_BYTES
class LoadRequest(BaseModel):
"""Request to load a model for inference"""
@ -54,8 +56,16 @@ class LoadRequest(BaseModel):
@field_validator("chat_template_override")
@classmethod
def normalize_blank_chat_template_override(cls, value: Optional[str]) -> Optional[str]:
if value is not None and value.strip() == "":
if value is None:
return None
# Char count is a lower bound on UTF-8 byte length: reject an oversized
# template before spending work encoding it.
if len(value) > MAX_CHAT_TEMPLATE_BYTES:
raise ValueError(f"Chat template exceeds the {MAX_CHAT_TEMPLATE_BYTES}-byte limit.")
if value.strip() == "":
return None
if len(value.encode("utf-8")) > MAX_CHAT_TEMPLATE_BYTES:
raise ValueError(f"Chat template exceeds the {MAX_CHAT_TEMPLATE_BYTES}-byte limit.")
return value
cache_type_kv: Optional[str] = Field(
@ -64,7 +74,7 @@ class LoadRequest(BaseModel):
)
gpu_ids: Optional[List[int]] = Field(
None,
description = "Physical GPU indices to use, for example [0, 1]. Omit or pass [] to use automatic selection. Explicit gpu_ids are unsupported when the parent CUDA_VISIBLE_DEVICES uses UUID/MIG entries. Not supported for GGUF models.",
description = "Physical GPU indices to use, for example [0, 1]. Omit or pass [] to use automatic selection. Explicit gpu_ids are unsupported when the parent CUDA_VISIBLE_DEVICES uses UUID/MIG entries. For GGUF models the picked devices are pinned via CUDA/HIP_VISIBLE_DEVICES.",
)
speculative_type: Optional[str] = Field(
None,
@ -100,6 +110,66 @@ class LoadRequest(BaseModel):
"No effect on a single GPU. Ignored for non-GGUF models."
),
)
gpu_memory_mode: Literal["auto", "manual"] = Field(
"auto",
description = (
"GPU memory strategy for GGUF models. 'auto' (default): Unsloth "
"selects GPUs and caps context to fit VRAM. 'manual': you own the "
"offload. Leave gpu_layers at -1 (Auto) to hand memory management to "
"llama.cpp's --fit (no device masking, no context auto-reduce, no "
"gpu-layer/tensor-split planning); set gpu_layers >= 0 to pin layers "
"and n_cpu_moe yourself (--fit off), with tensor_parallel still "
"applying (split by free VRAM unless tensor_split is set, no planner). "
"Ignored for non-GGUF."
),
)
gpu_layers: int = Field(
-1,
ge = -1,
description = (
"Manual mode only: number of layers to offload to the GPU "
"(--gpu-layers, with --fit off). A value >= the model's layer count "
"offloads all of them. -1 = Auto: hand layer + context sizing to "
"llama.cpp's --fit. Ignored unless gpu_memory_mode is 'manual'."
),
)
n_cpu_moe: int = Field(
0,
ge = 0,
description = (
"Manual mode only: keep the first N MoE expert layers on the CPU "
"(--n-cpu-moe) to save VRAM on MoE models. 0 = none, N = number of "
"MoE layers offloaded (the backend offsets past any leading dense "
"layers). Ignored unless gpu_memory_mode is 'manual' with gpu_layers >= 0."
),
)
tensor_split: Optional[List[float]] = Field(
None,
description = (
"Manual mode only: relative share of the model per GPU (--tensor-split), "
"in the order of the GPUs in use, e.g. [2, 1] for 2:1. Omit it to let "
"llama.cpp use its default, which splits by free VRAM. Any list given is "
"passed through as-is, so send [1, 1] to force an even split. Ignored "
"unless gpu_memory_mode is 'manual' with gpu_layers >= 0."
),
)
@field_validator("tensor_split")
@classmethod
def _reject_degenerate_tensor_split(cls, value: Optional[List[float]]) -> Optional[List[float]]:
# A negative / non-finite / all-zero split is silently dropped at launch
# (stored as None) yet still compared raw in the reload dedupe, so an
# identical Apply reloads forever. Reject it up front; [] = no split.
if not value:
return value
import math
if any((not math.isfinite(v)) or v < 0 for v in value):
raise ValueError("tensor_split entries must be finite and non-negative")
if sum(value) <= 0:
raise ValueError("tensor_split must have a positive total")
return value
llama_extra_args: Optional[List[str]] = Field(
None,
description = (
@ -133,11 +203,26 @@ class ValidateModelRequest(BaseModel):
max_seq_length: int = Field(0, ge = 0, le = 1048576)
load_in_4bit: bool = Field(True)
gpu_ids: Optional[List[int]] = Field(None)
gpu_memory_mode: Literal["auto", "manual"] = Field(
"auto",
description = (
"GGUF GPU-memory strategy intended for the follow-up load. Manual "
"placement bypasses the training coexistence estimate: Auto layers "
"delegate fitting to llama.cpp, while explicit layers are user-owned."
),
)
include_context_length: bool = Field(
False,
description = "Also read the native context length from the local GGUF header. "
"Opt-in so the normal load preflight doesn't pay for a cache scan it doesn't need.",
)
include_chat_template: bool = Field(
False,
description = "Also read the embedded chat template from the local GGUF header, so a "
"native (picked / drag-drop) file's default template can be shown before it is loaded. "
"Opt-in and, like include_context_length, a metadata-only probe that skips the training "
"guard. Only the leased file's own embedded template is read, never sibling sidecars.",
)
class TransformersUpgradeInfo(BaseModel):
@ -188,6 +273,21 @@ class ValidateModelResponse(BaseModel):
description = "Native training context length, read from the GGUF header when the file "
"is already downloaded locally; None for non-GGUF, gated, or not-yet-downloaded models.",
)
layer_count: Optional[int] = Field(
None,
description = "Total layer count (GGUF block_count), the manual gpu-layers ceiling, read "
"from the header alongside context_length; None when not read.",
)
moe_layer_count: Optional[int] = Field(
None,
description = "MoE expert-layer count (the manual --n-cpu-moe ceiling), read from the GGUF "
"header alongside context_length; 0 for dense models, None when not read.",
)
chat_template: Optional[str] = Field(
None,
description = "Embedded GGUF chat template, read from the header when include_chat_template "
"is set (native lease-backed picks); None for non-GGUF, over-cap, or not-read templates.",
)
# Additive fields; the consuming consent dialog ships in a follow-up frontend PR.
requires_transformers_upgrade: bool = Field(
False,
@ -333,6 +433,34 @@ class LoadResponse(BaseModel):
False,
description = "Whether tensor-parallel split (--split-mode tensor) is active.",
)
gpu_memory_mode: Literal["auto", "manual"] = Field(
"auto",
description = "Active GPU memory strategy ('auto' or 'manual').",
)
gpu_layers: int = Field(
-1,
description = "Manual mode: requested --gpu-layers value (-1 = Auto/--fit, or when not manual).",
)
n_cpu_moe: int = Field(
0,
description = "Manual mode: MoE expert layers pinned to CPU (--n-cpu-moe); 0 = none.",
)
tensor_split: Optional[List[float]] = Field(
None,
description = "Manual mode: relative model share per GPU (--tensor-split); None = default (split by free VRAM).",
)
n_layers: Optional[int] = Field(
None,
description = "Model's layer count (GGUF block_count), for the manual gpu-layers ceiling.",
)
n_moe_layers: int = Field(
0,
description = "Model's MoE expert-layer count (the n_cpu_moe ceiling); 0 if not an MoE model.",
)
gpu_ids: Optional[List[int]] = Field(
None,
description = "Physical GPU indices the model is pinned to, or None for automatic selection.",
)
class UnloadResponse(BaseModel):
@ -461,6 +589,42 @@ class InferenceStatusResponse(BaseModel):
False,
description = "Whether tensor-parallel split (--split-mode tensor) is active.",
)
gpu_memory_mode: Literal["auto", "manual"] = Field(
"auto",
description = "Active GPU memory strategy ('auto' or 'manual').",
)
gpu_layers: int = Field(
-1,
description = "Manual mode: requested --gpu-layers value (-1 = Auto/--fit, or when not manual).",
)
n_cpu_moe: int = Field(
0,
description = "Manual mode: MoE expert layers pinned to CPU (--n-cpu-moe); 0 = none.",
)
tensor_split: Optional[List[float]] = Field(
None,
description = "Manual mode: relative model share per GPU (--tensor-split); None = default (split by free VRAM).",
)
requested_context_length: Optional[int] = Field(
None,
description = (
"The n_ctx the active GGUF load was invoked with (0 = Auto). Lets the "
"UI re-seed a Manual + Auto-layers context pin on hydration, where "
"context_length only exposes the resolved value. None for non-GGUF."
),
)
n_layers: Optional[int] = Field(
None,
description = "Model's layer count (GGUF block_count), for the manual gpu-layers ceiling.",
)
n_moe_layers: int = Field(
0,
description = "Model's MoE expert-layer count (the n_cpu_moe ceiling); 0 if not an MoE model.",
)
gpu_ids: Optional[List[int]] = Field(
None,
description = "Physical GPU indices the model is pinned to, or None for automatic selection.",
)
llama_cpp_supports_mtp: bool = Field(
True,
description = (

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

@ -505,6 +505,13 @@ class TrainingStartRequest(BaseModel):
description = "S3 bucket configuration for loading datasets from AWS S3. Requires boto3 to be installed.",
)
@field_validator("target_modules", mode = "before")
@classmethod
def _normalize_target_modules(cls, value: Any) -> Any:
# Sanitized non-LoRA history stores the unused value as null; treat it as a
# fresh request's omitted/default empty list on resume.
return [] if value is None else value
@model_validator(mode = "after")
def _validate_streaming_splits(self) -> "TrainingStartRequest":
# Streaming load_dataset does not accept HF slice syntax (e.g. "train[:50%]"

2
studio/backend/picker/__init__.py Normal file
View file

@ -0,0 +1,2 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0

6
studio/backend/picker/routes/__init__.py Normal file
View file

@ -0,0 +1,6 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
from .templates import router as templates_router
__all__ = ["templates_router"]

45
studio/backend/picker/routes/templates.py Normal file
View file

@ -0,0 +1,45 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
from __future__ import annotations
import asyncio
from typing import Optional
from fastapi import APIRouter, Body, Depends, Query
from auth.authentication import get_current_subject
from hub.dependencies import get_hf_token
from ..schemas import (
MAX_CHAT_TEMPLATE_BYTES,
ModelTemplateResponse,
ValidateChatTemplateRequest,
ValidateChatTemplateResponse,
)
from ..service import read_default_chat_template, validate_chat_template
router = APIRouter()
@router.post("/validate-chat-template", response_model = ValidateChatTemplateResponse)
async def validate_chat_template_route(
body: ValidateChatTemplateRequest = Body(...),
current_subject: str = Depends(get_current_subject),
) -> ValidateChatTemplateResponse:
return await asyncio.to_thread(validate_chat_template, body.template)
@router.get("/chat-template/{model_name:path}", response_model = ModelTemplateResponse)
async def get_default_chat_template_route(
model_name: str,
gguf_variant: Optional[str] = Query(None),
hf_token: Optional[str] = Depends(get_hf_token),
current_subject: str = Depends(get_current_subject),
) -> ModelTemplateResponse:
template = await asyncio.to_thread(
read_default_chat_template, model_name, hf_token, gguf_variant
)
if template is not None and len(template.encode("utf-8")) > MAX_CHAT_TEMPLATE_BYTES:
template = None
return ModelTemplateResponse(model_name = model_name, chat_template = template)

32
studio/backend/picker/schemas.py Normal file
View file

@ -0,0 +1,32 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
from typing import Optional
from pydantic import BaseModel, Field, field_validator
# Mirror the frontend's 64 KiB chat-template contract (per-model-config.ts) at
# the API boundary so a direct caller cannot make Jinja parse an oversized
# template. MaxBodyMiddleware only caps the whole request body, not this field.
MAX_CHAT_TEMPLATE_BYTES = 65_536
class ValidateChatTemplateRequest(BaseModel):
template: str = Field(default = "")
@field_validator("template")
@classmethod
def _enforce_template_size(cls, value: str) -> str:
if len(value.encode("utf-8")) > MAX_CHAT_TEMPLATE_BYTES:
raise ValueError(f"Chat template exceeds the {MAX_CHAT_TEMPLATE_BYTES}-byte limit.")
return value
class ValidateChatTemplateResponse(BaseModel):
valid: bool
error: Optional[str] = None
class ModelTemplateResponse(BaseModel):
model_name: str
chat_template: Optional[str] = None

426
studio/backend/picker/service.py Normal file
View file

@ -0,0 +1,426 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
from __future__ import annotations
import json
import logging
import os
import re
from pathlib import Path
from typing import Optional
from hub.services.models.folder_browser import (
_build_browse_allowlist,
_is_path_inside_allowlist,
)
from hub.utils.gguf import extract_quant_label, iter_hf_cache_snapshots
from utils.models.gguf_metadata import read_gguf_chat_template
from utils.models.model_config import (
_extract_quant_label,
_is_big_endian_gguf_path,
_is_mmproj,
_is_mtp_drafter,
)
from utils.paths.path_utils import (
is_local_path,
normalize_path,
resolve_cached_repo_id_case,
)
from .schemas import MAX_CHAT_TEMPLATE_BYTES, ValidateChatTemplateResponse
logger = logging.getLogger(__name__)
_VALID_REPO_ID = re.compile(r"^[A-Za-z0-9._-]+/[A-Za-z0-9._-]+$")
def _is_valid_repo_id(repo_id: str) -> bool:
return bool(_VALID_REPO_ID.fullmatch(repo_id))
_TOKENIZER_CONFIG_PATHS = ("tokenizer_config.json", "LLM/tokenizer_config.json")
_JINJA_TEMPLATE_PATHS = ("chat_template.jinja", "LLM/chat_template.jinja")
_PROCESSOR_TEMPLATE_PATHS = ("chat_template.json", "LLM/chat_template.json")
# Cap sidecar reads so a malformed or hostile metadata file cannot exhaust memory
# before its template is size-checked. The JSON envelope may exceed a bare template
# (it carries other tokenizer metadata); the extracted template is still bounded by
# MAX_CHAT_TEMPLATE_BYTES downstream.
MAX_TEMPLATE_METADATA_BYTES = 4 * 1024 * 1024
def _read_bounded_text(path: Path, limit: int) -> Optional[str]:
"""Read at most `limit` bytes of UTF-8 text; None if larger or unreadable."""
try:
with path.open("rb") as f:
data = f.read(limit + 1)
except OSError:
return None
if len(data) > limit:
return None
try:
return data.decode("utf-8")
except UnicodeError:
return None
def _leaf_inside_allowlist(path: Path, allow_roots: Optional[list[Path]]) -> bool:
# Block symlinked children from escaping the validated directory (realpath-checked).
# None = trusted caller (HF cache / remote download).
return allow_roots is None or _is_path_inside_allowlist(path, allow_roots)
def validate_chat_template(template: str) -> ValidateChatTemplateResponse:
text = (template or "").strip()
if not text:
return ValidateChatTemplateResponse(valid = True, error = None)
# Import Jinja lazily: optional at runtime (e.g. GGUF-only installs), so a
# missing dependency must not crash API startup.
try:
from jinja2 import TemplateError
from jinja2.ext import Extension
from jinja2.sandbox import ImmutableSandboxedEnvironment
except ImportError:
return ValidateChatTemplateResponse(valid = True, error = None)
class _GenerationTag(Extension):
# Accept Transformers' {% generation %} assistant-mask tag so a pasted HF
# chat template validates (we only parse it).
tags = {"generation"}
def parse(self, parser):
next(parser.stream)
return parser.parse_statements(["name:endgeneration"], drop_needle = True)
try:
env = ImmutableSandboxedEnvironment(
trim_blocks = True,
lstrip_blocks = True,
extensions = ["jinja2.ext.loopcontrols", _GenerationTag],
)
env.parse(text)
return ValidateChatTemplateResponse(valid = True, error = None)
except TemplateError as exc:
message = getattr(exc, "message", None) or str(exc)
lineno = getattr(exc, "lineno", None)
if lineno:
message = f"Line {lineno}: {message}"
return ValidateChatTemplateResponse(valid = False, error = message)
except Exception as exc:
return ValidateChatTemplateResponse(valid = False, error = str(exc))
def _chat_template_from_tokenizer_config(config: dict) -> Optional[str]:
if not isinstance(config, dict):
return None
raw = config.get("chat_template")
if isinstance(raw, str) and raw.strip():
return raw
if isinstance(raw, list):
fallback: Optional[str] = None
for entry in raw:
if not isinstance(entry, dict):
continue
template = entry.get("template")
if not isinstance(template, str):
continue
if entry.get("name") == "default":
return template
if fallback is None:
fallback = template
return fallback
return None
def _chat_template_from_jinja_file(
dir_path: Path, allow_roots: Optional[list[Path]] = None
) -> Optional[str]:
for rel in _JINJA_TEMPLATE_PATHS:
template_file = dir_path / rel
if not template_file.exists() or not _leaf_inside_allowlist(template_file, allow_roots):
continue
try:
if template_file.stat().st_size > MAX_CHAT_TEMPLATE_BYTES:
continue
template = template_file.read_text(encoding = "utf-8")
except Exception:
continue
if template.strip():
return template
return None
def _chat_template_from_processor_payload(payload: object) -> Optional[str]:
# processor chat_template.json may be the template string itself or a
# {name: template} map, not only a tokenizer_config-shaped object.
if isinstance(payload, str):
return payload if payload.strip() else None
template = _chat_template_from_tokenizer_config(payload) # type: ignore[arg-type]
if template:
return template
if isinstance(payload, dict):
# Named-template map: prefer "default", else the first non-empty entry
# (mirrors the tokenizer-config list fallback).
default = payload.get("default")
if isinstance(default, str) and default.strip():
return default
for value in payload.values():
if isinstance(value, str) and value.strip():
return value
return None
def _chat_template_from_processor_json(
dir_path: Path, allow_roots: Optional[list[Path]] = None
) -> Optional[str]:
for rel in _PROCESSOR_TEMPLATE_PATHS:
config_file = dir_path / rel
if not config_file.exists() or not _leaf_inside_allowlist(config_file, allow_roots):
continue
raw = _read_bounded_text(config_file, MAX_TEMPLATE_METADATA_BYTES)
if raw is None:
continue
try:
payload = json.loads(raw)
except Exception:
continue
template = _chat_template_from_processor_payload(payload)
if template:
return template
return None
def _chat_template_from_tokenizer_dir(
dir_path: Path, allow_roots: Optional[list[Path]] = None
) -> Optional[str]:
jinja = _chat_template_from_jinja_file(dir_path, allow_roots)
if jinja:
return jinja
for rel in _TOKENIZER_CONFIG_PATHS:
config_file = dir_path / rel
if not config_file.exists() or not _leaf_inside_allowlist(config_file, allow_roots):
continue
raw = _read_bounded_text(config_file, MAX_TEMPLATE_METADATA_BYTES)
if raw is None:
continue
try:
config = json.loads(raw)
except Exception:
continue
template = _chat_template_from_tokenizer_config(config)
if template:
return template
return _chat_template_from_processor_json(dir_path, allow_roots)
_GGUF_SCAN_MAX_DEPTH = 2
def _iter_ggufs(dir_path: Path) -> list[Path]:
if dir_path == dir_path.parent:
return []
root = str(dir_path)
found: list[Path] = []
for current, dirs, files in os.walk(root, followlinks = False):
rel = os.path.relpath(current, root)
depth = 0 if rel == os.curdir else rel.count(os.sep) + 1
if depth >= _GGUF_SCAN_MAX_DEPTH:
dirs[:] = []
for name in files:
if not name.lower().endswith(".gguf") or _is_mmproj(name):
continue
path = Path(current) / name
try:
rel = path.relative_to(dir_path).as_posix()
except ValueError:
rel = name
quant = _extract_quant_label(rel)
if _is_mtp_drafter(rel) or _is_big_endian_gguf_path(rel, quant):
continue
found.append(path)
return found
def _variant_matches(relative_path: str, needle: str) -> bool:
quant = _extract_quant_label(relative_path).lower()
if quant == needle:
return True
if extract_quant_label(relative_path).lower() == needle:
return True
prefix = f"{needle}-"
if not quant.startswith(prefix):
return False
suffix = quant[len(prefix) :]
if not suffix.endswith("bpw"):
return False
value = suffix[:-3]
return bool(value) and value.replace(".", "", 1).isdigit()
_GGUF_SPLIT_INDEX_RE = re.compile(r"-(\d{3,})-of-\d{3,}$", re.IGNORECASE)
def _is_nonfirst_gguf_split(path: Path) -> bool:
match = _GGUF_SPLIT_INDEX_RE.search(path.stem)
return match is not None and int(match.group(1)) != 1
def _find_gguf_in_dir(dir_path: Path, gguf_variant: Optional[str]) -> Optional[Path]:
try:
ggufs = sorted(_iter_ggufs(dir_path))
except OSError:
return None
if not ggufs:
return None
needle = (gguf_variant or "").strip().lower()
if needle:
for path in ggufs:
try:
relative = path.relative_to(dir_path).as_posix()
except ValueError:
relative = path.name
if _variant_matches(relative, needle):
return path
return None
candidates = [path for path in ggufs if not _is_nonfirst_gguf_split(path)] or ggufs
try:
return max(candidates, key = lambda path: path.stat().st_size)
except OSError:
return candidates[0]
def _chat_template_from_dir(
dir_path: Path,
gguf_variant: Optional[str] = None,
allow_roots: Optional[list[Path]] = None,
) -> Optional[str]:
def from_gguf() -> Optional[str]:
gguf = _find_gguf_in_dir(dir_path, gguf_variant)
if gguf is None or not _leaf_inside_allowlist(gguf, allow_roots):
return None
return read_gguf_chat_template(str(gguf))
# Sidecar tokenizer files (chat_template.jinja / tokenizer_config.json) are the
# author's maintained template and supersede the GGUF's possibly-stale embedded
# copy. The variant only picks the GGUF fallback, so tokenizer-first precedence
# holds whether or not a variant is given.
return _chat_template_from_tokenizer_dir(dir_path, allow_roots) or from_gguf()
def read_default_chat_template(
model_name: str,
hf_token: Optional[str] = None,
gguf_variant: Optional[str] = None,
) -> Optional[str]:
if not isinstance(model_name, str) or not model_name.strip():
return None
name = model_name.strip()
if is_local_path(name):
try:
target = Path(normalize_path(name)).expanduser()
allow_roots = _build_browse_allowlist()
if not _is_path_inside_allowlist(target, allow_roots):
logger.debug("Refused chat template read outside allowed folders: %s", name)
return None
if name.lower().endswith(".gguf"):
# Prefer a maintained sidecar next to the file over the GGUF's
# embedded copy (tokenizer-first precedence, as elsewhere).
sidecar = _chat_template_from_tokenizer_dir(target.parent, allow_roots)
if sidecar:
return sidecar
return read_gguf_chat_template(str(target))
return _chat_template_from_dir(target, gguf_variant, allow_roots)
except Exception as exc:
logger.debug("Could not read local chat template for %s: %s", name, exc)
return None
if not _is_valid_repo_id(name):
return None
resolved = resolve_cached_repo_id_case(name)
try:
# Resolve within each cached revision, newest first. A revision's sidecar
# supersedes its own embedded GGUF copy, but must not override a newer
# revision, so precedence stays per-snapshot rather than global.
for snapshot in iter_hf_cache_snapshots(resolved):
template = _chat_template_from_dir(snapshot, gguf_variant)
if template:
return template
except Exception as exc:
logger.debug("Could not read cached chat template for %s: %s", resolved, exc)
try:
from huggingface_hub import HfApi, hf_hub_download
_api = HfApi()
def _remote_exceeds_cap(rel: str) -> bool:
# Best-effort: skip the download when the remote's advertised size
# exceeds the cap, so a maliciously large sidecar is never fetched.
try:
infos = _api.get_paths_info(resolved, [rel], repo_type = "model", token = hf_token)
except Exception:
return False
for info in infos:
size = getattr(info, "size", None)
if (
getattr(info, "path", None) == rel
and isinstance(size, int)
and size > MAX_TEMPLATE_METADATA_BYTES
):
return True
return False
def _download_text(rel: str) -> Optional[str]:
if _remote_exceeds_cap(rel):
return None
try:
path = hf_hub_download(resolved, rel, token = hf_token)
return _read_bounded_text(Path(path), MAX_TEMPLATE_METADATA_BYTES)
except Exception:
return None
for rel in _JINJA_TEMPLATE_PATHS:
template = _download_text(rel)
if not template or not template.strip():
continue
# A raw Jinja sidecar is the whole template, so it must fit the route's
# response cap (the local path skips oversized .jinja too). Download stays
# bounded at MAX_TEMPLATE_METADATA_BYTES so a large JSON embedding a small
# template still extracts below, but an over-cap Jinja is dropped so the
# search falls through to the tokenizer/processor template.
if len(template.encode("utf-8")) > MAX_CHAT_TEMPLATE_BYTES:
continue
return template
for rel in _TOKENIZER_CONFIG_PATHS:
raw = _download_text(rel)
if not raw:
continue
try:
config = json.loads(raw)
except Exception:
continue
template = _chat_template_from_tokenizer_config(config)
if template:
return template
for rel in _PROCESSOR_TEMPLATE_PATHS:
raw = _download_text(rel)
if not raw:
continue
try:
payload = json.loads(raw)
except Exception:
continue
template = _chat_template_from_processor_payload(payload)
if template:
return template
return None
except Exception as exc:
logger.debug("Could not fetch chat template for %s: %s", resolved, exc)
return None

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

@ -5,7 +5,7 @@
Chat history API routes backed by studio.db.
"""
from typing import Any, Literal, Optional
from typing import Annotated, Any, Literal, Optional
from fastapi import APIRouter, Depends, HTTPException, Query
from pydantic import BaseModel, ConfigDict, Field, ValidationError
@ -19,13 +19,16 @@ from storage.studio_db import (
clear_chat_history,
count_chat_threads,
count_forks_for_message,
delete_chat_attachment,
delete_chat_threads,
delete_chat_project,
ensure_chat_project_workspace,
fork_chat_thread,
get_chat_attachment,
get_chat_project,
get_chat_thread,
get_chat_message,
list_chat_attachments_page,
list_chat_projects,
list_chat_legacy_imports,
list_chat_settings,
@ -279,6 +282,131 @@ async def delete_threads(
return {"status": "deleted"}
@router.get("/attachments")
def list_attachments(
limit: Annotated[int, Query(ge = 1, le = 100)] = 50,
offset: Annotated[int, Query(ge = 0)] = 0,
current_subject: str = Depends(get_current_subject),
) -> dict:
"""One bounded page of chat uploads for the settings Data tab."""
attachments, next_offset = list_chat_attachments_page(limit = limit, offset = offset)
return {"attachments": attachments, "nextOffset": next_offset}
def _decode_attachment_base64(payload: str) -> bytes:
"""Strict base64 decode of a stored payload.
Normalizes first: strips whitespace, fixes padding, accepts the URL-safe
alphabet. validate=False would silently drop bad characters and serve
corrupted bytes instead of failing, so raise 422 on anything else.
"""
import base64
normalized = "".join(payload.split())
altchars = b"-_" if ("-" in normalized or "_" in normalized) else None
normalized += "=" * (-len(normalized) % 4)
try:
return base64.b64decode(normalized, altchars = altchars, validate = True)
except Exception as exc: # noqa: BLE001 - corrupt stored payload
raise HTTPException(status_code = 422, detail = "Attachment data is corrupt") from exc
_AUDIO_FORMAT_MEDIA_TYPES = {
"mp3": "audio/mpeg",
"wav": "audio/wav",
"ogg": "audio/ogg",
"flac": "audio/flac",
}
def _safe_image_media_type(media_type: str) -> str:
"""Clamp a data-URL media type to something inert to render.
Imported chats store image parts verbatim, so the embedded type can be
text/html or image/svg+xml; echoing those would execute markup with the
app origin when opened. Anything not a plain raster type downloads as
bytes instead.
"""
lowered = media_type.strip().lower()
if lowered.startswith("image/") and lowered != "image/svg+xml":
return lowered
return "application/octet-stream"
@router.get("/attachments/{message_id}/{attachment_id}/file")
def get_attachment_file(
message_id: str,
attachment_id: str,
current_subject: str = Depends(get_current_subject),
):
"""Serve one attachment's stored content: image or audio bytes, or
extracted text."""
import urllib.parse
from fastapi.responses import Response
attachment = get_chat_attachment(message_id, attachment_id)
if attachment is None:
raise HTTPException(status_code = 404, detail = "Attachment not found")
attachment_content_type = attachment.get("contentType")
texts: list[str] = []
for part in attachment.get("content") or []:
if not isinstance(part, dict):
continue
image = part.get("image")
if isinstance(image, str) and image[:5].lower() == "data:":
header, _, payload = image.partition(",")
media_type = _safe_image_media_type(
header[5:].split(";", 1)[0] or "application/octet-stream"
)
if "base64" not in header.lower():
# RFC 2397 non-base64 form stores percent-encoded bytes.
data = urllib.parse.unquote_to_bytes(payload)
return Response(content = data, media_type = media_type)
data = _decode_attachment_base64(payload)
return Response(content = data, media_type = media_type)
# Audio parts: the attachment adapter stores {data, format} with raw
# base64; compare chats store a bare base64 string.
audio = part.get("audio")
if isinstance(audio, dict) or (isinstance(audio, str) and audio):
if isinstance(audio, dict):
payload = audio.get("data")
audio_format = audio.get("format")
else:
payload = audio.rsplit(",", 1)[-1]
audio_format = None
if isinstance(payload, str) and payload:
data = _decode_attachment_base64(payload)
media_type = (
attachment_content_type
if isinstance(attachment_content_type, str)
and attachment_content_type.startswith("audio/")
else _AUDIO_FORMAT_MEDIA_TYPES.get(
str(audio_format or "").lower(), "application/octet-stream"
)
)
return Response(content = data, media_type = media_type)
text = part.get("text")
if isinstance(text, str) and text:
texts.append(text)
if texts:
return Response(content = "\n".join(texts), media_type = "text/plain; charset=utf-8")
raise HTTPException(status_code = 404, detail = "Attachment has no stored content")
@router.delete("/attachments/{message_id}/{attachment_id}")
def delete_attachment(
message_id: str,
attachment_id: str,
current_subject: str = Depends(get_current_subject),
) -> dict:
"""Remove one attachment from its chat message."""
if not delete_chat_attachment(message_id, attachment_id):
raise HTTPException(status_code = 404, detail = "Attachment not found")
return {"ok": True}
@router.get("/projects", response_model = ChatProjectListResponse)
async def list_projects(
include_archived: bool = Query(False), current_subject: str = Depends(get_current_subject)
@ -409,7 +537,7 @@ async def get_thread_message(
@router.put("/threads/{thread_id}/messages/{message_id}", response_model = ChatMessage)
async def save_thread_message(
def save_thread_message(
thread_id: str,
message_id: str,
payload: ChatMessage,
@ -432,7 +560,7 @@ async def save_thread_message(
@router.put("/threads/{thread_id}/messages", response_model = ChatMessageListResponse)
async def replace_thread_messages(
def replace_thread_messages(
thread_id: str,
payload: ChatMessageSyncRequest,
current_subject: str = Depends(get_current_subject),

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

@ -13,7 +13,7 @@ from pathlib import Path
from fastapi import APIRouter, Depends, HTTPException, Request, status
from fastapi.responses import StreamingResponse, JSONResponse, Response
from starlette.requests import ClientDisconnect
from typing import Any, Callable, List, Optional, Union
from typing import Any, Callable, List, Literal, Optional, Union
import json
import httpx
from loggers import get_logger
@ -1778,7 +1778,7 @@ from core.inference.providers import get_base_url
from core.inference.external_provider import ExternalProviderClient
from core.inference.chat_templates import resolve_effective_chat_template_override
from storage import providers_db
from utils.utils import safe_error_detail, log_and_http_error
from utils.utils import is_hf_authentication_error, safe_error_detail, log_and_http_error
import io
import base64
@ -3115,13 +3115,16 @@ def _normalise_settings_str(value: Optional[str]) -> Optional[str]:
def _should_strip_split_mode(request: LoadRequest, backend_extra: Optional[list[str]]) -> bool:
"""Whether an inherited --split-mode should be stripped on reload.
"""Whether an inherited --split-mode (and its coupled --tensor-split) should
be stripped on reload.
The binary Tensor Parallelism toggle can't carry --split-mode's row/none/
layer modes, so only strip when the toggle overrides it: tensor being turned
on, or the inherited mode is tensor (toggle turning it off). Non-tensor modes
survive. Shared by the inheritance strip and the already-loaded stale check
so they agree on what reload would do.
survive. A manual per-GPU ratio is handled by _should_strip_tensor_split,
which strips only --tensor-split so the inherited mode is kept. Shared by the
inheritance strip and the already-loaded stale check so they agree on what
reload would do.
"""
fields_set = getattr(request, "model_fields_set", set())
return "tensor_parallel" in fields_set and (
@ -3129,6 +3132,25 @@ def _should_strip_split_mode(request: LoadRequest, backend_extra: Optional[list[
)
def _should_strip_tensor_split(request: LoadRequest) -> bool:
"""Whether an inherited --tensor-split alone should be stripped on reload.
Manual explicit offload (gpu_layers >= 0) owns the per-GPU split: with a ratio
it emits its own --tensor-split (an inherited one, appended last, would
override it), and with the ratio cleared it wants llama.cpp's default
free-VRAM split. Either way an inherited --tensor-split must go, else the
cleared case silently keeps the stale ratio while status reports None.
Unlike _should_strip_split_mode this leaves --split-mode untouched, so a
user's row/none/layer mode survives a Studio split-ratio edit. When the
Tensor Parallelism toggle IS overriding the mode, _should_strip_split_mode
(called alongside this at every site) strips --split-mode anyway.
"""
return (
getattr(request, "gpu_memory_mode", "auto") == "manual"
and getattr(request, "gpu_layers", -1) >= 0
)
def _carry_preserved_tensor_intent(
*, preserved: bool, same_model: bool, explicit_drop: bool
) -> bool:
@ -3187,12 +3209,44 @@ def _request_matches_loaded_settings(
else strip_shadowing_flags(
backend_extra,
strip_split_mode = _should_strip_split_mode(request, backend_extra),
strip_tensor_split = _should_strip_tensor_split(request),
strip_offload = request.gpu_memory_mode == "manual",
)
)
if not _tensor_parallel_matches_loaded(
effective_extra, request.tensor_parallel, llama_backend.tensor_parallel
):
return False
# The diffusion runner is mode-agnostic (it always reports "auto" and ignores
# the layer/MoE/split knobs), so a standing manual preference in the request
# must not force a needless reload -- only the GPU pick matters.
if not llama_backend.is_diffusion:
if request.gpu_memory_mode != llama_backend.gpu_memory_mode:
return False
# Manual: a layer-count change always reloads; MoE/split only matter with
# an explicit offload (gpu_layers >= 0), so a leftover value under Auto
# must not force one. Mirrors LlamaCppBackend._already_in_target_state.
if request.gpu_memory_mode == "manual" and (
request.gpu_layers != llama_backend.gpu_layers
or (
request.gpu_layers >= 0
and (
request.n_cpu_moe != llama_backend.n_cpu_moe
or (request.tensor_split or None) != (llama_backend.tensor_split or None)
)
)
):
return False
# A changed GPU pick must reload. The diffusion runner collapses a multi-GPU
# request to its single lowest device (it drives one device only), so the
# backend records just that device; compare the request the same way, or a
# multi-GPU pick that resolves to the same device needlessly reloads.
if llama_backend.is_diffusion:
_req_gpu_ids = [sorted(request.gpu_ids)[0]] if request.gpu_ids else None
else:
_req_gpu_ids = sorted(request.gpu_ids) if request.gpu_ids else None
if _req_gpu_ids != llama_backend.gpu_ids:
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
@ -3235,14 +3289,17 @@ def _request_matches_loaded_settings(
# contain any shadow flag, so the reload path strips them rather than
# leaving a stale override in effect. (backend_extra computed above.)
if request.llama_extra_args is None:
# Mirror the reload's conditional split-mode strip, so a preserved
# non-tensor mode (row/none/layer) isn't seen as stale and doesn't
# trigger a needless reload of a healthy server.
# Mirror the reload's conditional strips, so a preserved non-tensor mode
# (row/none/layer) isn't seen as stale and doesn't trigger a needless
# reload of a healthy server, while an inherited offload/ratio flag that
# the reload *would* strip is correctly seen as stale.
if (
backend_extra
and strip_shadowing_flags(
backend_extra,
strip_split_mode = _should_strip_split_mode(request, backend_extra),
strip_tensor_split = _should_strip_tensor_split(request),
strip_offload = request.gpu_memory_mode == "manual",
)
!= backend_extra
):
@ -3861,6 +3918,46 @@ def _estimate_gguf_required_gb(
return None
def _classify_diffusion_gguf(config: ModelConfig) -> Optional[bool]:
"""Classify a GGUF as diffusion, normal, or unknown before it is loaded.
``None`` is important here: a remote GGUF whose header is not cached can
still be routed to the single-GPU diffusion runner after download. Treating
that case as normal would let Manual mode skip the training guard even
though the runner ignores Manual's llama-server placement controls.
"""
identity = " ".join(
str(getattr(config, attr, "") or "") for attr in ("identifier", "gguf_hf_repo", "gguf_file")
).lower()
if "diffusion" in identity:
return True
try:
main = getattr(config, "gguf_file", None)
if not (main and Path(main).is_file()):
repo = getattr(config, "gguf_hf_repo", None)
variant = getattr(config, "gguf_variant", None)
if repo and variant:
from hub.utils.gguf import resolve_local_gguf_path
main = resolve_local_gguf_path(repo, variant)
if not main or not Path(main).is_file():
return None
probe = LlamaCppBackend()
probe._read_gguf_metadata(str(main))
if probe.is_diffusion:
return True
# A successfully decoded architecture proves that this is a normal
# llama-server GGUF. No architecture means the lightweight probe could
# not establish the routing decision, so preserve the unknown state.
if getattr(probe, "_architecture", None):
return False
return None
except Exception as e:
logger.debug("Could not identify diffusion GGUF for training guard: %s", e)
return None
def _guard_chat_load_against_training(
config: ModelConfig,
*,
@ -3871,11 +3968,19 @@ def _guard_chat_load_against_training(
requested_gpu_ids: Optional[List[int]],
llama_extra_args: Optional[list[str]] = None,
n_parallel: int = 1,
gpu_memory_mode: Literal["auto", "manual"] = "auto",
) -> None:
"""Refuse loading a local chat model that would OOM an active training run.
"""Protect active training from automatically placed chat-model loads.
No-op when training is inactive or unknown. `load_in_4bit` must be the
effective quantization (see _effective_load_in_4bit). Raises HTTP 409 when the
model would not fit alongside training."""
effective quantization (see _effective_load_in_4bit). Manual chat-GGUF
placement is an explicit override: Auto layers delegate fitting to
llama.cpp's ``--fit`` and pinned layers are owned by the user, so neither is
estimated here. Diffusion is still guarded because its mode-agnostic runner
ignores those controls and uses one GPU. An unclassified GGUF is guarded as
potentially diffusion until its local header proves otherwise. Other loads
raise HTTP 409 when they would not fit beside training.
"""
from core.training import get_training_backend
from routes.training_vram import can_load_chat_during_training
@ -3887,6 +3992,19 @@ def _guard_chat_load_against_training(
return
is_gguf = bool(getattr(config, "is_gguf", False))
diffusion_kind = _classify_diffusion_gguf(config) if is_gguf else False
if is_gguf and gpu_memory_mode == "manual" and diffusion_kind is False:
return
diffusion_gpu = None
if is_gguf and diffusion_kind is not False:
# Use the same token selection as the runner: an explicit pick wins,
# followed by DG_GPU, the first parent-visible token, then GPU 0.
diffusion_gpu = LlamaCppBackend._diffusion_gpu_arg(
requested_gpu_ids,
cpu_only = LlamaCppBackend._effective_gpu_count() == 0,
)
required_override_gb = (
_estimate_gguf_required_gb(
config,
@ -3907,6 +4025,7 @@ def _guard_chat_load_against_training(
requested_gpu_ids = requested_gpu_ids,
is_gguf = is_gguf,
required_override_gb = required_override_gb,
single_device_gpu = diffusion_gpu,
)
if ok:
return
@ -3934,6 +4053,98 @@ def _guard_chat_load_against_training(
raise HTTPException(status_code = 409, detail = detail)
def _resolve_inherited_extra_args(
request,
config: ModelConfig,
model_identifier: str,
extra_llama_args: Optional[list[str]],
effective_chat_template_override: Optional[str] = None,
) -> Optional[list[str]]:
"""Effective pass-through extras for a GGUF request that omitted the field:
the previous same-model load's extras, shadow-stripped, so a settings-Apply
reload (which does not round-trip the extras field) keeps them (#5401)."""
if getattr(request, "llama_extra_args", None) is not None:
return extra_llama_args
if not getattr(config, "is_gguf", False):
return extra_llama_args
llama_backend = get_llama_cpp_backend()
if not llama_backend.extra_args:
return extra_llama_args
# Inherit the previous load's extras (the chat-settings Apply path doesn't
# round-trip them; an explicit [] still clears). Gated on (model_identifier,
# hf_variant) to refuse cross-model pickup, and shadowing flags are
# stripped so an inherited override can't win the last-wins CLI
# parse against a freshly-supplied first-class field.
source = llama_backend.extra_args_source
# Compare against the resolved variant, not the request field: callers
# commonly omit gguf_variant for local ``.gguf`` paths and HF auto-pick
# flows. ``config.gguf_variant`` is the variant load_model was actually
# invoked with, so both sides of the comparison key off the same string.
resolved_variant = (config.gguf_variant or "").lower()
request_variant = (request.gguf_variant or "").lower()
stored_variant = (source[1] or "").lower() if source else ""
same_model = bool(source and source[0] and source[0].lower() == model_identifier.lower())
if request.gguf_variant:
variant_mismatch = request_variant != stored_variant
else:
variant_mismatch = bool(stored_variant and resolved_variant != stored_variant)
same_source = same_model and not variant_mismatch
if not same_source:
logger.info(
"Not inheriting llama_extra_args: stored args came from %s, loading %s",
source,
(model_identifier, resolved_variant),
)
# Cross-model: clear explicitly so the backend doesn't
# inherit via "no opinion" semantics.
extra_llama_args = []
else:
# Strip only the groups whose first-class field was set by the caller, so
# an inherited --chat-template-file survives an Apply that omits
# chat_template_override. A bundled family template (e.g. gemma-4) counts as
# a first-class template even when the request omits chat_template_override,
# so strip the inherited --chat-template-file then too -- else the stale arg
# (appended last) shadows the bundled template while Studio reports its caps.
fields_set = getattr(request, "model_fields_set", set())
stripped = strip_shadowing_flags(
llama_backend.extra_args,
strip_context = "max_seq_length" in fields_set,
strip_cache = "cache_type_kv" in fields_set,
strip_spec = ("speculative_type" in fields_set or "spec_draft_n_max" in fields_set),
strip_template = (
"chat_template_override" in fields_set
or effective_chat_template_override is not None
),
strip_split_mode = _should_strip_split_mode(request, llama_backend.extra_args),
# manual + per-GPU ratio emits its own --tensor-split; drop
# an inherited one (appended last would override it) while
# keeping the user's --split-mode row/none/layer choice.
strip_tensor_split = _should_strip_tensor_split(request),
# manual emits its own --fit/--gpu-layers, so an inherited offload flag
# must not last-wins-override it. auto leaves a user's inherited -ngl
# alone. getattr: a validate request reuses this resolver, no offload fields.
strip_offload = getattr(request, "gpu_memory_mode", "auto") == "manual",
)
try:
extra_llama_args = validate_extra_args(stripped)
except ValueError:
# Shouldn't happen on already-validated args; degrade to
# no-extras rather than 400 if managed flags changed.
logger.warning(
"Stored llama_extra_args failed revalidation; loading without them: %s",
stripped,
)
extra_llama_args = []
else:
if extra_llama_args:
logger.info(
"Inheriting llama_extra_args from previous "
"load (same model, shadow-stripped): %s",
extra_llama_args,
)
return extra_llama_args
def _model_json_response(model, status_code: int = 200) -> Response:
"""Serialize a pydantic response once via pydantic-core.
@ -4040,6 +4251,35 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre
None if request.llama_extra_args is None else extra_llama_args
)
# Manual mode owns the offload flags: strip them from EXPLICIT extras
# too (the inherited path already does), or a last-wins --gpu-layers /
# --fit in extras re-enables GPU offload on a load status reports as
# CPU-only. Manual + per-GPU ratio owns --tensor-split the same way.
if request.gpu_memory_mode == "manual" and extra_llama_args:
_stripped_explicit = strip_shadowing_flags(
extra_llama_args,
strip_context = False,
strip_cache = False,
strip_spec = False,
strip_template = False,
strip_split_mode = False,
strip_tensor_split = _should_strip_tensor_split(request),
strip_offload = True,
)
if _stripped_explicit != extra_llama_args:
logger.info(
"Manual GPU memory owns the offload flags; stripping them "
"from explicit llama_extra_args: %s -> %s",
extra_llama_args,
_stripped_explicit,
)
extra_llama_args = _stripped_explicit
# Keep every downstream consumer on the normalized explicit list. In
# particular, the already-loaded comparator must not compare the raw
# request's managed offload flags against the stripped launch state.
request = request.model_copy(update = {"llama_extra_args": extra_llama_args})
model_identifier, model_log_label, native_grant_backed = (
_resolve_model_identifier_for_request(request, operation = "load-model")
)
@ -4121,6 +4361,13 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre
speculative_type = llama_backend.requested_spec_mode,
spec_draft_n_max = llama_backend.spec_draft_n_max,
tensor_parallel = llama_backend.tensor_parallel,
gpu_memory_mode = llama_backend.gpu_memory_mode,
gpu_layers = llama_backend.gpu_layers,
n_cpu_moe = llama_backend.n_cpu_moe,
tensor_split = llama_backend.tensor_split,
n_layers = llama_backend.n_layers,
n_moe_layers = llama_backend.n_moe_layers,
gpu_ids = llama_backend.gpu_ids,
)
else:
if (
@ -4187,12 +4434,41 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre
# Normalize gpu_ids: empty list means auto-selection, same as None
effective_gpu_ids = request.gpu_ids if request.gpu_ids else None
# Reject GGUF + gpu_ids first so the guard can't mask it with a VRAM 409.
# GGUF supports gpu_ids: validate the pick up front (before the training
# guard) so a bad pick is a clean 400, not masked by a VRAM 409. Rejects
# negative / out-of-range / duplicate ids and UUID/MIG parents. XPU hosts
# are rejected outright: the picker's indices are torch-xpu ordinals neither
# applicator speaks (CUDA/HIP masks don't apply, the Vulkan --device pin
# uses ggml's own Vulkan ordinals), so a pick could land on the wrong device.
if config.is_gguf and effective_gpu_ids is not None:
raise HTTPException(
status_code = 400,
detail = "gpu_ids is not supported for GGUF models yet.",
)
from utils.hardware import DeviceType, get_device
from utils.hardware.hardware import resolve_requested_gpu_ids
if get_device() == DeviceType.XPU:
raise HTTPException(
status_code = 400,
detail = (
"GPU selection (gpu_ids) is not supported on Intel XPU. "
"Omit gpu_ids to use all devices."
),
)
# Same reasoning for a Vulkan-only build: --device pins ggml's own
# Vulkan ordinals, so a physical pick can land on the wrong card on
# masked or non-contiguous hosts.
if LlamaCppBackend._is_vulkan_backend():
raise HTTPException(
status_code = 400,
detail = (
"GPU selection (gpu_ids) is not supported with a Vulkan "
"llama.cpp build: physical GPU ids have no defined "
"mapping to Vulkan device ordinals. Omit gpu_ids to use "
"all devices."
),
)
try:
resolve_requested_gpu_ids(effective_gpu_ids)
except ValueError as exc:
raise HTTPException(status_code = 400, detail = str(exc)) from exc
if not config.is_gguf and _mlx_distributed_launch_detected():
raise HTTPException(
status_code = 400,
@ -4222,8 +4498,20 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre
"architectures)"
)
# Refuse a load that would OOM active training, before the unload step below
# frees the resident model. Off-loop: guard does sync nvidia-smi / HF work.
# Inherit the previous same-model load's pass-through extras when this
# request omits the field (a settings-Apply reload doesn't round-trip
# them); shadow-stripped so an inherited flag can't override a
# first-class field the caller did set (#5401).
extra_llama_args = _resolve_inherited_extra_args(
request,
config,
model_identifier,
extra_llama_args,
effective_chat_template_override,
)
# Apply the training coexistence policy before the unload step below
# frees the resident model. Off-loop: the default-mode guard does sync work.
await asyncio.to_thread(
_guard_chat_load_against_training,
config,
@ -4234,6 +4522,7 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre
requested_gpu_ids = effective_gpu_ids,
llama_extra_args = extra_llama_args,
n_parallel = getattr(fastapi_request.app.state, "llama_parallel_slots", 1),
gpu_memory_mode = request.gpu_memory_mode,
)
# ── GGUF path: load via llama-server ──────────────────────
@ -4245,84 +4534,6 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre
from core.inference.llama_cpp import gguf_load_in_flight
gguf_load_stack.enter_context(gguf_load_in_flight(config.gguf_hf_repo))
# Inherit llama_extra_args from the previous load when the request
# omits the field (the chat-settings Apply path doesn't round-trip
# them; explicit [] still clears). Gated on (model_identifier,
# hf_variant) to refuse cross-model pickup, and shadowing flags are
# stripped so an inherited override can't win the last-wins CLI
# parse against a freshly-supplied first-class field.
if request.llama_extra_args is None and llama_backend.extra_args:
source = llama_backend.extra_args_source
# Compare against the resolved variant, not the request
# field: callers commonly omit gguf_variant for local
# ``.gguf`` paths and HF auto-pick flows. ``config.gguf_
# variant`` is the variant load_model was actually
# invoked with (see the HF / local branches below), so
# both sides of the comparison key off the same string.
resolved_variant = (config.gguf_variant or "").lower()
request_variant = (request.gguf_variant or "").lower()
stored_variant = (source[1] or "").lower() if source else ""
same_model = bool(
source and source[0] and source[0].lower() == model_identifier.lower()
)
if request.gguf_variant:
variant_mismatch = request_variant != stored_variant
else:
variant_mismatch = bool(stored_variant and resolved_variant != stored_variant)
same_source = same_model and not variant_mismatch
if not same_source:
logger.info(
"Not inheriting llama_extra_args: stored args came from %s, loading %s",
source,
(model_identifier, resolved_variant),
)
# Cross-model: clear explicitly so the backend doesn't
# inherit via "no opinion" semantics.
extra_llama_args = []
else:
# Strip only the groups whose first-class field was set by
# the caller, so an inherited --chat-template-file survives
# an Apply that omits chat_template_override. A bundled family
# template (e.g. the gemma-4 override) is an effective
# first-class template setting even when the raw request
# omits chat_template_override, so strip the inherited
# --chat-template-file in that case too -- otherwise the stale
# extra arg (appended last) shadows the bundled template while
# Unsloth reports the bundled template's capabilities.
fields_set = getattr(request, "model_fields_set", set())
stripped = strip_shadowing_flags(
llama_backend.extra_args,
strip_context = "max_seq_length" in fields_set,
strip_cache = "cache_type_kv" in fields_set,
strip_spec = (
"speculative_type" in fields_set or "spec_draft_n_max" in fields_set
),
strip_template = (
"chat_template_override" in fields_set
or effective_chat_template_override is not None
),
strip_split_mode = _should_strip_split_mode(
request, llama_backend.extra_args
),
)
try:
extra_llama_args = validate_extra_args(stripped)
except ValueError:
# Shouldn't happen on already-validated args; degrade to
# no-extras rather than 400 if managed flags changed.
logger.warning(
"Stored llama_extra_args failed revalidation; loading without them: %s",
stripped,
)
extra_llama_args = []
else:
if extra_llama_args:
logger.info(
"Inheriting llama_extra_args from previous "
"load (same model, shadow-stripped): %s",
extra_llama_args,
)
# Block cache writes that would race the download manager. This runs
# after pass-through argument inheritance so a carried --no-mmproj
# changes the companion requirement exactly as it does for the load.
@ -4370,6 +4581,11 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre
cache_type_kv = request.cache_type_kv,
speculative_type = request.speculative_type,
spec_draft_n_max = request.spec_draft_n_max,
gpu_memory_mode = request.gpu_memory_mode,
gpu_layers = request.gpu_layers,
n_cpu_moe = request.n_cpu_moe,
tensor_split = request.tensor_split,
gpu_ids = effective_gpu_ids,
n_parallel = _n_parallel,
)
if config.gguf_hf_repo:
@ -4494,7 +4710,7 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre
# Clear any idle-unload reload stash now, not only on the next poll.
from core.inference.llama_keepwarm import note_model_loaded
note_model_loaded()
await asyncio.to_thread(note_model_loaded, llama_backend)
# A plain load advertises its own identifier; auto-switch overwrites
# this with the repo id right after _load_model_impl returns.
llama_backend._openai_advertised_id = None
@ -4537,6 +4753,13 @@ async def _load_model_impl(request: LoadRequest, fastapi_request: Request, curre
speculative_type = llama_backend.requested_spec_mode,
spec_draft_n_max = llama_backend.spec_draft_n_max,
tensor_parallel = llama_backend.tensor_parallel,
gpu_memory_mode = llama_backend.gpu_memory_mode,
gpu_layers = llama_backend.gpu_layers,
n_cpu_moe = llama_backend.n_cpu_moe,
tensor_split = llama_backend.tensor_split,
n_layers = llama_backend.n_layers,
n_moe_layers = llama_backend.n_moe_layers,
gpu_ids = llama_backend.gpu_ids,
)
# ── Standard path: load via Unsloth/transformers ──────────
@ -4795,7 +5018,9 @@ def _requires_security_review_for_model(
@router.post("/validate", response_model = ValidateModelResponse)
async def validate_model(
request: ValidateModelRequest, current_subject: str = Depends(get_current_subject)
request: ValidateModelRequest,
fastapi_request: Request = None,
current_subject: str = Depends(get_current_subject),
):
"""
Lightweight validation endpoint for model identifiers.
@ -4823,15 +5048,39 @@ async def validate_model(
detail = f"Invalid model identifier: {model_log_label}",
)
# Refuse early (before the frontend unloads to load this) if it can't fit
# alongside training, using the same settings /load uses so they agree.
# Apply the same training coexistence policy as /load before the frontend
# unloads the current model.
effective_gpu_ids = request.gpu_ids if request.gpu_ids else None
# Mirror /load: reject GGUF + gpu_ids before the guard so both return 400.
# Mirror /load: GGUF supports gpu_ids, so validate the pick (a bad one is
# a clean 400) before the guard sizes the model against training VRAM.
# XPU-host picks are rejected like /load (no defined mapping from the
# picker's torch-xpu ordinals to the launcher's device spaces).
if config.is_gguf and effective_gpu_ids is not None:
raise HTTPException(
status_code = 400,
detail = "gpu_ids is not supported for GGUF models yet.",
)
from utils.hardware import DeviceType, get_device
from utils.hardware.hardware import resolve_requested_gpu_ids
if get_device() == DeviceType.XPU:
raise HTTPException(
status_code = 400,
detail = (
"GPU selection (gpu_ids) is not supported on Intel XPU. "
"Omit gpu_ids to use all devices."
),
)
if LlamaCppBackend._is_vulkan_backend():
raise HTTPException(
status_code = 400,
detail = (
"GPU selection (gpu_ids) is not supported with a Vulkan "
"llama.cpp build: physical GPU ids have no defined "
"mapping to Vulkan device ordinals. Omit gpu_ids to use "
"all devices."
),
)
try:
resolve_requested_gpu_ids(effective_gpu_ids)
except ValueError as exc:
raise HTTPException(status_code = 400, detail = str(exc)) from exc
effective_load_in_4bit = _effective_load_in_4bit(config, request.load_in_4bit)
# Both checks cover the [adapter, base] set (matching the scan route and workers):
@ -4895,16 +5144,32 @@ async def validate_model(
latest_tier_active_for, config.identifier, request.hf_token
):
effective_load_in_4bit = False
# Off-loop: guard does sync nvidia-smi / HF work.
await asyncio.to_thread(
_guard_chat_load_against_training,
config,
model_identifier = model_identifier,
hf_token = request.hf_token,
load_in_4bit = effective_load_in_4bit,
max_seq_length = request.max_seq_length,
requested_gpu_ids = effective_gpu_ids,
)
# A metadata-only probe reads the GGUF header and allocates no VRAM, so the
# training guard must not refuse it. Real loads omit include_context_length /
# include_chat_template, and /load applies the guard again.
if not (request.include_context_length or request.include_chat_template):
# Match /load's inherited llama.cpp extras and parallel slot count so
# validation cannot pass a smaller estimate than the subsequent load.
effective_extra_args = _resolve_inherited_extra_args(
request, config, model_identifier, None
)
# Off-loop: guard does sync nvidia-smi / HF work.
await asyncio.to_thread(
_guard_chat_load_against_training,
config,
model_identifier = model_identifier,
hf_token = request.hf_token,
load_in_4bit = effective_load_in_4bit,
max_seq_length = request.max_seq_length,
requested_gpu_ids = effective_gpu_ids,
llama_extra_args = effective_extra_args,
n_parallel = (
getattr(fastapi_request.app.state, "llama_parallel_slots", 1)
if fastapi_request is not None
else 1
),
gpu_memory_mode = request.gpu_memory_mode,
)
# A selected GGUF loads via llama.cpp: auto_map Python and root pickle weights in a
# mixed repo are inert for this load, so gating on them is a false positive. Only
@ -4918,10 +5183,21 @@ async def validate_model(
# Native context length, read from the local GGUF header when present.
# Lets the staged ("Load on selection" off) flow populate the context
# slider before the GPU load; None until the file is downloaded.
# Staged header dims (one read): native context, total layer count, and
# MoE expert-layer count -- let the staged flow size the context, GPU-
# layers and manual --n-cpu-moe sliders before the load.
context_length: Optional[int] = None
if request.include_context_length and is_gguf:
layer_count: Optional[int] = None
moe_layer_count: Optional[int] = None
chat_template: Optional[str] = None
# Both header probes read the same local GGUF, so resolve it once.
if (request.include_context_length or request.include_chat_template) and is_gguf:
from hub.utils.gguf import resolve_local_gguf_path
from utils.models.gguf_metadata import read_gguf_context_length
from picker.schemas import MAX_CHAT_TEMPLATE_BYTES
from utils.models.gguf_metadata import (
read_gguf_chat_template,
read_gguf_staged_dims,
)
# Best-effort: a header-read failure must never fail validation of an
# otherwise-valid model (the outer except turns it into a 400).
@ -4937,9 +5213,26 @@ async def validate_model(
model_identifier, request.gguf_variant
)
if local_gguf:
context_length = read_gguf_context_length(local_gguf)
if request.include_context_length:
# Header walk reads tokenizer arrays (tens of ms); keep it
# off the event loop.
dims = await asyncio.to_thread(read_gguf_staged_dims, local_gguf)
if dims:
context_length = dims["context_length"]
layer_count = dims["layer_count"]
moe_layer_count = dims["moe_layer_count"]
if request.include_chat_template:
# Read only the leased GGUF's own embedded template (the copy
# llama.cpp loads), never a sibling sidecar: the native grant
# authorizes just this path, so neighbours would be scope escalation.
raw_template = await asyncio.to_thread(read_gguf_chat_template, local_gguf)
if (
raw_template is not None
and len(raw_template.encode("utf-8")) <= MAX_CHAT_TEMPLATE_BYTES
):
chat_template = raw_template
except Exception as e:
logger.debug("Context-length probe failed for %s: %s", model_log_label, e)
logger.debug("Header probe failed for %s: %s", model_log_label, e)
return ValidateModelResponse(
valid = True,
@ -4954,6 +5247,9 @@ async def validate_model(
requires_trust_remote_code = requires_trust_remote_code,
requires_security_review = requires_security_review,
context_length = context_length,
layer_count = layer_count,
moe_layer_count = moe_layer_count,
chat_template = chat_template,
requires_transformers_upgrade = transformers_upgrade is not None,
transformers_upgrade = transformers_upgrade,
)
@ -4966,6 +5262,14 @@ async def validate_model(
raise HTTPException(status_code = 400, detail = str(e))
except Exception as e:
redacted_msg = redact_native_paths(str(e))
if is_hf_authentication_error(e):
raise HTTPException(
status_code = 400,
detail = (
"Hugging Face authentication failed. Check or clear the token "
"in Settings, and confirm access to this gated repository."
),
)
if _is_unsupported_nvfp4_inference_error(redacted_msg):
logger.warning(
"NVFP4 inference is not supported yet while validating '%s'",
@ -5593,6 +5897,14 @@ async def get_status(current_subject: str = Depends(get_current_subject)):
speculative_type = llama_backend.requested_spec_mode,
spec_draft_n_max = llama_backend.spec_draft_n_max,
tensor_parallel = llama_backend.tensor_parallel,
gpu_memory_mode = llama_backend.gpu_memory_mode,
gpu_layers = llama_backend.gpu_layers,
n_cpu_moe = llama_backend.n_cpu_moe,
tensor_split = llama_backend.tensor_split,
requested_context_length = llama_backend.requested_n_ctx,
n_layers = llama_backend.n_layers,
n_moe_layers = llama_backend.n_moe_layers,
gpu_ids = llama_backend.gpu_ids,
llama_cpp_supports_mtp = _supports_mtp,
spec_fallback_reason = llama_backend.spec_fallback_reason,
llama_cpp_prebuilt_stale = _stale,

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

@ -60,13 +60,52 @@ def _safe_is_dir(path) -> bool:
# Shared with the hub inventory scans; keep the private aliases so existing
# importers (core.inference.local_model_resolver, tests) stay valid.
# importers stay valid. ``_HF_REPO_ID_RE`` is the Hub repo id shape ("owner/name");
# anything else is treated as a local filesystem path.
from utils.hidden_models import (
_HF_REPO_ID_RE,
_existing_resolved_path,
_safe_resolve,
is_hidden_model as _is_hidden_model,
)
def hidden_model_matchers() -> tuple[list[str], list[str], list[str]]:
"""Substring needles, exact repo ids, and exact resolved paths identifying
infra models (the RAG embedder and the llama.cpp install validation probe)
that pickers hide. Served by the ``/api/hub/hidden-models`` endpoint. A
configured HF-repo embedder is published as its exact lowercased repo id
(mirroring ``utils.hidden_models.is_hidden_model``) and a local-path
embedder as its exact resolved path only: a generic basename like "model"
must not substring-hide unrelated chat models."""
from core.rag import config as rag_config
needles = [
# The validation probe's repo and its exact filename. The filename carries
# .gguf so it won't hide unrelated repos like ``user/stories260K-finetune-GGUF``.
"ggml-org/models",
"stories260k.gguf",
]
exact_ids: list[str] = []
exact_paths: list[str] = []
for model in (
rag_config.effective_embedding_model(),
rag_config.effective_gguf_repo(),
):
# Resolve an existing local path before the repo-id regex: a local embedder
# shaped like "models/embedder" is an exact path, not a Hub repo id.
existing_path = _existing_resolved_path(model)
if existing_path:
exact_paths.append(existing_path.lower())
elif _HF_REPO_ID_RE.match(model):
exact_ids.append(model.lower())
else:
resolved = _safe_resolve(Path(model).expanduser())
if resolved:
exact_paths.append(resolved.lower())
return needles, exact_ids, exact_paths
backend_path = Path(__file__).parent.parent.parent
if str(backend_path) not in sys.path:
sys.path.insert(0, str(backend_path))
@ -91,6 +130,7 @@ try:
_pick_best_gguf,
_extract_quant_label,
_is_big_endian_gguf_path,
_is_mtp_drafter,
is_audio_input_type,
)
from core.inference import get_inference_backend
@ -123,6 +163,7 @@ except ImportError:
_pick_best_gguf,
_extract_quant_label,
_is_big_endian_gguf_path,
_is_mtp_drafter,
is_audio_input_type,
)
from core.inference import get_inference_backend
@ -803,7 +844,7 @@ def collect_local_models(models_root: Path) -> List[LocalModelInfo]:
models = sorted(
deduped.values(),
key = lambda item: (item.updated_at or 0),
key = lambda item: item.updated_at or 0,
reverse = True,
)
return [m for m in models if not _is_hidden_model(m.id, m.model_id, m.path)]
@ -1750,9 +1791,11 @@ def _get_model_size_bytes(model_name: str, hf_token: Optional[str] = None) -> Op
async def get_model_config(
model_name: str,
hf_token: Optional[str] = Query(None),
header_hf_token: Optional[str] = Depends(get_hf_token),
current_subject: str = Depends(get_current_subject),
):
"""Get configuration for a specific model (wraps load_model_defaults)."""
hf_token = _normalize_hf_token(header_hf_token) or _normalize_hf_token(hf_token)
try:
if not is_local_path(model_name):
resolved = resolve_cached_repo_id_case(model_name)
@ -2471,6 +2514,7 @@ async def get_lora_base_model(lora_path: str, current_subject: str = Depends(get
async def check_vision_model(
model_name: str,
hf_token: Optional[str] = Query(None),
header_hf_token: Optional[str] = Depends(get_hf_token),
current_subject: str = Depends(get_current_subject),
):
"""
@ -2478,6 +2522,7 @@ async def check_vision_model(
This endpoint wraps the backend is_vision_model function.
"""
hf_token = _normalize_hf_token(header_hf_token) or _normalize_hf_token(hf_token)
try:
logger.info(f"Checking if vision model: {model_name}")
# Authenticate so a gated/private VLM classifies correctly (else 404 -> non-vision).
@ -2503,6 +2548,7 @@ async def check_vision_model(
async def check_embedding_model(
model_name: str,
hf_token: Optional[str] = Query(None),
header_hf_token: Optional[str] = Depends(get_hf_token),
current_subject: str = Depends(get_current_subject),
):
"""
@ -2510,6 +2556,7 @@ async def check_embedding_model(
This endpoint wraps the backend is_embedding_model function.
"""
hf_token = _normalize_hf_token(header_hf_token) or _normalize_hf_token(hf_token)
try:
logger.info(f"Checking if embedding model: {model_name}")
is_embedding = is_embedding_model(model_name, hf_token = hf_token)
@ -2573,12 +2620,6 @@ def _resolve_quant_gguf(repo_id: str, quant: str, is_local: bool) -> tuple[Optio
Q8_0 weights). Never raises.
"""
try:
from utils.models.model_config import (
_extract_quant_label,
_is_big_endian_gguf_path,
_is_mtp_drafter,
)
if is_local:
roots = [Path(repo_id)]
else:
@ -2595,25 +2636,19 @@ def _resolve_quant_gguf(repo_id: str, quant: str, is_local: bool) -> tuple[Optio
if snaps.is_dir():
roots.extend(s for s in snaps.iterdir() if s.is_dir())
want = quant.lower().replace("-", "").replace("_", "")
want = _normalized_quant_label(quant)
best_total = 0
best_first: Optional[str] = None
for root in roots:
matches: list[tuple[str, Path]] = []
total = 0
for f in _iter_gguf_paths(root):
if _is_mmproj_filename(f.name):
continue
try:
rel = f.relative_to(root).as_posix()
except ValueError:
rel = f.name
if _is_mtp_drafter(rel):
continue
q = _extract_quant_label(rel)
if _is_big_endian_gguf_path(rel, q):
continue
if q.lower().replace("-", "").replace("_", "") != want:
q = _main_variant_gguf_label(rel)
if q is None or _normalized_quant_label(q) != want:
continue
try:
total += f.stat().st_size
@ -2731,7 +2766,11 @@ async def get_gguf_variants(
],
has_vision = response.has_vision,
default_variant = response.default_variant,
context_length = _read_native_context_length(repo_id, is_local = local),
# The header walk reads tokenizer arrays on dense models (tens of
# ms per uncached file); keep it off the event loop.
context_length = await asyncio.to_thread(
_read_native_context_length, repo_id, is_local = local
),
)
except HTTPException:
raise
@ -3031,6 +3070,22 @@ def _is_main_gguf_filename(name: str) -> bool:
return _is_gguf_filename(name) and not _is_mmproj_filename(name)
def _main_variant_gguf_label(rel_path: str) -> Optional[str]:
name = rel_path.rsplit("/", 1)[-1]
if not _is_main_gguf_filename(name):
return None
if _is_mtp_drafter(rel_path):
return None
label = _extract_quant_label(rel_path)
if _is_big_endian_gguf_path(rel_path, label):
return None
return label
def _normalized_quant_label(label: str) -> str:
return label.lower().replace("-", "").replace("_", "")
def _repo_has_mmproj(repo_info) -> bool:
"""True if the repo ships a GGUF vision adapter (mmproj), so it can
take image inputs. Cheap: scans already-listed file names only."""
@ -3358,6 +3413,170 @@ async def delete_cached_model(
)
def _resolve_cached_model_path(repo_id: str, variant: Optional[str]) -> Path:
"""Absolute path of a cached repo (newest snapshot dir) or, with *variant*,
that quant's main GGUF file (first split of a sharded quant). Paths come
from the HF cache scan only, so callers can't probe arbitrary paths."""
cache_scans = _all_hf_cache_scans()
matching_repos = []
for hf_cache in cache_scans:
for repo_info in hf_cache.repos:
if repo_info.repo_type != "model":
continue
if repo_info.repo_id.lower() == repo_id.lower():
matching_repos.append(repo_info)
if not matching_repos:
raise HTTPException(status_code = 404, detail = "Model not found in cache")
if variant:
want = _normalized_quant_label(variant)
candidate_revisions = sorted(
(rev for repo_info in matching_repos for rev in repo_info.revisions),
key = lambda rev: getattr(rev, "last_modified", 0) or 0,
reverse = True,
)
for rev in candidate_revisions:
snapshot = getattr(rev, "snapshot_path", None)
matches = []
for f in rev.files:
p = Path(f.file_path)
rel = f.file_name
if snapshot:
try:
rel = p.relative_to(snapshot).as_posix()
except ValueError:
pass
label = _main_variant_gguf_label(rel)
if label is None or _normalized_quant_label(label) != want:
continue
if p.exists() or p.is_symlink():
matches.append((rel, p))
if matches:
# Path-sorted so a sharded quant deterministically yields its first split.
return sorted(matches, key = lambda m: m[0].lower())[0][1]
raise HTTPException(
status_code = 404,
detail = f"Variant {variant} not found in cache for {repo_id}",
)
def repo_size(repo_info) -> int:
gguf_size = _repo_gguf_size_bytes(repo_info)
if gguf_size > 0:
return gguf_size
return sum(
(getattr(f, "size_on_disk", None) or 0)
for rev in repo_info.revisions
for f in rev.files
)
def repo_last_modified(repo_info) -> float:
return max(
(getattr(rev, "last_modified", 0) or 0 for rev in repo_info.revisions),
default = 0,
)
target_repo = max(
matching_repos,
key = lambda repo_info: (repo_size(repo_info), repo_last_modified(repo_info)),
)
# Whole repo: the newest revision's snapshot dir holds the visible files.
revisions = sorted(
(rev for rev in target_repo.revisions if getattr(rev, "snapshot_path", None)),
key = lambda rev: getattr(rev, "last_modified", 0) or 0,
reverse = True,
)
for rev in revisions:
p = Path(rev.snapshot_path)
if p.exists():
return p
p = Path(target_repo.repo_path)
if p.exists():
return p
raise HTTPException(status_code = 404, detail = "Cached model path not found")
def _wsl_reveal_in_explorer(path: Path) -> bool:
import subprocess
from utils.paths.path_utils import _IS_WSL
if not _IS_WSL:
return False
try:
windows_path = subprocess.run(
["wslpath", "-w", str(path)],
capture_output = True,
text = True,
check = True,
timeout = 10,
).stdout.strip()
if not windows_path:
return False
argument = f"/select,{windows_path}" if path.is_file() else windows_path
subprocess.Popen(["explorer.exe", argument])
return True
except (OSError, subprocess.SubprocessError):
return False
def _reveal_in_file_manager(path: Path) -> None:
"""Open the OS file manager with *path* selected (best effort per platform)."""
import subprocess
target = str(path)
if sys.platform == "darwin":
cmd = ["open", "-R", target] if path.is_file() else ["open", target]
subprocess.Popen(cmd)
elif os.name == "nt":
if path.is_file():
subprocess.Popen(["explorer", f"/select,{target}"])
else:
os.startfile(target) # noqa: S606 - local user's own file manager
elif not _wsl_reveal_in_explorer(path):
# No cross-desktop "select file" standard on Linux; open the directory.
directory = target if path.is_dir() else str(path.parent)
subprocess.Popen(["xdg-open", directory])
class CachedModelPathResponse(BaseModel):
path: str
is_dir: bool
@router.get("/cached-model-path", response_model = CachedModelPathResponse)
async def get_cached_model_path(
repo_id: str = Query(..., description = "HuggingFace repo ID"),
variant: str = Query("", description = "Quantization variant (empty for whole repo)"),
current_subject: str = Depends(get_current_subject),
):
"""Absolute on-disk path of a cached repo or one of its GGUF variants."""
if not _is_valid_repo_id(repo_id):
raise HTTPException(status_code = 400, detail = "Invalid repo_id format")
path = await asyncio.to_thread(_resolve_cached_model_path, repo_id, variant.strip() or None)
return {"path": str(path), "is_dir": path.is_dir()}
@router.post("/reveal-cached-model")
async def reveal_cached_model(
repo_id: str = Body(...),
variant: Optional[str] = Body(None),
current_subject: str = Depends(get_current_subject),
):
"""Reveal a cached repo (or one GGUF variant's file) in the OS file manager."""
if not _is_valid_repo_id(repo_id):
raise HTTPException(status_code = 400, detail = "Invalid repo_id format")
variant = (variant or "").strip() or None
path = await asyncio.to_thread(_resolve_cached_model_path, repo_id, variant)
try:
await asyncio.to_thread(_reveal_in_file_manager, path)
except Exception as e:
logger.error(f"Failed to reveal {path}: {e}")
raise HTTPException(status_code = 500, detail = "Failed to open file manager")
return {"status": "ok", "path": str(path)}
@router.get("/checkpoints", response_model = CheckpointListResponse)
async def list_checkpoints(
outputs_dir: str = Query(

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

@ -318,6 +318,39 @@ def list_project_documents(project_id: str, subject: str = Depends(get_current_s
conn.close()
@router.get("/documents")
def list_all_uploaded_documents(subject: str = Depends(get_current_subject)) -> dict:
"""Every uploaded file across chats, projects, and knowledge bases (settings
Data tab)."""
_require_rag()
conn = rag_db.get_connection()
try:
docs = store.list_all_documents(conn)
kb_names = {kb["id"]: kb["name"] for kb in store.list_kbs(conn)}
finally:
conn.close()
from storage.studio_db import list_chat_projects
project_names = {p["id"]: p["name"] for p in list_chat_projects(include_archived = True)}
out = []
for doc in docs:
view = _doc_view(doc)
stored_path = doc.get("stored_path")
size = None
if stored_path:
try:
size = os.path.getsize(stored_path)
except OSError:
size = None
view["sizeBytes"] = size
view["kbName"] = kb_names.get(doc.get("kb_id"))
view["projectName"] = project_names.get(doc.get("project_id"))
out.append(view)
return {"documents": out}
@router.delete("/documents/{document_id}")
def delete_document(document_id: str, subject: str = Depends(get_current_subject)) -> dict:
_require_rag()
@ -424,8 +457,10 @@ _CONTENT_TYPES = {
".txt": "text/plain; charset=utf-8",
".md": "text/markdown; charset=utf-8",
".markdown": "text/markdown; charset=utf-8",
".html": "text/html; charset=utf-8",
".htm": "text/html; charset=utf-8",
# Served as plain text, never text/html: an uploaded HTML document rendered
# same-origin would execute its scripts with access to the app's storage.
".html": "text/plain; charset=utf-8",
".htm": "text/plain; charset=utf-8",
".docx": "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
}

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

@ -36,9 +36,10 @@ from utils.helper_precache_settings import (
)
from utils.coding_agents import CODING_AGENTS, detect_installed_coding_agents
from utils.openai_auto_switch_settings import (
DEFAULT_AUTO_UNLOAD_IDLE_SECONDS,
DEFAULT_AUTO_UNLOAD_KEEP_KV,
DEFAULT_OPENAI_AUTO_SWITCH_ENABLED,
get_auto_unload_idle_seconds,
get_auto_unload_keep_kv,
get_model_overrides,
get_openai_auto_switch_enabled,
get_stored_auto_unload_idle_seconds,
@ -90,7 +91,9 @@ class HelperPrecacheResponse(BaseModel):
class OpenAIAutoSwitchPayload(BaseModel):
enabled: bool
auto_unload_idle_seconds: int = Field(default = DEFAULT_AUTO_UNLOAD_IDLE_SECONDS, ge = 0)
# None leaves the stored value untouched (partial updates can't clobber it).
auto_unload_idle_seconds: Optional[int] = Field(default = None, ge = 0)
auto_unload_keep_kv: Optional[bool] = None
class OpenAIAutoSwitchResponse(BaseModel):
@ -101,6 +104,7 @@ class OpenAIAutoSwitchResponse(BaseModel):
# UNSLOTH_MODEL_IDLE_TTL set and nothing stored, this is true even while enabled
# is false, so the UI can show idle-unload as active instead of "needs enable".
idle_unload_active: bool = False
auto_unload_keep_kv: bool = DEFAULT_AUTO_UNLOAD_KEEP_KV
class ModelOverridePayload(BaseModel):
@ -198,6 +202,7 @@ def get_openai_auto_switch(
enabled = get_openai_auto_switch_enabled(),
auto_unload_idle_seconds = get_stored_auto_unload_idle_seconds(),
idle_unload_active = get_auto_unload_idle_seconds() > 0,
auto_unload_keep_kv = get_auto_unload_keep_kv(),
)
@ -206,8 +211,8 @@ def update_openai_auto_switch(
payload: OpenAIAutoSwitchPayload, current_subject: str = Depends(get_current_subject)
) -> OpenAIAutoSwitchResponse:
try:
enabled, idle_seconds = set_openai_auto_switch(
payload.enabled, payload.auto_unload_idle_seconds
enabled, idle_seconds, keep_kv = set_openai_auto_switch(
payload.enabled, payload.auto_unload_idle_seconds, payload.auto_unload_keep_kv
)
except ValueError as exc:
raise log_and_http_error(
@ -217,10 +222,16 @@ def update_openai_auto_switch(
event = "settings.update_openai_auto_switch_failed",
log = logger,
) from exc
idle_unload_active = get_auto_unload_idle_seconds() > 0
if not keep_kv or not idle_unload_active:
# Keep-KV off or idle unload disabled: drop already-saved chat context too.
from core.inference.llama_keepwarm import purge_kv_resume
purge_kv_resume()
return OpenAIAutoSwitchResponse(
enabled = enabled,
auto_unload_idle_seconds = idle_seconds,
idle_unload_active = get_auto_unload_idle_seconds() > 0,
idle_unload_active = idle_unload_active,
auto_unload_keep_kv = keep_kv,
)

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

@ -196,6 +196,7 @@ async def start_training(
request.local_eval_datasets, "Local eval dataset"
)
resume_output_dir: Optional[str] = None
resume_run: Optional[dict] = None
if request.resume_from_checkpoint:
try:
resume_output_dir = normalize_resume_output_dir(request.resume_from_checkpoint)
@ -208,7 +209,7 @@ async def start_training(
if not resume_run or not can_resume_run(resume_run):
raise HTTPException(
status_code = 400,
detail = "Resume checkpoint must belong to a stopped run with saved trainer state.",
detail = "Resume checkpoint must belong to a stopped or errored run with complete saved trainer state.",
)
resume_checkpoint = get_resume_checkpoint_path(resume_output_dir)
if not resume_checkpoint:
@ -458,7 +459,10 @@ async def start_training(
try:
success = backend.start_training(
job_id = job_id, before_spawn = _free_vram_for_training, **training_kwargs
job_id = job_id,
before_spawn = _free_vram_for_training,
resume_source_run_id = resume_run["id"] if resume_run else None,
**training_kwargs,
)
except SidecarSwapInProgress as exc:
# Expected loss of the race against a sidecar install: a retryable
@ -521,7 +525,10 @@ async def stop_training(
status = "idle", message = "No training job is currently running"
)
backend.stop_training(save = body.save)
if not backend.stop_training(save = body.save):
return TrainingStopResponse(
status = "idle", message = "No training job is currently running"
)
return TrainingStopResponse(
status = "stopped",
@ -637,9 +644,9 @@ async def get_training_status(current_subject: str = Depends(get_current_subject
"loss": getattr(progress, "loss", None),
"learning_rate": getattr(progress, "learning_rate", None),
}
output_dir = getattr(backend, "_output_dir", None)
if output_dir:
details["output_dir"] = output_dir
# Always present: an explicit null tells the client to drop a cached
# path (stop without save clears the run's output_dir).
details["output_dir"] = getattr(backend, "_output_dir", None) or None
# Metric history for chart recovery after SSE reconnection.
metric_history = None

42
studio/backend/routes/training_vram.py
View file

@ -197,15 +197,18 @@ def can_load_chat_during_training(
requested_gpu_ids: Optional[List[int]],
is_gguf: bool = False,
required_override_gb: Optional[float] = None,
single_device_gpu: Optional[str] = None,
) -> Tuple[bool, Dict[str, Any]]:
"""Decide if a NEW chat model can load without OOMing active training (inverse
of can_keep_chat_during_training: training is already resident, so size the
chat model against the free VRAM that remains). Sizes/places it the same way
the loader will: HF auto reuses auto_select_gpu_ids; HF explicit requires an
even-share per-GPU floor for device_map="balanced"; GGUF sizes from
required_override_gb over the visible pool. `load_in_4bit` must be effective
(LoRA can flip 4-bit -> 16-bit). Non-CUDA allows the load; default-deny on any
CUDA case it can't size, so a load never OOMs training."""
required_override_gb over the visible pool. ``single_device_gpu`` is the
exact physical device token selected by a single-device runner.
`load_in_4bit` must be effective (LoRA can flip 4-bit -> 16-bit). Non-CUDA
allows the load; default-deny on any CUDA case it can't size, so a load never
OOMs training."""
try:
from utils.hardware import (
DeviceType,
@ -251,26 +254,49 @@ def can_load_chat_during_training(
}
# Explicit GPUs, or GGUF: size directly and check live free VRAM.
if single_device_gpu is not None:
mode = "single_device"
elif is_gguf:
mode = "gguf"
else:
mode = "explicit"
required_gb = required_override_gb
if required_gb is None:
required_gb, _meta = estimate_required_model_memory_gb(model_name, **est_kwargs)
if required_gb is None:
mode = "explicit" if requested_gpu_ids else "gguf"
return False, {"mode": mode, "reason": "estimate_unavailable"}
free_by_index = _free_vram_by_index(get_visible_gpu_utilization().get("devices", []))
if requested_gpu_ids:
if single_device_gpu is not None:
token = str(single_device_gpu).strip()
if not token:
# Empty token = a CPU-only single-device runner (e.g. a CPU
# diffusion GGUF): it uses no GPU VRAM, so it never threatens
# active training and can always load.
return True, {"mode": "single_device", "reason": "cpu_only"}
try:
selected_gpu = int(token)
if selected_gpu < 0:
raise ValueError
except (TypeError, ValueError):
# A non-numeric device token (e.g. a CUDA UUID / MIG handle)
# can't be mapped to a free-VRAM index, but the runner still
# drives ONE device. Size against the worst-case visible device
# (min free), never the aggregate pool, so a single-device load
# is never OK'd on capacity it can't use and OOMs training.
free_vals = [min(free_by_index.values())] if free_by_index else []
else:
free_vals = [free_by_index.get(selected_gpu, 0.0)]
elif requested_gpu_ids:
# Invalid ids -> load_model 400s first, so don't block; missing id = 0.
try:
resolved = resolve_requested_gpu_ids(requested_gpu_ids)
except ValueError:
return True, {"mode": "explicit", "reason": "invalid_gpu_ids"}
return True, {"mode": mode, "reason": "invalid_gpu_ids"}
free_vals = [free_by_index.get(i, 0.0) for i in resolved]
mode = "explicit"
else:
# GGUF: llama.cpp picks the GPU(s); any visible GPU is a candidate.
free_vals = list(free_by_index.values())
mode = "gguf"
if not free_vals:
return False, {"mode": mode, "reason": "no_visible_gpus"}

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

File diff suppressed because it is too large Load diff

634
studio/backend/tests/test_chat_attachments.py Normal file
View file

@ -0,0 +1,634 @@
# 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 base64
import json
import os
import sqlite3
import sys
import pytest
from fastapi import HTTPException
_backend = os.path.join(os.path.dirname(__file__), "..")
sys.path.insert(0, _backend)
from routes import chat_history
from storage import studio_db
from utils.paths import studio_db_path
PNG_BYTES = base64.b64decode(
"iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8BQDwAEhQGAhKmMIQAAAABJRU5ErkJggg=="
)
PNG_DATA_URL = "data:image/png;base64," + base64.b64encode(PNG_BYTES).decode("ascii")
def _reset_studio_db(tmp_path, monkeypatch):
monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path))
monkeypatch.setenv("UNSLOTH_STUDIO_PROJECTS_HOME", str(tmp_path / "Projects"))
monkeypatch.setattr(studio_db, "_schema_ready", False)
def _thread(
thread_id: str = "thread-1",
title: str = "Test Chat",
pair_id: str | None = None,
) -> dict:
return {
"id": thread_id,
"title": title,
"modelType": "base",
"modelId": "test-model",
"pairId": pair_id,
"archived": False,
"createdAt": 1_700_000_000_000,
}
def _message(
message_id: str,
created_at: int = 1_700_000_000_000,
attachments = None,
thread_id: str = "thread-1",
) -> dict:
message = {
"id": message_id,
"threadId": thread_id,
"parentId": None,
"role": "user",
"content": [{"type": "text", "text": "hello"}],
"createdAt": created_at,
}
if attachments is not None:
message["attachments"] = attachments
return message
def _image_attachment(attachment_id: str = "att-1", name: str = "photo.png") -> dict:
return {
"id": attachment_id,
"type": "image",
"name": name,
"contentType": "image/png",
"content": [{"type": "image", "image": PNG_DATA_URL}],
"status": {"type": "complete"},
}
def _seed(
tmp_path,
monkeypatch,
attachments,
message_id: str = "msg-1",
):
_reset_studio_db(tmp_path, monkeypatch)
studio_db.upsert_chat_thread(_thread())
studio_db.upsert_chat_message(_message(message_id, attachments = attachments))
def _set_raw_attachments_json(message_id: str, raw: str) -> None:
conn = sqlite3.connect(studio_db_path())
try:
conn.execute(
"UPDATE chat_messages SET attachments_json = ? WHERE id = ?",
(raw, message_id),
)
conn.commit()
finally:
conn.close()
def _raw_attachments_json(message_id: str):
conn = sqlite3.connect(studio_db_path())
try:
row = conn.execute(
"SELECT attachments_json FROM chat_messages WHERE id = ?",
(message_id,),
).fetchone()
return row[0] if row is not None else None
finally:
conn.close()
# ---------------------------------------------------------------------------
# Storage: list_chat_attachments
# ---------------------------------------------------------------------------
def test_list_chat_attachments_empty_db(tmp_path, monkeypatch):
_reset_studio_db(tmp_path, monkeypatch)
assert studio_db.list_chat_attachments() == []
def test_list_chat_attachments_round_trip(tmp_path, monkeypatch):
_seed(tmp_path, monkeypatch, [_image_attachment()])
records = studio_db.list_chat_attachments()
assert len(records) == 1
record = records[0]
assert record["id"] == "att-1"
assert record["messageId"] == "msg-1"
assert record["threadId"] == "thread-1"
assert record["threadTitle"] == "Test Chat"
assert record["name"] == "photo.png"
assert record["type"] == "image"
assert record["contentType"] == "image/png"
assert record["createdAt"] == 1_700_000_000_000
# Base64 length estimate is within padding error of the decoded size.
assert abs(record["sizeBytes"] - len(PNG_BYTES)) <= 2
def test_list_chat_attachments_counts_text_utf8(tmp_path, monkeypatch):
text = "héllo wörld é世界"
attachment = {
"id": "att-txt",
"type": "document",
"name": "notes.txt",
"content": [{"type": "text", "text": text}],
}
_seed(tmp_path, monkeypatch, [attachment])
records = studio_db.list_chat_attachments()
assert records[0]["sizeBytes"] == len(text.encode("utf-8"))
def test_list_chat_attachments_no_content_size_is_none(tmp_path, monkeypatch):
attachment = {"id": "att-empty", "name": "ghost.bin", "content": []}
_seed(tmp_path, monkeypatch, [attachment])
records = studio_db.list_chat_attachments()
assert records[0]["sizeBytes"] is None
assert records[0]["name"] == "ghost.bin"
def test_list_chat_attachments_defaults_missing_name(tmp_path, monkeypatch):
attachment = {"id": "att-noname", "content": []}
_seed(tmp_path, monkeypatch, [attachment])
assert studio_db.list_chat_attachments()[0]["name"] == "attachment"
def test_list_chat_attachments_sanitizes_structured_metadata(tmp_path, monkeypatch):
attachment = {
"id": "att-weird",
"name": {"nested": "name"},
"type": ["image"],
"contentType": {"mime": "image/png"},
"content": [],
}
_seed(tmp_path, monkeypatch, [attachment])
record = studio_db.list_chat_attachments()[0]
assert record["name"] == "attachment"
assert record["type"] is None
assert record["contentType"] is None
def test_list_chat_attachments_skips_malformed_rows(tmp_path, monkeypatch):
_reset_studio_db(tmp_path, monkeypatch)
studio_db.upsert_chat_thread(_thread())
for i, raw in enumerate(
[
"not json at all",
'{"id": "att-obj"}',
"null",
"[]",
'[{"noid": true}, "just a string", 42]',
'[{"id": ""}]',
]
):
message_id = f"msg-bad-{i}"
studio_db.upsert_chat_message(_message(message_id))
_set_raw_attachments_json(message_id, raw)
studio_db.upsert_chat_message(_message("msg-good", attachments = [_image_attachment("att-ok")]))
records = studio_db.list_chat_attachments()
assert [r["id"] for r in records] == ["att-ok"]
def test_list_chat_attachments_orders_newest_first(tmp_path, monkeypatch):
_reset_studio_db(tmp_path, monkeypatch)
studio_db.upsert_chat_thread(_thread())
studio_db.upsert_chat_message(
_message("msg-old", 1_700_000_000_000, [_image_attachment("att-old")])
)
studio_db.upsert_chat_message(
_message("msg-new", 1_700_000_100_000, [_image_attachment("att-new")])
)
assert [r["id"] for r in studio_db.list_chat_attachments()] == ["att-new", "att-old"]
def test_list_chat_attachments_survives_missing_thread_row(tmp_path, monkeypatch):
_reset_studio_db(tmp_path, monkeypatch)
studio_db.upsert_chat_thread(_thread())
studio_db.upsert_chat_message(_message("msg-1", attachments = [_image_attachment()]))
conn = sqlite3.connect(studio_db_path())
try:
conn.execute("DELETE FROM chat_threads WHERE id = 'thread-1'")
conn.commit()
finally:
conn.close()
records = studio_db.list_chat_attachments()
assert len(records) == 1
assert records[0]["threadTitle"] is None
def test_list_chat_attachments_includes_compare_pair_id(tmp_path, monkeypatch):
_reset_studio_db(tmp_path, monkeypatch)
studio_db.upsert_chat_thread(_thread(pair_id = "pair-1"))
studio_db.upsert_chat_message(_message("msg-compare", attachments = [_image_attachment()]))
record = studio_db.list_chat_attachments()[0]
assert record["threadId"] == "thread-1"
assert record["pairId"] == "pair-1"
def test_list_chat_attachments_gone_after_thread_delete(tmp_path, monkeypatch):
_seed(tmp_path, monkeypatch, [_image_attachment()])
studio_db.delete_chat_threads(["thread-1"])
assert studio_db.list_chat_attachments() == []
# ---------------------------------------------------------------------------
# Storage: get_chat_attachment / delete_chat_attachment
# ---------------------------------------------------------------------------
def test_get_chat_attachment_found_and_missing(tmp_path, monkeypatch):
_seed(tmp_path, monkeypatch, [_image_attachment()])
attachment = studio_db.get_chat_attachment("msg-1", "att-1")
assert attachment is not None
assert attachment["content"][0]["image"] == PNG_DATA_URL
assert studio_db.get_chat_attachment("msg-1", "att-missing") is None
assert studio_db.get_chat_attachment("msg-missing", "att-1") is None
def test_delete_chat_attachment_keeps_others(tmp_path, monkeypatch):
_seed(
tmp_path,
monkeypatch,
[_image_attachment("att-1"), _image_attachment("att-2", "other.png")],
)
assert studio_db.delete_chat_attachment("msg-1", "att-1") is True
assert studio_db.get_chat_attachment("msg-1", "att-1") is None
assert studio_db.get_chat_attachment("msg-1", "att-2") is not None
assert [r["id"] for r in studio_db.list_chat_attachments()] == ["att-2"]
def test_delete_last_chat_attachment_stores_empty_list(tmp_path, monkeypatch):
_seed(tmp_path, monkeypatch, [_image_attachment()])
assert studio_db.delete_chat_attachment("msg-1", "att-1") is True
# '[]' rather than NULL: a NULL attachments field reads back as missing
# and triggers the legacy IndexedDB backfill, resurrecting the deleted
# attachment on the next chat load.
assert _raw_attachments_json("msg-1") == "[]"
assert studio_db.list_chat_attachments() == []
# The message itself must survive with its content intact.
message = studio_db.get_chat_message("thread-1", "msg-1")
assert message is not None
assert message["content"] == [{"type": "text", "text": "hello"}]
assert message["attachments"] == []
def test_delete_chat_attachment_missing_targets(tmp_path, monkeypatch):
_seed(tmp_path, monkeypatch, [_image_attachment()])
assert studio_db.delete_chat_attachment("msg-missing", "att-1") is False
assert studio_db.delete_chat_attachment("msg-1", "att-missing") is False
_set_raw_attachments_json("msg-1", "not json")
assert studio_db.delete_chat_attachment("msg-1", "att-1") is False
# ---------------------------------------------------------------------------
# Routes: /attachments endpoints (real storage, direct calls)
# ---------------------------------------------------------------------------
def test_list_attachments_route(tmp_path, monkeypatch):
_seed(tmp_path, monkeypatch, [_image_attachment()])
result = chat_history.list_attachments(current_subject = "unsloth")
assert [a["id"] for a in result["attachments"]] == ["att-1"]
def test_attachment_file_serves_image_bytes(tmp_path, monkeypatch):
_seed(tmp_path, monkeypatch, [_image_attachment()])
response = chat_history.get_attachment_file("msg-1", "att-1", current_subject = "unsloth")
assert response.body == PNG_BYTES
assert response.media_type == "image/png"
def test_attachment_file_tolerates_whitespace_in_base64(tmp_path, monkeypatch):
encoded = base64.b64encode(PNG_BYTES).decode("ascii")
wrapped = "\n".join(encoded[i : i + 8] for i in range(0, len(encoded), 8))
attachment = _image_attachment()
attachment["content"] = [{"type": "image", "image": "data:image/png;base64," + wrapped}]
_seed(tmp_path, monkeypatch, [attachment])
response = chat_history.get_attachment_file("msg-1", "att-1", current_subject = "unsloth")
assert response.body == PNG_BYTES
def test_attachment_file_corrupt_base64_is_422(tmp_path, monkeypatch):
attachment = _image_attachment()
attachment["content"] = [{"type": "image", "image": "data:image/png;base64,%%%"}]
_seed(tmp_path, monkeypatch, [attachment])
with pytest.raises(HTTPException) as excinfo:
chat_history.get_attachment_file("msg-1", "att-1", current_subject = "unsloth")
assert excinfo.value.status_code == 422
def test_attachment_file_accepts_urlsafe_base64(tmp_path, monkeypatch):
data = bytes(range(251, 256)) * 3 # encodes to characters remapped by urlsafe
payload = base64.urlsafe_b64encode(data).decode("ascii")
assert "-" in payload or "_" in payload
attachment = _image_attachment()
attachment["content"] = [{"type": "image", "image": "data:image/png;base64," + payload}]
_seed(tmp_path, monkeypatch, [attachment])
response = chat_history.get_attachment_file("msg-1", "att-1", current_subject = "unsloth")
assert response.body == data
def test_attachment_file_accepts_missing_padding(tmp_path, monkeypatch):
payload = base64.b64encode(PNG_BYTES).decode("ascii").rstrip("=")
attachment = _image_attachment()
attachment["content"] = [{"type": "image", "image": "data:image/png;base64," + payload}]
_seed(tmp_path, monkeypatch, [attachment])
response = chat_history.get_attachment_file("msg-1", "att-1", current_subject = "unsloth")
assert response.body == PNG_BYTES
def test_attachment_file_serves_percent_encoded_data_url(tmp_path, monkeypatch):
attachment = _image_attachment()
attachment["content"] = [{"type": "image", "image": "data:text/plain,hello%20world"}]
_seed(tmp_path, monkeypatch, [attachment])
response = chat_history.get_attachment_file("msg-1", "att-1", current_subject = "unsloth")
assert response.body == b"hello world"
# Non-image data URL types are clamped so markup never renders same-origin.
assert response.media_type == "application/octet-stream"
def test_attachment_file_serves_text_parts(tmp_path, monkeypatch):
attachment = {
"id": "att-txt",
"type": "document",
"name": "notes.txt",
"content": [
{"type": "text", "text": "first"},
{"type": "text", "text": "second"},
],
}
_seed(tmp_path, monkeypatch, [attachment])
response = chat_history.get_attachment_file("msg-1", "att-txt", current_subject = "unsloth")
assert response.body.decode("utf-8") == "first\nsecond"
assert response.media_type.startswith("text/plain")
def test_attachment_file_no_content_is_404(tmp_path, monkeypatch):
_seed(tmp_path, monkeypatch, [{"id": "att-empty", "name": "ghost", "content": []}])
with pytest.raises(HTTPException) as excinfo:
chat_history.get_attachment_file("msg-1", "att-empty", current_subject = "unsloth")
assert excinfo.value.status_code == 404
def test_attachment_file_missing_message_is_404(tmp_path, monkeypatch):
_reset_studio_db(tmp_path, monkeypatch)
with pytest.raises(HTTPException) as excinfo:
chat_history.get_attachment_file("nope", "att-1", current_subject = "unsloth")
assert excinfo.value.status_code == 404
def test_attachment_file_non_data_url_image_is_404(tmp_path, monkeypatch):
attachment = _image_attachment()
attachment["content"] = [{"type": "image", "image": "https://example.com/a.png"}]
_seed(tmp_path, monkeypatch, [attachment])
with pytest.raises(HTTPException) as excinfo:
chat_history.get_attachment_file("msg-1", "att-1", current_subject = "unsloth")
assert excinfo.value.status_code == 404
def test_attachment_file_defaults_media_type(tmp_path, monkeypatch):
payload = base64.b64encode(b"raw-bytes").decode("ascii")
attachment = _image_attachment()
attachment["content"] = [{"type": "image", "image": "data:;base64," + payload}]
_seed(tmp_path, monkeypatch, [attachment])
response = chat_history.get_attachment_file("msg-1", "att-1", current_subject = "unsloth")
assert response.body == b"raw-bytes"
assert response.media_type == "application/octet-stream"
def test_attachment_file_svg_media_type(tmp_path, monkeypatch):
svg = b"<svg xmlns='http://www.w3.org/2000/svg'/>"
payload = base64.b64encode(svg).decode("ascii")
attachment = _image_attachment()
attachment["content"] = [{"type": "image", "image": "data:image/svg+xml;base64," + payload}]
_seed(tmp_path, monkeypatch, [attachment])
response = chat_history.get_attachment_file("msg-1", "att-1", current_subject = "unsloth")
assert response.body == svg
# SVG can carry scripts, so it downloads as bytes instead of rendering.
assert response.media_type == "application/octet-stream"
def test_delete_attachment_route_then_404(tmp_path, monkeypatch):
_seed(tmp_path, monkeypatch, [_image_attachment()])
result = chat_history.delete_attachment("msg-1", "att-1", current_subject = "unsloth")
assert result == {"ok": True}
with pytest.raises(HTTPException) as excinfo:
chat_history.delete_attachment("msg-1", "att-1", current_subject = "unsloth")
assert excinfo.value.status_code == 404
# ---------------------------------------------------------------------------
# Audio attachments (adapter {data, format} and compare-chat bare base64)
# ---------------------------------------------------------------------------
WAV_BYTES = b"RIFF$\x00\x00\x00WAVEfmt \x10\x00\x00\x00\x01\x00\x01\x00"
WAV_B64 = base64.b64encode(WAV_BYTES).decode("ascii")
def _audio_attachment(attachment_id: str = "att-audio") -> dict:
return {
"id": attachment_id,
"type": "file",
"name": "clip.wav",
"contentType": "audio/wav",
"content": [{"type": "audio", "audio": {"data": WAV_B64, "format": "wav"}}],
"status": {"type": "complete"},
}
def test_audio_attachment_lists_with_size(tmp_path, monkeypatch):
_seed(tmp_path, monkeypatch, [_audio_attachment()])
records = studio_db.list_chat_attachments()
assert len(records) == 1
assert records[0]["id"] == "att-audio"
assert abs(records[0]["sizeBytes"] - len(WAV_BYTES)) <= 2
def test_audio_attachment_file_serves_bytes(tmp_path, monkeypatch):
_seed(tmp_path, monkeypatch, [_audio_attachment()])
response = chat_history.get_attachment_file("msg-1", "att-audio", current_subject = "unsloth")
assert response.body == WAV_BYTES
assert response.media_type == "audio/wav"
def test_audio_attachment_media_type_from_format(tmp_path, monkeypatch):
attachment = _audio_attachment()
attachment["contentType"] = None
attachment["content"] = [{"type": "audio", "audio": {"data": WAV_B64, "format": "mp3"}}]
_seed(tmp_path, monkeypatch, [attachment])
response = chat_history.get_attachment_file("msg-1", "att-audio", current_subject = "unsloth")
assert response.media_type == "audio/mpeg"
def test_audio_attachment_corrupt_payload_is_422(tmp_path, monkeypatch):
attachment = _audio_attachment()
attachment["content"] = [{"type": "audio", "audio": {"data": "%%%", "format": "wav"}}]
_seed(tmp_path, monkeypatch, [attachment])
with pytest.raises(HTTPException) as excinfo:
chat_history.get_attachment_file("msg-1", "att-audio", current_subject = "unsloth")
assert excinfo.value.status_code == 422
# ---------------------------------------------------------------------------
# Compare-chat uploads stored as message content parts
# ---------------------------------------------------------------------------
def _compare_message(message_id: str = "msg-cmp") -> dict:
return {
"id": message_id,
"threadId": "thread-1",
"parentId": None,
"role": "user",
"content": [
{"type": "image", "image": PNG_DATA_URL},
{"type": "audio", "audio": WAV_B64},
{"type": "text", "text": "compare these"},
],
"createdAt": 1_700_000_000_000,
}
def _seed_compare(tmp_path, monkeypatch):
_reset_studio_db(tmp_path, monkeypatch)
studio_db.upsert_chat_thread(_thread())
studio_db.upsert_chat_message(_compare_message())
_CONTENT_PART_PREFIX = "content-part-sha256-"
def _content_part_id_for(message_id: str, kind: str) -> str:
"""Resolve the stable content-hash id for a message's stored blob.
Content-part ids are SHA-256 hashes of the blob payload, not array
indices, so tests look them up from the listing instead of hardcoding an
index that would shift when an earlier part is deleted.
"""
for record in studio_db.list_chat_attachments():
if record["messageId"] == message_id and record["type"] == kind:
return record["id"]
raise AssertionError(f"no {kind} content-part upload for {message_id}")
def test_content_part_uploads_are_listed(tmp_path, monkeypatch):
_seed_compare(tmp_path, monkeypatch)
records = studio_db.list_chat_attachments()
# Ids are stable content hashes, not array indices.
assert all(r["id"].startswith(_CONTENT_PART_PREFIX) for r in records)
assert {r["type"] for r in records} == {"image", "audio"}
image = next(r for r in records if r["type"] == "image")
assert image["contentType"] == "image/png"
assert abs(image["sizeBytes"] - len(PNG_BYTES)) <= 2
audio = next(r for r in records if r["type"] == "audio")
assert audio["type"] == "audio"
def test_content_part_file_serves_image_bytes(tmp_path, monkeypatch):
_seed_compare(tmp_path, monkeypatch)
image_id = _content_part_id_for("msg-cmp", "image")
response = chat_history.get_attachment_file("msg-cmp", image_id, current_subject = "unsloth")
assert response.body == PNG_BYTES
assert response.media_type == "image/png"
def test_content_part_delete_keeps_text(tmp_path, monkeypatch):
_seed_compare(tmp_path, monkeypatch)
image_id = _content_part_id_for("msg-cmp", "image")
assert studio_db.delete_chat_attachment("msg-cmp", image_id) is True
message = studio_db.get_chat_message("thread-1", "msg-cmp")
types = [p["type"] for p in message["content"]]
assert types == ["audio", "text"]
# The surviving audio blob keeps its own stable hash id after the delete.
remaining = studio_db.list_chat_attachments()
assert [r["type"] for r in remaining] == ["audio"]
assert remaining[0]["id"].startswith(_CONTENT_PART_PREFIX)
assert remaining[0]["id"] != image_id
def test_content_part_delete_rejects_non_blob(tmp_path, monkeypatch):
_seed_compare(tmp_path, monkeypatch)
# The text part is not a stored upload, so it never gets an id: only the
# image and audio blobs are addressable.
assert len(studio_db.list_chat_attachments()) == 2
# A well-formed but unknown content-hash id, and malformed ids, all no-op.
assert studio_db.delete_chat_attachment("msg-cmp", _CONTENT_PART_PREFIX + "0" * 64) is False
assert studio_db.delete_chat_attachment("msg-cmp", "content-part-99") is False
assert studio_db.delete_chat_attachment("msg-cmp", "content-part-x") is False
def test_text_only_messages_not_listed_as_uploads(tmp_path, monkeypatch):
_reset_studio_db(tmp_path, monkeypatch)
studio_db.upsert_chat_thread(_thread())
# The word "image" inside text must not create phantom upload rows.
message = _message("msg-txt")
message["content"] = [{"type": "text", "text": 'discussing an "image" and "audio" here'}]
studio_db.upsert_chat_message(message)
assert studio_db.list_chat_attachments() == []
def test_remote_image_urls_are_not_listed_as_uploads(tmp_path, monkeypatch):
_reset_studio_db(tmp_path, monkeypatch)
studio_db.upsert_chat_thread(_thread())
message = _message("msg-remote")
message["content"] = [
{"type": "image", "image": "https://example.com/cat.png"},
{"type": "text", "text": "look at this"},
]
studio_db.upsert_chat_message(message)
# No stored bytes: nothing to list, open, or delete.
assert studio_db.list_chat_attachments() == []
assert studio_db.get_chat_attachment("msg-remote", "content-part-0") is None
assert studio_db.delete_chat_attachment("msg-remote", "content-part-0") is False
stored = studio_db.get_chat_message("thread-1", "msg-remote")
assert [p["type"] for p in stored["content"]] == ["image", "text"]
def test_html_data_url_serves_as_octet_stream(tmp_path, monkeypatch):
_reset_studio_db(tmp_path, monkeypatch)
studio_db.upsert_chat_thread(_thread())
html_b64 = base64.b64encode(b"<script>alert(1)</script>").decode()
message = _message("msg-html")
message["content"] = [
{"type": "image", "image": f"data:text/html;base64,{html_b64}"},
]
studio_db.upsert_chat_message(message)
attachment_id = _content_part_id_for("msg-html", "image")
response = chat_history.get_attachment_file(
"msg-html", attachment_id, current_subject = "unsloth"
)
# Never echo a script-capable media type back under the app origin.
assert response.media_type == "application/octet-stream"
assert response.body == b"<script>alert(1)</script>"
def test_svg_data_url_serves_as_octet_stream(tmp_path, monkeypatch):
_reset_studio_db(tmp_path, monkeypatch)
studio_db.upsert_chat_thread(_thread())
svg_b64 = base64.b64encode(b"<svg onload='x'/>").decode()
message = _message("msg-svg")
message["content"] = [
{"type": "image", "image": f"data:image/svg+xml;base64,{svg_b64}"},
]
studio_db.upsert_chat_message(message)
attachment_id = _content_part_id_for("msg-svg", "image")
response = chat_history.get_attachment_file("msg-svg", attachment_id, current_subject = "unsloth")
assert response.media_type == "application/octet-stream"
def test_png_data_url_keeps_its_media_type(tmp_path, monkeypatch):
_seed_compare(tmp_path, monkeypatch)
image_id = _content_part_id_for("msg-cmp", "image")
response = chat_history.get_attachment_file("msg-cmp", image_id, current_subject = "unsloth")
assert response.media_type == "image/png"

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

@ -168,11 +168,14 @@ class TestCanLoadGGUF(_GpuCacheResetMixin, unittest.TestCase):
devices,
required_override = None,
estimate = None,
single_device_gpu = None,
gpu_ids = None,
):
with (
patch("utils.hardware.get_device", return_value = DeviceType.CUDA),
patch("utils.hardware.estimate_required_model_memory_gb", return_value = (estimate, {})),
patch("utils.hardware.get_visible_gpu_utilization", return_value = {"devices": devices}),
patch("utils.hardware.resolve_requested_gpu_ids", return_value = gpu_ids),
patch("utils.hardware.auto_select_gpu_ids") as auto_mock,
):
ok, info = tv.can_load_chat_during_training(
@ -180,9 +183,10 @@ class TestCanLoadGGUF(_GpuCacheResetMixin, unittest.TestCase):
hf_token = None,
load_in_4bit = True,
max_seq_length = 0,
requested_gpu_ids = None,
requested_gpu_ids = gpu_ids,
is_gguf = True,
required_override_gb = required_override,
single_device_gpu = single_device_gpu,
)
return ok, info, auto_mock
@ -198,6 +202,88 @@ class TestCanLoadGGUF(_GpuCacheResetMixin, unittest.TestCase):
ok, _, _ = self._run(devices = _devices((0, 80, 35), (1, 80, 70)), required_override = 20.0)
self.assertTrue(ok)
def test_no_per_gpu_floor_for_gguf_with_explicit_gpu_ids(self):
# gpu_ids narrows llama.cpp's candidate pool but does not turn its
# self-placement into HF device_map="balanced". The uneven selected
# pair therefore keeps the aggregate GGUF check without an even-share
# floor on the nearly-full card.
ok, info, _ = self._run(
devices = _devices((0, 80, 35), (1, 80, 70), (2, 80, 0)),
required_override = 20.0,
gpu_ids = [0, 1],
)
self.assertTrue(ok)
self.assertEqual(info["mode"], "gguf")
def test_single_device_uses_selected_gpu(self):
# The model needs 27 GB with headroom. GPU 0 has 45 GB free, while an
# unrelated training-heavy GPU 1 has only 10 GB free.
ok, info, _ = self._run(
devices = _devices((0, 80, 35), (1, 80, 70)),
required_override = 20.0,
single_device_gpu = "0",
)
self.assertTrue(ok)
self.assertEqual(info["usable_gb"], 45.0)
blocked, blocked_info, _ = self._run(
devices = _devices((0, 80, 35), (1, 80, 70)),
required_override = 20.0,
single_device_gpu = "1",
)
self.assertFalse(blocked)
self.assertEqual(blocked_info["usable_gb"], 10.0)
def test_single_device_unresolved_token_sizes_against_worst_device(self):
# A non-numeric device token (a CUDA UUID / MIG handle) can't map to a
# free-VRAM index. The runner still drives ONE device, so size against the
# worst-case visible device (min free), not the aggregate pool: one GPU
# with 80 GB free vs a 20 GB model -> allow.
ok, info, _ = self._run(
devices = _devices((0, 80, 0)),
required_override = 20.0,
single_device_gpu = "GPU-uuid",
)
self.assertTrue(ok)
self.assertEqual(info["mode"], "single_device")
self.assertNotIn("reason", info)
def test_single_device_unresolved_token_refuses_when_worst_device_full(self):
# Same UUID fallback, worst-case device nearly full (2 GB for a 20 GB
# model) -> refuse (default-deny), not on an unresolved-token technicality.
ok, info, _ = self._run(
devices = _devices((0, 80, 78)),
required_override = 20.0,
single_device_gpu = "GPU-uuid",
)
self.assertFalse(ok)
self.assertNotEqual(info.get("reason"), "unresolved_gpu_id")
def test_single_device_unresolved_token_uses_min_free_not_aggregate(self):
# The single-device runner uses ONE device but we can't tell which from a
# UUID token. Sizing against the aggregate pool would let a 20 GB model
# "fit" 160 GB of pooled free VRAM while landing on a 2 GB card and OOMing
# training. Min-free (2 GB) is the safe worst case -> refuse.
ok, info, _ = self._run(
devices = _devices((0, 80, 78), (1, 80, 0), (2, 80, 0)),
required_override = 20.0,
single_device_gpu = "GPU-uuid",
)
self.assertFalse(ok)
self.assertEqual(info["mode"], "single_device")
def test_single_device_cpu_token_allows(self):
# An empty device token = a CPU-only single-device runner (CPU diffusion
# GGUF): it uses no GPU VRAM, so it never threatens training -> allow
# regardless of how full the GPUs are.
ok, info, _ = self._run(
devices = _devices((0, 80, 78)),
required_override = 20.0,
single_device_gpu = "",
)
self.assertTrue(ok)
self.assertEqual(info["reason"], "cpu_only")
def test_estimate_unavailable_refuses(self):
# No override and the estimator can't size it -> default-deny.
ok, info, _ = self._run(devices = _devices((0, 80, 0)), required_override = None, estimate = None)
@ -309,6 +395,8 @@ class TestChatLoadGuardRoute(unittest.TestCase):
captured = None,
training_active,
decision,
gpu_memory_mode = "auto",
requested_gpu_ids = None,
):
config = config or SimpleNamespace(is_gguf = False, is_lora = False, path = None)
with _stub_guard_deps(
@ -320,7 +408,8 @@ class TestChatLoadGuardRoute(unittest.TestCase):
hf_token = None,
load_in_4bit = True,
max_seq_length = 0,
requested_gpu_ids = None,
requested_gpu_ids = requested_gpu_ids,
gpu_memory_mode = gpu_memory_mode,
)
def test_noop_when_training_inactive(self):
@ -332,6 +421,141 @@ class TestChatLoadGuardRoute(unittest.TestCase):
def test_allows_when_fits(self):
self._guard(training_active = True, decision = (True, {"mode": "auto"}))
def test_diffusion_detection_uses_name_before_download(self):
config = SimpleNamespace(
identifier = "unsloth/DiffusionGemma-GGUF",
gguf_hf_repo = "unsloth/DiffusionGemma-GGUF",
gguf_file = None,
)
self.assertTrue(self.route._classify_diffusion_gguf(config))
def test_uncached_gguf_classification_remains_unknown(self):
config = SimpleNamespace(
identifier = "owner/renamed-model",
gguf_hf_repo = "owner/renamed-model",
gguf_variant = "Q4_K_M",
gguf_file = None,
)
self.assertIsNone(self.route._classify_diffusion_gguf(config))
def test_diffusion_detection_reuses_loader_metadata_probe(self):
import tempfile
seen = []
class _Probe:
is_diffusion = False
_architecture = None
def _read_gguf_metadata(self, path):
seen.append(path)
self.is_diffusion = True
with tempfile.TemporaryDirectory() as d:
model = Path(d) / "renamed.gguf"
model.write_bytes(b"GGUF")
config = SimpleNamespace(identifier = "local", gguf_file = str(model))
with patch.object(self.route, "LlamaCppBackend", _Probe):
self.assertTrue(self.route._classify_diffusion_gguf(config))
self.assertEqual(seen, [str(model)])
def test_local_chat_gguf_classification_is_definitive(self):
import tempfile
class _Probe:
is_diffusion = False
_architecture = "llama"
def _read_gguf_metadata(self, _path):
pass
with tempfile.TemporaryDirectory() as d:
model = Path(d) / "renamed.gguf"
model.write_bytes(b"GGUF")
config = SimpleNamespace(identifier = "local", gguf_file = str(model))
with patch.object(self.route, "LlamaCppBackend", _Probe):
self.assertFalse(self.route._classify_diffusion_gguf(config))
def test_manual_known_normal_gguf_bypasses_training_estimate(self):
captured = []
config = SimpleNamespace(is_gguf = True)
with patch.object(self.route, "_classify_diffusion_gguf", return_value = False):
self._guard(
config = config,
captured = captured,
training_active = True,
decision = (False, {"reason": "must not run"}),
gpu_memory_mode = "manual",
)
self.assertEqual(captured, [])
def test_manual_unknown_gguf_keeps_single_device_training_guard(self):
captured = []
config = SimpleNamespace(is_gguf = True)
with (
patch.object(self.route, "_classify_diffusion_gguf", return_value = None),
patch.object(self.route, "_estimate_gguf_required_gb", return_value = 12.5),
patch.object(
self.route.LlamaCppBackend,
"_diffusion_gpu_arg",
return_value = "2",
),
):
self._guard(
config = config,
captured = captured,
training_active = True,
decision = (True, {"mode": "single_device"}),
gpu_memory_mode = "manual",
)
self.assertEqual(len(captured), 1)
self.assertEqual(captured[0]["single_device_gpu"], "2")
def test_manual_diffusion_uses_single_device_guard(self):
captured = []
config = SimpleNamespace(is_gguf = True)
with (
patch.object(self.route, "_classify_diffusion_gguf", return_value = True),
patch.object(self.route, "_estimate_gguf_required_gb", return_value = 12.5),
):
self._guard(
config = config,
captured = captured,
training_active = True,
decision = (True, {"mode": "gguf"}),
gpu_memory_mode = "manual",
requested_gpu_ids = [3, 1],
)
self.assertEqual(len(captured), 1)
self.assertEqual(captured[0]["single_device_gpu"], "1")
self.assertEqual(captured[0]["requested_gpu_ids"], [3, 1])
def test_unpinned_diffusion_uses_runner_default_gpu(self):
captured = []
config = SimpleNamespace(is_gguf = True)
with (
patch.object(self.route, "_classify_diffusion_gguf", return_value = True),
patch.object(self.route, "_estimate_gguf_required_gb", return_value = 12.5),
patch.object(
self.route.LlamaCppBackend,
"_effective_gpu_count",
return_value = 2,
),
patch.object(
self.route.LlamaCppBackend,
"_diffusion_gpu_arg",
return_value = "3",
) as gpu_arg,
):
self._guard(
config = config,
captured = captured,
training_active = True,
decision = (True, {"mode": "single_device"}),
gpu_memory_mode = "manual",
)
gpu_arg.assert_called_once_with(None, cpu_only = False)
self.assertEqual(captured[0]["single_device_gpu"], "3")
def test_refuses_with_headroom_number(self):
info = {"required_gb": 30.0, "usable_gb": 6.0, "needed_gb": 39.0, "mode": "auto"}
with self.assertRaises(HTTPException) as exc:
@ -467,36 +691,189 @@ class TestValidateRefusesDuringTraining(unittest.TestCase):
self.assertEqual(captured[0]["load_in_4bit"], False)
self.assertEqual(captured[0]["max_seq_length"], 4096)
def test_rejects_gguf_with_gpu_ids_before_guard(self):
# /validate must mirror /load's GGUF + gpu_ids 400, before the VRAM guard.
def test_validate_forwards_manual_gpu_memory_mode_to_guard(self):
from models.inference import ValidateModelRequest
request = ValidateModelRequest(model_path = "x.gguf", gpu_ids = [0])
request = ValidateModelRequest(
model_path = "unsloth/model-GGUF",
gguf_variant = "Q4_K_M",
gpu_memory_mode = "manual",
)
cfg = SimpleNamespace(
identifier = "x.gguf",
display_name = "x",
identifier = "unsloth/model-GGUF",
display_name = "model-GGUF",
is_gguf = True,
is_lora = False,
is_vision = False,
path = None,
base_model = None,
)
captured = []
captured = {}
with (
patch.object(
self.route,
"_resolve_model_identifier_for_request",
return_value = ("x.gguf", "x.gguf", False),
return_value = ("unsloth/model-GGUF", "unsloth/model-GGUF", False),
),
patch.object(self.route.ModelConfig, "from_identifier", return_value = cfg),
patch.object(self.route, "load_inference_config", return_value = {}),
_stub_guard_deps(training_active = True, decision = (True, {}), captured = captured),
patch.object(
self.route,
"_guard_chat_load_against_training",
lambda config, **kw: captured.update(kw),
),
):
with self.assertRaises(HTTPException) as exc:
asyncio.run(self.route.validate_model(request, current_subject = "u"))
self.assertEqual(exc.exception.status_code, 400)
self.assertIn("gpu_ids is not supported for GGUF", exc.exception.detail)
self.assertEqual(captured, []) # guard never reached
asyncio.run(self.route.validate_model(request, current_subject = "u"))
self.assertEqual(captured.get("gpu_memory_mode"), "manual")
def test_validate_forwards_inherited_extras_and_parallel_to_guard(self):
# Regression: /load resolves inherited same-model extras and passes the
# real slot count to the guard; validate must do the same, else it sizes
# a smaller estimate (no inherited -c/--model-draft, n_parallel=1) and
# /load then 409s after the frontend has already unloaded.
from models.inference import ValidateModelRequest
request = ValidateModelRequest(model_path = "unsloth/Qwen3-1.7B", max_seq_length = 4096)
cfg = SimpleNamespace(
identifier = "unsloth/Qwen3-1.7B",
display_name = "Qwen3-1.7B",
is_gguf = False,
is_lora = False,
is_vision = False,
path = None,
base_model = None,
)
captured = {}
with (
patch.object(
self.route,
"_resolve_model_identifier_for_request",
return_value = ("unsloth/Qwen3-1.7B", "unsloth/Qwen3-1.7B", False),
),
patch.object(self.route.ModelConfig, "from_identifier", return_value = cfg),
patch.object(self.route, "load_inference_config", return_value = {}),
patch.object(self.route, "_resolve_inherited_extra_args", return_value = ["-c", "32768"]),
patch.object(
self.route,
"_guard_chat_load_against_training",
lambda config, **kw: captured.update(kw),
),
):
asyncio.run(self.route.validate_model(request, current_subject = "u"))
self.assertEqual(captured.get("llama_extra_args"), ["-c", "32768"])
self.assertIn("n_parallel", captured)
def test_metadata_probe_skips_training_guard(self):
# A header-only probe (include_context_length) allocates no VRAM, so the
# training guard must not run -- else the staging GPU-layers / MoE sliders
# it feeds are hidden exactly when a during-training user needs them.
from models.inference import ValidateModelRequest
request = ValidateModelRequest(
model_path = "unsloth/Qwen3-1.7B",
max_seq_length = 4096,
include_context_length = True,
)
cfg = SimpleNamespace(
identifier = "unsloth/Qwen3-1.7B",
display_name = "Qwen3-1.7B",
is_gguf = False,
is_lora = False,
is_vision = False,
path = None,
base_model = None,
)
guard_called = []
with (
patch.object(
self.route,
"_resolve_model_identifier_for_request",
return_value = ("unsloth/Qwen3-1.7B", "unsloth/Qwen3-1.7B", False),
),
patch.object(self.route.ModelConfig, "from_identifier", return_value = cfg),
patch.object(self.route, "load_inference_config", return_value = {}),
patch.object(
self.route,
"_guard_chat_load_against_training",
lambda *a, **kw: guard_called.append(True),
),
):
asyncio.run(self.route.validate_model(request, current_subject = "u"))
self.assertEqual(guard_called, [])
def _validate_gguf_template(
self,
*,
template,
canonical_path = "/picked/model.gguf",
):
# Drive validate_model for a native lease-backed GGUF template probe and
# capture what the embedded-template reader was called with.
from models.inference import ValidateModelRequest
request = ValidateModelRequest(
model_path = "model.gguf",
gguf_variant = "Q4_K_M",
native_path_lease = "signed-lease",
include_chat_template = True,
)
cfg = SimpleNamespace(
identifier = canonical_path,
display_name = "model.gguf",
is_gguf = True,
is_lora = False,
is_vision = False,
gguf_file = canonical_path,
path = None,
base_model = None,
)
import utils.models.gguf_metadata as gguf_meta
seen = {}
def _fake_read(path):
seen["path"] = path
return template
guard_called = []
with (
patch.object(
self.route,
"_resolve_model_identifier_for_request",
return_value = (canonical_path, "model.gguf", True),
),
patch.object(self.route.ModelConfig, "from_identifier", return_value = cfg),
patch.object(self.route, "load_inference_config", return_value = {}),
patch.object(gguf_meta, "read_gguf_chat_template", _fake_read),
patch.object(
self.route,
"_guard_chat_load_against_training",
lambda *a, **kw: guard_called.append(True),
),
):
resp = asyncio.run(self.route.validate_model(request, current_subject = "u"))
return resp, seen, guard_called
def test_include_chat_template_reads_leased_gguf_embedded_template(self):
# The picker chat-template GET has no lease plumbing, so a native picked
# GGUF surfaces its default template through this lease-aware probe: the
# embedded template is read from the granted canonical path and returned.
resp, seen, _ = self._validate_gguf_template(template = "{{ messages }}")
self.assertEqual(resp.chat_template, "{{ messages }}")
# Read strictly the leased file's own embedded template, never a sibling
# sidecar: the grant authorizes just this one path.
self.assertEqual(seen["path"], "/picked/model.gguf")
def test_include_chat_template_skips_training_guard(self):
# A template-only probe allocates no VRAM, so like include_context_length
# it must not be refused by the training guard.
_, _, guard_called = self._validate_gguf_template(template = "{{ messages }}")
self.assertEqual(guard_called, [])
def test_include_chat_template_over_cap_is_dropped(self):
from picker.schemas import MAX_CHAT_TEMPLATE_BYTES
resp, _, _ = self._validate_gguf_template(template = "a" * (MAX_CHAT_TEMPLATE_BYTES + 1))
self.assertIsNone(resp.chat_template)
# ── _estimate_gguf_required_gb (sizes the same weights the loader loads) ──────

73
studio/backend/tests/test_cuda_torch_spec.py Normal file
View file

@ -0,0 +1,73 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Tests for _CUDA_TORCH_PKG_SPEC in install_python_stack.py.
The CUDA repair path installs the torch trio from an exclusive --index-url (no
PyPI fallback), so these pinned ranges decide which torch the venv gets. The
upper bound is locked to the 2.11.x family to match the base image and rocm7.2
spec and to keep the companions off a torch-2.12 wheel that would ABI-mismatch.
"""
from __future__ import annotations
import sys
from pathlib import Path
import pytest
from packaging.requirements import Requirement
# install_python_stack.py lives at repo_root/studio/install_python_stack.py
_INSTALL_SCRIPT = Path(__file__).resolve().parents[2] / "install_python_stack.py"
def _load_module(monkeypatch):
"""(Re-)import and return install_python_stack (mirrors test_torchao_select)."""
sys.modules.pop("install_python_stack", None)
monkeypatch.syspath_prepend(str(_INSTALL_SCRIPT.parent))
import install_python_stack
return install_python_stack
def _spec_of(pkg_spec: str):
"""Parse 'torch>=2.4,<2.12.0' into a packaging SpecifierSet."""
return Requirement(pkg_spec).specifier
@pytest.mark.parametrize(
"index, allowed, rejected",
[
# torch: 2.11.x allowed (matches base image); 2.12.x excluded.
(0, ["2.11.0", "2.11.2", "2.10.0", "2.4.0"], ["2.12.0", "2.3.0", "1.13.1"]),
# torchvision: 0.26.x (torch 2.11 companion) allowed; 0.27.x (torch 2.12) out.
(1, ["0.26.0", "0.26.1", "0.19.0"], ["0.27.0", "0.18.0"]),
# torchaudio: same 2.11.x window as torch.
(2, ["2.11.0", "2.10.0", "2.4.0"], ["2.12.0", "2.3.0"]),
],
)
def test_cuda_spec_bounds(monkeypatch, index, allowed, rejected):
mod = _load_module(monkeypatch)
spec = _spec_of(mod._CUDA_TORCH_PKG_SPEC[index])
for v in allowed:
assert spec.contains(v, prereleases = True), f"{v} should satisfy {spec}"
for v in rejected:
assert not spec.contains(v, prereleases = True), f"{v} should not satisfy {spec}"
def test_cuda_spec_matches_rocm72_upper_bound(monkeypatch):
"""CUDA and rocm7.2 target the same torch 2.11.x family, so their upper
bounds must stay in lockstep (bump both together at 2.12.x)."""
mod = _load_module(monkeypatch)
rocm72 = mod._ROCM_TORCH_PKG_SPECS["rocm7.2"]
def _upper(pkg_spec: str) -> str:
for clause in _spec_of(pkg_spec):
if clause.operator == "<":
return clause.version
raise AssertionError(f"no upper bound in {pkg_spec!r}")
for cuda_pkg, rocm_pkg in zip(mod._CUDA_TORCH_PKG_SPEC, rocm72, strict = True):
assert _upper(cuda_pkg) == _upper(
rocm_pkg
), f"CUDA {cuda_pkg!r} upper bound must match rocm7.2 {rocm_pkg!r}"

1
studio/backend/tests/test_export_absolute_paths.py
View file

@ -158,6 +158,7 @@ def _install_lightweight_backend_stubs(monkeypatch):
utils_model_config._pick_best_gguf = lambda variants: variants[0] if variants else None
utils_model_config._extract_quant_label = lambda value: value
utils_model_config._is_big_endian_gguf_path = lambda *args, **kwargs: False
utils_model_config._is_mtp_drafter = lambda *args, **kwargs: False
utils_model_config.is_audio_input_type = lambda *args, **kwargs: None
monkeypatch.setitem(
sys.modules,

6
studio/backend/tests/test_gguf_load_cache_reuse.py
View file

@ -728,9 +728,13 @@ class TestLoadHubDownloadExclusion:
source = (Path(__file__).resolve().parent.parent / "routes" / "inference.py").read_text()
gguf_branch = source[source.index("if config.is_gguf:") :]
# The gguf_load_in_flight marker must be entered before the hub-download
# guard and the unload so a concurrent load can't race the download
# manager. The llama_extra_args inheritance that used to sit between the
# marker and the guard now runs in _guard_chat_load_against_training, ahead
# of the GGUF branch, so it is no longer a landmark inside this slice.
assert (
gguf_branch.index("enter_context(gguf_load_in_flight")
< gguf_branch.index("if request.llama_extra_args is None")
< gguf_branch.index("_hub_download_blocks_gguf_load")
< gguf_branch.index("unsloth_backend.unload_model")
)

73
studio/backend/tests/test_gguf_metadata.py
View file

@ -15,6 +15,7 @@ from utils.models.gguf_metadata import (
pairing_score,
read_gguf_context_length,
read_gguf_general_metadata,
read_gguf_staged_dims,
read_mmproj_audio_capability,
)
@ -153,6 +154,78 @@ def test_context_length_ignores_foreign_arch_key(tmp_path: Path):
assert read_gguf_context_length(str(p)) is None
# --- read_gguf_staged_dims (one pass: context + layer + moe counts) ----
def test_staged_dims_none_for_missing_or_non_gguf(tmp_path: Path):
assert read_gguf_staged_dims(str(tmp_path / "nope.gguf")) is None
p = tmp_path / "garbage.gguf"
p.write_bytes(b"not a gguf at all")
assert read_gguf_staged_dims(str(p)) is None
def test_staged_dims_moe_with_leading_dense(tmp_path: Path):
# GLM-4.7-Flash shape: context + total layers + MoE layers in one read.
p = _write_synthetic_gguf(
tmp_path / "glm.gguf",
{"general.architecture": "deepseek2"},
extra_uint32 = {
"deepseek2.context_length": 202752,
"deepseek2.block_count": 47,
"deepseek2.expert_count": 64,
"deepseek2.leading_dense_block_count": 1,
},
)
assert read_gguf_staged_dims(str(p)) == {
"context_length": 202752,
"layer_count": 47,
"moe_layer_count": 46,
}
def test_staged_dims_dense_model(tmp_path: Path):
# Dense: layer_count present, moe_layer_count 0 (slider hidden).
p = _write_synthetic_gguf(
tmp_path / "dense.gguf",
{"general.architecture": "qwen3"},
extra_uint32 = {"qwen3.context_length": 40960, "qwen3.block_count": 36},
)
assert read_gguf_staged_dims(str(p)) == {
"context_length": 40960,
"layer_count": 36,
"moe_layer_count": 0,
}
def test_staged_dims_all_moe_no_leading_dense(tmp_path: Path):
# Experts present, no leading_dense key -> every block is a MoE layer.
p = _write_synthetic_gguf(
tmp_path / "moe.gguf",
{"general.architecture": "qwen35moe"},
extra_uint32 = {"qwen35moe.block_count": 40, "qwen35moe.expert_count": 256},
)
assert read_gguf_staged_dims(str(p)) == {
"context_length": None,
"layer_count": 40,
"moe_layer_count": 40,
}
def test_staged_dims_uint64_block_count(tmp_path: Path):
# block_count stored as uint64 (vtype 10) still parses; moe == block_count.
p = _write_synthetic_gguf(
tmp_path / "moe64.gguf",
{"general.architecture": "gpt-oss"},
extra_uint32 = {"gpt-oss.expert_count": 32},
extra_uint64 = {"gpt-oss.block_count": 24},
)
assert read_gguf_staged_dims(str(p)) == {
"context_length": None,
"layer_count": 24,
"moe_layer_count": 24,
}
def test_context_length_read_from_uint64(tmp_path: Path):
# Some models store context_length as a uint64 (vtype 10).
p = _write_synthetic_gguf(

879
studio/backend/tests/test_gpu_memory_mode.py Normal file
View file

@ -0,0 +1,879 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Backend contract for the GPU Memory mode dropdown.
The dropdown threads a single ``gpu_memory_mode`` ("auto" | "manual") from the
chat UI through the load request. "manual" lets the user own the offload: with
``gpu_layers < 0`` (Auto, the default) it hands all memory management to
llama.cpp's ``--fit on`` (no CUDA/HIP device masking, no context auto-reduce, no
gpu-layer or tensor-split planning); with ``gpu_layers >= 0`` it pins the layers
and MoE offload itself (``--fit off``). These tests pin:
* the pydantic request/response/status contract (snake_case key, default
"auto", unknown values rejected),
* the backend ``gpu_memory_mode`` property and its reset on unload,
* the ``_already_in_target_state`` reload-detection branch, and
* that the manual + Auto-layers branch in ``load_model`` empties the probed
GPU set and drops tensor parallelism so the selection below no-ops, while
the explicit-offload branch emits ``--gpu-layers`` / ``--fit off``.
"""
from __future__ import annotations
import inspect
import sys
import types as _types
from pathlib import Path
import pytest
_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
if _BACKEND_DIR not in sys.path:
sys.path.insert(0, _BACKEND_DIR)
# Same external-dep stubs as the other llama_cpp unit tests so importing
# the backend doesn't drag in structlog / httpx / loggers.
_loggers_stub = _types.ModuleType("loggers")
_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name)
sys.modules.setdefault("loggers", _loggers_stub)
_structlog_stub = _types.ModuleType("structlog")
_structlog_stub.get_logger = lambda *a, **k: __import__("logging").getLogger("stub")
sys.modules.setdefault("structlog", _structlog_stub)
# httpx is a real, installed backend dependency: import it so the genuine module
# is in sys.modules. A hand-rolled stub here is inevitably incomplete and, since
# setdefault installs it before real httpx loads, would poison a combined pytest
# run -- routes/inference references httpx.Response (and other attrs) at def time.
import httpx # noqa: F401
from core.inference import llama_cpp as llama_cpp_module
from core.inference.llama_cpp import LlamaCppBackend
from models.inference import (
InferenceStatusResponse,
LoadRequest,
LoadResponse,
)
# ── Pydantic contract (snake_case key, default "auto") ───────────────
def test_load_request_defaults_gpu_memory_mode_auto():
assert LoadRequest(model_path = "owner/repo").gpu_memory_mode == "auto"
def test_load_request_round_trips_json_key():
req = LoadRequest.model_validate({"model_path": "owner/repo", "gpu_memory_mode": "manual"})
assert req.gpu_memory_mode == "manual"
assert req.model_dump()["gpu_memory_mode"] == "manual"
def test_load_request_rejects_unknown_mode():
with pytest.raises(ValueError):
LoadRequest(model_path = "owner/repo", gpu_memory_mode = "bogus")
@pytest.mark.parametrize("model_cls", [LoadResponse, InferenceStatusResponse])
def test_response_models_emit_gpu_memory_mode(model_cls):
if model_cls is LoadResponse:
default = model_cls(
status = "loaded",
model = "owner/repo",
display_name = "repo",
inference = {},
)
manual = model_cls(
status = "loaded",
model = "owner/repo",
display_name = "repo",
inference = {},
gpu_memory_mode = "manual",
)
else:
default = model_cls()
manual = model_cls(gpu_memory_mode = "manual")
assert default.model_dump()["gpu_memory_mode"] == "auto"
assert manual.model_dump()["gpu_memory_mode"] == "manual"
# ── Backend property + reset ─────────────────────────────────────────
class _FakeProcess:
"""Stand-in for subprocess.Popen so _kill_process is a no-op."""
def terminate(self):
pass
def wait(self, timeout = None):
return 0
def kill(self):
pass
def poll(self):
return 0
def test_gpu_memory_mode_property_defaults_auto():
assert LlamaCppBackend().gpu_memory_mode == "auto"
def test_gpu_memory_mode_property_reflects_field():
backend = LlamaCppBackend()
backend._gpu_memory_mode = "manual"
assert backend.gpu_memory_mode == "manual"
def test_unload_resets_gpu_memory_mode():
backend = LlamaCppBackend()
backend._process = _FakeProcess()
backend._gpu_memory_mode = "manual"
backend.unload_model()
assert backend.gpu_memory_mode == "auto"
# ── _already_in_target_state reload-detection branch ─────────────────
def _loaded_backend(gpu_memory_mode: str) -> LlamaCppBackend:
backend = LlamaCppBackend()
backend._process = _FakeProcess() # is_loaded only checks "is not None"
backend._healthy = True
backend._model_identifier = "owner/repo"
backend._hf_variant = "Q4_K_M"
backend._requested_n_ctx = 8192
backend._cache_type_kv = None
backend._requested_spec_mode = "auto"
backend._chat_template_override = None
backend._is_vision = False
backend._extra_args = None
backend._gguf_path = None
backend._gpu_memory_mode = gpu_memory_mode
return backend
def _target_state(backend: LlamaCppBackend, gpu_memory_mode: str) -> bool:
return backend._already_in_target_state(
gguf_path = None,
model_identifier = "owner/repo",
hf_variant = "Q4_K_M",
n_ctx = 8192,
cache_type_kv = None,
speculative_type = "auto",
chat_template_override = None,
extra_args = None,
is_vision = False,
gpu_memory_mode = gpu_memory_mode,
)
@pytest.mark.parametrize("mode", ["auto", "manual"])
def test_already_in_target_state_matches_same_mode(mode):
assert _target_state(_loaded_backend(mode), mode) is True
@pytest.mark.parametrize("loaded,requested", [("auto", "manual"), ("manual", "auto")])
def test_already_in_target_state_reloads_on_mode_change(loaded, requested):
# Flipping the dropdown either direction must force a reload so the command
# is rebuilt with/without the Unsloth GPU masking.
assert _target_state(_loaded_backend(loaded), requested) is False
def test_already_in_target_state_ignores_mode_for_diffusion():
# The diffusion runner is mode-agnostic (always "auto"), so a standing manual
# preference must not force a needless reload.
backend = _loaded_backend("auto")
backend._is_diffusion = True
assert _target_state(backend, "manual") is True
# ── load_model: manual + Auto layers bypasses Unsloth GPU management ──
def _load_model_source() -> str:
return inspect.getsource(llama_cpp_module.LlamaCppBackend.load_model)
def test_auto_layers_branch_empties_gpus_and_drops_tensor_parallel():
# Emptying the probed set makes the selection / TP planning below no-op, so
# gpu_indices stays None and use_fit True (--fit on).
src = _load_model_source()
gate = src.find('if gpu_memory_mode == "manual" and gpu_layers < 0:')
assert gate != -1, "load_model must branch on manual + Auto layers (gpu_layers < 0)"
block = src[gate : gate + 1400]
assert "gpus = []" in block, "Auto-layers branch must empty the probed GPU set"
# --fit aborts under --split-mode tensor, so a raw-extras split-mode is stripped.
assert "strip_split_mode_only(extra_args)" in block
assert "requested_ctx if requested_ctx > 0 else 0" in block
# The branch sits before GPU selection assigns gpu_indices; --fit on is its emission.
assert gate < src.find("gpu_indices, use_fit = None, True")
assert 'cmd.extend(["--fit", "on"])' in src
# TP drops for this path, but at a guard BEFORE the quantized-KV cache-drop, so
# a requested quantized cache survives into the --fit load.
tp_drop = src.find('if tensor_parallel and gpu_memory_mode == "manual" and gpu_layers < 0:')
assert tp_drop != -1, "manual + Auto layers must drop tensor_parallel"
assert "tensor_parallel = False" in src[tp_drop : tp_drop + 400]
cache_drop = src.find("Tensor parallelism requires a non-quantized KV cache")
assert cache_drop != -1
assert (
tp_drop < cache_drop
), "TP must drop before the cache-drop so a quantized KV survives --fit"
def test_auto_layers_never_sends_ctx_size_zero():
# Sending "-c 0" sets fit_params_min_ctx = UINT32_MAX in llama.cpp, pinning
# the full native context and disabling --fit's reduction. So the base cmd
# must never carry -c, "-c 0" is emitted only outside the Auto-layers (--fit)
# case, and a positive context is passed through (which --fit optimizes
# layers around).
src = _load_model_source()
base_start = src.find("cmd = [")
base_end = src.find("\n ]", base_start)
base_block = src[base_start:base_end]
assert '"-c"' not in base_block, "-c must be conditional, not in the base cmd list"
assert 'cmd.extend(["-c", str(effective_ctx)])' in src, "positive ctx must pass -c"
assert 'auto_fit = gpu_memory_mode == "manual" and gpu_layers < 0' in src
zero = src.find('cmd.extend(["-c", "0"])')
assert zero != -1, '"-c 0" emission must exist outside the Auto-layers case'
guard = src.rfind("elif not auto_fit:", 0, zero)
assert guard != -1 and zero - guard < 120, '"-c 0" must sit under the not-auto_fit guard'
def test_manual_mode_clears_inherited_main_model_placement_env():
env = {name: "inherited" for name in LlamaCppBackend._MANUAL_PLACEMENT_ENV_VARS}
env["LLAMA_ARG_N_GPU_LAYERS_DRAFT"] = "7"
env["UNRELATED"] = "kept"
LlamaCppBackend._clear_manual_placement_env(env)
assert not (set(env) & set(LlamaCppBackend._MANUAL_PLACEMENT_ENV_VARS))
assert env["LLAMA_ARG_N_GPU_LAYERS_DRAFT"] == "7"
assert env["UNRELATED"] == "kept"
def test_load_model_sanitizes_manual_env_after_building_child_env():
src = _load_model_source()
env_build = src.find("env = self._llama_server_env_for_binary(binary)")
env_clear = src.find("self._clear_manual_placement_env(env)", env_build)
launch = src.find("subprocess.Popen", env_build)
assert env_build != -1
assert env_build < env_clear < launch
# ── Manual offload (--gpu-layers + --fit off + --n-cpu-moe) ───────────
def test_load_request_accepts_manual():
req = LoadRequest(
model_path = "owner/repo",
gpu_memory_mode = "manual",
gpu_layers = 20,
n_cpu_moe = 8,
tensor_split = [2, 1],
)
assert req.gpu_memory_mode == "manual"
assert req.gpu_layers == 20
assert req.n_cpu_moe == 8
assert req.tensor_split == [2, 1]
def test_load_request_manual_defaults():
req = LoadRequest(model_path = "owner/repo")
assert req.gpu_layers == -1
assert req.n_cpu_moe == 0
assert req.tensor_split is None
@pytest.mark.parametrize("bad", [[0, 0], [-1, 2], [float("inf"), 1], [float("nan"), 1]])
def test_load_request_rejects_degenerate_tensor_split(bad):
# A negative/non-finite/all-zero split is dropped at launch but compared raw
# in the reload dedupe, so it would reload forever -- reject it up front.
with pytest.raises(ValueError):
LoadRequest(model_path = "owner/repo", tensor_split = bad)
@pytest.mark.parametrize("good", [[2, 1], [1, 1], [], None])
def test_load_request_accepts_valid_tensor_split(good):
assert LoadRequest(model_path = "owner/repo", tensor_split = good).tensor_split == good
def test_route_normalizes_explicit_extras_before_reload_dedupe():
route_src = (Path(_BACKEND_DIR) / "routes" / "inference.py").read_text(encoding = "utf-8")
load_impl = route_src[route_src.index("async def _load_model_impl") :]
strip = load_impl.index("_stripped_explicit = strip_shadowing_flags")
normalize = load_impl.index(
'request = request.model_copy(update = {"llama_extra_args": extra_llama_args})'
)
dedupe = load_impl.index("and _request_matches_loaded_settings(")
assert strip < normalize < dedupe
@pytest.mark.parametrize("model_cls", [LoadResponse, InferenceStatusResponse])
def test_response_models_emit_manual_fields(model_cls):
if model_cls is LoadResponse:
obj = model_cls(
status = "loaded",
model = "owner/repo",
display_name = "repo",
inference = {},
gpu_memory_mode = "manual",
gpu_layers = 20,
n_cpu_moe = 8,
tensor_split = [2, 1],
n_layers = 32,
n_moe_layers = 32,
)
else:
obj = model_cls(
gpu_memory_mode = "manual",
gpu_layers = 20,
n_cpu_moe = 8,
tensor_split = [2, 1],
n_layers = 32,
n_moe_layers = 32,
)
dumped = obj.model_dump()
assert dumped["gpu_memory_mode"] == "manual"
assert dumped["gpu_layers"] == 20
assert dumped["n_cpu_moe"] == 8
assert dumped["tensor_split"] == [2, 1]
assert dumped["n_layers"] == 32
assert dumped["n_moe_layers"] == 32
def test_manual_properties_default_and_reflect_and_reset():
backend = LlamaCppBackend()
assert backend.gpu_layers == -1 and backend.n_cpu_moe == 0
assert backend.tensor_split is None
backend._gpu_layers = 20
backend._n_cpu_moe = 8
backend._tensor_split = [2, 1]
assert backend.gpu_layers == 20 and backend.n_cpu_moe == 8
assert backend.tensor_split == [2, 1]
backend._process = _FakeProcess()
backend.unload_model()
assert backend.gpu_layers == -1 and backend.n_cpu_moe == 0
assert backend.tensor_split is None
def test_n_moe_layers_property():
# 0 for a dense model (hides the slider); block_count for all-MoE;
# block_count - leading_dense otherwise (GLM-4.7-Flash: 47 - 1 -> 46).
b = LlamaCppBackend()
b._n_layers = 36
b._n_experts = None
assert b.n_moe_layers == 0
b._n_experts = 128
b._leading_dense_block_count = None
assert b.n_moe_layers == 36
b._n_layers = 47
b._leading_dense_block_count = 1
assert b.n_moe_layers == 46
def _target_state_manual(
backend,
*,
gpu_layers,
n_cpu_moe,
tensor_split = None,
):
return backend._already_in_target_state(
gguf_path = None,
model_identifier = "owner/repo",
hf_variant = "Q4_K_M",
n_ctx = 8192,
cache_type_kv = None,
speculative_type = "auto",
chat_template_override = None,
extra_args = None,
is_vision = False,
gpu_memory_mode = "manual",
gpu_layers = gpu_layers,
n_cpu_moe = n_cpu_moe,
tensor_split = tensor_split,
)
def test_manual_reloads_on_gpu_layers_or_n_cpu_moe_or_split_change():
backend = _loaded_backend("manual")
backend._gpu_layers = 20
backend._n_cpu_moe = 0
backend._tensor_split = None
# Same knobs -> no reload.
assert _target_state_manual(backend, gpu_layers = 20, n_cpu_moe = 0) is True
# Changed layer count -> reload.
assert _target_state_manual(backend, gpu_layers = 16, n_cpu_moe = 0) is False
# Changed MoE offload -> reload.
assert _target_state_manual(backend, gpu_layers = 20, n_cpu_moe = 8) is False
# Added a GPU split -> reload.
assert _target_state_manual(backend, gpu_layers = 20, n_cpu_moe = 0, tensor_split = [2, 1]) is False
# Same GPU split -> no reload.
backend._tensor_split = [2, 1]
assert _target_state_manual(backend, gpu_layers = 20, n_cpu_moe = 0, tensor_split = [2, 1]) is True
def test_auto_layers_reload_tracks_only_gpu_layers():
# Under Auto (gpu_layers < 0) the MoE/split knobs don't apply, so a leftover
# request value must not reload -- only a gpu_layers change (Auto -> pinned) does.
backend = _loaded_backend("manual")
backend._gpu_layers = -1
backend._n_cpu_moe = 0
backend._tensor_split = None
# Same Auto, leftover MoE/split in the request -> still no reload.
assert _target_state_manual(backend, gpu_layers = -1, n_cpu_moe = 8, tensor_split = [2, 1]) is True
# Auto -> explicit offload reloads.
assert _target_state_manual(backend, gpu_layers = 20, n_cpu_moe = 0) is False
def test_manual_offload_emits_gpu_layers_fit_off_and_n_cpu_moe():
src = _load_model_source()
gate = src.find('elif gpu_memory_mode == "manual":')
assert gate != -1, "load_model must have an explicit-offload manual branch"
block = src[gate : gate + 700]
# Empties the probed set (skips the planner) but keeps the user's TP choice
# (only the Auto-layers branch above drops TP).
assert "gpus = []" in block
assert "tensor_parallel = False" not in block
# The cmd emits the layer count with fit disabled, gated on gpu_layers >= 0.
assert 'if gpu_memory_mode == "manual" and gpu_layers >= 0:' in src
assert 'cmd.extend(["--gpu-layers", str(gpu_layers), "--fit", "off"])' in src
# MoE offload uses --n-cpu-moe via _resolve_cpu_moe_flag (tested behaviorally below).
assert "_resolve_cpu_moe_flag(" in src
assert 'cmd.extend(["--n-cpu-moe", str(moe_flag)])' in src
# A count requested on a dense model is never emitted, so it must also be
# dropped from the recorded state -- else /status and /load report a count
# llama-server never received (same rule as the tensor-split drop below).
moe_emit = src.find('cmd.extend(["--n-cpu-moe", str(moe_flag)])')
assert "elif n_cpu_moe:" in src[moe_emit : moe_emit + 300]
assert "self._n_cpu_moe = 0" in src[moe_emit : moe_emit + 300]
# The offload path forces use_fit False so --fit-ctx is never added under --fit off.
emit = src.find('cmd.extend(["--gpu-layers", str(gpu_layers), "--fit", "off"])')
assert "use_fit = False" in src[src.rfind("\n", 0, emit) - 200 : emit + 80]
def test_status_reports_requested_context_length():
# The hydration path re-seeds a Manual+Auto context pin from the REQUESTED
# n_ctx (0 = Auto); context_length only exposes the resolved value.
assert "requested_context_length" in InferenceStatusResponse.model_fields
s = InferenceStatusResponse(requested_context_length = 8192)
assert s.model_dump()["requested_context_length"] == 8192
assert InferenceStatusResponse().model_dump()["requested_context_length"] is None
# The /status route must actually wire it from the backend (a declared-but-
# never-populated field would leave hydration silently reverting the pin).
from pathlib import Path as _P
route_src = (_P(_BACKEND_DIR) / "routes" / "inference.py").read_text(encoding = "utf-8")
assert "requested_context_length = llama_backend.requested_n_ctx" in route_src
def test_manual_offload_emits_tensor_split():
# The offload path emits --tensor-split from the per-GPU shares, only when
# provided, with >1 GPU in use, AND matching that count (a stale ratio on a
# narrowed picker or a mismatched direct-API list must not emit -- llama-
# server aborts on a split/GPU-count mismatch).
src = _load_model_source()
assert "if tensor_split and _split_gpus > 1:" in src
# Emit only on a length match AND a positive sanitized total: a mismatched
# or all-zero split aborts llama-server / assigns nothing, so it's dropped.
# The emitted list is the sanitized one (clamping tested behaviorally below).
assert "_sanitized_split = self._sanitize_tensor_split(tensor_split)" in src
assert "if len(_sanitized_split) == _split_gpus and _split_total > 0:" in src
assert '"--tensor-split"' in src
# Joined as a comma list (e.g. "2,1") within the explicit-offload cmd branch.
gate = src.find('if gpu_memory_mode == "manual" and gpu_layers >= 0:')
nxt = src.find("elif use_fit:", gate)
assert '","' in src[gate:nxt] and "tensor_split" in src[gate:nxt]
# A split with a single effective GPU is never emitted, so it must also be
# dropped from the recorded state -- else /status and /load report a ratio
# llama-server never received and the dedupe baseline preserves it.
assert "elif tensor_split:" in src[gate:nxt]
drop = src.find("elif tensor_split:", gate, nxt)
assert "self._tensor_split = None" in src[drop : drop + 250]
def test_sanitize_tensor_split_clamps_negative_and_non_finite():
# Negative entries would launch a placement different from the ratio the
# UI showed; inf passes a plain > 0 total gate and would emit
# "--tensor-split inf,..." (llama.cpp normalizes shares by the running
# total, so an inf poisons the shares from that entry on). Both clamp to 0.
sanitize = LlamaCppBackend._sanitize_tensor_split
assert sanitize([2, 1]) == [2.0, 1.0]
assert sanitize([-1, 2]) == [0.0, 2.0]
assert sanitize([float("inf"), 1]) == [0.0, 1.0]
assert sanitize([float("nan"), 1]) == [0.0, 1.0]
# All-zero survives sanitization; the call site's total gate drops it.
assert sanitize([0, 0]) == [0.0, 0.0]
# Unreadable input -> []; the call site's length gate drops it.
assert sanitize(["x", 1]) == []
assert sanitize([10**400, 1]) == []
def test_zero_offload_mask_honors_device_pin_spellings():
# A user device pin must keep the GPUs visible: llama-server aborts on a
# pin it can't see ('error: invalid device'). The pin can arrive as
# --device or its -dev alias, as the draft forms (parsed even with no
# drafter loaded), or as an inherited LLAMA_ARG_DEVICE env var.
load_src = _load_model_source()
assert "self._zero_offload_keeps_gpu_visible(cmd, env)" in load_src
block = inspect.getsource(LlamaCppBackend._cmd_has_gpu_device_pin)
for flag in (
'"--device"',
'"-dev"',
'"--spec-draft-device"',
'"-devd"',
'"--device-draft"',
):
assert flag in block
assert '"LLAMA_ARG_DEVICE"' in block
def test_resolve_cpu_moe_flag():
# Clamp the requested MoE-layer count to the model's MoE layers, then offset
# past leading dense layers (--n-cpu-moe counts from layer 0).
R = LlamaCppBackend._resolve_cpu_moe_flag
assert R(0, 40, 0) is None # nothing requested
assert R(8, 0, 0) is None # dense model (no MoE layers)
assert R(8, 40, 0) == 8 # all-MoE: direct
assert R(100, 40, 0) == 40 # clamp to the MoE layer count
# GLM-4.7-Flash (deepseek2): block_count 47, leading_dense 1, n_moe 46.
assert R(5, 46, 1) == 6 # offset past the 1 dense layer
assert R(46, 46, 1) == 47 # all MoE on CPU == block_count
def test_manual_allows_tensor_parallel_via_split_mode():
# Manual offload keeps the user's TP choice but skips the memory-based planner
# (plan_tp excludes manual, so its empty gpu set can't downgrade TP). The
# --split-mode tensor emission gates on tensor_parallel alone, so manual
# reaches it -- with tp_tensor_split None it's an even split (no
# --tensor-split). --fit off means no fit/tensor abort.
src = _load_model_source()
assert 'plan_tp = tensor_parallel and gpu_memory_mode != "manual"' in src
assert "if plan_tp:" in src
assert "if plan_tp and len(tp_gpus) < 2:" in src
sm = src.find('cmd.extend(["--split-mode", "tensor"])')
assert sm != -1, "TP must emit --split-mode tensor"
guard = src.rfind("if tensor_parallel:", 0, sm)
assert guard != -1 and sm - guard < 200, "split-mode gates on tensor_parallel"
# The tensor-split is only emitted for a planned (non-even) split, which
# manual never produces, so manual stays an even split.
assert "if tp_tensor_split and len(tp_tensor_split) > 1:" in src
def test_fit_sets_target_margin():
# Manual + Auto (auto_fit) tightens the per-device VRAM margin to 512 MiB.
caps = {"supports_fit_target": True}
flags = LlamaCppBackend._ctx_integrity_flags(1, True, True, 0, 0, caps)
assert flags[flags.index("--fit-target") + 1] == "512"
# Not emitted on the legacy auto path (fit on but not auto_fit): -c 0 pins
# native there, so the tighter margin must not ride along.
assert "--fit-target" not in LlamaCppBackend._ctx_integrity_flags(1, True, False, 0, 0, caps)
# Not emitted when fit is off.
assert "--fit-target" not in LlamaCppBackend._ctx_integrity_flags(1, False, False, 0, 0, caps)
# Not emitted when the binary lacks support.
assert "--fit-target" not in LlamaCppBackend._ctx_integrity_flags(
1, True, True, 0, 0, {"supports_fit_target": False}
)
# ── GPU picker (gpu_ids -> CUDA_VISIBLE_DEVICES) ─────────────────────
def test_load_request_accepts_gpu_ids():
req = LoadRequest(model_path = "owner/repo", gpu_ids = [1, 0])
assert req.gpu_ids == [1, 0]
assert LoadRequest(model_path = "owner/repo").gpu_ids is None
@pytest.mark.parametrize("model_cls", [LoadResponse, InferenceStatusResponse])
def test_response_models_emit_gpu_ids(model_cls):
if model_cls is LoadResponse:
obj = model_cls(status = "loaded", model = "m", display_name = "m", inference = {}, gpu_ids = [1])
else:
obj = model_cls(gpu_ids = [1])
assert obj.model_dump()["gpu_ids"] == [1]
def test_gpu_ids_property_default_and_reset():
backend = LlamaCppBackend()
assert backend.gpu_ids is None
backend._gpu_ids = [0, 1]
assert backend.gpu_ids == [0, 1]
backend._process = _FakeProcess()
backend.unload_model()
assert backend.gpu_ids is None
def _target_state_gpu_ids(backend, gpu_ids):
return backend._already_in_target_state(
gguf_path = None,
model_identifier = "owner/repo",
hf_variant = "Q4_K_M",
n_ctx = 8192,
cache_type_kv = None,
speculative_type = "auto",
chat_template_override = None,
extra_args = None,
is_vision = False,
gpu_ids = gpu_ids,
)
def test_gpu_ids_reload_detection_is_order_insensitive():
backend = _loaded_backend("auto")
backend._gpu_ids = [0, 1]
# Same set, different order -> no reload.
assert _target_state_gpu_ids(backend, [1, 0]) is True
# Different set -> reload.
assert _target_state_gpu_ids(backend, [0]) is False
# Dropping the pick (auto) -> reload.
assert _target_state_gpu_ids(backend, None) is False
def test_gpu_ids_reload_detection_collapses_diffusion_to_single_device():
# The diffusion runner drives only its single lowest device, so the backend
# records [lowest]. A later multi-GPU request that still resolves to that
# same lowest device must dedupe (no needless reload); a request whose lowest
# device moves, or that drops the pick, must reload.
backend = _loaded_backend("auto")
backend._is_diffusion = True
backend._gpu_ids = [1] # loaded on the lowest of an earlier [3, 1] pick
assert _target_state_gpu_ids(backend, [3, 1]) is True
assert _target_state_gpu_ids(backend, [1]) is True
# Lowest device changes (2, not 1) -> reload.
assert _target_state_gpu_ids(backend, [3, 2]) is False
# Dropping the pick (auto) -> reload.
assert _target_state_gpu_ids(backend, None) is False
def test_start_diffusion_server_resets_tensor_parallel():
# A prior tensor-parallel chat load leaves self._tensor_parallel True (load_model
# phase 1 only kills the process, it skips the unload reset). Diffusion is never
# TP, so startup must clear it -- else /status misreports TP and an identical
# diffusion re-Apply reloads against stale tensor-parallel state.
src = inspect.getsource(llama_cpp_module.LlamaCppBackend._start_diffusion_server)
assert "self._tensor_parallel = False" in src
def test_route_matches_loaded_settings_collapses_diffusion_gpu_ids():
# The route-level reload dedupe mirrors the backend: for a loaded diffusion
# model it compares the request against the single recorded device, not the
# full requested list, or a same-device multi-GPU pick reloads needlessly.
route_src = (Path(_BACKEND_DIR) / "routes" / "inference.py").read_text(encoding = "utf-8")
match_impl = route_src[route_src.index("def _request_matches_loaded_settings") :]
guard = match_impl.index("if llama_backend.is_diffusion:")
collapse = match_impl.index("[sorted(request.gpu_ids)[0]] if request.gpu_ids else None")
compare = match_impl.index("if _req_gpu_ids != llama_backend.gpu_ids:")
assert guard < collapse < compare
# ── Manual tensor split: child enumeration pinned to the picker's order ──────
def _patch_split_pin_env(monkeypatch, *, inherited, reported):
"""Point the pin helper at a fake inherited mask and picker report.
``reported`` None = enumeration unavailable (falls back to ascending)."""
import utils.hardware as hw
monkeypatch.setattr(
LlamaCppBackend, "_resolve_visible_physical_ids", staticmethod(lambda: inherited)
)
info = (
{"available": False}
if reported is None
else {
"available": True,
"index_kind": "physical",
"devices": [{"index": i} for i in reported],
}
)
monkeypatch.setattr(hw, "get_backend_visible_gpu_info", lambda: info)
def test_split_pin_reorders_inherited_numeric_mask(monkeypatch):
# Parent CUDA_VISIBLE_DEVICES=3,1 makes the child enumerate dev0=phys3, but
# nvidia-smi reported the picker's list ascending -- the mask must be
# re-emitted in that order or the per-GPU shares land on the wrong cards.
_patch_split_pin_env(monkeypatch, inherited = [3, 1], reported = [1, 3])
env = {"CUDA_VISIBLE_DEVICES": "3,1"}
LlamaCppBackend._pin_visible_gpu_order_for_split(env)
assert env["CUDA_DEVICE_ORDER"] == "PCI_BUS_ID"
assert env["CUDA_VISIBLE_DEVICES"] == "1,3"
def test_split_pin_keeps_mask_order_when_picker_reported_it(monkeypatch):
# Torch-fallback enumeration (no nvidia-smi) reports devices in inherited
# mask order, so the picker's split list follows the mask -- the pin must
# keep that order, not re-sort it into a mismatch.
_patch_split_pin_env(monkeypatch, inherited = [3, 1], reported = [3, 1])
env = {"CUDA_VISIBLE_DEVICES": "3,1"}
LlamaCppBackend._pin_visible_gpu_order_for_split(env)
assert env["CUDA_VISIBLE_DEVICES"] == "3,1"
def test_split_pin_falls_back_to_ascending_without_report(monkeypatch):
# Enumeration unavailable: ascending physical is the best guess (it matches
# the dominant nvidia-smi report order).
_patch_split_pin_env(monkeypatch, inherited = [3, 1], reported = None)
env = {"CUDA_VISIBLE_DEVICES": "3,1"}
LlamaCppBackend._pin_visible_gpu_order_for_split(env)
assert env["CUDA_VISIBLE_DEVICES"] == "1,3"
def test_split_pin_without_mask_only_sets_pci_order(monkeypatch):
# No inherited mask (or a UUID/MIG one resolving to None): enumeration order
# is fully fixed by CUDA_DEVICE_ORDER, so no mask is written.
_patch_split_pin_env(monkeypatch, inherited = None, reported = None)
env = {}
LlamaCppBackend._pin_visible_gpu_order_for_split(env)
assert env == {"CUDA_DEVICE_ORDER": "PCI_BUS_ID"}
def test_split_pin_mirrors_hip_mask_on_rocm(monkeypatch):
# ROCm: the pin must land in HIP_VISIBLE_DEVICES too, and an inherited ROCR
# mask is cleared so the mask can't apply twice (ROCR re-indexes, then HIP
# would index into the already-reduced set).
_patch_split_pin_env(monkeypatch, inherited = [3, 1], reported = [1, 3])
torch_stub = _types.ModuleType("torch")
torch_stub.version = _types.SimpleNamespace(hip = "6.0")
monkeypatch.setitem(sys.modules, "torch", torch_stub)
env = {"CUDA_VISIBLE_DEVICES": "3,1", "ROCR_VISIBLE_DEVICES": "3,1"}
LlamaCppBackend._pin_visible_gpu_order_for_split(env)
assert env["CUDA_VISIBLE_DEVICES"] == "1,3"
assert env["HIP_VISIBLE_DEVICES"] == "1,3"
assert "ROCR_VISIBLE_DEVICES" not in env
# ── Diffusion single-device selection ───────────────────────────────────────
def test_diffusion_gpu_arg_uses_lowest_explicit_physical_id(monkeypatch):
monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "3,1")
monkeypatch.setenv("DG_GPU", "7")
assert LlamaCppBackend._diffusion_gpu_arg([3, 1]) == "1"
def test_diffusion_gpu_arg_preserves_parent_mask_order(monkeypatch):
monkeypatch.delenv("DG_GPU", raising = False)
monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "3,1")
assert LlamaCppBackend._diffusion_gpu_arg(None) == "3"
def test_diffusion_gpu_arg_honors_override_and_cpu_mask(monkeypatch):
monkeypatch.setenv("DG_GPU", "GPU-abc")
assert LlamaCppBackend._diffusion_gpu_arg(None) == "GPU-abc"
assert LlamaCppBackend._diffusion_gpu_arg(None, cpu_only = True) == ""
# ── Deliberate zero-offload (manual gpu_layers=0): training-skip flag ─────────
def test_zero_offload_flag_false_without_companions():
# CPU-only by construction: False lets training skip unloading a server that
# holds no VRAM.
cmd = ["llama-server", "-m", "model.gguf", "--gpu-layers", "0", "--fit", "off"]
assert LlamaCppBackend._zero_offload_gpu_flag(cmd, [(0, 8000, 24000)], {}) is False
@pytest.mark.parametrize(
"companion",
["--mmproj", "--model-draft", "-md", "--spec-draft-model", "-hfd"],
)
def test_zero_offload_flag_true_with_companion(companion):
# mmproj / a drafter offload to GPU regardless of --gpu-layers, so the
# server still holds VRAM and training must unload it. Drafter detection
# reuses the extras parser, so pass-through aliases count too.
cmd = ["llama-server", "-m", "model.gguf", "--gpu-layers", "0", companion, "x.gguf"]
assert LlamaCppBackend._zero_offload_gpu_flag(cmd, [(0, 8000, 24000)], {}) is True
def test_zero_offload_flag_true_with_inline_companion_forms():
cmd = ["llama-server", "-m", "model.gguf", "--spec-draft-model=x.gguf"]
assert LlamaCppBackend._zero_offload_gpu_flag(cmd, [(0, 8000, 24000)], {}) is True
cmd = ["llama-server", "-m", "model.gguf", "--mmproj=proj.gguf"]
assert LlamaCppBackend._zero_offload_gpu_flag(cmd, [(0, 8000, 24000)], {}) is True
def test_zero_offload_flag_true_with_env_drafter():
cmd = ["llama-server", "-m", "model.gguf", "--gpu-layers", "0"]
env = {"LLAMA_ARG_SPEC_DRAFT_MODEL": "x.gguf"}
assert LlamaCppBackend._zero_offload_gpu_flag(cmd, [(0, 8000, 24000)], env) is True
@pytest.mark.parametrize(
"device_args",
[
["--device", "CUDA0"],
["--device=CUDA0"],
["-dev", "CUDA0"],
["--spec-draft-device", "CUDA0"],
["--device-draft=CUDA0"],
],
)
def test_zero_offload_flag_true_with_device_pin(device_args):
cmd = ["llama-server", "-m", "model.gguf", "--gpu-layers", "0", *device_args]
assert LlamaCppBackend._zero_offload_gpu_flag(cmd, [(0, 8000, 24000)], {}) is True
def test_zero_offload_flag_true_with_env_device_pin():
cmd = ["llama-server", "-m", "model.gguf", "--gpu-layers", "0"]
env = {"LLAMA_ARG_DEVICE": "CUDA0"}
assert LlamaCppBackend._zero_offload_gpu_flag(cmd, [(0, 8000, 24000)], env) is True
@pytest.mark.parametrize(
("device_args", "env"),
[
(["--device", "cpu"], {}),
(["--device=none"], {}),
(["--spec-draft-device", "cpu"], {}),
([], {"LLAMA_ARG_DEVICE": "none"}),
(["--device", "CUDA0", "--device", "cpu"], {}),
],
)
def test_zero_offload_flag_false_with_cpu_device_pin(device_args, env):
cmd = ["llama-server", "-m", "model.gguf", "--gpu-layers", "0", *device_args]
assert LlamaCppBackend._zero_offload_gpu_flag(cmd, [(0, 8000, 24000)], env) is False
def test_zero_offload_flag_true_with_surviving_tensor_mode():
cmd = ["llama-server", "-m", "model.gguf", "--gpu-layers", "0", "--split-mode", "tensor"]
assert LlamaCppBackend._zero_offload_gpu_flag(cmd, [(0, 8000, 24000)], {}) is True
def test_zero_offload_flag_true_for_unmasked_vulkan(monkeypatch):
monkeypatch.setattr(LlamaCppBackend, "_is_vulkan_backend", staticmethod(lambda: True))
cmd = ["llama-server", "-m", "model.gguf", "--gpu-layers", "0"]
assert LlamaCppBackend._zero_offload_gpu_flag(cmd, [(0, 8000, 24000)], {}) is True
def test_zero_offload_flag_none_without_gpus():
cmd = ["llama-server", "-m", "model.gguf", "--gpu-layers", "0"]
assert LlamaCppBackend._zero_offload_gpu_flag(cmd, [], {}) is None
def test_cmd_has_gpu_companion_detection():
# The env mask for CPU-only zero-offload loads keys off this scan: any
# --mmproj form or a drafter (flag aliases / env) keeps the GPUs visible.
has = LlamaCppBackend._cmd_has_gpu_companion
assert has(["llama-server", "-m", "m.gguf"], {}) is False
assert has(["llama-server", "--mmproj", "p.gguf"], {}) is True
assert has(["llama-server", "--mmproj=p.gguf"], {}) is True
assert has(["llama-server", "-md", "d.gguf"], {}) is True
assert has(["llama-server"], {"LLAMA_ARG_SPEC_DRAFT_MODEL": "d.gguf"}) is True
def test_cmd_companion_ignores_cpu_forced_drafter():
# A CPU-pinned drafter holds no VRAM: the zero-offload mask may hide the GPUs
# and training may leave the server alone.
has = LlamaCppBackend._cmd_has_gpu_companion
cmd = ["llama-server", "-md", "d.gguf", "--spec-draft-ngl", "0"]
assert has(cmd, {}) is False
cmd = ["llama-server", "-md", "d.gguf", "--spec-draft-device", "cpu"]
assert has(cmd, {}) is False
# mmproj still counts even alongside a CPU drafter.
cmd = ["llama-server", "-md", "d.gguf", "--spec-draft-ngl", "0", "--mmproj", "p.gguf"]
assert has(cmd, {}) is True

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

@ -853,7 +853,13 @@ class TestRouteErrors(unittest.TestCase):
self.assertIn("only supported on CUDA devices", str(exc_info.exception))
def test_inference_route_rejects_gpu_ids_for_gguf(self):
def test_inference_route_validates_gpu_ids_for_gguf(self):
# gpu_ids is now SUPPORTED for GGUF (the GPU picker), but still
# validated: a rejected pick surfaces as a clean 400, not the old
# "not supported for GGUF" rejection. Patch the validator so the test
# is deterministic regardless of the host's (or a prior test's) GPU env.
import utils.hardware.hardware as hardware_mod
inference_route = _load_route_module(
"inference_route_module_for_gguf_gpu_ids_test",
"routes/inference.py",
@ -887,6 +893,11 @@ class TestRouteErrors(unittest.TestCase):
),
patch.object(inference_route.asyncio, "to_thread", new = _inline_to_thread),
patch.object(inference_route, "_hf_offline_if_dns_dead", nullcontext),
patch.object(
hardware_mod,
"resolve_requested_gpu_ids",
side_effect = ValueError("Invalid gpu_ids [0, 1]: rejected by test"),
),
):
with self.assertRaises(HTTPException) as exc_info:
asyncio.run(
@ -901,8 +912,11 @@ class TestRouteErrors(unittest.TestCase):
)
)
# The validator's ValueError becomes a clean 400 (not the removed
# "not supported for GGUF" rejection).
self.assertEqual(exc_info.exception.status_code, 400)
self.assertIn("GGUF", exc_info.exception.detail)
self.assertIn("gpu_ids", exc_info.exception.detail.lower())
self.assertNotIn("not supported", exc_info.exception.detail.lower())
def test_training_route_returns_400_for_invalid_gpu_ids(self):
training_route = _load_route_module(

165
studio/backend/tests/test_hf_token_validation.py Normal file
View file

@ -0,0 +1,165 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Focused coverage for cached, rate-limited HF token validation."""
from __future__ import annotations
from pathlib import Path
import sys
import httpx
import pytest
_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
if _BACKEND_DIR not in sys.path:
sys.path.insert(0, _BACKEND_DIR)
import utils.hf_token_validation as validation
@pytest.fixture(autouse = True)
def _reset_validation_state():
validation.reset_hf_token_validation_state()
yield
validation.reset_hf_token_validation_state()
def test_cached_token_does_not_spend_another_attempt(monkeypatch):
calls = []
def _check(token):
calls.append(token)
return validation.TokenValidationResult(status = "valid")
monkeypatch.setattr(validation, "_check_remote", _check)
first = validation.validate_hf_token("hf_valid", rate_key = "user:ip")
second = validation.validate_hf_token("hf_valid", rate_key = "user:ip")
assert first.status == second.status == "valid"
assert calls == ["hf_valid"]
def test_three_uncached_attempts_per_hour(monkeypatch):
monkeypatch.setattr(
validation,
"_check_remote",
lambda _token: validation.TokenValidationResult(status = "invalid"),
)
for index in range(3):
result = validation.validate_hf_token(f"hf_bad_{index}", rate_key = "user:ip")
assert result.status == "invalid"
limited = validation.validate_hf_token("hf_bad_4", rate_key = "user:ip")
assert limited.status == "rate_limited"
assert limited.retry_after_seconds is not None
assert limited.retry_after_seconds > 0
other_user = validation.validate_hf_token("hf_other", rate_key = "other:ip")
assert other_user.status == "invalid"
def test_window_rolls_forward(monkeypatch):
clock = {"now": 100.0}
monkeypatch.setattr(validation.time, "monotonic", lambda: clock["now"])
monkeypatch.setattr(validation, "_MAX_ATTEMPTS", 1)
monkeypatch.setattr(validation, "_WINDOW_SECONDS", 10.0)
monkeypatch.setattr(
validation,
"_check_remote",
lambda _token: validation.TokenValidationResult(status = "invalid"),
)
assert validation.validate_hf_token("hf_a", rate_key = "user:ip").status == "invalid"
assert validation.validate_hf_token("hf_b", rate_key = "user:ip").status == "rate_limited"
clock["now"] += 11.0
assert validation.validate_hf_token("hf_b", rate_key = "user:ip").status == "invalid"
@pytest.mark.parametrize(
("status_code", "expected"),
[(200, "valid"), (401, "invalid"), (429, "rate_limited"), (500, "unavailable")],
)
def test_remote_status_classification(monkeypatch, status_code, expected):
response = httpx.Response(
status_code,
request = httpx.Request("GET", "https://huggingface.co/api/whoami-v2"),
headers = {"Retry-After": "42"} if status_code == 429 else None,
)
class _Session:
def get(self, url, *, headers, timeout):
assert url == "https://huggingface.co/api/whoami-v2"
assert headers["authorization"] == "Bearer hf_test"
assert timeout == validation._REMOTE_TIMEOUT_SECONDS
return response
monkeypatch.setattr(validation, "get_session", lambda: _Session())
result = validation._check_remote("hf_test")
assert result.status == expected
if status_code == 429:
assert result.retry_after_seconds == 42
def test_wrapped_http_401_is_invalid(monkeypatch):
response = httpx.Response(
401,
request = httpx.Request("GET", "https://huggingface.co/api/whoami-v2"),
)
class _Session:
def get(self, _url, **_kwargs):
error = RuntimeError("Invalid user token.")
error.response = response
raise error
monkeypatch.setattr(validation, "get_session", lambda: _Session())
assert validation._check_remote("hf_test").status == "invalid"
def test_remote_timeout_is_bounded_and_unavailable(monkeypatch):
class _Session:
def get(self, _url, *, headers, timeout):
assert headers["authorization"] == "Bearer hf_test"
assert timeout == validation._REMOTE_TIMEOUT_SECONDS
raise TimeoutError("timed out")
monkeypatch.setattr(validation, "get_session", lambda: _Session())
assert validation._check_remote("hf_test").status == "unavailable"
def test_raw_token_is_not_retained(monkeypatch):
monkeypatch.setattr(
validation,
"_check_remote",
lambda _token: validation.TokenValidationResult(status = "valid"),
)
token = "hf_do_not_store_this_value"
validation.validate_hf_token(token, rate_key = "user:ip")
assert token not in repr(validation._cache)
assert token not in repr(validation._attempts)
def test_unexpected_remote_exception_releases_singleflight(monkeypatch):
calls = 0
monkeypatch.setattr(validation, "_INFLIGHT_WAIT_SECONDS", 0.0)
def _check(_token):
nonlocal calls
calls += 1
if calls == 1:
raise RuntimeError("unexpected failure")
return validation.TokenValidationResult(status = "valid")
monkeypatch.setattr(validation, "_check_remote", _check)
with pytest.raises(RuntimeError, match = "unexpected failure"):
validation.validate_hf_token("hf_test", rate_key = "user:ip")
result = validation.validate_hf_token("hf_test", rate_key = "user:ip")
assert result.status == "valid"
assert calls == 2
assert validation._inflight == {}

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

@ -445,6 +445,101 @@ def test_route_to_vulkan_prebuilt_cpu_fallback_wins():
assert routed is host
@pytest.mark.parametrize("cpu_flag", ["--cpu-fallback", "--force-cpu"])
def test_resolve_prebuilt_cpu_fallback_overrides_intel_vulkan(monkeypatch, capsys, cpu_flag):
"""Either CPU flag via CLI must suppress Vulkan even on an Intel GPU host: both
drop GPU detection (--force-cpu additionally persists, on the install path)."""
monkeypatch.setattr(
ilp,
"detect_host",
lambda: _host(is_linux = True, is_x86_64 = True, has_intel_gpu = True),
)
seen = {}
def _resolver(tag, host, repo, published_release_tag):
seen["host"] = host
seen["repo"] = repo
raise ilp.PrebuiltFallback("no asset")
monkeypatch.setattr(ilp, "resolve_simple_install_release_plans", _resolver)
monkeypatch.setattr(
sys,
"argv",
[
"install_llama_prebuilt.py",
"--resolve-prebuilt",
"latest",
cpu_flag,
"--output-format",
"json",
],
)
assert ilp.main() == ilp.EXIT_SUCCESS
# The CPU flag must suppress Intel GPU, route to fork (not upstream Vulkan)
assert seen["host"].has_intel_gpu is False
assert seen["repo"] == FORK
@pytest.mark.parametrize(
"flags, expect_force, expect_persist",
[
([], False, False),
# Automatic/transient last resort (arm64 GPU-build recovery): drops GPU but
# does NOT persist, so a later update heals to a GPU bundle (#6097).
(["--cpu-fallback"], True, False),
# Deliberate CPU-only (UNSLOTH_LLAMA_CPP_BACKEND=cpu): drops GPU AND persists so
# the updater re-asserts it and never revives the Intel iGPU crash (#7213).
(["--force-cpu"], True, True),
(["--cpu-fallback", "--force-cpu"], True, True),
],
)
def test_cli_cpu_flags_thread_force_and_persist(
monkeypatch, tmp_path, flags, expect_force, expect_persist
):
captured = {}
monkeypatch.setattr(ilp, "install_prebuilt", lambda **kw: captured.update(kw))
monkeypatch.setattr(
sys,
"argv",
["install_llama_prebuilt.py", "--install-dir", str(tmp_path / "llama.cpp"), *flags],
)
assert ilp.main() == ilp.EXIT_SUCCESS
assert captured["force_cpu"] is expect_force
assert captured["persist_force_cpu"] is expect_persist
@pytest.mark.parametrize(
"existing, requested, expected",
[
# A deliberate --force-cpu on top of a naturally-installed CPU bundle (same
# asset, install skipped) must still flip the marker to true (#7213).
(False, True, True),
(None, True, True),
# No spurious writes when already in sync, and a released force syncs down.
(True, True, True),
(False, False, False),
(True, False, False),
],
)
def test_sync_marker_force_cpu(tmp_path, existing, requested, expected):
marker = {"tag": "b9585", "asset": "llama-b9585-bin-ubuntu-x64.tar.gz"}
if existing is not None:
marker["force_cpu"] = existing
marker_path = tmp_path / "UNSLOTH_PREBUILT_INFO.json"
marker_path.write_text(json.dumps(marker))
ilp.sync_marker_force_cpu(tmp_path, requested)
written = json.loads(marker_path.read_text())
assert written["force_cpu"] is expected
# Unrelated fields are preserved.
assert written["asset"] == "llama-b9585-bin-ubuntu-x64.tar.gz"
def test_sync_marker_force_cpu_missing_marker_is_noop(tmp_path):
# No marker (or unreadable) must not crash the reuse path.
ilp.sync_marker_force_cpu(tmp_path, True)
assert not (tmp_path / "UNSLOTH_PREBUILT_INFO.json").exists()
def test_route_to_vulkan_prebuilt_hidden_nvidia_not_rerouted():
# A mixed NVIDIA+Intel host that hid NVIDIA (CUDA_VISIBLE_DEVICES=""/-1):
# physical NVIDIA present but not usable. Must NOT auto-route to Vulkan, or

74
studio/backend/tests/test_llama_cpp_mtp_detection.py
View file

@ -9,6 +9,7 @@ the _already_in_target_state mirror that prevents needless reloads.
from __future__ import annotations
import ast
import inspect
import os
import struct
@ -345,10 +346,62 @@ def test_windows_full_offload_flags_use_current_llama_server_args():
stale_checkpoint_flag = "--checkpoint-" + "every-n-tokens"
assert '"--cache-ram"' in src
assert '"--ctx-checkpoints"' in src
assert '"--no-cache-prompt"' in src
# Prompt caching stays on (in-VRAM prefix reuse); #5692 only needed the host-RAM
# checkpoints (--cache-ram / --ctx-checkpoints) disabled, not prompt reuse.
assert '"--no-cache-prompt"' not in src
assert stale_checkpoint_flag not in src
# Backend-wide guard: Unsloth must never inject --no-cache-prompt into a llama-server
# command. It disables in-VRAM prompt-prefix reuse, re-prefilling every repeated prompt
# (#5692 only needed --cache-ram / --ctx-checkpoints off; #7260 dropped the stray flag).
# Detecting it (_is_real) or honouring a user-supplied one (_prompt_cache_off) is fine.
_NO_CACHE_PROMPT_FLAG = "--no-cache-prompt"
_LIST_MUTATORS = frozenset({"append", "extend", "insert"})
def _has_flag_literal(node: ast.AST) -> bool:
return any(
isinstance(n, ast.Constant) and n.value == _NO_CACHE_PROMPT_FLAG for n in ast.walk(node)
)
def _no_cache_prompt_injections(source: str, filename: str) -> list[tuple[str, int]]:
"""(file, lineno) for each spot adding --no-cache-prompt to a list."""
hits: list[tuple[str, int]] = []
for node in ast.walk(ast.parse(source, filename = filename)):
# cmd.append/extend/insert(... flag ...) or cmd += [... flag ...]
if (
isinstance(node, ast.Call)
and isinstance(node.func, ast.Attribute)
and node.func.attr in _LIST_MUTATORS
and any(_has_flag_literal(a) for a in node.args)
) or (
isinstance(node, ast.AugAssign)
and isinstance(node.op, ast.Add)
and _has_flag_literal(node.value)
):
hits.append((filename, node.lineno))
return hits
def test_unsloth_never_injects_no_cache_prompt_into_any_command():
root = Path(_BACKEND_DIR)
files = [p for p in root.rglob("*.py") if "tests" not in p.relative_to(root).parts]
violations: list[tuple[str, int]] = []
for path in files:
try:
violations += _no_cache_prompt_injections(path.read_text(encoding = "utf-8"), str(path))
except (OSError, UnicodeDecodeError, SyntaxError):
continue
assert files, "no backend source files were scanned"
assert violations == [], (
"Unsloth must never add --no-cache-prompt to a llama-server command "
"(it disables prompt-prefix reuse); detecting or honouring a user-supplied "
f"one is fine. Offending sites: {violations}"
)
def test_load_model_sets_threads_once():
src = inspect.getsource(LlamaCppBackend.load_model)
assert src.count('cmd.extend(["--threads", str(') == 1
@ -741,6 +794,25 @@ def test_probe_reports_windows_cache_flags_absent_for_older_binary(tmp_path):
assert caps["supports_no_cache_prompt"] is False
@_NEEDS_BASH
def test_probe_detects_slot_save_path(tmp_path):
fake = _make_fake_llama_server(
tmp_path / "llama-server",
"--slot-save-path PATH path to save slot kv cache\n--threads N\n",
)
_clear_caps_cache()
caps = LlamaCppBackend.probe_server_capabilities(str(fake))
assert caps["supports_slot_save"] is True
@_NEEDS_BASH
def test_probe_reports_slot_save_absent_for_older_binary(tmp_path):
fake = _make_fake_llama_server(tmp_path / "llama-server", "--threads N\n")
_clear_caps_cache()
caps = LlamaCppBackend.probe_server_capabilities(str(fake))
assert caps["supports_slot_save"] is False
def test_build_ngram_mod_flags_new():
flags = _build_ngram_mod_flags({"ngram_mod_flavor": "new"})
assert flags == [

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

@ -118,9 +118,17 @@ def test_flag_sits_inside_the_base_cmd_list():
"conditional branch -- otherwise some code paths would still "
"run with silent context shift enabled."
)
# Pin that it sits next to -c / --ctx so the grouping makes sense.
assert '"-c"' in block
assert '"--flash-attn"' in block
# -c is emitted in the conditional right after the base list, not inside
# it: auto-fit (--fit on with no pinned context) must omit -c entirely,
# because "-c 0" pins the full native context and disables --fit's
# VRAM-based sizing. Pin that it still sits next to the base block so the
# context grouping stays intact.
after = rest[end_rel : end_rel + 1000]
assert '"-c"' in after, (
"-c must still be emitted in the conditional immediately after the "
"base cmd list (omitted only in auto-fit, where --fit sizes context)."
)
def _iter_lines_with_offset(text: str):

31
studio/backend/tests/test_llama_cpp_props_readback.py
View file

@ -225,31 +225,46 @@ def test_kv_unified_added_for_multi_slot():
"""Explicit --parallel N disables llama-server's auto-slots kv-unified
default, splitting -c into per-slot windows of -c/N; Unsloth must restore
the shared pool so one request can use the full advertised context."""
flags = LlamaCppBackend._ctx_integrity_flags(4, False, 98304, 98304, _CAPS_ALL)
flags = LlamaCppBackend._ctx_integrity_flags(4, False, False, 98304, 98304, _CAPS_ALL)
assert "--kv-unified" in flags
def test_kv_unified_skipped_for_single_slot_or_old_build():
assert "--kv-unified" not in LlamaCppBackend._ctx_integrity_flags(
1, False, 98304, 98304, _CAPS_ALL
1, False, False, 98304, 98304, _CAPS_ALL
)
assert "--kv-unified" not in LlamaCppBackend._ctx_integrity_flags(
4, False, 98304, 98304, _CAPS_NONE
4, False, False, 98304, 98304, _CAPS_NONE
)
def test_fit_ctx_floors_explicit_request_under_fit():
flags = LlamaCppBackend._ctx_integrity_flags(1, True, 98304, 98304, _CAPS_ALL)
# An explicit requested ctx floors --fit-ctx at that value on any --fit
# path, including legacy auto (auto_fit False).
flags = LlamaCppBackend._ctx_integrity_flags(1, True, False, 98304, 98304, _CAPS_ALL)
assert flags[flags.index("--fit-ctx") + 1] == "98304"
def test_fit_ctx_skipped_without_fit_or_explicit_ctx_or_support():
def test_fit_ctx_skipped_without_fit_or_support():
# No --fit on -> no --fit-ctx.
assert "--fit-ctx" not in LlamaCppBackend._ctx_integrity_flags(
1, False, 98304, 98304, _CAPS_ALL
1, False, False, 98304, 98304, _CAPS_ALL
)
assert "--fit-ctx" not in LlamaCppBackend._ctx_integrity_flags(1, True, 0, 262144, _CAPS_ALL)
# --fit on but the binary doesn't support --fit-ctx.
assert "--fit-ctx" not in LlamaCppBackend._ctx_integrity_flags(
1, True, 98304, 98304, _CAPS_NONE
1, True, True, 98304, 98304, _CAPS_NONE
)
def test_fit_ctx_floors_auto_request_at_8192_only_under_auto_fit():
# Manual + Auto (auto_fit) floors the auto window at 8192 so --fit can't
# shrink it to a tiny size.
flags = LlamaCppBackend._ctx_integrity_flags(1, True, True, 0, 262144, _CAPS_ALL)
assert flags[flags.index("--fit-ctx") + 1] == "8192"
# Legacy auto (fit on but not auto_fit) emits -c 0 to pin native, so the
# 8192 floor must NOT ride along and override that pin.
assert "--fit-ctx" not in LlamaCppBackend._ctx_integrity_flags(
1, True, False, 0, 262144, _CAPS_ALL
)

494
studio/backend/tests/test_llama_cpp_slot_resume.py Normal file
View file

@ -0,0 +1,494 @@
# 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 os
from types import SimpleNamespace
import core.inference.llama_cpp as llama_cpp
from core.inference.llama_cpp import LlamaCppBackend
def _resume_backend(tmp_path, n_slots = 1):
backend = LlamaCppBackend()
backend._healthy = True
# No-op lifecycle methods so the atexit cleanup can kill the fake quietly.
backend._process = SimpleNamespace(
poll = lambda: None,
terminate = lambda: None,
wait = lambda *a, **k: 0,
kill = lambda: None,
pid = 0,
)
backend._port = 8081
backend._slot_save_dir = str(tmp_path)
backend._slot_save_binary = ("/bin/llama-server", 1)
(tmp_path / "model.gguf").write_bytes(b"gguf")
backend._gguf_path = str(tmp_path / "model.gguf")
backend._effective_parallel_slots = n_slots
backend._estimate_kv_cache_bytes = lambda *a, **k: 0
return backend
def _fake_disk(monkeypatch, free = 1 << 40):
monkeypatch.setattr(llama_cpp.shutil, "disk_usage", lambda _p: SimpleNamespace(free = free))
class _Resp:
def __init__(
self,
status_code = 200,
body = None,
):
self.status_code = status_code
self._body = body or {}
def json(self):
return self._body
def test_save_returns_none_when_slot_save_disabled(monkeypatch, tmp_path):
backend = _resume_backend(tmp_path)
backend._slot_save_dir = None
monkeypatch.setattr(
llama_cpp.httpx,
"post",
lambda *a, **k: (_ for _ in ()).throw(AssertionError),
raising = False,
)
assert backend.save_slots_for_resume() is None
def test_save_skipped_when_prompt_cache_disabled(monkeypatch, tmp_path):
backend = _resume_backend(tmp_path)
backend._prompt_cache_disabled = True
monkeypatch.setattr(
llama_cpp.httpx,
"post",
lambda *a, **k: (_ for _ in ()).throw(AssertionError),
raising = False,
)
assert backend.save_slots_for_resume() is None
def test_save_skipped_when_insufficient_free_disk(monkeypatch, tmp_path):
backend = _resume_backend(tmp_path)
backend._estimate_kv_cache_bytes = lambda *a, **k: 1 << 40
_fake_disk(monkeypatch, free = 1 << 20)
monkeypatch.setattr(
llama_cpp.httpx,
"post",
lambda *a, **k: (_ for _ in ()).throw(AssertionError),
raising = False,
)
assert backend.save_slots_for_resume() is None
def test_save_collects_manifest_across_slots(monkeypatch, tmp_path):
backend = _resume_backend(tmp_path, n_slots = 2)
_fake_disk(monkeypatch)
calls = []
def fake_post(url, **kwargs):
calls.append((url, kwargs["params"], kwargs["json"]))
return _Resp(200, {"n_saved": 40, "n_written": 100})
monkeypatch.setattr(llama_cpp.httpx, "post", fake_post, raising = False)
manifest = backend.save_slots_for_resume()
assert manifest is not None
assert manifest["dir"] == str(tmp_path)
assert manifest["binary"] == ("/bin/llama-server", 1)
assert manifest["gguf"] == str(tmp_path / "model.gguf")
st = os.stat(manifest["gguf"])
assert manifest["gguf_stat"] == ((st.st_size, st.st_mtime_ns),)
assert manifest["launch"] == backend._slot_launch_fingerprint()
assert [e["id"] for e in manifest["slots"]] == [0, 1]
assert all(e["n_saved"] == 40 for e in manifest["slots"])
assert [c[1] for c in calls] == [{"action": "save"}] * 2
assert "/slots/0" in calls[0][0] and "/slots/1" in calls[1][0]
def test_save_unlinks_empty_slot_and_returns_none(monkeypatch, tmp_path):
backend = _resume_backend(tmp_path)
_fake_disk(monkeypatch)
def fake_post(url, **kwargs):
(tmp_path / kwargs["json"]["filename"]).write_bytes(b"")
return _Resp(200, {"n_saved": 0, "n_written": 0})
monkeypatch.setattr(llama_cpp.httpx, "post", fake_post, raising = False)
assert backend.save_slots_for_resume() is None
assert list(tmp_path.glob("resume-*.bin")) == [] # empty-slot file removed
def test_save_cap_breach_discards_all_files(monkeypatch, tmp_path):
backend = _resume_backend(tmp_path, n_slots = 2)
_fake_disk(monkeypatch)
monkeypatch.setattr(llama_cpp, "_SLOT_SAVE_MAX_BYTES", 150)
def fake_post(url, **kwargs):
(tmp_path / kwargs["json"]["filename"]).write_bytes(b"x" * 100)
return _Resp(200, {"n_saved": 40, "n_written": 100})
monkeypatch.setattr(llama_cpp.httpx, "post", fake_post, raising = False)
assert backend.save_slots_for_resume() is None # 200 bytes > 150 cap
assert list(tmp_path.glob("resume-*.bin")) == []
def test_save_transport_error_aborts_remaining_slots(monkeypatch, tmp_path):
backend = _resume_backend(tmp_path, n_slots = 3)
_fake_disk(monkeypatch)
calls = []
def fake_post(url, **kwargs):
calls.append(url)
raise OSError("connection refused")
monkeypatch.setattr(llama_cpp.httpx, "post", fake_post, raising = False)
assert backend.save_slots_for_resume() is None
assert len(calls) == 1 # no retries against a dead server
def test_save_transport_error_unlinks_partial_file(monkeypatch, tmp_path):
backend = _resume_backend(tmp_path)
_fake_disk(monkeypatch)
def fake_post(url, **kwargs):
(tmp_path / kwargs["json"]["filename"]).write_bytes(b"partial")
raise OSError("timed out")
monkeypatch.setattr(llama_cpp.httpx, "post", fake_post, raising = False)
assert backend.save_slots_for_resume() is None
assert list(tmp_path.glob("resume-*.bin")) == []
def test_fingerprint_tracks_lora_sidecar_rewrite(tmp_path):
backend = _resume_backend(tmp_path)
adapter = tmp_path / "adapter.gguf"
adapter.write_bytes(b"v1")
backend._extra_args = ["--lora", str(adapter)]
before = backend._slot_launch_fingerprint()
adapter.write_bytes(b"v2-different") # re-exported adapter, same path
assert backend._slot_launch_fingerprint() != before
backend._extra_args = [f"--lora={adapter}"]
assert backend._sidecar_weight_files() == [str(adapter)]
backend._extra_args = ["--lora-scaled", str(adapter), "0.5"]
assert backend._sidecar_weight_files() == [str(adapter)]
backend._extra_args = ["--control-vector", str(adapter), "--threads", "4"]
assert backend._sidecar_weight_files() == [str(adapter)]
def test_sidecar_files_parse_csv_and_colon_scale(tmp_path):
backend = _resume_backend(tmp_path)
a, b = tmp_path / "a.gguf", tmp_path / "b.gguf"
backend._extra_args = ["--lora", f"{a},{b}"]
files = backend._sidecar_weight_files()
assert str(a) in files and str(b) in files
backend._extra_args = ["--lora-scaled", f"{a}:0.5"]
assert str(a) in backend._sidecar_weight_files()
backend._extra_args = ["--control-vector-scaled", f"{a}:1.0,{b}:2.0"]
files = backend._sidecar_weight_files()
assert str(a) in files and str(b) in files
# Windows drive letter must not be mistaken for a scale separator.
backend._extra_args = ["--lora-scaled", "C:\\adapters\\a.gguf:0.75"]
assert "C:\\adapters\\a.gguf" in backend._sidecar_weight_files()
backend._extra_args = ["--lora", "C:\\adapters\\a.gguf"]
assert backend._sidecar_weight_files() == ["C:\\adapters\\a.gguf"]
def test_fingerprint_tracks_colon_scaled_adapter_rewrite(tmp_path):
backend = _resume_backend(tmp_path)
adapter = tmp_path / "adapter.gguf"
adapter.write_bytes(b"v1")
backend._extra_args = ["--lora-scaled", f"{adapter}:0.5"]
before = backend._slot_launch_fingerprint()
adapter.write_bytes(b"v2-different") # re-exported adapter, same path
assert backend._slot_launch_fingerprint() != before
def test_fingerprint_tracks_effective_context_length(tmp_path):
backend = _resume_backend(tmp_path)
backend._effective_context_length = 8192
before = backend._slot_launch_fingerprint()
backend._effective_context_length = 4096 # auto-fit landed smaller on reload
assert backend._slot_launch_fingerprint() != before
def test_gguf_file_identity_covers_split_shards(tmp_path):
backend = _resume_backend(tmp_path)
first = tmp_path / "m-00001-of-00002.gguf"
second = tmp_path / "m-00002-of-00002.gguf"
first.write_bytes(b"a")
second.write_bytes(b"bb")
before = backend._gguf_file_identity(str(first))
st1, st2 = os.stat(first), os.stat(second)
assert before == ((st1.st_size, st1.st_mtime_ns), (st2.st_size, st2.st_mtime_ns))
second.write_bytes(b"rewritten") # sibling changes, primary untouched
after = backend._gguf_file_identity(str(first))
assert after is not None and after != before
assert after[0] == before[0] # primary shard unchanged
second.unlink()
assert backend._gguf_file_identity(str(first)) is None # missing shard
def test_save_skipped_when_user_disabled_prompt_cache(monkeypatch, tmp_path):
backend = _resume_backend(tmp_path)
backend._extra_args = ["--no-cache-prompt"]
monkeypatch.setattr(
llama_cpp.httpx,
"post",
lambda *a, **k: (_ for _ in ()).throw(AssertionError),
raising = False,
)
assert backend.save_slots_for_resume() is None
def test_save_skipped_when_env_disables_prompt_cache(monkeypatch, tmp_path):
backend = _resume_backend(tmp_path)
monkeypatch.setenv("LLAMA_ARG_CACHE_PROMPT", "0")
monkeypatch.setattr(
llama_cpp.httpx,
"post",
lambda *a, **k: (_ for _ in ()).throw(AssertionError),
raising = False,
)
assert backend.save_slots_for_resume() is None
monkeypatch.delenv("LLAMA_ARG_CACHE_PROMPT")
monkeypatch.setenv("LLAMA_ARG_NO_CACHE_PROMPT", "1") # legacy negative form
assert backend.save_slots_for_resume() is None
def test_explicit_cache_prompt_flag_overrides_env(monkeypatch, tmp_path):
backend = _resume_backend(tmp_path)
monkeypatch.setenv("LLAMA_ARG_CACHE_PROMPT", "0")
backend._extra_args = ["--cache-prompt"] # CLI wins over env in llama.cpp
_fake_disk(monkeypatch)
monkeypatch.setattr(
llama_cpp.httpx,
"post",
lambda *a, **k: _Resp(200, {"n_saved": 1, "n_written": 1}),
raising = False,
)
assert backend.save_slots_for_resume() is not None
def test_user_cache_prompt_overrides_studio_no_cache_flag(monkeypatch, tmp_path):
# User extras follow Studio's flags, so an explicit --cache-prompt wins.
backend = _resume_backend(tmp_path)
backend._prompt_cache_disabled = True
backend._extra_args = ["--cache-prompt"]
_fake_disk(monkeypatch)
monkeypatch.setattr(
llama_cpp.httpx,
"post",
lambda *a, **k: _Resp(200, {"n_saved": 1, "n_written": 1}),
raising = False,
)
assert backend.save_slots_for_resume() is not None
# Last flag wins when both appear in extras.
backend._extra_args = ["--cache-prompt", "--no-cache-prompt"]
assert backend.save_slots_for_resume() is None
def test_save_stops_writing_once_cap_exceeded(monkeypatch, tmp_path):
backend = _resume_backend(tmp_path, n_slots = 3)
_fake_disk(monkeypatch)
monkeypatch.setattr(llama_cpp, "_SLOT_SAVE_MAX_BYTES", 150)
calls = []
def fake_post(url, **kwargs):
calls.append(url)
(tmp_path / kwargs["json"]["filename"]).write_bytes(b"x" * 100)
return _Resp(200, {"n_saved": 1, "n_written": 100})
monkeypatch.setattr(llama_cpp.httpx, "post", fake_post, raising = False)
assert backend.save_slots_for_resume() is None
assert len(calls) == 2 # cap blown after slot 1; slot 2 never attempted
assert list(tmp_path.glob("resume-*.bin")) == []
def test_save_aborts_between_slots_when_no_longer_idle(monkeypatch, tmp_path):
backend = _resume_backend(tmp_path, n_slots = 3)
_fake_disk(monkeypatch)
calls = []
def fake_post(url, **kwargs):
calls.append(url)
return _Resp(200, {"n_saved": 5, "n_written": 10})
monkeypatch.setattr(llama_cpp.httpx, "post", fake_post, raising = False)
aborts = iter([False, True, True])
manifest = backend.save_slots_for_resume(should_abort = lambda: next(aborts))
assert len(calls) == 1 # slots 1 and 2 skipped
assert manifest is not None
assert [e["id"] for e in manifest["slots"]] == [0]
def test_save_non_200_slot_is_skipped_but_others_kept(monkeypatch, tmp_path):
backend = _resume_backend(tmp_path, n_slots = 2)
_fake_disk(monkeypatch)
def fake_post(url, **kwargs):
if "/slots/0" in url:
return _Resp(500)
return _Resp(200, {"n_saved": 5, "n_written": 10})
monkeypatch.setattr(llama_cpp.httpx, "post", fake_post, raising = False)
manifest = backend.save_slots_for_resume()
assert manifest is not None
assert [e["id"] for e in manifest["slots"]] == [1]
def test_restore_posts_each_slot_and_tolerates_failures(monkeypatch, tmp_path):
backend = _resume_backend(tmp_path)
calls = []
def fake_post(url, **kwargs):
calls.append((url, kwargs["params"], kwargs["json"]))
return _Resp(500 if "/slots/0" in url else 200, {"n_restored": 5})
monkeypatch.setattr(llama_cpp.httpx, "post", fake_post, raising = False)
backend.restore_slots_for_resume(
{
"slots": [
{"id": 0, "filename": "resume-a-slot0.bin", "n_saved": 5},
{"id": 1, "filename": "resume-a-slot1.bin", "n_saved": 5},
]
}
)
assert [c[1] for c in calls] == [{"action": "restore"}] * 2
assert calls[0][2] == {"filename": "resume-a-slot0.bin"}
def test_restore_transport_error_stops_early(monkeypatch, tmp_path):
backend = _resume_backend(tmp_path)
calls = []
def fake_post(url, **kwargs):
calls.append(url)
raise OSError("connection refused")
monkeypatch.setattr(llama_cpp.httpx, "post", fake_post, raising = False)
backend.restore_slots_for_resume(
{"slots": [{"id": 0, "filename": "a.bin"}, {"id": 1, "filename": "b.bin"}]}
)
assert len(calls) == 1
def test_save_deletes_orphan_on_malformed_response(monkeypatch, tmp_path):
# A 200 that writes a file but returns a non-numeric counter must be cleaned
# up like any other save failure, not left orphaned holding chat KV.
backend = _resume_backend(tmp_path)
_fake_disk(monkeypatch)
def fake_post(url, **kwargs):
(tmp_path / kwargs["json"]["filename"]).write_bytes(b"chat-kv")
return _Resp(200, {"n_saved": "not-an-int"})
monkeypatch.setattr(llama_cpp.httpx, "post", fake_post, raising = False)
assert backend.save_slots_for_resume() is None
assert list(tmp_path.glob("resume-*.bin")) == []
def test_save_deletes_orphan_on_non_dict_response(monkeypatch, tmp_path):
backend = _resume_backend(tmp_path)
_fake_disk(monkeypatch)
def fake_post(url, **kwargs):
(tmp_path / kwargs["json"]["filename"]).write_bytes(b"chat-kv")
return _Resp(200, ["unexpected", "list"])
monkeypatch.setattr(llama_cpp.httpx, "post", fake_post, raising = False)
assert backend.save_slots_for_resume() is None
assert list(tmp_path.glob("resume-*.bin")) == []
def test_save_cap_uses_actual_file_size_not_reported_bytes(monkeypatch, tmp_path):
# A binary under-reporting n_written must not slip past the disk cap: the
# cap is enforced against the bytes actually on disk.
backend = _resume_backend(tmp_path)
_fake_disk(monkeypatch)
monkeypatch.setattr(llama_cpp, "_SLOT_SAVE_MAX_BYTES", 150)
def fake_post(url, **kwargs):
(tmp_path / kwargs["json"]["filename"]).write_bytes(b"x" * 200)
return _Resp(200, {"n_saved": 5, "n_written": 1}) # under-reported
monkeypatch.setattr(llama_cpp.httpx, "post", fake_post, raising = False)
assert backend.save_slots_for_resume() is None # 200 real bytes > 150 cap
assert list(tmp_path.glob("resume-*.bin")) == []
def test_save_skipped_when_estimate_exceeds_cap(monkeypatch, tmp_path):
# An estimate over the cap skips before writing any slot at all.
backend = _resume_backend(tmp_path)
backend._estimate_kv_cache_bytes = lambda *a, **k: 1 << 40
monkeypatch.setattr(llama_cpp, "_SLOT_SAVE_MAX_BYTES", 1 << 20)
_fake_disk(monkeypatch)
monkeypatch.setattr(
llama_cpp.httpx,
"post",
lambda *a, **k: (_ for _ in ()).throw(AssertionError),
raising = False,
)
assert backend.save_slots_for_resume() is None
def test_save_skipped_when_model_file_changed_since_load(monkeypatch, tmp_path):
# The GGUF/sidecars were swapped on disk after the server loaded them, so the
# live KV belongs to the old weights: refuse to persist it (no POST at all).
backend = _resume_backend(tmp_path)
backend._slot_loaded_identity = ((("stale", 0),), ()) # != current identity
_fake_disk(monkeypatch)
monkeypatch.setattr(
llama_cpp.httpx,
"post",
lambda *a, **k: (_ for _ in ()).throw(AssertionError),
raising = False,
)
assert backend.save_slots_for_resume() is None
def test_save_proceeds_when_load_identity_matches(monkeypatch, tmp_path):
# Matching load-time snapshot: the save runs normally.
backend = _resume_backend(tmp_path)
backend._slot_loaded_identity = (
backend._gguf_file_identity(backend._gguf_path),
backend._slot_launch_fingerprint(),
)
_fake_disk(monkeypatch)
def fake_post(url, **kwargs):
(tmp_path / kwargs["json"]["filename"]).write_bytes(b"kv")
return _Resp(200, {"n_saved": 5, "n_written": 2})
monkeypatch.setattr(llama_cpp.httpx, "post", fake_post, raising = False)
manifest = backend.save_slots_for_resume()
assert manifest is not None
assert [e["id"] for e in manifest["slots"]] == [0]
def test_save_skipped_when_estimate_unavailable_and_low_disk(monkeypatch, tmp_path):
# A 0 estimate means metadata was insufficient, not a zero-byte cache: the save
# must demand room for the whole cap, not just 1 GiB, on a low-disk host.
backend = _resume_backend(tmp_path)
backend._estimate_kv_cache_bytes = lambda *a, **k: 0 # metadata unavailable
monkeypatch.setattr(llama_cpp, "_SLOT_SAVE_MAX_BYTES", 8 << 30) # 8 GiB cap
_fake_disk(monkeypatch, free = 2 << 30) # 2 GiB free < 8 + 1 GiB required
monkeypatch.setattr(
llama_cpp.httpx,
"post",
lambda *a, **k: (_ for _ in ()).throw(AssertionError),
raising = False,
)
assert backend.save_slots_for_resume() is None

125
studio/backend/tests/test_llama_cpp_stall_timeout.py Normal file
View file

@ -0,0 +1,125 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Regression test for the post-first-token stall timeout in the cancel-aware read.
httpcore snapshots ``request.extensions["timeout"]["read"]`` once at body start, so
when ``_iter_text_cancellable`` lowers it after the first token, a one-token-then-silent
server hangs for the full prefill window. The fix re-reads the live extensions timeout
per call; a fake clock and always-silent stream check the read gives up after the live
stall timeout, not the stale prefill one.
"""
from __future__ import annotations
import inspect
import sys
import threading
import types as _types
from pathlib import Path
import pytest
_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
if _BACKEND_DIR not in sys.path:
sys.path.insert(0, _BACKEND_DIR)
# Mirror sibling tests' stubbing so the module imports without fastapi.
_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"))
import httpcore # noqa: E402
from core.inference import llama_cpp as llama_cpp_mod # noqa: E402
from core.inference.llama_cpp import LlamaCppBackend # noqa: E402
_PREFILL_TIMEOUT = 1200.0 # what httpcore snapshots from the prefill timeout
_STALL_TIMEOUT = 120.0 # the post-first-token stall timeout the wrapper must honor
class _Obj:
pass
def _install(response, clock, silent_stream):
"""Wire fake client/pool so _install_cancel_aware_read finds the stream; return the wrapped stream.read."""
inner = _Obj()
inner._network_stream = silent_stream
connection = _Obj()
connection._connection = inner
pool = _Obj()
pool._connections = [connection]
transport = _Obj()
transport._pool = pool
client = _Obj()
client._transport = transport
cancel_event = threading.Event() # never set: we test the stall path, not cancel
sig = inspect.signature(LlamaCppBackend._install_cancel_aware_read)
if "response" in sig.parameters:
# Fixed signature: wrapper reads the live extensions timeout.
LlamaCppBackend._install_cancel_aware_read(client, cancel_event, response)
else:
# Pre-fix signature: no response, so the stall assertion fails (proves the bug).
LlamaCppBackend._install_cancel_aware_read(client, cancel_event)
return silent_stream.read
def test_stall_timeout_honored_after_first_token(monkeypatch):
clock = {"t": 0.0}
monkeypatch.setattr(llama_cpp_mod.time, "monotonic", lambda: clock["t"])
# One token then silence: every read times out, advancing fake time by its timeout.
def silent_read(max_bytes, timeout = None):
clock["t"] += timeout if timeout is not None else 0.0
raise httpcore.ReadTimeout("slice timed out on silence")
stream = _Obj()
stream.read = silent_read
# First token seen: the live read timeout is lowered to the stall timeout.
request = _Obj()
request.extensions = {"timeout": {"read": _STALL_TIMEOUT}}
response = _Obj()
response.request = request
wrapped_read = _install(response, clock, stream)
# httpcore still passes the stale prefill timeout it snapshotted at body start.
with pytest.raises(httpcore.ReadTimeout):
wrapped_read(65536, timeout = _PREFILL_TIMEOUT)
# Must give up ~stall timeout after the last token, not the prefill window.
assert clock["t"] <= _STALL_TIMEOUT * 1.5, (
f"stall timeout not honored: waited {clock['t']}s "
f"(expected ~{_STALL_TIMEOUT}s, not {_PREFILL_TIMEOUT}s)"
)
assert clock["t"] >= _STALL_TIMEOUT * 0.5
def test_prefill_timeout_used_when_no_live_override(monkeypatch):
"""Without a lowered live timeout, the wrapper honors the passed prefill timeout, so the normal first-token wait is unchanged."""
clock = {"t": 0.0}
monkeypatch.setattr(llama_cpp_mod.time, "monotonic", lambda: clock["t"])
def silent_read(max_bytes, timeout = None):
clock["t"] += timeout if timeout is not None else 0.0
raise httpcore.ReadTimeout("slice timed out on silence")
stream = _Obj()
stream.read = silent_read
# No timeout extension: wrapper falls back to httpcore's passed timeout.
request = _Obj()
request.extensions = {}
response = _Obj()
response.request = request
wrapped_read = _install(response, clock, stream)
with pytest.raises(httpcore.ReadTimeout):
wrapped_read(65536, timeout = _PREFILL_TIMEOUT)
assert clock["t"] >= _PREFILL_TIMEOUT * 0.9

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

@ -83,6 +83,7 @@ def _write_install(
repo: str = "unslothai/llama.cpp",
asset: str | None = None,
release_tag: str | None = None,
force_cpu: bool | None = None,
) -> str:
"""Create a fake prebuilt install and return the llama-server path."""
bin_dir = dir_ / "build" / "bin"
@ -99,6 +100,8 @@ def _write_install(
}
if asset is not None:
marker["asset"] = asset
if force_cpu is not None:
marker["force_cpu"] = force_cpu
(dir_ / MARKER).write_text(json.dumps(marker))
return str(binary)
@ -493,6 +496,47 @@ def test_start_update_preserves_vulkan_via_env(monkeypatch, tmp_path):
assert popen_kwargs["env"]["UNSLOTH_FORCE_VULKAN"] == "1"
@pytest.mark.parametrize(
"force_cpu, expect_flag",
[
# A deliberate CPU install (marker force_cpu=True) re-asserts --force-cpu on
# update so detect_host on a GPU host cannot re-route and revive the crash
# (#7213); --force-cpu also re-persists the flag for the next update.
(True, True),
# A transient fallback (or a legacy marker without the flag) stays free to
# heal to a GPU bundle (#6097).
(False, False),
(None, False),
],
)
def test_start_update_cpu_fallback_preserved_by_flag(monkeypatch, tmp_path, force_cpu, expect_flag):
asset = "llama-b9493-bin-ubuntu-x64.tar.gz"
install_dir = tmp_path / "llama.cpp"
binary = _write_install(install_dir, "b9493", asset = asset, force_cpu = force_cpu)
monkeypatch.setattr(upd, "_find_binary", lambda: binary)
monkeypatch.setattr(upd, "_installer_script", lambda: tmp_path / "install_llama_prebuilt.py")
monkeypatch.setattr(freshness, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9518")
captured: dict = {}
def _on_start(cmd):
captured["cmd"] = cmd
_write_install(install_dir, "b9518", asset = asset, force_cpu = force_cpu)
_patch_installer_popen(monkeypatch, lines = ["installed\n"], on_start = _on_start)
assert upd.start_update()["started"] is True
deadline = time.time() + 10
while time.time() < deadline:
job = upd.get_update_status()["job"]
if job["state"] in ("success", "error"):
break
time.sleep(0.05)
assert job["state"] == "success", job
assert ("--force-cpu" in captured["cmd"]) is expect_flag
assert "--cpu-fallback" not in captured["cmd"]
def test_start_update_reports_full_release_tag(monkeypatch, tmp_path):
install_dir = tmp_path / "llama.cpp"
binary = _write_install(install_dir, "b9595")
@ -676,7 +720,7 @@ def test_install_cmd_rocm_marker_forwards_gfx(monkeypatch, tmp_path):
assert "--rocm-gfx" in cmd
assert cmd[cmd.index("--rocm-gfx") + 1] == "gfx110x"
assert "--has-rocm" not in cmd
assert "--cpu-fallback" not in cmd
assert "--force-cpu" not in cmd
assert "--simple-policy" not in cmd
assert "--published-repo" in cmd and "unslothai/llama.cpp" in cmd
@ -690,17 +734,17 @@ def test_install_cmd_fork_rocm_marker_forwards_has_rocm(monkeypatch, tmp_path):
def test_install_cmd_ggml_cpu_marker_has_no_cpu_fallback(monkeypatch, tmp_path):
# Legacy CPU installs recorded a ggml-org marker (new installs use the fork).
# Re-running into the same install-dir/repo reproduces the same CPU bundle;
# --cpu-fallback (which force-drops GPU detection) is reserved for setup.sh's
# arm64 rescue and must not appear here.
# Legacy CPU installs recorded a ggml-org marker (new installs use the fork) with
# no force_cpu field. Re-running into the same install-dir/repo reproduces the same
# CPU bundle; --force-cpu (the persisted-CPU re-assert) must not appear for a marker
# that never recorded a deliberate CPU choice, so it can still heal to GPU (#6097).
cmd = _capture_install_cmd(
monkeypatch,
tmp_path,
repo = "ggml-org/llama.cpp",
asset = "llama-b9334-bin-ubuntu-x64.tar.gz",
)
assert "--cpu-fallback" not in cmd
assert "--force-cpu" not in cmd
assert "--rocm-gfx" not in cmd
assert "--has-rocm" not in cmd
assert "--simple-policy" not in cmd
@ -714,7 +758,7 @@ def test_install_cmd_cuda_marker_minimal_and_backward_compatible(monkeypatch, tm
assert "--simple-policy" not in cmd
assert "--rocm-gfx" not in cmd
assert "--has-rocm" not in cmd
assert "--cpu-fallback" not in cmd
assert "--force-cpu" not in cmd
def test_install_cmd_pins_offered_release_tag(monkeypatch, tmp_path):

57
studio/backend/tests/test_llama_server_args.py
View file

@ -183,6 +183,8 @@ def test_non_flag_token_passes_through():
"--reranking",
# llama-server's own --tools clashes with Unsloth's tool policy.
"--tools",
# Slot-state dir: Studio owns it for KV persistence across idle unload.
"--slot-save-path",
],
)
def test_denylist_rejects_all_aliases(denied):
@ -224,6 +226,16 @@ def test_denylist_rejects_equals_form():
validate_extra_args(["--port=9000"])
def test_slot_save_path_is_managed_in_all_forms():
for args in (["--slot-save-path", "/tmp/x"], ["--slot-save-path=/tmp/x"], ["--slot-save-path"]):
with pytest.raises(ValueError, match = "--slot-save-path"):
validate_extra_args(args)
assert is_managed_flag("--slot-save-path") is True
assert is_managed_flag("--slot-save-path=/tmp/x") is True
# --slots (read-only diagnostics endpoint) stays a user choice.
assert is_managed_flag("--slots") is False
@pytest.mark.parametrize(
"padded",
[" --parallel", "--parallel ", "\t--parallel", " -np", "-np \n", "-np\t"],
@ -747,6 +759,34 @@ def test_strip_shadowing_flags_defaults_strip_split_mode_too():
assert strip_shadowing_flags(["--split-mode", "tensor"]) == []
def test_strip_offload_is_opt_in_and_covers_moe():
base = dict(
strip_context = False,
strip_cache = False,
strip_spec = False,
strip_template = False,
strip_split_mode = False,
)
# Default: offload (incl. MoE) flags are NOT stripped.
assert strip_shadowing_flags(["--n-cpu-moe", "8", "--top-k", "20"], **base) == [
"--n-cpu-moe",
"8",
"--top-k",
"20",
]
# Opt-in strips layer AND MoE offload flags (value-aware), keeps the rest.
assert strip_shadowing_flags(
["--n-cpu-moe", "8", "--gpu-layers", "33", "--fit", "off", "--top-k", "20"],
**base,
strip_offload = True,
) == ["--top-k", "20"]
# Boolean --cpu-moe drops the flag only, not the following value.
assert strip_shadowing_flags(["--cpu-moe", "--seed", "-1"], **base, strip_offload = True) == [
"--seed",
"-1",
]
@pytest.mark.parametrize(
"args",
[
@ -796,6 +836,23 @@ def test_strip_split_mode_only_drops_tensor_split_too():
assert strip_split_mode_only(["-sm=tensor", "-ts=3,1"]) == []
def test_strip_tensor_split_alone_preserves_split_mode():
# Manual mode emits its own --tensor-split, so an inherited ratio is dropped
# -- but the user's --split-mode row/none/layer choice (which the manual
# ratio toggle can't express) must survive. strip_tensor_split removes only
# the ratio, unlike strip_split_mode which removes the whole group.
out = strip_shadowing_flags(
["--split-mode", "row", "--tensor-split", "1,1", "--top-k", "20"],
strip_context = False,
strip_cache = False,
strip_spec = False,
strip_template = False,
strip_split_mode = False,
strip_tensor_split = True,
)
assert out == ["--split-mode", "row", "--top-k", "20"]
def test_strip_shadowing_flags_keeps_model_draft_without_spec():
out = strip_shadowing_flags(
["--model-draft", "/custom/mtp.gguf"],

66
studio/backend/tests/test_middleware.py
View file

@ -14,6 +14,7 @@ import pytest
from fastapi import FastAPI, HTTPException, Request
from fastapi.responses import Response
from fastapi.testclient import TestClient
from starlette.middleware.gzip import GZipMiddleware
_BACKEND_ROOT = Path(__file__).resolve().parents[1]
@ -471,6 +472,71 @@ class TestSecurityHeadersMiddleware:
assert b"server" in names
class TestFrontendAssets:
def test_hashed_assets_are_compressed_and_cached(self, tmp_path, main_module):
content = b"export const value = 'responsive';\n" * 200
(tmp_path / "page-abc123.js").write_bytes(content)
app = FastAPI()
assets_app = GZipMiddleware(
main_module.ImmutableStaticFiles(directory = tmp_path),
minimum_size = 1024,
compresslevel = 6,
)
app.mount("/assets", assets_app, name = "assets")
response = TestClient(app).get(
"/assets/page-abc123.js",
headers = {"Accept-Encoding": "gzip"},
)
assert response.status_code == 200
assert response.content == content
assert response.headers["content-encoding"] == "gzip"
assert response.headers["cache-control"] == (main_module._IMMUTABLE_ASSET_CACHE_CONTROL)
assert "accept-encoding" in response.headers["vary"].lower()
def test_asset_revalidation_keeps_immutable_cache_header(self, tmp_path, main_module):
(tmp_path / "page-abc123.js").write_text("export {};", encoding = "utf-8")
app = FastAPI()
app.mount(
"/assets",
main_module.ImmutableStaticFiles(directory = tmp_path),
name = "assets",
)
client = TestClient(app)
first = client.get("/assets/page-abc123.js")
response = client.get(
"/assets/page-abc123.js",
headers = {"If-None-Match": first.headers["etag"]},
)
assert response.status_code == 304
assert response.headers["cache-control"] == (main_module._IMMUTABLE_ASSET_CACHE_CONTROL)
def test_range_request_is_not_compressed(self, tmp_path, main_module):
content = b"export const value = 'responsive';\n" * 200
(tmp_path / "page-abc123.js").write_bytes(content)
app = FastAPI()
assets_app = main_module._AssetGZipMiddleware(
main_module.ImmutableStaticFiles(directory = tmp_path),
minimum_size = 1024,
compresslevel = 6,
)
app.mount("/assets", assets_app, name = "assets")
response = TestClient(app).get(
"/assets/page-abc123.js",
headers = {"Accept-Encoding": "gzip", "Range": "bytes=0-99"},
)
assert response.status_code == 206
assert response.headers.get("content-encoding") != "gzip"
assert response.headers["content-range"] == f"bytes 0-99/{len(content)}"
assert response.content == content[:100]
assert response.headers["cache-control"] == (main_module._IMMUTABLE_ASSET_CACHE_CONTROL)
# /api/health auth gate

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

@ -0,0 +1,137 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Regression tests for MLX stop-and-save checkpoint handling."""
import importlib.util
import json
import sys
import types
from pathlib import Path
import numpy as np
from safetensors.numpy import save_file
_BACKEND = Path(__file__).resolve().parents[1]
def _load_worker_module():
spec = importlib.util.spec_from_file_location(
"training_worker_under_test",
_BACKEND / "core" / "training" / "worker.py",
)
module = importlib.util.module_from_spec(spec)
assert spec.loader is not None
spec.loader.exec_module(module)
return module
worker = _load_worker_module()
class _FakeTrainer:
def __init__(self, step: int):
self._global_step = step
self._train_loss_history = []
self.model = object()
def _write_checkpoint(out: Path, step: int) -> Path:
checkpoint = out / f"checkpoint-{step}"
checkpoint.mkdir(parents = True, exist_ok = True)
(checkpoint / "trainer_state.json").write_text(
json.dumps({"global_step": step}), encoding = "utf-8"
)
save_file({"weight": np.ones(1, dtype = np.float32)}, checkpoint / "adapters.safetensors")
save_file(
{"state": np.ones(1, dtype = np.float32)},
checkpoint / "optimizer_state.safetensors",
)
return checkpoint
def test_mlx_has_checkpoint_at_step_requires_complete_state(tmp_path):
out = tmp_path / "outputs" / "run_x"
_write_checkpoint(out, 5)
assert worker._mlx_has_checkpoint_at_step(out, 5) is True
def test_write_mlx_stop_checkpoint_returns_true_when_current_step_checkpoint_exists(tmp_path):
out = tmp_path / "outputs" / "run_x"
_write_checkpoint(out, 5)
assert worker._write_mlx_stop_checkpoint(_FakeTrainer(step = 5), object(), out) is True
def test_write_mlx_stop_checkpoint_writes_current_step_when_only_older_checkpoint_exists(
tmp_path, monkeypatch
):
out = tmp_path / "outputs" / "run_x"
_write_checkpoint(out, 5)
saved_steps: list[int] = []
def _save_state(_value, path, name):
save_file({"state": np.ones(1, dtype = np.float32)}, Path(path, name))
def _save_trainer_state(state, ckpt_dir, **_kwargs):
Path(ckpt_dir, "trainer_state.json").write_text(json.dumps(state), encoding = "utf-8")
saved_steps.append(int(state["global_step"]))
fake_utils = types.SimpleNamespace(
save_trainable_adapters = lambda model, path: _save_state(
model, path, "adapters.safetensors"
),
save_optimizer_state = lambda optimizer, path: _save_state(
optimizer, path, "optimizer_state.safetensors"
),
save_trainer_state = _save_trainer_state,
)
monkeypatch.setitem(sys.modules, "unsloth_zoo.mlx.utils", fake_utils)
assert worker._write_mlx_stop_checkpoint(_FakeTrainer(step = 10), object(), out) is True
assert saved_steps == [10]
assert (out / "checkpoint-10" / "trainer_state.json").is_file()
def test_write_mlx_stop_checkpoint_returns_false_without_optimizer(tmp_path):
out = tmp_path / "outputs" / "run_x"
out.mkdir(parents = True)
assert worker._write_mlx_stop_checkpoint(_FakeTrainer(step = 5), None, out) is False
def test_write_mlx_stop_checkpoint_rejects_incomplete_current_checkpoint(tmp_path):
out = tmp_path / "outputs" / "run_x"
ckpt = out / "checkpoint-5"
ckpt.mkdir(parents = True)
(ckpt / "trainer_state.json").write_text('{"global_step": 5}', encoding = "utf-8")
assert worker._write_mlx_stop_checkpoint(_FakeTrainer(step = 5), None, out) is False
def test_write_mlx_stop_checkpoint_ignores_stale_checkpoint_without_optimizer(tmp_path):
# An older checkpoint does not cover the current step, so this still fails.
out = tmp_path / "outputs" / "run_x"
_write_checkpoint(out, 5)
assert worker._write_mlx_stop_checkpoint(_FakeTrainer(step = 10), None, out) is False
def test_write_mlx_stop_checkpoint_returns_false_when_save_fails(tmp_path, monkeypatch):
out = tmp_path / "outputs" / "run_x"
out.mkdir(parents = True)
def _boom(*_args, **_kwargs):
raise RuntimeError("save failed")
fake_utils = types.SimpleNamespace(
save_trainable_adapters = _boom,
save_optimizer_state = lambda *_a, **_k: None,
save_trainer_state = lambda *_a, **_k: None,
)
monkeypatch.setitem(sys.modules, "unsloth_zoo.mlx.utils", fake_utils)
assert worker._write_mlx_stop_checkpoint(_FakeTrainer(step = 5), object(), out) is False

232
studio/backend/tests/test_model_picker_regression.py Normal file
View file

@ -0,0 +1,232 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Regression guards for the model-picker per-model-config feature (the set of
bugs that got the predecessor PR reverted). Pure-function / validation checks
only, so they run on CPU in the backend pytest job with no model download.
Covers, at the backend layer:
- infra-model hiding: the RAG embedder (bge-small-en-v1.5) and the llama.cpp
install-validation probe (ggml-org/models / stories260K) stay hidden, while
normal chat repos are not hidden;
- the HF token is honored from the dedicated header with the query string as a
fallback, never the other way around;
- the chat-template byte caps reject oversized overrides (both the char-count
fast path and the UTF-8 byte path) and the sidecar reader is size-bounded.
"""
from __future__ import annotations
import sys
import types
import pytest
# Keep this test runnable without the optional structlog dependency (mirrors
# tests/test_cached_gguf_routes.py), since importing routes.models pulls it in.
if "structlog" not in sys.modules:
class _DummyLogger:
def __getattr__(self, _name):
return lambda *args, **kwargs: None
sys.modules["structlog"] = types.SimpleNamespace(
BoundLogger = _DummyLogger,
get_logger = lambda *args, **kwargs: _DummyLogger(),
)
import routes.models as models_route
from core.rag import config as rag_config
from hub.dependencies import get_hf_token
from models.inference import LoadRequest
from picker.schemas import MAX_CHAT_TEMPLATE_BYTES
from picker.service import _read_bounded_text
from utils.hidden_models import is_hidden_model
@pytest.fixture(autouse = True)
def _pin_default_embedder(monkeypatch):
"""Pin the effective embedder to Studio's static default so hiding is
deterministic and cannot depend on ambient RAG config / env."""
default = "unsloth/bge-small-en-v1.5"
monkeypatch.setattr(rag_config, "EMBEDDING_MODEL", default, raising = False)
monkeypatch.setattr(rag_config, "effective_embedding_model", lambda: default)
monkeypatch.setattr(rag_config, "effective_gguf_repo", lambda: default)
monkeypatch.setattr(rag_config, "default_gguf_repo", lambda: default)
# --------------------------------------------------------------------------- #
# Infra-model hiding (the "infra models resurfaced in the picker" regression) #
# --------------------------------------------------------------------------- #
@pytest.mark.parametrize(
"value",
[
"ggml-org/models", # the probe repo id
"unsloth/bge-small-en-v1.5", # the RAG embedder repo
"unsloth/bge-small-en-v1.5-GGUF", # its GGUF companion
"/root/.cache/huggingface/hub/x/stories260K.gguf", # probe on disk
"/root/.cache/x/Stories260K.GGUF", # case-insensitive
r"C:\\models\\stories260K.gguf", # windows-style path
"/opt/models/bge-small-en-v1.5", # embedder basename folder
"/opt/models/bge-small-en-v1.5-Q8_0.gguf", # suffixed local weight
],
)
def test_infra_models_are_hidden(value):
assert is_hidden_model(value) is True
@pytest.mark.parametrize(
"value",
[
"unsloth/gemma-3-270m-it-GGUF", # a normal small chat GGUF
"unsloth/Qwen3-0.6B", # a normal non-GGUF chat model
"user/stories260K-finetune-GGUF", # repo id merely contains "stories260k"
"user/model-chat", # generic repo must not be hidden
"meta-llama/Llama-3.1-8B-Instruct",
],
)
def test_normal_models_are_not_hidden(value):
assert is_hidden_model(value) is False
def test_is_hidden_model_ignores_empty_values():
assert is_hidden_model(None) is False
assert is_hidden_model("") is False
assert is_hidden_model(None, "", "unsloth/gemma-3-270m-it-GGUF") is False
def test_hidden_model_matchers_expose_probe_needles():
needles, exact_ids, _exact_paths = models_route.hidden_model_matchers()
lowered = [n.lower() for n in needles]
assert "ggml-org/models" in lowered
assert "stories260k.gguf" in lowered
# The configured embedder is exposed as an exact repo id, never as a
# basename needle that would substring-hide unrelated chat models.
assert "bge-small-en-v1.5" not in lowered
assert "unsloth/bge-small-en-v1.5" in exact_ids
def test_hidden_model_matchers_custom_repo_publishes_exact_ids(monkeypatch):
monkeypatch.setattr(rag_config, "effective_embedding_model", lambda: "org/model")
monkeypatch.setattr(rag_config, "effective_gguf_repo", lambda: "org/model-GGUF")
needles, exact_ids, exact_paths = models_route.hidden_model_matchers()
assert needles == ["ggml-org/models", "stories260k.gguf"]
assert "org/model" in exact_ids
assert "org/model-gguf" in exact_ids
assert exact_paths == []
def test_hidden_model_matchers_local_owner_name_path_is_exact_path(monkeypatch, tmp_path):
# A local embedder shaped like owner/name that exists on disk must be an
# exact resolved path, not a Hub repo id (mirroring is_hidden_model), so the
# local row stays hidden instead of showing as a chat model.
(tmp_path / "models" / "embedder").mkdir(parents = True)
monkeypatch.chdir(tmp_path)
monkeypatch.setattr(rag_config, "effective_embedding_model", lambda: "models/embedder")
monkeypatch.setattr(rag_config, "effective_gguf_repo", lambda: "ggml-org/models")
_needles, exact_ids, exact_paths = models_route.hidden_model_matchers()
resolved = str((tmp_path / "models" / "embedder").resolve()).lower()
assert resolved in exact_paths
assert "models/embedder" not in exact_ids
# --------------------------------------------------------------------------- #
# HF token via header, query string only as a fallback (the token-leak fix) #
# --------------------------------------------------------------------------- #
def test_get_hf_token_strips_and_returns():
assert get_hf_token(" hf_abc ") == "hf_abc"
@pytest.mark.parametrize("value", [None, "", " ", "\n\t"])
def test_get_hf_token_blank_is_none(value):
assert get_hf_token(value) is None
@pytest.mark.parametrize(
"value,expected",
[(" hf_x ", "hf_x"), ("", None), (" ", None), (None, None), (1234, None)],
)
def test_normalize_hf_token(value, expected):
assert models_route._normalize_hf_token(value) == expected
def test_header_token_wins_over_query():
header, query = "hf_header", "hf_query"
resolved = models_route._normalize_hf_token(header) or models_route._normalize_hf_token(query)
assert resolved == "hf_header"
def test_query_token_is_fallback_when_header_absent():
resolved = models_route._normalize_hf_token(None) or models_route._normalize_hf_token(
"hf_query"
)
assert resolved == "hf_query"
# --------------------------------------------------------------------------- #
# Chat-template byte caps (the unbounded-template hardening) #
# --------------------------------------------------------------------------- #
def _load_request(**overrides):
data = {"model_path": "unsloth/test-model-GGUF", "gguf_variant": "Q4_K_M"}
data.update(overrides)
return LoadRequest.model_validate(data)
def test_blank_chat_template_override_normalizes_to_none():
assert _load_request(chat_template_override = " \n\t").chat_template_override is None
def test_nonblank_chat_template_override_preserved_verbatim():
template = " {{ messages }} "
assert _load_request(chat_template_override = template).chat_template_override == template
def test_chat_template_at_byte_limit_is_accepted():
template = "a" * MAX_CHAT_TEMPLATE_BYTES # exactly the limit, 1 byte/char
assert (
len(_load_request(chat_template_override = template).chat_template_override)
== MAX_CHAT_TEMPLATE_BYTES
)
def test_chat_template_over_char_limit_is_rejected():
with pytest.raises(Exception): # pydantic ValidationError wrapping ValueError
_load_request(chat_template_override = "a" * (MAX_CHAT_TEMPLATE_BYTES + 1))
def test_chat_template_over_byte_limit_is_rejected():
# Char count stays under the limit but UTF-8 bytes exceed it (3 bytes/char),
# so only the byte-count branch can catch this.
multibyte = "€" * (MAX_CHAT_TEMPLATE_BYTES // 2) # euro sign, 3 bytes each
assert len(multibyte) <= MAX_CHAT_TEMPLATE_BYTES
assert len(multibyte.encode("utf-8")) > MAX_CHAT_TEMPLATE_BYTES
with pytest.raises(Exception):
_load_request(chat_template_override = multibyte)
def test_read_bounded_text_reads_within_limit(tmp_path):
p = tmp_path / "t.json"
p.write_text("hello", encoding = "utf-8")
assert _read_bounded_text(p, 16) == "hello"
def test_read_bounded_text_rejects_over_limit(tmp_path):
p = tmp_path / "big.json"
p.write_bytes(b"x" * 100)
assert _read_bounded_text(p, 50) is None
def test_read_bounded_text_at_limit_is_read(tmp_path):
p = tmp_path / "exact.json"
p.write_bytes(b"x" * 50)
assert _read_bounded_text(p, 50) == "x" * 50
def test_read_bounded_text_missing_file_is_none(tmp_path):
assert _read_bounded_text(tmp_path / "nope.json", 50) is None

61
studio/backend/tests/test_model_update_robustness.py
View file

@ -314,6 +314,7 @@ def test_cached_model_scan_keeps_local_safetensors_repo(monkeypatch, tmp_path):
file_name = "model.safetensors",
size_on_disk = 100,
blob_path = str(repo_path / "blobs" / "modelsha"),
blob_last_modified = 3_000.0,
),
]
)
@ -336,6 +337,51 @@ def test_cached_model_scan_keeps_local_safetensors_repo(monkeypatch, tmp_path):
assert rows[0]["repo_id"] == "Org/SafeTensorRepo"
assert rows[0]["model_format"] == "safetensors"
assert rows[0]["size_bytes"] == 100
assert rows[0]["last_modified"] == 3_000.0
def test_cached_gguf_scan_keeps_download_timestamp(monkeypatch, tmp_path):
repo_path = tmp_path / "models--Org--GgufRepo"
repo = SimpleNamespace(
repo_id = "Org/GgufRepo",
repo_type = "model",
repo_path = repo_path,
revisions = [
SimpleNamespace(
files = [
SimpleNamespace(
file_name = "model-Q4_K_M.gguf",
size_on_disk = 100,
blob_path = None,
blob_last_modified = 5_000.0,
),
]
)
],
)
monkeypatch.setattr(
CI,
"all_hf_cache_scans",
lambda: [SimpleNamespace(repos = [repo])],
)
monkeypatch.setattr(
CI.hf_cache_scan,
"is_gguf_repo_partial",
lambda *args, **kwargs: False,
)
monkeypatch.setattr(
CI,
"_gguf_variant_state_summary",
lambda _repo_id: (False, 0),
)
rows = CI._scan_cached_gguf()
assert len(rows) == 1
assert rows[0]["repo_id"] == "Org/GgufRepo"
assert rows[0]["model_format"] == "gguf"
assert rows[0]["size_bytes"] == 100
assert rows[0]["last_modified"] == 5_000.0
# ── hf_hub_download_with_xet_fallback force_download bypass (X2/F2) ───
@ -636,3 +682,18 @@ def test_reclaim_replaced_gguf_variant_keeps_no_symlink_current_file(monkeypatch
assert snap.exists() is True # the current file must survive
assert result["removed_snapshots"] == 0
assert result["deleted_blobs"] == 0
def _mmproj_repo(*file_names: str):
return SimpleNamespace(
revisions = [SimpleNamespace(files = [SimpleNamespace(file_name = n) for n in file_names])]
)
def test_repo_has_mmproj_requires_gguf_projector():
# A non-GGUF sidecar whose name merely contains "mmproj" must NOT mark the
# repo vision-capable; the runtime's projector detection is GGUF-only.
assert CI._repo_has_mmproj(_mmproj_repo("model-Q4_K_M.gguf", "mmproj_config.json")) is False
assert CI._repo_has_mmproj(_mmproj_repo("model-Q4_K_M.gguf", "README-mmproj.md")) is False
# A real GGUF projector still marks the repo vision-capable.
assert CI._repo_has_mmproj(_mmproj_repo("model-Q4_K_M.gguf", "mmproj-F16.gguf")) is True

60
studio/backend/tests/test_offline_gguf_cache_fallback.py
View file

@ -119,6 +119,13 @@ def _build_cache(
return snap
def _symlink_or_skip(link: Path, target: Path) -> None:
try:
link.symlink_to(target)
except OSError as exc:
pytest.skip(f"symlinks unavailable: {exc}")
@pytest.fixture
def hf_cache(tmp_path, monkeypatch):
"""Point ``huggingface_hub.constants.HF_HUB_CACHE`` at a temp dir."""
@ -1084,7 +1091,7 @@ class TestListLocalGgufVariantsSubdir:
target.write_bytes(b"\0" * 20)
out = _find_local_gguf_by_variant(str(tmp_path), "Q4_K_M")
assert out == str(target.resolve())
assert out == str(target.absolute())
def test_find_local_gguf_by_variant_skips_big_endian_only_match(self, tmp_path):
from utils.models.model_config import _find_local_gguf_by_variant
@ -1094,6 +1101,57 @@ class TestListLocalGgufVariantsSubdir:
assert _find_local_gguf_by_variant(str(tmp_path), "Q4_K_M") is None
def test_find_local_gguf_by_variant_keeps_split_symlink_name(self, tmp_path):
from utils.models.model_config import _find_local_gguf_by_variant
blobs = tmp_path / "blobs"
blobs.mkdir()
snap = tmp_path / "snapshots" / "rev" / "BF16"
snap.mkdir(parents = True)
(tmp_path / "snapshots" / "rev" / "config.json").write_text("{}")
for i, sha in enumerate(("aa" * 32, "bb" * 32), start = 1):
(blobs / sha).write_bytes(b"\0" * 10)
_symlink_or_skip(snap / f"model-BF16-0000{i}-of-00002.gguf", blobs / sha)
out = _find_local_gguf_by_variant(str(tmp_path / "snapshots" / "rev"), "BF16")
assert out is not None
assert Path(out).name == "model-BF16-00001-of-00002.gguf"
def test_detect_gguf_model_keeps_split_symlink_name(self, tmp_path):
from utils.models.model_config import detect_gguf_model
blobs = tmp_path / "blobs"
blobs.mkdir()
snap = tmp_path / "snapshots" / "rev"
snap.mkdir(parents = True)
for i, (sha, size) in enumerate((("cc" * 32, 10), ("dd" * 32, 20)), start = 1):
(blobs / sha).write_bytes(b"\0" * size)
_symlink_or_skip(snap / f"model-BF16-0000{i}-of-00002.gguf", blobs / sha)
out = detect_gguf_model(str(snap))
assert out is not None
assert Path(out).name == "model-BF16-00001-of-00002.gguf"
def test_lone_split_symlink_uses_colocated_target_shards(self, tmp_path):
from utils.models.model_config import _find_local_gguf_by_variant, detect_gguf_model
target_dir = tmp_path / "external" / "BF16"
target_dir.mkdir(parents = True)
target = target_dir / "model-BF16-00001-of-00002.gguf"
target.write_bytes(b"\0" * 10)
(target_dir / "model-BF16-00002-of-00002.gguf").write_bytes(b"\0" * 10)
local = tmp_path / "local"
local.mkdir()
(local / "config.json").write_text("{}")
link = local / target.name
_symlink_or_skip(link, target)
expected = str(target.absolute())
assert _find_local_gguf_by_variant(str(local), "BF16") == expected
assert detect_gguf_model(str(local)) == expected
assert detect_gguf_model(str(link)) == expected
def test_model_config_variant_ignores_big_endian_sibling(self, tmp_path):
from utils.models.model_config import ModelConfig

588
studio/backend/tests/test_openai_auto_switch.py
View file

@ -8,6 +8,7 @@ tests/test_gguf_completion_usage.py.
"""
import asyncio
import os
import pytest
@ -18,6 +19,10 @@ from utils import openai_auto_switch_settings as settings
class _FakeBackend:
effective_parallel_slots = 1
_slot_save_binary = None
_gguf_path = None
def __init__(
self,
loaded_id = None,
@ -29,6 +34,22 @@ class _FakeBackend:
self.hf_variant = hf_variant
self._openai_advertised_id = advertised_id
def save_slots_for_resume(self, should_abort = None):
return None
def restore_slots_for_resume(self, manifest):
return None
def _slot_launch_fingerprint(self):
return ((), None, None, 1)
def _gguf_file_identity(self, path):
try:
st = os.stat(path)
except OSError:
return None
return ((st.st_size, st.st_mtime_ns),)
class _LoadRecorder:
"""Stand-in for the load route: records calls and simulates a load."""
@ -53,10 +74,15 @@ class _LoadRecorder:
from fastapi import HTTPException
raise HTTPException(status_code = 503, detail = "load failed")
self.backend.model_identifier = request.model_path
self.backend.hf_variant = getattr(request, "gguf_variant", None)
self.backend._gguf_path = request.model_path
self.backend.is_loaded = True
# Mirror _load_model_impl: a load advertises its own id until the
# auto-switch caller overwrites it with the repo id.
self.backend._openai_advertised_id = None
from core.inference import llama_keepwarm as kw
kw.note_model_loaded(self.backend)
return None
@ -446,6 +472,75 @@ def test_idle_loop_unloads_after_ttl_and_stashes_for_reload(monkeypatch):
assert stash is not None and stash[0] == "unsloth/Idle-GGUF" and stash[1] == "Q4_K_M"
def test_idle_loop_deletes_saved_kv_when_unload_fails(monkeypatch, tmp_path):
import time
from core.inference import llama_keepwarm as kw
monkeypatch.setattr(settings, "get_auto_unload_idle_seconds", lambda: 0.005)
monkeypatch.setattr(settings, "get_auto_unload_keep_kv", lambda: True)
kw._inflight = 0
kw._pending = 0
kw._last_active = time.monotonic() - 3600
kw._last_unloaded_model = None
kw._kv_resume = None
saved = tmp_path / "resume-abc-slot0.bin"
backend = _FakeBackend("unsloth/Idle-GGUF")
manifests = []
def _save(should_abort = None):
if manifests:
return None
saved.write_bytes(b"kv")
manifest = {"dir": str(tmp_path), "slots": [{"id": 0, "filename": saved.name}]}
manifests.append(manifest)
return manifest
def _unload():
raise RuntimeError("cuda teardown failed")
backend.save_slots_for_resume = _save
backend.unload_model = _unload
monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: backend)
async def _drive():
task = asyncio.create_task(kw.idle_unload_loop(poll_seconds = 0.01))
for _ in range(200):
await asyncio.sleep(0.01)
if manifests and not saved.exists():
break
task.cancel()
try:
await task
except asyncio.CancelledError:
pass
asyncio.run(_drive())
assert manifests and not saved.exists()
assert kw._kv_resume is None
def test_disabling_idle_unload_purges_saved_kv(monkeypatch, tmp_path):
# PUT leaves keep-KV on but makes idle unload inactive: saved KV must go too.
import routes.settings as settings_route
from core.inference import llama_keepwarm as kw
saved = tmp_path / "resume-abc-slot0.bin"
saved.write_bytes(b"kv")
kw._kv_resume = {
"identity": ("m", None, "m"),
"dir": str(tmp_path),
"slots": [{"id": 0, "filename": saved.name}],
}
monkeypatch.setattr(settings_route, "set_openai_auto_switch", lambda *a: (False, 300, True))
monkeypatch.setattr(settings_route, "get_auto_unload_idle_seconds", lambda: 0)
payload = settings_route.OpenAIAutoSwitchPayload(enabled = False)
resp = settings_route.update_openai_auto_switch(payload, "tester")
assert resp.idle_unload_active is False and resp.auto_unload_keep_kv is True
assert kw._kv_resume is None and not saved.exists()
def test_audio_generate_is_tracked_as_inference_path():
# Direct GGUF TTS uses the llama backend and can outlive the idle TTL, so
# the keep-warm middleware must count it as in-flight inference.
@ -2912,8 +3007,10 @@ def test_non_gguf_load_clears_reload_stash():
# A non-GGUF (Transformers/Unsloth) load must clear the stash like the GGUF
# branch, so it never lingers until the idle poll (or forever, idle-unload off).
import inspect
src = inspect.getsource(inference_route._load_model_impl)
assert src.count("note_model_loaded()") >= 2
assert src.count("note_model_loaded()") >= 1 # non-GGUF branch
assert "to_thread(note_model_loaded, llama_backend)" in src # GGUF branch
def test_chat_rejects_malformed_tool_choice_before_switch(monkeypatch):
@ -3121,6 +3218,495 @@ def test_responses_stream_hint_matches_toggle_regardless_of_active_model(monkeyp
assert "Model auto-switch" in non_gguf_loaded
# ── idle-unload KV persistence (slot save/restore) ──────────────────
def _seed_kv_manifest(
tmp_path,
identity = ("unsloth/A-GGUF", "Q4_K_M", "unsloth/A-GGUF"),
gguf = None,
):
if gguf is None:
gguf_file = tmp_path / "model.gguf"
gguf_file.write_bytes(b"gguf")
gguf = str(gguf_file)
st = os.stat(gguf)
state_file = tmp_path / "resume-abc-slot0.bin"
state_file.write_bytes(b"kv")
return state_file, {
"identity": identity,
"dir": str(tmp_path),
"binary": ("/bin/llama-server", 111),
"gguf": gguf,
"gguf_stat": ((st.st_size, st.st_mtime_ns),),
"launch": ((), None, None, 1),
"slots": [{"id": 0, "filename": state_file.name, "n_saved": 42}],
}
def _drive_idle_loop(
kw,
poll_seconds = 0.02,
run_for = 0.2,
):
async def _drive():
task = asyncio.create_task(kw.idle_unload_loop(poll_seconds = poll_seconds))
await asyncio.sleep(run_for)
task.cancel()
try:
await task
except asyncio.CancelledError:
pass
asyncio.run(_drive())
def test_idle_unload_saves_slots_before_unload_and_stashes_manifest(monkeypatch, tmp_path):
import time
from core.inference import llama_keepwarm as kw
monkeypatch.setattr(settings, "get_auto_unload_idle_seconds", lambda: 0.005)
monkeypatch.setattr(settings, "get_auto_unload_keep_kv", lambda: True)
kw._inflight = 0
kw._pending = 0
kw._last_active = time.monotonic() - 3600
kw._last_unloaded_model = None
kw._kv_resume = None
events = []
backend = _FakeBackend("unsloth/Idle-GGUF", hf_variant = "Q4_K_M")
manifest = {
"dir": str(tmp_path),
"binary": ("bin", 1),
"slots": [{"id": 0, "filename": "f.bin", "n_saved": 42}],
}
def _save(should_abort = None):
events.append("save")
return manifest
def _unload():
events.append("unload")
backend.is_loaded = False
backend.save_slots_for_resume = _save
backend.unload_model = _unload
monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: backend)
_drive_idle_loop(kw)
# KV must be saved while the server is still alive, then exactly one unload.
assert events == ["save", "unload"]
assert kw.get_last_unloaded_model()[:2] == ("unsloth/Idle-GGUF", "Q4_K_M")
resume = kw.take_kv_resume()
assert resume is not None
assert resume["identity"][:2] == ("unsloth/Idle-GGUF", "Q4_K_M")
assert resume["slots"][0]["filename"] == "f.bin"
def test_idle_save_failure_still_unloads_plain(monkeypatch):
import time
from core.inference import llama_keepwarm as kw
monkeypatch.setattr(settings, "get_auto_unload_idle_seconds", lambda: 0.005)
monkeypatch.setattr(settings, "get_auto_unload_keep_kv", lambda: True)
kw._inflight = 0
kw._pending = 0
kw._last_active = time.monotonic() - 3600
kw._last_unloaded_model = None
kw._kv_resume = None
unloads = []
backend = _FakeBackend("unsloth/Idle-GGUF", hf_variant = "Q4_K_M")
def _save(should_abort = None):
raise RuntimeError("slot save exploded")
def _unload():
unloads.append(1)
backend.is_loaded = False
backend.save_slots_for_resume = _save
backend.unload_model = _unload
monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: backend)
_drive_idle_loop(kw)
assert unloads == [1] # the save failure must not skip the unload
assert kw.get_last_unloaded_model() is not None
assert kw.take_kv_resume() is None
def test_keep_kv_setting_off_skips_save(monkeypatch):
import time
from core.inference import llama_keepwarm as kw
monkeypatch.setattr(settings, "get_auto_unload_idle_seconds", lambda: 0.005)
monkeypatch.setattr(settings, "get_auto_unload_keep_kv", lambda: False)
kw._inflight = 0
kw._pending = 0
kw._last_active = time.monotonic() - 3600
kw._last_unloaded_model = None
kw._kv_resume = None
saves, unloads = [], []
backend = _FakeBackend("unsloth/Idle-GGUF")
def _unload():
unloads.append(1)
backend.is_loaded = False
backend.save_slots_for_resume = lambda *a, **k: saves.append(1)
backend.unload_model = _unload
monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: backend)
_drive_idle_loop(kw)
assert saves == []
assert unloads == [1]
assert kw.take_kv_resume() is None
def test_keep_kv_disabled_mid_save_discards_manifest(monkeypatch, tmp_path):
import time
from core.inference import llama_keepwarm as kw
keep = {"on": True}
monkeypatch.setattr(settings, "get_auto_unload_idle_seconds", lambda: 0.005)
monkeypatch.setattr(settings, "get_auto_unload_keep_kv", lambda: keep["on"])
kw._inflight = 0
kw._pending = 0
kw._last_active = time.monotonic() - 3600
kw._last_unloaded_model = None
kw._kv_resume = None
unloads = []
backend = _FakeBackend("unsloth/Idle-GGUF", hf_variant = "Q4_K_M")
state_file = tmp_path / "resume-mid-slot0.bin"
state_file.write_bytes(b"kv")
manifest = {
"dir": str(tmp_path),
"binary": ("bin", 1),
"slots": [{"id": 0, "filename": state_file.name, "n_saved": 1}],
}
def _save(should_abort = None):
keep["on"] = False # user flips the toggle while the save runs
return manifest
def _unload():
unloads.append(1)
backend.is_loaded = False
backend.save_slots_for_resume = _save
backend.unload_model = _unload
monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: backend)
_drive_idle_loop(kw)
assert unloads == [1] # still unloads; only the stash is dropped
assert kw.take_kv_resume() is None
assert not state_file.exists()
def test_idle_ttl_disabled_mid_save_skips_unload(monkeypatch, tmp_path):
import time
from core.inference import llama_keepwarm as kw
ttl = {"v": 0.005}
monkeypatch.setattr(settings, "get_auto_unload_idle_seconds", lambda: ttl["v"])
monkeypatch.setattr(settings, "get_auto_unload_keep_kv", lambda: True)
kw._inflight = 0
kw._pending = 0
kw._last_active = time.monotonic() - 3600
kw._last_unloaded_model = None
kw._kv_resume = None
unloads = []
backend = _FakeBackend("unsloth/Idle-GGUF", hf_variant = "Q4_K_M")
state_file = tmp_path / "resume-mid-slot0.bin"
state_file.write_bytes(b"kv")
manifest = {
"dir": str(tmp_path),
"binary": ("bin", 1),
"slots": [{"id": 0, "filename": state_file.name, "n_saved": 1}],
}
def _save(should_abort = None):
ttl["v"] = 0 # user turns idle unload off while the save runs
return manifest
backend.save_slots_for_resume = _save
backend.unload_model = lambda: unloads.append(1)
monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: backend)
_drive_idle_loop(kw)
assert unloads == [] # the unload was cancelled by the setting change
assert kw.take_kv_resume() is None
assert not state_file.exists()
def test_alias_reload_restores_slots_and_deletes_files(monkeypatch, tmp_path):
from core.inference import llama_keepwarm as kw
backend = _FakeBackend(None) # idle-unload emptied the backend
backend._slot_save_binary = ("/bin/llama-server", 111)
restored = []
backend.restore_slots_for_resume = lambda manifest: restored.append(manifest)
rec = _LoadRecorder(backend)
_wire(monkeypatch, enabled = True, resolves_to = None, backend = backend, recorder = rec)
monkeypatch.setattr(kw, "_inflight", 0)
state_file, manifest = _seed_kv_manifest(tmp_path)
monkeypatch.setattr(kw, "_last_unloaded_model", (manifest["gguf"], "Q4_K_M"))
monkeypatch.setattr(kw, "_kv_resume", manifest)
_run_hook("gpt-4o-mini")
assert len(rec.calls) == 1
assert len(restored) == 1 # same model + binary: restore ran
assert not state_file.exists() # state file deleted after the restore
assert kw._kv_resume is None
def test_no_restore_when_different_model_loads(monkeypatch, tmp_path):
from core.inference import llama_keepwarm as kw
backend = _FakeBackend(None)
backend._slot_save_binary = ("/bin/llama-server", 111)
restored = []
backend.restore_slots_for_resume = lambda manifest: restored.append(manifest)
rec = _LoadRecorder(backend)
_wire(
monkeypatch,
enabled = True,
resolves_to = ("unsloth/B-GGUF", None, "unsloth/B-GGUF"),
backend = backend,
recorder = rec,
)
monkeypatch.setattr(kw, "_inflight", 0)
state_file, manifest = _seed_kv_manifest(tmp_path) # manifest is for model A
monkeypatch.setattr(kw, "_kv_resume", manifest)
_run_hook("unsloth/B-GGUF")
assert len(rec.calls) == 1
assert restored == [] # different model: never restored
assert not state_file.exists() # but the stale files are gone
assert kw._kv_resume is None
def test_restore_skipped_when_binary_changed(monkeypatch, tmp_path):
from core.inference import llama_keepwarm as kw
state_file, manifest = _seed_kv_manifest(tmp_path)
backend = _FakeBackend("unsloth/A-GGUF", hf_variant = "Q4_K_M")
backend._gguf_path = manifest["gguf"]
backend._slot_save_binary = ("/bin/llama-server", 222) # newer mtime
restored = []
backend.restore_slots_for_resume = lambda manifest: restored.append(manifest)
kw.restore_kv_resume(backend, manifest)
assert restored == []
assert not state_file.exists()
def test_restore_skipped_when_launch_config_changed(tmp_path):
from core.inference import llama_keepwarm as kw
state_file, manifest = _seed_kv_manifest(tmp_path)
backend = _FakeBackend("unsloth/A-GGUF", hf_variant = "Q4_K_M")
backend._gguf_path = manifest["gguf"]
backend._slot_save_binary = ("/bin/llama-server", 111)
backend._slot_launch_fingerprint = lambda: (("--rope-freq-scale", "0.5"), None, None, 1)
restored = []
backend.restore_slots_for_resume = lambda manifest: restored.append(manifest)
kw.restore_kv_resume(backend, manifest)
assert restored == []
assert not state_file.exists()
def test_restore_skipped_when_gguf_rewritten_in_place(tmp_path):
from core.inference import llama_keepwarm as kw
state_file, manifest = _seed_kv_manifest(tmp_path)
with open(manifest["gguf"], "wb") as fh:
fh.write(b"different weights") # same path, new content
backend = _FakeBackend("unsloth/A-GGUF", hf_variant = "Q4_K_M")
backend._gguf_path = manifest["gguf"]
backend._slot_save_binary = ("/bin/llama-server", 111)
restored = []
backend.restore_slots_for_resume = lambda manifest: restored.append(manifest)
kw.restore_kv_resume(backend, manifest)
assert restored == []
assert not state_file.exists()
def test_note_model_unloaded_purges_manifest_and_files(tmp_path):
from core.inference import llama_keepwarm as kw
state_file, manifest = _seed_kv_manifest(tmp_path)
kw._set_last_unloaded(("org/A-GGUF", "Q4_K_M"))
kw._set_kv_resume(manifest)
kw.note_model_unloaded()
assert kw.get_last_unloaded_model() is None
assert kw.take_kv_resume() is None
assert not state_file.exists()
def test_note_model_loaded_purges_manifest_and_files(tmp_path):
from core.inference import llama_keepwarm as kw
state_file, manifest = _seed_kv_manifest(tmp_path)
kw._set_last_unloaded(("org/A-GGUF", "Q4_K_M"))
kw._set_kv_resume(manifest)
kw.note_model_loaded()
assert kw.get_last_unloaded_model() is None
assert kw.take_kv_resume() is None
assert not state_file.exists()
def test_new_idle_save_purges_previous_manifest_files(tmp_path):
from core.inference import llama_keepwarm as kw
old_file, old_manifest = _seed_kv_manifest(tmp_path)
kw._set_kv_resume(old_manifest)
new_file = tmp_path / "resume-def-slot0.bin"
new_file.write_bytes(b"kv2")
kw._set_kv_resume(
{
"identity": ("unsloth/B-GGUF", None, "unsloth/B-GGUF"),
"dir": str(tmp_path),
"binary": ("/bin/llama-server", 111),
"slots": [{"id": 0, "filename": new_file.name, "n_saved": 7}],
}
)
assert not old_file.exists() # replaced manifest's files purged
assert new_file.exists()
assert kw.take_kv_resume()["slots"][0]["filename"] == new_file.name
def test_sweep_slot_save_dir_removes_only_resume_files(monkeypatch, tmp_path):
from core.inference import llama_keepwarm as kw
from utils.paths import storage_roots
monkeypatch.setattr(storage_roots, "llama_slot_cache_root", lambda: tmp_path)
stale = tmp_path / "resume-old-slot0.bin"
stale.write_bytes(b"kv")
other = tmp_path / "unrelated.txt"
other.write_text("keep")
kw.sweep_slot_save_dir()
assert not stale.exists()
assert other.exists()
def test_keep_kv_setting_roundtrip_and_default(monkeypatch):
import storage.studio_db as db
store = {}
monkeypatch.setattr(db, "upsert_app_settings", lambda m: store.update(m))
monkeypatch.setattr(settings, "_cached_setting", lambda k, d = None: store.get(k, d))
assert settings.get_auto_unload_keep_kv() is True # default when never stored
assert settings.set_openai_auto_switch(True, 60, False)[2] is False
assert store[settings.AUTO_UNLOAD_KEEP_KV_SETTING_KEY] is False
assert settings.get_auto_unload_keep_kv() is False
# None leaves the stored value untouched (older clients can't reset it).
assert settings.set_openai_auto_switch(True, 60, None)[2] is False
assert store[settings.AUTO_UNLOAD_KEEP_KV_SETTING_KEY] is False
with pytest.raises(ValueError, match = "true or false"):
settings.set_openai_auto_switch(True, 60, "garbage")
def test_stale_stash_cleanup_waits_for_lifecycle_gate(monkeypatch, tmp_path):
# The loop's stale-stash purge must wait on the gate a mid-reload holds.
import time
from core.inference import llama_keepwarm as kw
monkeypatch.setattr(settings, "get_auto_unload_idle_seconds", lambda: 3600)
kw._inflight = 0
kw._pending = 0
kw._last_active = time.monotonic()
backend = _FakeBackend("unsloth/New-GGUF")
monkeypatch.setattr(inference_route, "get_llama_cpp_backend", lambda: backend)
state_file, manifest = _seed_kv_manifest(tmp_path)
kw._kv_resume = manifest
kw._last_unloaded_model = ("unsloth/A-GGUF", "Q4_K_M")
assert kw._lifecycle_lock.acquire(blocking = False) # simulate in-flight reload
try:
_drive_idle_loop(kw)
assert kw._kv_resume is manifest # purge deferred while the gate is held
assert state_file.exists()
finally:
kw._lifecycle_lock.release()
_drive_idle_loop(kw)
assert kw._kv_resume is None # gate freed: genuinely stale stash purged
assert not state_file.exists()
def test_put_route_disabling_keep_kv_purges_saved_state(monkeypatch, tmp_path):
import routes.settings as settings_route
import storage.studio_db as db
from core.inference import llama_keepwarm as kw
store = {}
monkeypatch.setattr(db, "upsert_app_settings", lambda m: store.update(m))
monkeypatch.setattr(settings, "_cached_setting", lambda k, d = None: store.get(k, d))
state_file, manifest = _seed_kv_manifest(tmp_path)
monkeypatch.setattr(kw, "_kv_resume", manifest)
payload = settings_route.OpenAIAutoSwitchPayload(enabled = True, auto_unload_keep_kv = False)
resp = settings_route.update_openai_auto_switch(payload, "tester")
assert resp.auto_unload_keep_kv is False
assert kw._kv_resume is None
assert not state_file.exists()
def test_keep_kv_only_update_leaves_env_idle_ttl_active(monkeypatch):
# A keep-KV-only update must not materialize the env TTL as a stored value.
import routes.settings as settings_route
import storage.studio_db as db
store = {}
monkeypatch.setattr(db, "upsert_app_settings", lambda m: store.update(m))
monkeypatch.setattr(settings, "_cached_setting", lambda k, d = None: store.get(k, d))
monkeypatch.setenv(settings.MODEL_IDLE_TTL_ENV_VAR, "600")
assert settings_route.OpenAIAutoSwitchPayload(enabled = False).auto_unload_idle_seconds is None
enabled, idle, keep_kv = settings.set_openai_auto_switch(False, None, False)
assert settings.AUTO_UNLOAD_IDLE_SETTING_KEY not in store # idle untouched
assert settings.get_auto_unload_idle_seconds() == 600 # env TTL still active
assert (enabled, idle, keep_kv) == (False, 600, False)
def test_load_impl_notes_loaded_with_backend_off_loop():
import inspect
src = inspect.getsource(inference_route._load_model_impl)
assert "to_thread(note_model_loaded, llama_backend)" in src
def test_restore_matches_gguf_realpath_across_naming(tmp_path):
from core.inference import llama_keepwarm as kw
blob = tmp_path / "blob.gguf"
blob.write_bytes(b"gguf")
link = tmp_path / "snapshot.gguf"
try:
link.symlink_to(blob)
except OSError:
pytest.skip("symlinks unsupported on this host")
backend = _FakeBackend("/hf/snapshots/d7f5", hf_variant = None)
backend._gguf_path = str(link) # reload resolved the symlink spelling
backend._slot_save_binary = ("/bin/llama-server", 111)
restored = []
backend.restore_slots_for_resume = lambda manifest: restored.append(manifest)
state_file, manifest = _seed_kv_manifest(
tmp_path, identity = ("unsloth/A-GGUF", None, "unsloth/A-GGUF"), gguf = str(blob)
)
kw.restore_kv_resume(backend, manifest)
assert len(restored) == 1 # names differ, file identical: restore ran
assert not state_file.exists()
def test_setter_rejects_idle_below_floor(monkeypatch):
import storage.studio_db as db

266
studio/backend/tests/test_picker_service.py Normal file
View file

@ -0,0 +1,266 @@
# 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
from types import SimpleNamespace
from picker.service import (
MAX_TEMPLATE_METADATA_BYTES,
_chat_template_from_dir,
_chat_template_from_processor_json,
_chat_template_from_tokenizer_config,
_chat_template_from_tokenizer_dir,
_find_gguf_in_dir,
_iter_ggufs,
read_default_chat_template,
validate_chat_template,
)
def test_iter_ggufs_skips_gguf_companions(tmp_path):
mtp_dir = tmp_path / "MTP"
mtp_dir.mkdir()
main = tmp_path / "model-Q8_0.gguf"
main.write_bytes(b"")
(tmp_path / "mmproj-F16.gguf").write_bytes(b"")
(tmp_path / "mtp-model-Q8_0.gguf").write_bytes(b"")
(mtp_dir / "model-Q8_0-MTP.gguf").write_bytes(b"")
(tmp_path / "model-Q8_0-be.gguf").write_bytes(b"")
assert _iter_ggufs(tmp_path) == [main]
def test_find_gguf_in_dir_matches_quant_label(tmp_path):
mtp_dir = tmp_path / "MTP"
mtp_dir.mkdir()
main = tmp_path / "model-Q8_0.gguf"
main.write_bytes(b"")
(mtp_dir / "model-Q8_0-MTP.gguf").write_bytes(b"")
(tmp_path / "model-Q4_K_M.gguf").write_bytes(b"")
assert _find_gguf_in_dir(tmp_path, "Q8_0") == main
assert _find_gguf_in_dir(tmp_path, "Q4_K") is None
def test_find_gguf_in_dir_without_variant_prefers_largest_model(tmp_path):
smaller = tmp_path / "a-model-Q4_K_M.gguf"
larger = tmp_path / "z-model-Q8_0.gguf"
smaller.write_bytes(b"0")
larger.write_bytes(b"00")
assert _find_gguf_in_dir(tmp_path, None) == larger
def test_find_gguf_in_dir_without_variant_prefers_first_split(tmp_path):
first = tmp_path / "model-Q4_K_M-00001-of-00003.gguf"
second = tmp_path / "model-Q4_K_M-00002-of-00003.gguf"
third = tmp_path / "model-Q4_K_M-00003-of-00003.gguf"
first.write_bytes(b"0")
second.write_bytes(b"000")
third.write_bytes(b"00")
assert _find_gguf_in_dir(tmp_path, None) == first
first.unlink()
assert _find_gguf_in_dir(tmp_path, None) == second
def test_find_gguf_in_dir_matches_bpw_variant_base_label(tmp_path):
target = tmp_path / "model-IQ4_XS-3.53bpw.gguf"
target.write_bytes(b"")
(tmp_path / "model-Q4_K_M.gguf").write_bytes(b"")
assert _find_gguf_in_dir(tmp_path, "IQ4_XS") == target
assert _find_gguf_in_dir(tmp_path, "IQ4_XS-3.53bpw") == target
assert _find_gguf_in_dir(tmp_path, "Q4_K") is None
def test_validate_chat_template_accepts_valid_and_empty():
assert validate_chat_template("{{ messages[0].content }}").valid is True
assert validate_chat_template("").valid is True
assert validate_chat_template(" ").valid is True
def test_validate_chat_template_reports_syntax_error_with_line():
result = validate_chat_template("{% if %}{% endif %}")
assert result.valid is False
assert result.error is not None
assert result.error.startswith("Line ")
def test_chat_template_from_tokenizer_config_reads_string():
assert _chat_template_from_tokenizer_config({"chat_template": "HELLO"}) == "HELLO"
assert _chat_template_from_tokenizer_config({"chat_template": " "}) is None
assert _chat_template_from_tokenizer_config({}) is None
def test_chat_template_from_tokenizer_config_prefers_named_default():
config = {
"chat_template": [
{"name": "tool_use", "template": "TOOL"},
{"name": "default", "template": "DEFAULT"},
]
}
assert _chat_template_from_tokenizer_config(config) == "DEFAULT"
def test_chat_template_from_tokenizer_config_falls_back_to_first_entry():
config = {
"chat_template": [
{"name": "tool_use", "template": "TOOL"},
{"name": "other", "template": "OTHER"},
]
}
assert _chat_template_from_tokenizer_config(config) == "TOOL"
def test_chat_template_from_tokenizer_dir_prefers_jinja_file(tmp_path):
(tmp_path / "chat_template.jinja").write_text("FROM_JINJA", encoding = "utf-8")
(tmp_path / "tokenizer_config.json").write_text(
json.dumps({"chat_template": "FROM_CONFIG"}), encoding = "utf-8"
)
assert _chat_template_from_tokenizer_dir(tmp_path) == "FROM_JINJA"
def test_chat_template_from_tokenizer_dir_reads_tokenizer_config(tmp_path):
(tmp_path / "tokenizer_config.json").write_text(
json.dumps({"chat_template": "FROM_CONFIG"}), encoding = "utf-8"
)
assert _chat_template_from_tokenizer_dir(tmp_path) == "FROM_CONFIG"
def test_chat_template_from_dir_without_variant_prefers_tokenizer(tmp_path):
(tmp_path / "tokenizer_config.json").write_text(
json.dumps({"chat_template": "FROM_CONFIG"}), encoding = "utf-8"
)
assert _chat_template_from_dir(tmp_path) == "FROM_CONFIG"
def test_chat_template_from_dir_with_variant_still_prefers_tokenizer(tmp_path, monkeypatch):
(tmp_path / "tokenizer_config.json").write_text(
json.dumps({"chat_template": "FROM_CONFIG"}), encoding = "utf-8"
)
(tmp_path / "model-Q4_K_M.gguf").write_bytes(b"")
monkeypatch.setattr("picker.service.read_gguf_chat_template", lambda _path: "FROM_GGUF")
# Selecting a variant must not flip precedence to the embedded GGUF template.
assert _chat_template_from_dir(tmp_path, "Q4_K_M") == "FROM_CONFIG"
def test_chat_template_from_dir_with_variant_falls_back_to_gguf(tmp_path, monkeypatch):
(tmp_path / "model-Q4_K_M.gguf").write_bytes(b"")
monkeypatch.setattr("picker.service.read_gguf_chat_template", lambda _path: "FROM_GGUF")
# With no tokenizer sidecar, the embedded GGUF template is still the fallback.
assert _chat_template_from_dir(tmp_path, "Q4_K_M") == "FROM_GGUF"
def test_chat_template_from_dir_returns_none_when_absent(tmp_path):
assert _chat_template_from_dir(tmp_path) is None
def test_read_default_chat_template_direct_gguf_prefers_sidecar(tmp_path, monkeypatch):
gguf = tmp_path / "model-Q4_K_M.gguf"
gguf.write_bytes(b"")
(tmp_path / "tokenizer_config.json").write_text(
json.dumps({"chat_template": "FROM_CONFIG"}), encoding = "utf-8"
)
monkeypatch.setattr("picker.service._build_browse_allowlist", lambda: [tmp_path])
monkeypatch.setattr("picker.service.read_gguf_chat_template", lambda _path: "FROM_GGUF")
# A directly selected .gguf must prefer a maintained sidecar over its embedded copy.
assert read_default_chat_template(str(gguf)) == "FROM_CONFIG"
def test_read_default_chat_template_direct_gguf_falls_back_to_embedded(tmp_path, monkeypatch):
gguf = tmp_path / "model-Q4_K_M.gguf"
gguf.write_bytes(b"")
monkeypatch.setattr("picker.service._build_browse_allowlist", lambda: [tmp_path])
monkeypatch.setattr("picker.service.read_gguf_chat_template", lambda _path: "FROM_GGUF")
# With no sidecar next to the file, the embedded GGUF template is the fallback.
assert read_default_chat_template(str(gguf)) == "FROM_GGUF"
def test_tokenizer_config_over_size_limit_is_skipped_not_parsed(tmp_path):
# An oversized tokenizer_config.json must be skipped before json.loads so a
# hostile sidecar cannot exhaust memory.
padding = "x" * (MAX_TEMPLATE_METADATA_BYTES + 1024)
(tmp_path / "tokenizer_config.json").write_text(
json.dumps({"chat_template": "HELLO", "_pad": padding}), encoding = "utf-8"
)
assert _chat_template_from_tokenizer_dir(tmp_path) is None
def test_processor_json_over_size_limit_is_skipped_not_parsed(tmp_path):
padding = "x" * (MAX_TEMPLATE_METADATA_BYTES + 1024)
(tmp_path / "chat_template.json").write_text(
json.dumps({"default": "HELLO", "_pad": padding}), encoding = "utf-8"
)
assert _chat_template_from_processor_json(tmp_path) is None
def test_tokenizer_config_at_size_limit_is_still_read(tmp_path):
# A normal-sized config is unaffected by the bound (regression guard).
(tmp_path / "tokenizer_config.json").write_text(
json.dumps({"chat_template": "FROM_CONFIG"}), encoding = "utf-8"
)
assert _chat_template_from_tokenizer_dir(tmp_path) == "FROM_CONFIG"
def test_remote_template_over_size_limit_is_skipped_before_download(monkeypatch):
# An uncached Hub repo whose template exceeds the cap must be skipped via the
# remote size pre-check, never downloaded.
import huggingface_hub
monkeypatch.setattr("picker.service.resolve_cached_repo_id_case", lambda name: name)
monkeypatch.setattr("picker.service.iter_hf_cache_snapshots", lambda resolved: [])
def _fail_download(*args, **kwargs):
raise AssertionError("oversized remote template must not be downloaded")
def _fake_get_paths_info(self, repo_id, paths, **kwargs):
return [SimpleNamespace(path = p, size = MAX_TEMPLATE_METADATA_BYTES + 1) for p in paths]
monkeypatch.setattr(huggingface_hub, "hf_hub_download", _fail_download)
monkeypatch.setattr(huggingface_hub.HfApi, "get_paths_info", _fake_get_paths_info)
assert read_default_chat_template("org/oversized-model") is None
def test_remote_oversized_jinja_falls_through_to_tokenizer_template(tmp_path, monkeypatch):
# A raw chat_template.jinja between the response cap (MAX_CHAT_TEMPLATE_BYTES)
# and the download bound (MAX_TEMPLATE_METADATA_BYTES) must not be returned: the
# route drops it, so the remote path must skip the oversized Jinja and fall
# through to the smaller tokenizer_config.json.
import huggingface_hub
from picker.schemas import MAX_CHAT_TEMPLATE_BYTES
big_jinja = tmp_path / "chat_template.jinja"
big_jinja.write_text("{{ x }}" * (MAX_CHAT_TEMPLATE_BYTES // 4), encoding = "utf-8")
assert MAX_CHAT_TEMPLATE_BYTES < big_jinja.stat().st_size < MAX_TEMPLATE_METADATA_BYTES
tokenizer_config = tmp_path / "tokenizer_config.json"
tokenizer_config.write_text(json.dumps({"chat_template": "SMALL_TEMPLATE"}), encoding = "utf-8")
files = {
"chat_template.jinja": big_jinja,
"tokenizer_config.json": tokenizer_config,
}
monkeypatch.setattr("picker.service.resolve_cached_repo_id_case", lambda name: name)
monkeypatch.setattr("picker.service.iter_hf_cache_snapshots", lambda resolved: [])
def _fake_download(repo_id, rel, **kwargs):
target = files.get(rel)
if target is None:
raise FileNotFoundError(rel)
return str(target)
def _fake_get_paths_info(self, repo_id, paths, **kwargs):
return [
SimpleNamespace(
path = p,
size = files[p].stat().st_size if p in files else 0,
)
for p in paths
]
monkeypatch.setattr(huggingface_hub, "hf_hub_download", _fake_download)
monkeypatch.setattr(huggingface_hub.HfApi, "get_paths_info", _fake_get_paths_info)
assert read_default_chat_template("org/big-jinja-model") == "SMALL_TEMPLATE"

361
studio/backend/tests/test_rocm_windows_vram_7072.py Normal file
View file

@ -0,0 +1,361 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Regression tests for issue #7072 -- "VRAM Usage in System Tab is wrong".
Reporter: dual AMD (Radeon PRO W7900 ~48GB + W7500 8GB), Windows 10, ROCm 7.13,
torch 2.11.0+rocm7.13. On Windows without a HIP SDK, amd-smi is permanently
disabled (avoids a UAC/DiskPart prompt) and hipMemGetInfo returns free==total
(used 0). Two symptoms followed:
* System tab (/api/system -> get_visible_gpu_utilization) showed ~0 VRAM used
on every GPU (torch mem_get_info free==total quirk; ROCm/ROCm#1909).
* get_gpu_utilization()'s Windows fallback SUMMED "GPU Adapter Memory\\Dedicated
Usage" across all adapters into ONE fake device with only GPU 0's total, so
the second GPU never appeared.
The fix reads the per-adapter (LUID-instanced) Dedicated Usage performance
counter -- Task Manager's source -- for per-GPU used, takes per-GPU total from
torch device properties, and guards the free==total mem_get_info quirk. CI has no
AMD GPU/Windows, so torch, the performance counter, and platform are all mocked.
"""
from __future__ import annotations
import subprocess
import sys
import types
import pytest
from utils.hardware import hardware as hw
GB = 1024**3
MiB = 1024**2
# ----------------------------------------------------------------------------- #
# Fakes
# ----------------------------------------------------------------------------- #
def _fake_torch(
devices,
*,
free_equals_total = False,
used_per_device = None,
):
"""Build a fake `torch` module. devices: list of (name, total_bytes)."""
dev = list(devices)
class _Props:
def __init__(self, name, total):
self.name = name
self.total_memory = total
def get_device_properties(i):
name, total = dev[i]
return _Props(name, total)
def mem_get_info(i):
_, total = dev[i]
if free_equals_total:
return (total, total)
used = used_per_device[i] if used_per_device is not None else 0
return (total - used, total)
t = types.ModuleType("torch")
t.__version__ = "2.11.0+rocm7.13"
t.version = types.SimpleNamespace(hip = "7.13", cuda = None)
t.cuda = types.SimpleNamespace(
is_available = lambda: len(dev) > 0,
device_count = lambda: len(dev),
current_device = lambda: 0,
get_device_properties = get_device_properties,
mem_get_info = mem_get_info,
memory_allocated = lambda i: 0,
memory_reserved = lambda i: 0,
)
return t
def _adapter_output(adapters):
if not adapters:
return "__NONE__\n"
return "".join(f"{name}|{int(used)}\n" for name, used in adapters)
def _subprocess_run(*, adapter_output = "__NONE__\n", util_output = "12.0\n"):
def fake_run(cmd, *a, **k):
joined = " ".join(cmd) if isinstance(cmd, list) else str(cmd)
if "GPU Adapter Memory" in joined and "InstanceName" in joined:
out = adapter_output
elif "engtype_3D" in joined or "GPU Engine" in joined:
out = util_output
else:
out = "-1\n"
return subprocess.CompletedProcess(args = cmd, returncode = 0, stdout = out, stderr = "")
return fake_run
@pytest.fixture
def win_rocm(monkeypatch):
"""Configure the hardware module as a Windows ROCm host with 2 visible GPUs."""
monkeypatch.setattr(hw, "get_device", lambda: hw.DeviceType.CUDA)
monkeypatch.setattr(hw, "IS_ROCM", True)
monkeypatch.setattr(hw.platform, "system", lambda: "Windows")
monkeypatch.setattr(hw.sys, "platform", "win32")
monkeypatch.setattr(hw, "_smi_query", lambda *a, **k: None) # amd-smi disabled
# Visible set via HIP mask so we don't shell out to amd-smi for the count.
monkeypatch.setenv("HIP_VISIBLE_DEVICES", "0,1")
monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising = False)
monkeypatch.delenv("ROCR_VISIBLE_DEVICES", raising = False)
return monkeypatch
REPORTER_ADAPTERS = [
("luid_0x00000000_0x0000d1e2_phys_0", 40.0 * GB), # W7900, model loaded
("luid_0x00000000_0x0000e34a_phys_0", 0.5 * GB), # W7500, idle
("luid_0x00000000_0x0000f001_phys_0", 3 * MiB), # Basic Render Driver
]
DEVICES = [("AMD Radeon PRO W7900", 48 * GB), ("AMD Radeon PRO W7500", 8 * GB)]
# ----------------------------------------------------------------------------- #
# System tab (get_visible_gpu_utilization) -- the reporter's screenshot
# ----------------------------------------------------------------------------- #
def test_system_tab_shows_per_gpu_used(win_rocm, monkeypatch):
monkeypatch.setitem(sys.modules, "torch", _fake_torch(DEVICES, free_equals_total = True))
monkeypatch.setattr(
hw.subprocess, "run", _subprocess_run(adapter_output = _adapter_output(REPORTER_ADAPTERS))
)
devices = hw.get_visible_gpu_utilization()["devices"]
by_idx = {d["index"]: d for d in devices}
assert len(devices) == 2
assert by_idx[0]["vram_total_gb"] == 48.0
assert by_idx[0]["vram_used_gb"] == pytest.approx(40.0, abs = 0.01) # not 0
assert by_idx[1]["vram_total_gb"] == 8.0 # own total
# The 3 MiB Basic Render Driver counter makes this a hidden-adapter case: only
# the 40 GiB is forced onto the 48 GiB card; the idle card reads Unknown.
assert by_idx[1]["vram_used_gb"] is None
assert by_idx[1]["vram_utilization_pct"] is None
assert all(
d["vram_used_gb"] <= d["vram_total_gb"] for d in devices if d["vram_used_gb"] is not None
)
def test_gpu_utilization_does_not_collapse(win_rocm, monkeypatch):
monkeypatch.setitem(sys.modules, "torch", _fake_torch(DEVICES, free_equals_total = True))
monkeypatch.setattr(
hw.subprocess, "run", _subprocess_run(adapter_output = _adapter_output(REPORTER_ADAPTERS))
)
result = hw.get_gpu_utilization()
devices = result["devices"]
assert sorted(d["index"] for d in devices) == [0, 1] # both GPUs, no collapse
assert {d["vram_total_gb"] for d in devices} == {48.0, 8.0}
assert result["vram_total_gb"] == 48.0 # legacy primary mirror preserved
def test_localized_counter_reports_unknown_not_zero(win_rocm, monkeypatch):
monkeypatch.setitem(sys.modules, "torch", _fake_torch(DEVICES, free_equals_total = True))
monkeypatch.setattr(hw.subprocess, "run", _subprocess_run(adapter_output = "__NONE__\n"))
devices = hw.get_visible_gpu_utilization()["devices"]
assert len(devices) == 2 # both still shown with correct totals
assert {d["vram_total_gb"] for d in devices} == {48.0, 8.0}
assert all(d["vram_used_gb"] is None for d in devices) # unknown, not fake 0
assert all(d["vram_utilization_pct"] is None for d in devices)
# ----------------------------------------------------------------------------- #
# mem_get_info free==total guard scoping
# ----------------------------------------------------------------------------- #
def test_mem_get_info_guard_scopes_to_windows_rocm(monkeypatch):
torch_mod = _fake_torch(DEVICES, free_equals_total = True)
monkeypatch.setattr(hw, "get_device", lambda: hw.DeviceType.CUDA)
monkeypatch.setitem(sys.modules, "torch", torch_mod)
# Windows ROCm -> used unknown (None), total kept.
monkeypatch.setattr(hw, "IS_ROCM", True)
monkeypatch.setattr(hw.sys, "platform", "win32")
win = hw._torch_get_per_device_info([0, 1])
assert [d["used_gb"] for d in win] == [None, None]
assert [d["total_gb"] for d in win] == [48.0, 8.0]
# Linux ROCm -> unchanged numeric used.
monkeypatch.setattr(hw.sys, "platform", "linux")
assert [d["used_gb"] for d in hw._torch_get_per_device_info([0, 1])] == [0.0, 0.0]
# Windows NVIDIA -> guard must not fire.
monkeypatch.setattr(hw, "IS_ROCM", False)
monkeypatch.setattr(hw.sys, "platform", "win32")
assert [d["used_gb"] for d in hw._torch_get_per_device_info([0, 1])] == [0.0, 0.0]
# ----------------------------------------------------------------------------- #
# Per-adapter attribution helpers (pure unit)
# ----------------------------------------------------------------------------- #
def test_match_adapter_pairs_and_clamps():
assert hw._match_adapter_used_to_devices([40 * GB, 0.5 * GB], [48 * GB, 8 * GB]) == [
40 * GB,
0.5 * GB,
]
assert hw._match_adapter_used_to_devices([100 * GB], [48 * GB]) == [48 * GB] # clamp
assert hw._match_adapter_used_to_devices([40 * GB], [48 * GB, 8 * GB]) == [40 * GB, None]
def test_match_adapter_reports_unknown_when_more_active_than_visible():
# More adapters actively using VRAM than are visible (a GPU outside the mask):
# attribution would fabricate a value, so report unknown for every device.
assert hw._match_adapter_used_to_devices([40 * GB, 0.5 * GB], [8 * GB]) == [None]
def test_match_adapter_reports_unknown_when_hidden_high_use_adapter_survives_filter():
# Idle 8 GiB card (10 MiB noise) beside a hidden 48 GiB card at 40 GiB: the
# 40 GiB can't fit the 8 GiB device, so clamping there would fabricate. Unknown.
assert hw._match_adapter_used_to_devices([40 * GB, 10 * MiB], [8 * GB]) == [None]
# Order of the counters must not matter.
assert hw._match_adapter_used_to_devices([10 * MiB, 40 * GB], [8 * GB]) == [None]
def test_match_adapter_reports_unknown_for_placeholder_fallback():
# Every counter below the 64 MiB floor plus a placeholder: no LUID-to-ordinal
# mapping tells placeholder from idle GPU, so report unknown, not fabricate.
# Single visible 8 GiB card idle (10 MiB) beside a 50 MiB placeholder counter.
assert hw._match_adapter_used_to_devices([50 * MiB, 10 * MiB], [8 * GB]) == [None]
# Order of the counters must not matter.
assert hw._match_adapter_used_to_devices([10 * MiB, 50 * MiB], [8 * GB]) == [None]
# Two idle visible GPUs plus a placeholder: all three counters below the floor.
assert hw._match_adapter_used_to_devices([50 * MiB, 10 * MiB, 5 * MiB], [48 * GB, 8 * GB]) == [
None,
None,
]
def test_match_adapter_reports_unknown_when_usage_not_capacity_ordered():
# 8 GiB card at 7 GiB beside a 48 GiB card at 5 GiB: the bigger usage still fits
# the smaller card, so both pairings are feasible -> unknown.
assert hw._match_adapter_used_to_devices([7 * GB, 5 * GB], [8 * GB, 48 * GB]) == [None, None]
# Device order must not matter (same physical situation, ordinals flipped).
assert hw._match_adapter_used_to_devices([7 * GB, 5 * GB], [48 * GB, 8 * GB]) == [None, None]
# Same-capacity cards with unequal usage are equally unattributable.
assert hw._match_adapter_used_to_devices([12 * GB, 8 * GB], [24 * GB, 24 * GB]) == [None, None]
# A single usage that fits both cards can sit on either -> unknown.
assert hw._match_adapter_used_to_devices([5 * GB], [48 * GB, 8 * GB]) == [None, None]
# But a capacity-forced assignment (usage exceeds the smaller card) is kept:
# 40 GiB can only be the 48 GiB card, so it is not fabrication.
assert hw._match_adapter_used_to_devices([40 * GB], [48 * GB, 8 * GB]) == [40 * GB, None]
def test_match_adapter_reports_unknown_when_hidden_usage_fits_visible_card():
# A survivor that merely *fits* a visible card must not be pinned onto it. Two
# cards (48/8 GiB) at 40 GiB / 10 MiB beside a hidden 6 GiB adapter: the 6 GiB
# fits the idle 8 GiB card but isn't forced -> Unknown; only 40 GiB is forced.
assert hw._match_adapter_used_to_devices([40 * GB, 10 * MiB, 6 * GB], [48 * GB, 8 * GB]) == [
40 * GB,
None,
]
# Counter order must not matter.
assert hw._match_adapter_used_to_devices([6 * GB, 40 * GB, 10 * MiB], [48 * GB, 8 * GB]) == [
40 * GB,
None,
]
# A single visible card with a hidden adapter is never attributable: a fitting
# survivor could be the hidden GPU's while the visible card is idle.
assert hw._match_adapter_used_to_devices([6 * GB, 10 * MiB], [8 * GB]) == [None]
def test_match_adapter_capacity_forced_matrix():
"""Exhaustive hidden-adapter matrix for the capacity-forced rule.
A value is emitted only when the supra-threshold counters number exactly the
visible devices AND a device's ranked usage strictly exceeds every smaller
card's capacity. Otherwise (a visible card idle, a merely-fitting usage, or the
smallest card) every device reports unknown.
"""
m = hw._match_adapter_used_to_devices
# -- exactly-n supra-threshold counters, capacity-forced survivors are kept - #
# Both visible cards have a real reading (the 3 MiB is a placeholder): 40 GiB
# forced onto the 48 GiB card, 0.5 GiB not forced -> None.
assert m([40 * GB, 0.5 * GB, 3 * MiB], [48 * GB, 8 * GB]) == [40 * GB, None]
# Three visible cards all active (supra-threshold) + placeholder: 40 > 24 and
# 20 > 8, both forced; the 8 GiB card is not forced -> None.
assert m([40 * GB, 20 * GB, 5 * GB, 3 * MiB], [48 * GB, 24 * GB, 8 * GB]) == [
40 * GB,
20 * GB,
None,
]
# -- fewer supra-threshold counters than visible cards -> all unknown ------ #
# A visible card is idle, so even a "forced" 40 could be the hidden GPU's.
assert m([40 * GB, 3 * MiB, 3 * MiB], [48 * GB, 8 * GB]) == [None, None]
assert m([40 * GB, 10 * MiB, 10 * MiB], [48 * GB, 8 * GB]) == [None, None]
assert m([40 * GB, 20 * GB, 3 * MiB, 3 * MiB], [48 * GB, 24 * GB, 8 * GB]) == [
None,
None,
None,
]
# Middle usage (6 GiB) fits both the 24 and 8 GiB cards, and only two cards are
# active for three visible -> not a bijection -> all unknown.
assert m([40 * GB, 6 * GB, 3 * MiB, 3 * MiB], [48 * GB, 24 * GB, 8 * GB]) == [
None,
None,
None,
]
# -- hidden larger than every visible card -> all unknown ----------------- #
assert m([40 * GB, 10 * MiB], [8 * GB]) == [None]
assert m([48 * GB, 3 * MiB, 3 * MiB], [24 * GB, 8 * GB]) == [None, None]
# -- more active adapters than visible cards -> all unknown --------------- #
assert m([40 * GB, 7 * GB, 6 * GB, 3 * MiB], [48 * GB, 8 * GB]) == [None, None]
assert m([40 * GB, 7 * GB, 6 * GB, 3 * MiB, 3 * MiB], [48 * GB, 8 * GB]) == [None, None]
# -- every counter below the noise floor (placeholder fallback) -> unknown - #
assert m([50 * MiB, 10 * MiB], [8 * GB]) == [None]
assert m([50 * MiB, 10 * MiB, 5 * MiB], [48 * GB, 8 * GB]) == [None, None]
# -- equal-capacity cards with a hidden adapter: nothing is forced -------- #
assert m([40 * GB, 40 * GB, 3 * MiB], [48 * GB, 48 * GB]) == [None, None]
assert m([40 * GB, 30 * GB, 3 * MiB], [48 * GB, 48 * GB]) == [None, None]
def test_perf_counter_parser_and_sentinel(monkeypatch):
monkeypatch.setattr(hw.platform, "system", lambda: "Windows")
monkeypatch.setattr(
hw.subprocess, "run", _subprocess_run(adapter_output = _adapter_output(REPORTER_ADAPTERS))
)
parsed = hw._rocm_windows_perf_counter_vram_by_adapter()
assert parsed is not None and len(parsed) == 3
assert parsed[0][0].startswith("luid_")
monkeypatch.setattr(hw.subprocess, "run", _subprocess_run(adapter_output = "__NONE__\n"))
assert hw._rocm_windows_perf_counter_vram_by_adapter() is None
# ----------------------------------------------------------------------------- #
# Unified-memory (Strix Halo APU) total reconciliation (Codex #7238)
# ----------------------------------------------------------------------------- #
def test_unified_memory_adopts_torch_total_even_when_used_unknown():
"""Windows ROCm unified-memory APU: torch's used is None but its total (the full
GTT pool) is authoritative. The correction must still adopt the larger total;
used stays at amd-smi's figure when torch's is unknown."""
metrics = {"vram_total_gb": 8.0, "vram_used_gb": 2.0, "vram_utilization_pct": 25.0}
hw._apply_unified_memory_correction(metrics, {"total_gb": 124.0, "used_gb": None, "index": 0})
assert metrics["vram_total_gb"] == 124.0 # full unified pool, not the 8 GB carve-out
assert metrics["vram_used_gb"] == 2.0 # amd-smi used preserved (torch's was None)
assert metrics["vram_utilization_pct"] == pytest.approx(round(2.0 / 124.0 * 100, 1))
def test_unified_memory_overwrites_used_when_torch_used_known():
"""When torch reports both a larger total and a known used, both are adopted
and utilization is recomputed against the corrected total (unchanged path)."""
metrics = {"vram_total_gb": 8.0, "vram_used_gb": 2.0, "vram_utilization_pct": 25.0}
hw._apply_unified_memory_correction(metrics, {"total_gb": 124.0, "used_gb": 40.0, "index": 0})
assert metrics["vram_total_gb"] == 124.0
assert metrics["vram_used_gb"] == 40.0
assert metrics["vram_utilization_pct"] == pytest.approx(round(40.0 / 124.0 * 100, 1))
def test_unified_memory_no_op_when_torch_total_not_larger():
"""A discrete GPU where torch total does not exceed amd-smi's is left untouched."""
metrics = {"vram_total_gb": 48.0, "vram_used_gb": 10.0, "vram_utilization_pct": 20.8}
hw._apply_unified_memory_correction(metrics, {"total_gb": 48.0, "used_gb": None, "index": 0})
assert metrics["vram_total_gb"] == 48.0
assert metrics["vram_used_gb"] == 10.0
assert metrics["vram_utilization_pct"] == 20.8

154
studio/backend/tests/test_setup_llama_cpp_backend.py Normal file
View file

@ -0,0 +1,154 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""setup.sh and setup.ps1 must map UNSLOTH_LLAMA_CPP_BACKEND=cpu to
install_llama_prebuilt.py's --force-cpu so users can force the CPU-only prebuilt
on GPU hosts (#7213). The match is case-insensitive and whitespace-trimmed, an
unrecognized value warns instead of silently falling back, and macOS warns (no
CPU-only bundle). Runs the real block extracted from each script so the tests
track the shipped logic.
"""
import os
import re
import shutil
import subprocess
from pathlib import Path
import pytest
_STUDIO = Path(__file__).resolve().parents[2]
_SETUP_SH = _STUDIO / "setup.sh"
_SETUP_PS1 = _STUDIO / "setup.ps1"
_SKIP_NO_BASH = pytest.mark.skipif(shutil.which("bash") is None, reason = "bash unavailable")
_SKIP_NO_PWSH = pytest.mark.skipif(shutil.which("pwsh") is None, reason = "pwsh unavailable")
def _backend_block() -> str:
text = _SETUP_SH.read_text(encoding = "utf-8")
m = re.search(r"_llama_backend=.*?esac", text, re.DOTALL)
assert m, "UNSLOTH_LLAMA_CPP_BACKEND block not found in setup.sh"
return m.group(0)
def _run(value: str | None, system: str = "Linux") -> tuple[list[str], str]:
# Pass the value through env (not the script text) so whitespace survives, and
# stub the setup.sh logging helpers the unknown-value branch calls. system sets
# _HOST_SYSTEM so the macOS (Darwin) no-op branch can be exercised.
env = {k: v for k, v in os.environ.items() if k != "UNSLOTH_LLAMA_CPP_BACKEND"}
if value is not None:
env["UNSLOTH_LLAMA_CPP_BACKEND"] = value
harness = (
f'_PREBUILT_CMD=()\nC_WARN=""\n_HOST_SYSTEM="{system}"\n'
'step() { printf "STEP: %s\\n" "$*" >&2; }\n'
f"{_backend_block()}\n"
'printf "%s\\n" "${_PREBUILT_CMD[@]}"'
)
out = subprocess.run(
["bash", "-c", harness], capture_output = True, text = True, env = env, check = True
)
return out.stdout.split(), out.stderr
@_SKIP_NO_BASH
@pytest.mark.parametrize("value", ["cpu", "CPU", "Cpu", " cpu ", "CPU\t"])
def test_backend_cpu_appends_flag(value):
# A deliberate CPU choice persists, so it uses --force-cpu (not the transient
# --cpu-fallback the arm64 GPU-build recovery uses).
args, stderr = _run(value)
assert "--force-cpu" in args
assert "--cpu-fallback" not in args
assert "Ignoring" not in stderr
@_SKIP_NO_BASH
@pytest.mark.parametrize("value", ["cpu", "CPU", " cpu "])
def test_backend_cpu_macos_warns_no_flag(value):
# macOS has no CPU-only bundle (the universal build already runs on CPU), so the
# override warns instead of writing a misleading forced-CPU marker.
args, stderr = _run(value, system = "Darwin")
assert "--force-cpu" not in args
assert "--cpu-fallback" not in args
assert "macOS" in stderr
@_SKIP_NO_BASH
@pytest.mark.parametrize("value", [None, "", "auto", "AUTO", " "])
def test_backend_auto_no_flag_no_warn(value):
args, stderr = _run(value)
assert "--force-cpu" not in args
assert "Ignoring" not in stderr
@_SKIP_NO_BASH
@pytest.mark.parametrize("value", ["vulkan", "gpu", "cuda"])
def test_backend_unknown_warns_and_no_flag(value):
args, stderr = _run(value)
assert "--force-cpu" not in args
assert "Ignoring" in stderr
@_SKIP_NO_BASH
def test_arm64_recovery_uses_transient_cpu_fallback():
# The arm64 Linux GPU-build recovery must stay transient (--cpu-fallback), never
# the persisted --force-cpu, so a later update can still heal to a GPU bundle (#6097).
text = _SETUP_SH.read_text(encoding = "utf-8")
m = re.search(r"_ARM64_CPU_CMD=\((.*?)\)", text, re.DOTALL)
assert m, "arm64 CPU recovery command not found in setup.sh"
block = m.group(1)
assert "--cpu-fallback" in block
assert "--force-cpu" not in block
def _ps1_search(pattern: str, flags = 0) -> str:
m = re.search(pattern, _SETUP_PS1.read_text(encoding = "utf-8"), flags)
assert m, f"setup.ps1 block not found: {pattern}"
return m.group(0)
def _run_ps1(value: str | None) -> str:
# The override is normalized (assign + warn) at the top of the prebuilt block and
# applied to $prebuiltArgs lower down; compose both real snippets.
normalize = _ps1_search(
r'\$llamaBackend = "\$\(\$env:UNSLOTH_LLAMA_CPP_BACKEND\)".*?Write-Host.*?\n\s*\}',
re.DOTALL,
)
apply_flag = _ps1_search(
r'if \(\$llamaBackend -eq "cpu"\) \{\s*\$prebuiltArgs \+= "--force-cpu"\s*\}'
)
env = {k: v for k, v in os.environ.items() if k != "UNSLOTH_LLAMA_CPP_BACKEND"}
if value is not None:
env["UNSLOTH_LLAMA_CPP_BACKEND"] = value
harness = f'$prebuiltArgs = @()\n{normalize}\n{apply_flag}\n"ARGS:" + ($prebuiltArgs -join ",")'
out = subprocess.run(
["pwsh", "-NoProfile", "-Command", harness],
capture_output = True,
text = True,
env = env,
check = True,
)
return out.stdout
@_SKIP_NO_PWSH
@pytest.mark.parametrize("value", ["cpu", "CPU", "Cpu", " cpu ", "CPU\t"])
def test_ps1_backend_cpu_appends_flag(value):
out = _run_ps1(value)
assert "--force-cpu" in out
assert "Ignoring" not in out
@_SKIP_NO_PWSH
@pytest.mark.parametrize("value", [None, "", "auto", "AUTO", " "])
def test_ps1_backend_auto_no_flag_no_warn(value):
out = _run_ps1(value)
assert "--force-cpu" not in out
assert "Ignoring" not in out
@_SKIP_NO_PWSH
@pytest.mark.parametrize("value", ["vulkan", "gpu", "cuda"])
def test_ps1_backend_unknown_warns_and_no_flag(value):
out = _run_ps1(value)
assert "--force-cpu" not in out
assert "Ignoring" in out

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

@ -262,9 +262,12 @@ def test_proportional_tensor_split_is_emitted_in_tensor_mode():
src = _load_model_source()
assert '"--tensor-split"' in src
gate = src.find("if tensor_parallel:")
ts = src.find('"--tensor-split"')
# Find the TP block's emission (after the gate); manual mode emits its own
# --tensor-split earlier in the source from the user's per-GPU shares.
ts = src.find('"--tensor-split"', gate)
nxt_else = src.find("self._tensor_parallel = False")
assert 0 <= gate < ts < nxt_else, "--tensor-split must be emitted under `if tensor_parallel:`"
assert "tp_tensor_split" in src[gate:nxt_else]
def test_mtp_decode_probe_wired_under_tensor_parallel():

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

@ -126,10 +126,21 @@ _ALLOWED_TP_DROP_GUARDS = {
# Capability: --split-mode tensor aborted for this (binary, model) (#6415).
# Self-healing -- tried by default, skipped only after a real abort (vs #6416).
"tensor_parallel and self._tensor_split_aborts(binary, model_identifier)",
# Capacity: tensor needs >= 2 GPUs clearing the compute-buffer reserve.
"tensor_parallel and len(tp_gpus) < 2",
# Capacity: tensor needs >= 2 GPUs clearing the compute-buffer reserve. Gated
# on plan_tp (not raw tensor_parallel) so manual mode skips this planner (#6414).
"plan_tp and len(tp_gpus) < 2",
# Capacity: pooled usable VRAM can't hold weights + MTP reserve -> layer split.
"_tp_weight_budget_mib <= _tp_required_mib",
# Manual mode, Auto layers: --fit owns memory and is incompatible with a
# tensor split, so TP is dropped (surfaced via logger.info) before the
# cache-drop, so a quantized KV survives into the --fit load (#6414).
"tensor_parallel and gpu_memory_mode == 'manual' and (gpu_layers < 0)",
# Manual mode, explicit layers: a tensor split still needs >= 2 GPUs in use.
"tensor_parallel and gpu_memory_mode == 'manual' and (gpu_layers >= 0) and (self._effective_gpu_count(sorted(gpu_ids) if gpu_ids else None) < 2)",
# Manual mode, zero layers: nothing to split on the GPU, and a tensor-mode
# launch under the CPU-only GPU mask (no visible devices) aborts the server
# instead of the intended CPU-only load (#6414).
"gpu_memory_mode == 'manual' and gpu_layers == 0",
}
@ -364,7 +375,7 @@ def test_compute_buffer_downgrade_preserves_multi_gpu_intent():
full GPU set too, so it is symmetric with the budget/geometry downgrades and
doesn't collapse a multi-GPU layer load to one card (reviewer.py P1 on #6659)."""
src = inspect.getsource(LlamaCppBackend.load_model)
gate = src.find("tensor_parallel and len(tp_gpus) < 2")
gate = src.find("plan_tp and len(tp_gpus) < 2")
assert gate != -1
# Bound to exactly this block: from its gate to the next (budget) downgrade.
nxt = src.find("_tp_weight_budget_mib <= _tp_required_mib", gate)

74
studio/backend/tests/test_training_pump_resilience.py
View file

@ -310,6 +310,80 @@ def test_pump_finalizes_when_read_keeps_raising_on_dead_worker(monkeypatch):
assert b._pump_running is False
def test_interrupted_cancel_clears_in_memory_output_dir(monkeypatch):
# Stop-without-save interrupted before its complete event: /status must not
# keep serving the cleared run's output_dir.
b = TrainingBackend()
finalized: dict = {}
monkeypatch.setattr(b, "_ensure_db_run_created", lambda: None)
monkeypatch.setattr(b, "_finalize_run_in_db", lambda **kw: finalized.update(kw))
b._proc = _FakeProc(alive = False)
b._event_queue = _IdleQueue()
b._progress.is_training = True
b._should_stop = True
b._cancel_requested = True
b._output_dir = "/out/x"
b._pump_loop()
assert b._output_dir is None
assert finalized.get("status") == "stopped"
assert finalized.get("output_dir") is None
assert finalized.get("clear_output_dir") is True
def test_worker_exit_reuses_terminal_stop_save_error(monkeypatch):
b = TrainingBackend()
finalized: dict = {}
monkeypatch.setattr(b, "_ensure_db_run_created", lambda: None)
monkeypatch.setattr(b, "_finalize_run_in_db", lambda **kw: finalized.update(kw))
b._proc = _FakeProc(alive = False)
b._event_queue = _IdleQueue()
b._progress.is_training = True
b._should_stop = True
b._cancel_requested = False
b._output_dir = "/out/x"
b.current_job_id = "job-x"
b._terminal_finalize_payload = {
"status": "error",
"error_message": "checkpoint failed",
"output_dir": "/out/x",
"clear_output_dir": False,
"resume_blocked": True,
"expected_job_id": "job-x",
}
b._pump_loop()
assert b._output_dir == "/out/x"
assert finalized.get("status") == "error"
assert finalized.get("output_dir") == "/out/x"
assert finalized.get("clear_output_dir") is False
assert finalized.get("resume_blocked") is True
def test_dead_worker_crash_preserves_output_dir(monkeypatch):
# A crash (no stop requested) after output_dir was emitted must keep the dir
# in the error finalize: checkpoints under it may still exist.
b = TrainingBackend()
finalized: dict = {}
monkeypatch.setattr(b, "_ensure_db_run_created", lambda: None)
monkeypatch.setattr(b, "_finalize_run_in_db", lambda **kw: finalized.update(kw))
b._proc = _FakeProc(alive = False)
b._event_queue = _IdleQueue()
b._progress.is_training = True
b._output_dir = "/out/x"
b._pump_loop()
assert finalized.get("status") == "error"
assert finalized.get("output_dir") == "/out/x"
assert finalized.get("clear_output_dir") is False
def test_start_training_clears_stale_pump_running_flag():
# A prior pump that died abnormally leaves _pump_running True. The next
# start_training must clear it during reset so the start-time watchdog can't

498
studio/backend/tests/test_training_resume.py
View file

@ -7,6 +7,9 @@ import importlib.util
import json
from pathlib import Path
import pytest
import torch
_BACKEND = Path(__file__).resolve().parents[1]
@ -25,6 +28,30 @@ def _load_resume_module():
resume = _load_resume_module()
def test_resume_request_accepts_sanitized_null_target_modules():
from models.training import TrainingStartRequest
request = TrainingStartRequest(
model_name = "unsloth/Qwen3-0.6B",
training_type = "Full Finetuning",
format_type = "alpaca",
target_modules = None,
)
assert request.target_modules == []
def _write_checkpoint(out: Path, step: int) -> Path:
checkpoint = out / f"checkpoint-{step}"
checkpoint.mkdir(parents = True, exist_ok = True)
(checkpoint / "trainer_state.json").write_text(
json.dumps({"global_step": step}), encoding = "utf-8"
)
torch.save({"weight": torch.ones(1)}, checkpoint / "adapter_model.bin")
torch.save({"state": {0: torch.ones(1)}}, checkpoint / "optimizer.pt")
torch.save({"last_epoch": step}, checkpoint / "scheduler.pt")
return checkpoint
def _stopped_run(**overrides):
run = {
"status": "stopped",
@ -44,6 +71,36 @@ def test_can_resume_run_allows_checkpointed_non_s3_run(monkeypatch):
assert resume.can_resume_run(_stopped_run()) is True
def test_can_resume_run_allows_errored_run_with_checkpoint(monkeypatch):
monkeypatch.setattr(resume, "has_resume_state", lambda _path: True)
assert resume.can_resume_run(_stopped_run(status = "error")) is True
def test_can_resume_run_rejects_errored_run_without_checkpoint(monkeypatch):
monkeypatch.setattr(resume, "has_resume_state", lambda _path: False)
assert resume.can_resume_run(_stopped_run(status = "error")) is False
def test_can_resume_run_allows_errored_run_at_final_step(monkeypatch):
# A save-time crash records final_step == total_steps; resuming re-runs the
# final-save path from the checkpoint.
monkeypatch.setattr(resume, "has_resume_state", lambda _path: True)
run = _stopped_run(status = "error", final_step = 10, total_steps = 10)
assert resume.can_resume_run(run) is True
def test_can_resume_run_rejects_stopped_run_at_final_step(monkeypatch):
monkeypatch.setattr(resume, "has_resume_state", lambda _path: True)
run = _stopped_run(final_step = 10, total_steps = 10)
assert resume.can_resume_run(run) is False
def test_can_resume_run_rejects_s3_dataset_source(monkeypatch):
monkeypatch.setattr(resume, "has_resume_state", lambda _path: True)
@ -91,3 +148,444 @@ def test_list_runs_includes_config_json_for_resume_policy(monkeypatch, tmp_path)
result = studio_db.list_runs()
assert result["runs"][0]["config_json"] == config_json
def test_crashed_run_with_persisted_output_dir_is_resumable(monkeypatch, tmp_path):
from storage import studio_db
monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path))
monkeypatch.setattr(studio_db, "_schema_ready", False)
out = tmp_path / "outputs" / "run_x"
_write_checkpoint(out, 10)
studio_db.create_run(
id = "run-crash",
model_name = "m",
dataset_name = "d",
config_json = "{}",
started_at = "2026-01-01T00:00:00Z",
total_steps = 20,
)
studio_db.update_run_output_dir("run-crash", str(out))
conn = studio_db.get_connection()
conn.execute("UPDATE training_runs SET status = 'error' WHERE id = 'run-crash'")
conn.commit()
conn.close()
run = studio_db.get_run("run-crash")
assert run["output_dir"] == str(out)
assert resume.can_resume_run(run) is True
def test_checkpoint_discovery_skips_malformed_newest(monkeypatch, tmp_path):
monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path))
out = tmp_path / "outputs" / "run_x"
valid = _write_checkpoint(out, 5)
(_write_checkpoint(out, 8) / "scheduler.pt").unlink()
malformed = out / "checkpoint-10"
malformed.mkdir()
(malformed / "trainer_state.json").write_text(json.dumps({"global_step": 10}), encoding = "utf-8")
(malformed / "adapter_model.bin").write_bytes(b"not a torch archive")
(malformed / "optimizer.pt").write_bytes(b"not a torch archive")
assert resume.get_resume_checkpoint_path(str(out)) == str(valid)
def test_completed_run_keeps_output_dir_and_rejects_stale_cancel(monkeypatch, tmp_path):
from storage import studio_db
monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path))
monkeypatch.setattr(studio_db, "_schema_ready", False)
studio_db.create_run(
id = "r",
model_name = "m",
dataset_name = "d",
config_json = "{}",
started_at = "2026-01-01T00:00:00Z",
total_steps = 10,
)
studio_db.update_run_output_dir("r", "/out/x")
studio_db.finish_run(
id = "r",
status = "completed",
ended_at = "t",
final_step = 2,
final_loss = None,
duration_seconds = 1,
loss_sparkline = "[]",
output_dir = "/out/x",
error_message = None,
)
assert studio_db.get_run("r")["output_dir"] == "/out/x"
assert studio_db.mark_run_cancel_requested("r") is False
assert studio_db.get_run("r")["output_dir"] == "/out/x"
assert studio_db.get_run("r")["resume_blocked"] == 0
def test_finish_run_clears_output_dir_for_stop_without_save(monkeypatch, tmp_path):
from storage import studio_db
monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path))
monkeypatch.setattr(studio_db, "_schema_ready", False)
studio_db.create_run(
id = "r",
model_name = "m",
dataset_name = "d",
config_json = "{}",
started_at = "2026-01-01T00:00:00Z",
total_steps = 10,
)
studio_db.update_run_output_dir("r", "/out/x")
studio_db.finish_run(
id = "r",
status = "stopped",
ended_at = "t",
final_step = 2,
final_loss = None,
duration_seconds = 1,
loss_sparkline = "[]",
output_dir = None,
error_message = None,
clear_output_dir = True,
)
assert studio_db.get_run("r")["output_dir"] is None
conn = studio_db.get_connection()
conn.execute(
"UPDATE training_runs SET status = 'running', output_dir = '/out/x', resume_blocked = 0 WHERE id = 'r'"
)
conn.commit()
conn.close()
studio_db.mark_run_cancel_requested("r")
studio_db.cleanup_orphaned_runs()
assert studio_db.get_run("r")["status"] == "stopped"
assert studio_db.get_run("r")["output_dir"] is None
def test_finish_run_clears_output_dir_on_cancel_error_finalize(monkeypatch, tmp_path):
from storage import studio_db
monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path))
monkeypatch.setattr(studio_db, "_schema_ready", False)
studio_db.create_run(
id = "r",
model_name = "m",
dataset_name = "d",
config_json = "{}",
started_at = "2026-01-01T00:00:00Z",
total_steps = 10,
)
studio_db.update_run_output_dir("r", "/out/x")
studio_db.finish_run(
id = "r",
status = "stopped",
ended_at = "t",
final_step = 2,
final_loss = None,
duration_seconds = 1,
loss_sparkline = "[]",
output_dir = "/out/x",
error_message = "worker failed during cancel",
clear_output_dir = True,
)
assert studio_db.get_run("r")["output_dir"] is None
def test_finish_run_preserves_output_dir_for_interrupted_stop_and_save(monkeypatch, tmp_path):
from storage import studio_db
monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path))
monkeypatch.setattr(studio_db, "_schema_ready", False)
studio_db.create_run(
id = "r",
model_name = "m",
dataset_name = "d",
config_json = "{}",
started_at = "2026-01-01T00:00:00Z",
total_steps = 10,
)
studio_db.update_run_output_dir("r", "/out/x")
studio_db.finish_run(
id = "r",
status = "stopped",
ended_at = "t",
final_step = 2,
final_loss = None,
duration_seconds = 1,
loss_sparkline = "[]",
output_dir = None,
error_message = None,
)
assert studio_db.get_run("r")["output_dir"] == "/out/x"
def test_resumed_errored_run_is_not_offered_again(monkeypatch, tmp_path):
from storage import studio_db
monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path))
monkeypatch.setattr(studio_db, "_schema_ready", False)
out = tmp_path / "outputs" / "run_x"
_write_checkpoint(out, 10)
studio_db.create_run(
id = "run-old",
model_name = "m",
dataset_name = "d",
config_json = "{}",
started_at = "2026-01-01T00:00:00Z",
total_steps = 20,
)
studio_db.update_run_output_dir("run-old", str(out))
studio_db.finish_run(
id = "run-old",
status = "error",
ended_at = "2026-01-01T00:05:00Z",
final_step = 10,
final_loss = None,
duration_seconds = 1,
loss_sparkline = "[]",
output_dir = None,
error_message = "killed",
)
studio_db.create_run(
id = "run-new",
model_name = "m",
dataset_name = "d",
config_json = "{}",
started_at = "2026-01-02T00:00:00Z",
total_steps = 20,
output_dir = str(out),
resumed_from_run_id = "run-old",
)
with pytest.raises(RuntimeError, match = "no longer available"):
studio_db.create_run(
id = "run-duplicate",
model_name = "m",
dataset_name = "d",
config_json = "{}",
started_at = "2026-01-02T00:00:01Z",
total_steps = 20,
output_dir = str(out),
resumed_from_run_id = "run-old",
)
assert studio_db.get_run("run-duplicate") is None
studio_db.finish_run(
id = "run-new",
status = "error",
ended_at = "2026-01-02T00:05:00Z",
final_step = 15,
final_loss = None,
duration_seconds = 1,
loss_sparkline = "[]",
output_dir = None,
error_message = "killed again",
)
old_run = studio_db.get_run("run-old")
new_run = studio_db.get_run("run-new")
assert old_run["resumed_later"] == 1
assert resume.can_resume_run(old_run) is False
assert new_run["resumed_later"] == 0
assert resume.can_resume_run(new_run) is True
assert studio_db.get_resumable_run_by_output_dir(str(out))["id"] == "run-new"
def test_running_continuation_blocks_older_resume(monkeypatch, tmp_path):
from storage import studio_db
monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path))
monkeypatch.setattr(studio_db, "_schema_ready", False)
out = tmp_path / "outputs" / "run_x"
_write_checkpoint(out, 10)
studio_db.create_run(
id = "run-old",
model_name = "m",
dataset_name = "d",
config_json = "{}",
started_at = "2026-01-01T00:00:00Z",
total_steps = 20,
)
studio_db.update_run_output_dir("run-old", str(out))
studio_db.finish_run(
id = "run-old",
status = "error",
ended_at = "2026-01-01T00:05:00Z",
final_step = 10,
final_loss = None,
duration_seconds = 1,
loss_sparkline = "[]",
output_dir = None,
error_message = "killed",
)
studio_db.create_run(
id = "run-new",
model_name = "m",
dataset_name = "d",
config_json = "{}",
started_at = "2026-01-02T00:00:00Z",
total_steps = 20,
output_dir = str(out),
resumed_from_run_id = "run-old",
)
old_run = studio_db.get_run("run-old")
assert old_run["resumed_later"] == 1
assert resume.can_resume_run(old_run) is False
assert studio_db.get_resumable_run_by_output_dir(str(out)) is None
def test_stop_save_checkpoint_failure_keeps_error_status(monkeypatch, tmp_path):
# A stop-and-save whose checkpoint write failed must finalize as an error so
# history explains the missing resume state (keep_error_status flag).
from core.training.training import TrainingBackend
from storage import studio_db
monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path))
monkeypatch.setattr(studio_db, "_schema_ready", False)
studio_db.create_run(
id = "run-failed-save",
model_name = "m",
dataset_name = "d",
config_json = "{}",
started_at = "2026-01-01T00:00:00Z",
total_steps = 10,
)
backend = TrainingBackend()
backend.current_job_id = "run-failed-save"
backend._db_run_created = True
backend._should_stop = True
backend._handle_event(
{
"type": "error",
"error": "Failed to save a resumable checkpoint after stop.",
"keep_error_status": True,
}
)
run = studio_db.get_run("run-failed-save")
assert run["status"] == "error"
assert "resumable checkpoint" in run["error_message"]
def test_can_resume_run_rejects_resume_blocked_run(monkeypatch):
monkeypatch.setattr(resume, "has_resume_state", lambda _path: True)
assert resume.can_resume_run(_stopped_run(status = "error", resume_blocked = 1)) is False
def test_stop_save_checkpoint_failure_with_stale_checkpoint_is_not_resumable(monkeypatch, tmp_path):
# A failed stop-and-save must not offer Resume from an older periodic checkpoint;
# that would roll back past the recorded final step.
from core.training.training import TrainingBackend
from storage import studio_db
monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path))
monkeypatch.setattr(studio_db, "_schema_ready", False)
out = tmp_path / "outputs" / "run_x"
_write_checkpoint(out, 10)
studio_db.create_run(
id = "run-stale-ckpt",
model_name = "m",
dataset_name = "d",
config_json = "{}",
started_at = "2026-01-01T00:00:00Z",
total_steps = 20,
)
studio_db.update_run_output_dir("run-stale-ckpt", str(out))
backend = TrainingBackend()
backend.current_job_id = "run-stale-ckpt"
backend._db_run_created = True
backend._should_stop = True
backend._output_dir = str(out)
backend._handle_event(
{
"type": "error",
"error": "Failed to save a resumable checkpoint after stop.",
"keep_error_status": True,
"resume_blocked": True,
}
)
run = studio_db.get_run("run-stale-ckpt")
assert run["status"] == "error"
assert run["resume_blocked"] == 1
assert run["output_dir"] == str(out)
assert resume.can_resume_run(run) is False
def test_user_stop_error_without_checkpoint_ack_is_blocked(monkeypatch, tmp_path):
from core.training.training import TrainingBackend
from storage import studio_db
monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path))
monkeypatch.setattr(studio_db, "_schema_ready", False)
studio_db.create_run(
id = "run-user-stop",
model_name = "m",
dataset_name = "d",
config_json = "{}",
started_at = "2026-01-01T00:00:00Z",
total_steps = 10,
)
backend = TrainingBackend()
backend.current_job_id = "run-user-stop"
backend._db_run_created = True
backend._should_stop = True
backend._handle_event({"type": "error", "error": "interrupted"})
run = studio_db.get_run("run-user-stop")
assert run["status"] == "error" and run["resume_blocked"] == 1
def test_terminal_fallback_keeps_resumable_when_current_checkpoint_landed(monkeypatch, tmp_path):
# Worker died before its terminal event, but a valid current-step checkpoint
# is on disk: the fallback must keep the run resumable, not block it.
from core.training.training import TrainingBackend
monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path))
out = tmp_path / "outputs" / "run_ok"
_write_checkpoint(out, 7)
backend = TrainingBackend()
backend.current_job_id = "run-ok"
backend._should_stop = True
backend._output_dir = str(out)
backend._progress.step = 7
kwargs = backend._terminal_finalize_kwargs()
assert kwargs["status"] == "stopped"
assert kwargs["resume_blocked"] is False
def test_terminal_fallback_blocks_when_no_current_checkpoint(monkeypatch, tmp_path):
# Same path, but only a stale (older-step) checkpoint exists: must block.
from core.training.training import TrainingBackend
monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path))
out = tmp_path / "outputs" / "run_stale"
_write_checkpoint(out, 5)
backend = TrainingBackend()
backend.current_job_id = "run-stale"
backend._should_stop = True
backend._output_dir = str(out)
backend._progress.step = 7
kwargs = backend._terminal_finalize_kwargs()
assert kwargs["status"] == "error"
assert kwargs["resume_blocked"] is True

91
studio/backend/tests/test_training_stop_watchdog.py
View file

@ -353,7 +353,7 @@ def test_finalize_after_escalation_clears_state(monkeypatch):
# stopped so the UI leaves "Stopping..." and a new run can start.
b = TrainingBackend()
finstop: list = []
monkeypatch.setattr(b, "_finish_stopped_run", lambda *a: finstop.append(a))
monkeypatch.setattr(b, "_finish_stopped_run", lambda *a, **k: finstop.append(a))
b._proc = _FakeProc(alive = True) # wedged: still reports alive
b._should_stop = True
@ -365,7 +365,7 @@ def test_finalize_after_escalation_clears_state(monkeypatch):
assert b._proc is None, "the wedged handle must be dropped so is_training_active clears"
assert b._progress.is_training is False
assert b._progress.status_message == "Training stopped."
assert "valid current-step checkpoint" in b._progress.status_message
assert finstop and finstop[0][0] == "job_c", "the captured run must be finalized by id"
assert b.is_training_active() is False
@ -375,7 +375,7 @@ def test_finalize_after_escalation_preserves_output_dir(monkeypatch):
# must record it even if the watchdog wins the finalize race against the pump.
b = TrainingBackend()
finstop: list = []
monkeypatch.setattr(b, "_finish_stopped_run", lambda *a: finstop.append(a))
monkeypatch.setattr(b, "_finish_stopped_run", lambda *a, **k: finstop.append(a))
b._proc = _FakeProc(alive = True)
b._should_stop = True
@ -390,6 +390,28 @@ def test_finalize_after_escalation_preserves_output_dir(monkeypatch):
assert finstop[0][1] == "/tmp/outputs/run-123"
def test_finalize_after_escalation_clears_output_dir_on_cancel(monkeypatch):
# Stop-without-saving promises no resume: a cancel that escalates through the
# watchdog clears the persisted output_dir, not a checkpoint path.
b = TrainingBackend()
finstop: list = []
monkeypatch.setattr(b, "_finish_stopped_run", lambda *a, **k: finstop.append((a, k)))
b._proc = _FakeProc(alive = True)
b._should_stop = True
b._cancel_requested = True
b.current_job_id = "job_c"
b._db_run_created = True
b._output_dir = "/tmp/outputs/run-123"
b._finalize_stopped_after_escalation(watched_job_id = "job_c")
assert finstop and finstop[0][0][0] == "job_c"
assert finstop[0][0][1] is None, "a cancelled run must not record a checkpoint path"
assert finstop[0][1].get("clear_output_dir") is True
assert b._output_dir is None, "/status must stop exposing the cancelled run's dir"
def test_stop_training_starts_watchdog_only_when_worker_alive(monkeypatch):
# No worker -> nothing to escalate; the watchdog must not spawn.
b = TrainingBackend()
@ -409,7 +431,7 @@ def test_finalize_after_escalation_no_ops_when_superseded(monkeypatch):
# The escalation finalize must then leave the NEW run untouched, not drop its handle.
b = TrainingBackend()
finstop: list = []
monkeypatch.setattr(b, "_finish_stopped_run", lambda *a: finstop.append(a))
monkeypatch.setattr(b, "_finish_stopped_run", lambda *a, **k: finstop.append(a))
old_proc = _FakeProc(alive = False) # force-terminated worker we were watching
new_proc = _FakeProc(alive = True) # a new run already took over
@ -430,7 +452,7 @@ def test_finalize_after_escalation_runs_for_its_own_worker(monkeypatch):
# finalizes the captured run by id.
b = TrainingBackend()
finstop: list = []
monkeypatch.setattr(b, "_finish_stopped_run", lambda *a: finstop.append(a))
monkeypatch.setattr(b, "_finish_stopped_run", lambda *a, **k: finstop.append(a))
proc = _FakeProc(alive = False)
b._proc = proc
@ -451,7 +473,7 @@ def test_finalize_after_escalation_no_ops_on_job_change_during_startup(monkeypat
# catch this even though the proc-only guard would not.
b = TrainingBackend()
finstop: list = []
monkeypatch.setattr(b, "_finish_stopped_run", lambda *a: finstop.append(a))
monkeypatch.setattr(b, "_finish_stopped_run", lambda *a, **k: finstop.append(a))
old_proc = _FakeProc(alive = False) # old worker, dead; new _proc not installed yet
b._proc = old_proc # still the old handle (== target), so proc guard would pass
@ -509,6 +531,7 @@ def _install_fake_db(monkeypatch):
recs["insert_ids"].append(job_id),
)
fake_db.update_run_progress = lambda **kw: recs["progress_ids"].append(kw.get("id"))
fake_db.mark_run_cancel_requested = lambda _run_id: True
fake_storage.studio_db = fake_db
monkeypatch.setitem(sys.modules, "storage", fake_storage)
monkeypatch.setitem(sys.modules, "storage.studio_db", fake_db)
@ -518,9 +541,48 @@ def _install_fake_db(monkeypatch):
return recs
def test_stop_without_save_creates_missing_row_before_signal(monkeypatch):
recs = _install_fake_db(monkeypatch)
b = TrainingBackend()
b.current_job_id, b._db_config = "job_missing", {"model_name": "m"}
b._stop_queue = queue.Queue()
assert b.stop_training(save = False) is True
assert [run["id"] for run in recs["created"]] == ["job_missing"]
assert b._stop_queue.get_nowait() == {"type": "stop", "save": False}
b._cancel_requested = b._should_stop = False
sys.modules["storage.studio_db"].mark_run_cancel_requested = lambda _run_id: False
assert b.stop_training(save = False) is False
assert not b._cancel_requested and b._stop_queue.empty()
new_queue = queue.Queue()
b.current_job_id, b._db_run_created = "job_old", True
b._cancel_requested = b._should_stop = False
def _supersede(_run_id):
b.current_job_id = "job_new"
b._stop_queue = new_queue
return True
sys.modules["storage.studio_db"].mark_run_cancel_requested = _supersede
assert b.stop_training(save = False) is False
assert not b._cancel_requested and new_queue.empty()
def test_finalize_run_in_db_single_winner_under_concurrency(monkeypatch):
# The watchdog and pump can both finalize; only one call may reach finish_run.
recs = _install_fake_db(monkeypatch)
monkeypatch.setitem(_G, "_DB_FINALIZE_RETRY_S", 0.0)
attempts = 0
def flaky_finish(**kw):
nonlocal attempts
attempts += 1
if attempts < 3:
raise RuntimeError("database is locked")
recs["finished"].append(kw)
sys.modules["storage.studio_db"].finish_run = flaky_finish
b = TrainingBackend()
b.current_job_id = "job_x"
b._db_run_created = True
@ -539,6 +601,7 @@ def test_finalize_run_in_db_single_winner_under_concurrency(monkeypatch):
t.join(timeout = 5)
assert len(recs["finished"]) == 1, f"finalize must run once, got {len(recs['finished'])}"
assert attempts == 3
assert b._run_finalized is True
@ -646,7 +709,13 @@ def test_ensure_db_run_created_publishes_only_after_insert(monkeypatch):
monkeypatch.setitem(sys.modules, "storage", fake_storage)
monkeypatch.setitem(sys.modules, "storage.studio_db", fake_db)
b._ensure_db_run_created()
b._run_intent_lock.acquire()
creator = threading.Thread(target = b._ensure_db_run_created)
creator.start()
time.sleep(0.02)
assert b._db_create_in_progress is False
b._run_intent_lock.release()
creator.join(timeout = 5)
assert observed["flag_during_create"] is False, "flag must not be published before insert"
assert observed["in_progress_during_create"] is True
@ -718,6 +787,7 @@ def test_escalation_finalizes_watched_run_by_id_end_to_end(monkeypatch):
b = TrainingBackend()
b.current_job_id = "job_old"
b._db_run_created = True
b._should_stop = True
b._proc = _FakeProc(alive = False)
b._progress.is_training = True
b._progress.step = 42
@ -726,7 +796,8 @@ def test_escalation_finalizes_watched_run_by_id_end_to_end(monkeypatch):
b._finalize_stopped_after_escalation(target_proc = b._proc, watched_job_id = "job_old")
assert [f["id"] for f in recs["finished"]] == ["job_old"], "must finish the captured run by id"
assert recs["finished"][0]["status"] == "stopped"
assert recs["finished"][0]["status"] == "error"
assert recs["finished"][0]["resume_blocked"] is True
assert recs["insert_ids"] == ["job_old"], "buffered metrics must land on the captured run"
assert b._metric_buffer == [], "the captured batch must be drained"
@ -737,7 +808,7 @@ def test_escalation_defers_when_row_cannot_be_created_here(monkeypatch):
# so the pump's create-then-finalize records the run. Parent state still clears.
b = TrainingBackend()
called: list = []
monkeypatch.setattr(b, "_finish_stopped_run", lambda *a: called.append(a))
monkeypatch.setattr(b, "_finish_stopped_run", lambda *a, **k: called.append(a))
b._proc = _FakeProc(alive = False)
b.current_job_id = "job_q"
@ -785,7 +856,7 @@ def test_escalation_does_not_drop_a_new_runs_handle(monkeypatch):
new_proc = _FakeProc(alive = True)
b._proc = old_proc
def hijack(*a):
def hijack(*a, **k):
b._proc = new_proc # a new run takes over during the finalize
monkeypatch.setattr(b, "_finish_stopped_run", hijack)

16
studio/backend/tests/test_utils.py
View file

@ -38,7 +38,7 @@ from utils.hardware import (
DeviceType,
)
import utils.hardware.hardware as _hw_module
from utils.utils import format_error_message
from utils.utils import format_error_message, is_hf_authentication_error
# ========== Helpers ==========
@ -439,6 +439,20 @@ class TestFormatErrorMessage:
msg = format_error_message(err, "any/model")
assert "invalid" in msg.lower()
def test_hf_authentication_error_follows_wrapped_401(self):
response = type("Response", (), {"status_code": 401})()
auth_error = Exception("request failed")
auth_error.response = response
wrapper = RuntimeError("model validation failed")
wrapper.__cause__ = auth_error
assert is_hf_authentication_error(wrapper) is True
def test_hf_authentication_error_does_not_treat_429_as_invalid(self):
response = type("Response", (), {"status_code": 429})()
rate_error = Exception("too many requests")
rate_error.response = response
assert is_hf_authentication_error(rate_error) is False
# --- OOM on CUDA ---
@needs_torch

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

@ -538,21 +538,31 @@ def _torch_get_physical_gpu_count() -> Optional[int]:
def _torch_get_per_device_info(device_indices: list[int]) -> list[Dict[str, Any]]:
"""Query torch for per-GPU name, total VRAM, and used VRAM."""
"""Query torch for per-GPU name, total VRAM, and used VRAM.
``used_gb`` is ``None`` on Windows ROCm when ``hipMemGetInfo`` reports
``free == total`` (ROCm/ROCm#1909): that 0 means unknown, not empty.
"""
mod, _ = _torch_get_device_module()
if mod is None:
return []
# free==total is a Windows-ROCm-only quirk.
_win_rocm = sys.platform == "win32" and IS_ROCM
devices = []
for ordinal, phys_idx in enumerate(device_indices):
try:
# torch ordinals are 0-based relative to CUDA_VISIBLE_DEVICES.
props = mod.get_device_properties(ordinal)
total_bytes = props.total_memory
used_bytes: Optional[int]
# Prefer mem_get_info (system-wide) so auto-select sees other consumers.
if hasattr(mod, "mem_get_info"):
free_bytes, total_bytes = mod.mem_get_info(ordinal)
used_bytes = total_bytes - free_bytes
# free==total is the broken-API sentinel, not an idle GPU.
if _win_rocm and free_bytes == total_bytes:
used_bytes = None
else:
used_bytes = mod.memory_allocated(ordinal)
devices.append(
@ -561,7 +571,7 @@ def _torch_get_per_device_info(device_indices: list[int]) -> list[Dict[str, Any]
"visible_ordinal": ordinal,
"name": props.name,
"total_gb": round(total_bytes / (1024**3), 2),
"used_gb": round(used_bytes / (1024**3), 2),
"used_gb": round(used_bytes / (1024**3), 2) if used_bytes is not None else None,
}
)
except Exception as e:
@ -724,20 +734,30 @@ def _rocm_linux_sysfs_vram_gb() -> tuple[Optional[float], Optional[float]]:
return None, None
def _rocm_windows_perf_counter_vram_gb() -> tuple[Optional[float], Optional[float]]:
"""Query system-wide dedicated GPU VRAM via Windows Performance Counters.
# ── Windows AMD/ROCm per-adapter VRAM (issue #7072) ──────────────────────────
# amd-smi is disabled and hipMemGetInfo reports free==total, so read used from the
# per-LUID "GPU Adapter Memory" perf counters and take each total from torch, so
# every GPU shows instead of one fake device with GPU 0's total.
# Placeholder adapters (Basic Render Driver / idle iGPU) drop only when they would
# outnumber the real torch devices.
_ROCM_WIN_ADAPTER_MIN_BYTES = 64 * 1024 * 1024 # 64 MiB
Same data source as Task Manager, so cross-process usage is accurate.
Works for any GPU vendor without amd-smi or nvidia-smi.
Returns (used_gb, total_gb) or (None, None) on failure.
def _rocm_windows_perf_counter_vram_by_adapter() -> Optional[list[tuple[str, float]]]:
"""Per-adapter dedicated VRAM usage on Windows via Performance Counters.
Returns ``[(instance_name, used_bytes)]`` (one per LUID-named adapter), or
``None`` when the counter is unavailable/localized/empty so callers fall back.
"""
if platform.system() != "Windows":
return None, None
return None
try:
# Emit "<InstanceName>|<CookedValue>" per sample, or a __NONE__ sentinel.
ps = (
"$s=(Get-Counter '\\GPU Adapter Memory(*)\\Dedicated Usage'"
" -ErrorAction SilentlyContinue).CounterSamples;"
"if($s){($s|Measure-Object CookedValue -Sum).Sum}else{-1}"
"if($s){$s|ForEach-Object{'{0}|{1}' -f $_.InstanceName,[int64]$_.CookedValue}}"
"else{'__NONE__'}"
)
r = subprocess.run(
["powershell", "-NoProfile", "-NonInteractive", "-Command", ps],
@ -746,16 +766,167 @@ def _rocm_windows_perf_counter_vram_gb() -> tuple[Optional[float], Optional[floa
timeout = 5,
)
if r.returncode != 0 or not r.stdout.strip():
return None, None
used_bytes = float(r.stdout.strip())
if used_bytes < 0:
return None, None
import torch as _torch
total_bytes = _torch.cuda.get_device_properties(0).total_memory
return round(used_bytes / (1024**3), 2), round(total_bytes / (1024**3), 2)
return None
adapters: list[tuple[str, float]] = []
for line in r.stdout.splitlines():
line = line.strip()
if not line or line == "__NONE__" or "|" not in line:
continue
instance, _, raw = line.rpartition("|")
try:
used = float(raw.strip())
except (ValueError, TypeError):
continue
if used < 0:
continue
adapters.append((instance.strip(), used))
return adapters or None
except Exception:
return None, None
return None
def _match_adapter_used_to_devices(
adapter_useds: list[float], device_totals: list[float]
) -> list[Optional[float]]:
"""Attribute per-adapter used bytes to torch devices by capacity ranking.
Windows shares no key between LUID counters and torch ordinals, so usages are
ranked against device totals and each is trusted only when capacity *forces* it
(it exceeds every smaller device); an ambiguous ranking reports unknown
(``None``) rather than fabricate a per-index free.
Extra counters mean a hidden/display adapter, and the noise filter may have
dropped a real reading, so values are emitted only when the supra-threshold
counters number EXACTLY the visible devices AND capacity forces the mapping;
otherwise every device is unknown. Best-effort but correct for the common
loaded-card case (#7072). Returns a list aligned to ``device_totals``.
"""
n = len(device_totals)
if n == 0:
return []
useds = sorted(adapter_useds, reverse = True)
ranked_positions = sorted(range(n), key = lambda i: -device_totals[i])
ranked_totals = [device_totals[pos] for pos in ranked_positions]
assigned: list[Optional[float]]
# More counters than devices -> a hidden/display adapter (check before noise filter).
if len(useds) > n:
non_trivial = [u for u in useds if u >= _ROCM_WIN_ADAPTER_MIN_BYTES]
if len(non_trivial) != n:
# Not a clean bijection (a masked GPU is busy or a visible card idle):
# no counter maps to a specific card, so report unknown.
return [None] * n
# Exactly n supra-threshold counters: extras were placeholders, so a
# capacity-ranked bijection is plausible.
useds = non_trivial
ranked_useds = [useds[rank] for rank in range(n)]
# A usage above its ranked capacity is a hidden larger GPU; clamping onto the
# smaller card would fabricate a fully-used reading.
for rank in range(n):
if ranked_useds[rank] > ranked_totals[rank]:
return [None] * n
# Capacity forces the mapping only when the usage exceeds the next-smaller
# capacity; the smallest card and merely-fitting usages stay unknown.
# Keeps 40 GiB over 48/8 GiB -> [40, None].
assigned = [None] * n
for rank, pos in enumerate(ranked_positions):
if rank + 1 < n and ranked_useds[rank] > ranked_totals[rank + 1]:
assigned[pos] = min(ranked_useds[rank], device_totals[pos])
return assigned
# No hidden adapters: every counter is a visible card, so ranking is a permutation.
ranked_useds = [useds[rank] if rank < len(useds) else 0.0 for rank in range(n)]
# Ambiguous if a strictly larger usage also fits the next smaller card: the two
# could be swapped without breaking capacity, so ranking can't tell them apart.
for rank in range(n - 1):
upper, lower = ranked_useds[rank], ranked_useds[rank + 1]
if upper > lower and upper <= ranked_totals[rank + 1]:
return [None] * n
assigned = [None] * n
for rank, pos in enumerate(ranked_positions):
if rank < len(useds):
assigned[pos] = min(useds[rank], device_totals[pos])
return assigned
def _rocm_windows_per_device_vram(device_indices: list[int]) -> list[Dict[str, Any]]:
"""Per-GPU VRAM on Windows AMD/ROCm: total from torch properties (reliable),
used from the per-adapter Dedicated Usage counter.
Returns ``{index, visible_ordinal, name, used_gb, total_gb}`` per visible GPU
(``used_gb`` may be ``None`` when the counter is unavailable), or ``[]`` when
torch can't enumerate devices so callers fall through to the torch last resort.
"""
if platform.system() != "Windows":
return []
mod, _ = _torch_get_device_module()
if mod is None:
return []
# Totals/names from torch properties (mem_get_info's free==total quirk zeroes used).
dev_meta: list[Dict[str, Any]] = []
for ordinal, phys_idx in enumerate(device_indices):
try:
props = mod.get_device_properties(ordinal)
dev_meta.append(
{
"index": phys_idx,
"visible_ordinal": ordinal,
"name": props.name,
"total_bytes": int(props.total_memory),
}
)
except Exception as e:
logger.debug("torch property probe failed for ordinal %d: %s", ordinal, e)
if not dev_meta:
return []
adapters = _rocm_windows_perf_counter_vram_by_adapter()
if adapters:
assigned = _match_adapter_used_to_devices(
[used for _, used in adapters],
[d["total_bytes"] for d in dev_meta],
)
else:
# Counter unavailable: show every GPU with a correct total, used unknown.
assigned = [None] * len(dev_meta)
devices: list[Dict[str, Any]] = []
for meta, used_bytes in zip(dev_meta, assigned):
total_gb = round(meta["total_bytes"] / (1024**3), 2)
used_gb = round(used_bytes / (1024**3), 2) if used_bytes is not None else None
devices.append(
{
"index": meta["index"],
"visible_ordinal": meta["visible_ordinal"],
"name": meta["name"],
"used_gb": used_gb,
"total_gb": total_gb,
}
)
return devices
def _rocm_windows_device_payload_entry(
device: DeviceType, dev: Dict[str, Any], gpu_util_pct: Optional[float]
) -> Dict[str, Any]:
"""Build a ``get_gpu_utilization`` device entry from a per-device VRAM dict."""
total_gb = dev["total_gb"]
used_gb = dev["used_gb"]
return {
"available": True,
"backend": _backend_label(device),
"index": dev["index"],
"visible_ordinal": dev["visible_ordinal"],
"name": dev.get("name", "Unknown"),
"gpu_utilization_pct": gpu_util_pct,
"temperature_c": None,
"vram_used_gb": used_gb,
"vram_total_gb": total_gb,
"vram_utilization_pct": round((used_gb / total_gb) * 100, 1)
if total_gb and total_gb > 0 and used_gb is not None
else None,
"power_draw_w": None,
"power_limit_w": None,
"power_utilization_pct": None,
}
def _gpu_utilization_payload(
@ -821,30 +992,24 @@ def get_gpu_utilization() -> Dict[str, Any]:
index_kind = result.get("index_kind"),
)
# Fallback Windows ROCm
# Fallback Windows ROCm: per-adapter VRAM attribution (issue #7072), so
# every visible GPU is shown instead of a sum collapsed onto one device.
if IS_ROCM and platform.system() == "Windows":
_win_used, _win_total = _rocm_windows_perf_counter_vram_gb()
if _win_used is not None and _win_total is not None:
_win_util = _rocm_windows_perf_counter_gpu_util_pct()
_win_ids = _get_parent_visible_gpu_spec().get("numeric_ids")
if not _win_ids:
_win_ids = list(range(_torch_get_physical_gpu_count() or 0))
_win_devices = _rocm_windows_per_device_vram(_win_ids)
if _win_devices:
# A single visible GPU can own the aggregate 3D-engine utilization;
# across several GPUs the sum isn't per-device, so leave it unset.
_win_util = (
_rocm_windows_perf_counter_gpu_util_pct() if len(_win_devices) == 1 else None
)
return _gpu_utilization_payload(
device,
[
{
"available": True,
"backend": _backend_label(device),
"index": 0,
"visible_ordinal": 0,
"gpu_utilization_pct": _win_util,
"temperature_c": None,
"vram_used_gb": _win_used,
"vram_total_gb": _win_total,
"vram_utilization_pct": round((_win_used / _win_total) * 100, 1)
if _win_total > 0
else None,
"power_draw_w": None,
"power_limit_w": None,
"power_utilization_pct": None,
}
_rocm_windows_device_payload_entry(device, _wd, _win_util)
for _wd in _win_devices
],
)
@ -901,7 +1066,7 @@ def get_gpu_utilization() -> Dict[str, Any]:
"vram_used_gb": _used,
"vram_total_gb": _total,
"vram_utilization_pct": round((_used / _total) * 100, 1)
if _total > 0
if _total > 0 and _used is not None
else None,
"power_draw_w": None,
"power_limit_w": None,
@ -995,19 +1160,27 @@ def _apply_unified_memory_correction(
endpoints stay in sync on AMD iGPUs with unified memory.
"""
torch_total_gb = torch_info["total_gb"]
torch_used_gb = torch_info.get("used_gb")
smi_total_gb = device_metrics.get("vram_total_gb") or 0.0
# torch sees the full unified (GTT) pool; amd-smi only the dedicated carve-out.
# Adopt torch's larger total regardless of used: on Windows ROCm torch_used is
# None (free==total sentinel) but its total stays authoritative. Overwrite used
# only when torch's is known, then recompute utilization against whatever remains.
if torch_total_gb > smi_total_gb:
torch_used_gb = torch_info["used_gb"]
device_metrics["vram_total_gb"] = torch_total_gb
device_metrics["vram_used_gb"] = torch_used_gb
if torch_used_gb is not None:
device_metrics["vram_used_gb"] = torch_used_gb
_used_for_pct = device_metrics.get("vram_used_gb")
device_metrics["vram_utilization_pct"] = (
round((torch_used_gb / torch_total_gb) * 100, 1) if torch_total_gb > 0 else None
round((_used_for_pct / torch_total_gb) * 100, 1)
if torch_total_gb > 0 and _used_for_pct is not None
else None
)
logger.debug(
"ROCm unified memory: replaced amd-smi VRAM (%.2f GB) with "
"torch mem_get_info total (%.2f GB) for device %s",
smi_total_gb,
"ROCm unified memory: adopted torch mem_get_info total (%.2f GB) over "
"amd-smi (%.2f GB) for device %s",
torch_total_gb,
smi_total_gb,
torch_info.get("index"),
)
@ -1067,6 +1240,49 @@ def get_visible_gpu_utilization() -> Dict[str, Any]:
_reconcile_rocm_unified_memory(result, numeric_ids)
return result
# Windows AMD/ROCm (issue #7072): the System tab's VRAM source. The torch
# fallback below would report used==0 (free==total), so read per-adapter
# Dedicated Usage instead; total from torch properties.
if IS_ROCM and platform.system() == "Windows":
win_numeric_ids = parent_visible_spec.get("numeric_ids")
if win_numeric_ids:
win_ids = win_numeric_ids
win_index_kind = "physical"
else:
win_ids = list(range(_torch_get_physical_gpu_count() or 0))
win_index_kind = "relative"
win_devices = _rocm_windows_per_device_vram(win_ids)
if win_devices:
devices = []
for wd in win_devices:
total = wd["total_gb"]
used = wd["used_gb"]
devices.append(
{
"index": wd["index"],
"index_kind": win_index_kind,
"visible_ordinal": wd["visible_ordinal"],
"name": wd.get("name"),
"gpu_utilization_pct": None,
"temperature_c": None,
"vram_used_gb": used,
"vram_total_gb": total,
"vram_utilization_pct": round((used / total) * 100, 1)
if total and total > 0 and used is not None
else None,
"power_draw_w": None,
"power_limit_w": None,
"power_utilization_pct": None,
}
)
return {
"available": True,
"backend": _backend_label(device),
"parent_visible_gpu_ids": win_numeric_ids or [],
"devices": devices,
"index_kind": win_index_kind,
}
# Torch-based fallback for CUDA (nvidia-smi unavailable, AMD ROCm) and XPU (Intel)
if device in (DeviceType.CUDA, DeviceType.XPU):
parent_ids = get_parent_visible_gpu_ids()
@ -1094,7 +1310,7 @@ def get_visible_gpu_utilization() -> Dict[str, Any]:
"vram_used_gb": used,
"vram_total_gb": total,
"vram_utilization_pct": round((used / total) * 100, 1)
if total > 0
if total > 0 and used is not None
else None,
"power_draw_w": None,
"power_limit_w": None,

208
studio/backend/utils/hf_token_validation.py Normal file
View file

@ -0,0 +1,208 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Cached, rate-limited Hugging Face token validation."""
from __future__ import annotations
import hashlib
import threading
import time
from collections import deque
from dataclasses import dataclass
from typing import Literal
from huggingface_hub import HfApi
from huggingface_hub.utils import build_hf_headers, get_session
TokenValidationStatus = Literal["valid", "invalid", "rate_limited", "unavailable"]
@dataclass(frozen = True)
class TokenValidationResult:
status: TokenValidationStatus
retry_after_seconds: int | None = None
_WINDOW_SECONDS = 3600.0
_MAX_ATTEMPTS = 3
_CACHE_TTL_SECONDS = 3600.0
_TEMPORARY_CACHE_TTL_SECONDS = 15.0
_MAX_BUCKETS = 4096
_MAX_CACHE_ENTRIES = 4096
_INFLIGHT_WAIT_SECONDS = 30.0
_REMOTE_TIMEOUT_SECONDS = 10.0
_attempts: dict[str, deque[float]] = {}
_cache: dict[str, tuple[float, TokenValidationResult]] = {}
_inflight: dict[str, threading.Event] = {}
_lock = threading.Lock()
def _fingerprint(token: str) -> str:
return hashlib.sha256(token.encode("utf-8")).hexdigest()
def _prune_attempts(bucket: deque[float], now: float) -> None:
while bucket and now - bucket[0] >= _WINDOW_SECONDS:
bucket.popleft()
def _prune_locked(now: float) -> None:
for key in list(_attempts):
bucket = _attempts[key]
_prune_attempts(bucket, now)
if not bucket:
del _attempts[key]
for key, (expires_at, _result) in list(_cache.items()):
if expires_at <= now:
del _cache[key]
def _cached_locked(fingerprint: str, now: float) -> TokenValidationResult | None:
cached = _cache.get(fingerprint)
if cached is None:
return None
expires_at, result = cached
if expires_at <= now:
del _cache[fingerprint]
return None
return result
def _retry_after(bucket: deque[float], now: float) -> int:
return max(1, int(_WINDOW_SECONDS - (now - bucket[0])) + 1)
def _reserve_attempt_locked(rate_key: str, now: float) -> TokenValidationResult | None:
bucket = _attempts.get(rate_key)
if bucket is None:
if len(_attempts) >= _MAX_BUCKETS:
_prune_locked(now)
if len(_attempts) >= _MAX_BUCKETS:
return TokenValidationResult(
status = "rate_limited",
retry_after_seconds = max(1, int(_WINDOW_SECONDS)),
)
bucket = _attempts[rate_key] = deque()
_prune_attempts(bucket, now)
if len(bucket) >= _MAX_ATTEMPTS:
return TokenValidationResult(
status = "rate_limited",
retry_after_seconds = _retry_after(bucket, now),
)
bucket.append(now)
return None
def _http_status(response: object | None) -> int | None:
status = getattr(response, "status_code", None)
try:
return int(status) if status is not None else None
except (TypeError, ValueError):
return None
def _remote_retry_after(response: object | None) -> int | None:
headers = getattr(response, "headers", None)
if not headers:
return None
raw = headers.get("Retry-After")
try:
return max(1, int(float(raw))) if raw is not None else None
except (TypeError, ValueError):
return None
def _classify_response(response: object | None) -> TokenValidationResult:
status = _http_status(response)
if status is not None and 200 <= status < 300:
return TokenValidationResult(status = "valid")
if status == 401:
return TokenValidationResult(status = "invalid")
if status == 429:
return TokenValidationResult(
status = "rate_limited",
retry_after_seconds = _remote_retry_after(response),
)
return TokenValidationResult(status = "unavailable")
def _check_remote(token: str) -> TokenValidationResult:
api = HfApi()
try:
# HfApi.whoami has no timeout parameter in the pinned Hub client.
# Use its session and headers against the same whoami endpoint.
response = get_session().get(
f"{api.endpoint}/api/whoami-v2",
headers = build_hf_headers(token = token),
timeout = _REMOTE_TIMEOUT_SECONDS,
)
except Exception as exc:
# huggingface-hub 0.36.x can wrap a 401 as requests.HTTPError.
return _classify_response(getattr(exc, "response", None))
return _classify_response(response)
def validate_hf_token(token: str, *, rate_key: str) -> TokenValidationResult:
"""Validate ``token`` without retaining it, sharing results across callers.
Cached checks do not consume the caller's three-per-hour network budget. A
single-flight event also prevents simultaneously mounted UI surfaces from
sending duplicate ``whoami`` requests for the same token.
"""
normalized = token.strip()
if not normalized:
return TokenValidationResult(status = "invalid")
token_fingerprint = _fingerprint(normalized)
owner_event: threading.Event | None = None
try:
while True:
now = time.monotonic()
with _lock:
cached = _cached_locked(token_fingerprint, now)
if cached is not None:
return cached
waiting = _inflight.get(token_fingerprint)
if waiting is None:
limited = _reserve_attempt_locked(rate_key, now)
if limited is not None:
return limited
owner_event = threading.Event()
_inflight[token_fingerprint] = owner_event
break
if not waiting.wait(_INFLIGHT_WAIT_SECONDS):
return TokenValidationResult(status = "unavailable")
result = _check_remote(normalized)
now = time.monotonic()
ttl = (
_CACHE_TTL_SECONDS
if result.status in ("valid", "invalid")
else max(_TEMPORARY_CACHE_TTL_SECONDS, float(result.retry_after_seconds or 0))
)
with _lock:
if len(_cache) >= _MAX_CACHE_ENTRIES:
_prune_locked(now)
if len(_cache) < _MAX_CACHE_ENTRIES:
_cache[token_fingerprint] = (now + ttl, result)
return result
finally:
if owner_event is not None:
with _lock:
event = _inflight.get(token_fingerprint)
if event is owner_event:
_inflight.pop(token_fingerprint, None)
event.set()
def reset_hf_token_validation_state() -> None:
"""Clear process state for test isolation."""
with _lock:
for event in _inflight.values():
event.set()
_inflight.clear()
_attempts.clear()
_cache.clear()

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

@ -479,6 +479,7 @@ def _run_update(
asset: Optional[str],
script: Path,
pin_release_tag: Optional[str] = None,
force_cpu: bool = False,
) -> None:
"""Worker: put the backend into a maintenance state, run the installer for
the latest prebuilt, then refresh caches so the next load uses the new build.
@ -522,6 +523,12 @@ def _run_update(
if pin_release_tag:
cmd.extend(["--published-release-tag", pin_release_tag])
cmd.extend(_rocm_install_args(asset))
# Re-assert a deliberate CPU install (--force-cpu) so detect_host on a GPU host
# does not re-route to a GPU/Vulkan bundle and revive the crash (#7213). --force-cpu
# (not --cpu-fallback) also re-persists force_cpu, keeping the choice across future
# updates. A natural fallback (or a legacy marker without the flag) heals to GPU (#6097).
if force_cpu:
cmd.append("--force-cpu")
logger.info("llama update: installing", cmd = " ".join(cmd))
# Stream progress lines into job["progress"].
env = dict(os.environ, UNSLOTH_PROGRESS_PERCENT_STEP = "5")
@ -671,6 +678,7 @@ def start_update() -> dict:
repo = marker.get("published_repo") or DEFAULT_PUBLISHED_REPO
from_tag = marker.get("tag") or marker.get("release_tag")
asset = marker.get("asset")
force_cpu = bool(marker.get("force_cpu"))
# Install exactly the release the banner offered: the installer's own
# "latest" is commit-date ordered and can lag the published_at pick
# above, reinstalling the current build in a loop (the #6219 class).
@ -705,6 +713,8 @@ def start_update() -> dict:
repo = (res or {}).get("repo") or DEFAULT_PUBLISHED_REPO
from_tag = None
asset = (res or {}).get("asset")
# Source builds carry no forced-CPU marker, so nothing to preserve here.
force_cpu = False
# No pin: source-build detection resolves via --resolve-prebuilt latest,
# the same resolver the unpinned apply uses, so the two already agree.
pin_release_tag = None
@ -735,7 +745,7 @@ def start_update() -> dict:
thread = threading.Thread(
target = _run_update,
args = (install_dir, repo, asset, script, pin_release_tag),
args = (install_dir, repo, asset, script, pin_release_tag, force_cpu),
name = "llama-cpp-update",
daemon = True,
)

190
studio/backend/utils/models/gguf_metadata.py
View file

@ -50,9 +50,15 @@ _CACHE_MAX_ENTRIES = 4096
# keyed by (file cache key, wanted key). None = key absent / file unreadable.
_BOOL_CACHE: Dict[Tuple[_CacheKey, str], Optional[bool]] = {}
# Native training context length (``{arch}.context_length``). None = absent /
# unreadable. Lets the UI show the real context ceiling before a model loads.
_CONTEXT_CACHE: Dict[_CacheKey, Optional[int]] = {}
_STRING_CACHE: Dict[Tuple[_CacheKey, str], Optional[str]] = {}
# GGUF header dims for the staged/deferred-load UI: context_length, layer_count
# (block_count), and moe_layer_count (block_count minus leading dense layers; 0
# if not MoE). One cached pass fills all three so the staged sheet can size every
# slider before the model loads. None = unreadable / not a GGUF. The native
# training context length (``{arch}.context_length``) the UI shows before a model
# loads is read from here via read_gguf_context_length.
_DIMS_CACHE: Dict[_CacheKey, Optional[Dict[str, Optional[int]]]] = {}
def _cache_key(path: str) -> Optional[_CacheKey]:
@ -142,32 +148,45 @@ def _parse_gguf_header(path: str) -> Optional[Dict[str, str]]:
return out
def read_gguf_context_length(path: str) -> Optional[int]:
"""Return the GGUF's native training context length (``{arch}.context_length``),
or ``None`` if missing/unreadable/not a GGUF. Cached by (path, mtime, size).
Lets the UI populate the context slider before the model is loaded."""
def read_gguf_staged_dims(path: str) -> Optional[Dict[str, Optional[int]]]:
"""GGUF header dims for the staged-load UI in one cached pass:
``{"context_length", "layer_count", "moe_layer_count"}``. Each may be None
when absent (moe_layer_count is 0 for a dense model). Returns ``None`` if not
a GGUF / unreadable. Cached by (path, mtime, size). Lets the staged sheet size
the context, GPU-layers and MoE sliders before the model loads."""
key = _cache_key(path)
if key is None:
return None
with _CACHE_LOCK:
if key in _CONTEXT_CACHE:
return _CONTEXT_CACHE[key]
result = _parse_gguf_context_length(path)
if key in _DIMS_CACHE:
return _DIMS_CACHE[key]
result = _parse_gguf_staged_dims(path)
with _CACHE_LOCK:
while len(_CONTEXT_CACHE) >= _CACHE_MAX_ENTRIES:
while len(_DIMS_CACHE) >= _CACHE_MAX_ENTRIES:
try:
_CONTEXT_CACHE.pop(next(iter(_CONTEXT_CACHE)))
_DIMS_CACHE.pop(next(iter(_DIMS_CACHE)))
except StopIteration:
break
_CONTEXT_CACHE[key] = result
_DIMS_CACHE[key] = result
return result
def _parse_gguf_context_length(path: str) -> Optional[int]:
# The context key is architecture-namespaced (``llama.context_length`` etc.),
# so we learn the key only after reading ``general.architecture``. GGUF writes
# general.* before arch.* keys, matching the loader's own parser.
ctx_key: Optional[str] = None
def read_gguf_context_length(path: str) -> Optional[int]:
"""Native training context length (``{arch}.context_length``), or ``None``.
Thin accessor over read_gguf_staged_dims."""
dims = read_gguf_staged_dims(path)
return dims["context_length"] if dims else None
def _parse_gguf_arch_uints(path: str, wanted_suffixes: frozenset[str]) -> Optional[Dict[str, int]]:
"""Walk a GGUF header once and return the requested architecture-namespaced
uint (vtype 4/10) keys, e.g. ``{"block_count": 32}``. Keys are
``{arch}.<suffix>``; the arch is learned from ``general.architecture`` (GGUF
writes general.* before arch.* keys, matching the loader's own parser).
Returns ``None`` if not a GGUF / unreadable, else a dict (possibly empty or
partial when some keys are absent)."""
arch: Optional[str] = None
found: Dict[str, int] = {}
try:
with open(path, "rb") as f:
head = f.read(24)
@ -204,28 +223,68 @@ def _parse_gguf_context_length(path: str) -> Optional[int]:
sbytes = f.read(slen)
if len(sbytes) < slen:
break
ctx_key = f"{sbytes.decode('utf-8', 'replace')}.context_length"
elif ctx_key is not None and key == ctx_key and vtype in (4, 10):
arch = sbytes.decode("utf-8", "replace")
elif (
arch is not None
and vtype in (4, 10)
and key.startswith(f"{arch}.")
and key[len(arch) + 1 :] in wanted_suffixes
):
width = 4 if vtype == 4 else 8
n_bytes = f.read(width)
if len(n_bytes) < width:
break
value = struct.unpack("<I" if vtype == 4 else "<Q", n_bytes)[0]
# A real context length is positive; treat 0/garbage as
# absent so the UI never builds a slider with max < min.
return value if value > 0 else None
found[key[len(arch) + 1 :]] = struct.unpack(
"<I" if vtype == 4 else "<Q", n_bytes
)[0]
if len(found) == len(wanted_suffixes):
break
else:
if not _skip_gguf_value(f, vtype):
break
except (struct.error, UnicodeDecodeError):
break
except OSError as e:
logger.debug(f"read_gguf_context_length: cannot open {path}: {e}")
logger.debug(f"_parse_gguf_arch_uints: cannot open {path}: {e}")
return None
except Exception as e:
logger.debug(f"read_gguf_context_length: parse failure on {path}: {e}")
logger.debug(f"_parse_gguf_arch_uints: parse failure on {path}: {e}")
return None
return None
return found
def _parse_gguf_staged_dims(path: str) -> Optional[Dict[str, Optional[int]]]:
vals = _parse_gguf_arch_uints(
path,
frozenset(
{
"context_length",
"block_count",
"expert_count",
"leading_dense_block_count",
}
),
)
if vals is None:
return None
ctx = vals.get("context_length")
block = vals.get("block_count")
# A real context/layer count is positive; treat 0/garbage as absent so the
# UI never builds a slider with max < min.
context_length = ctx if ctx and ctx > 0 else None
layer_count = block if block and block > 0 else None
# MoE layer count = block_count - leading dense layers, only when experts
# exist; else 0 (dense -> slider hidden). Mirrors n_moe_layers in
# core/inference/llama_cpp.py.
if not vals.get("expert_count") or not block:
moe_layer_count: Optional[int] = 0
else:
moe_layer_count = max(0, block - (vals.get("leading_dense_block_count") or 0))
return {
"context_length": context_length,
"layer_count": layer_count,
"moe_layer_count": moe_layer_count,
}
# Strings (8) and arrays (9) are handled inline.
@ -353,6 +412,83 @@ def _read_gguf_bool(path: str, wanted_key: str) -> Optional[bool]:
return result
def _parse_gguf_string(path: str, wanted_key: str) -> Optional[str]:
try:
with open(path, "rb") as f:
head = f.read(24)
if len(head) < 24:
return None
magic, _version, _tcount, kv_count = struct.unpack("<IIQQ", head)
if magic != _GGUF_MAGIC:
return None
for _ in range(kv_count):
try:
klen_bytes = f.read(8)
if len(klen_bytes) < 8:
break
klen = struct.unpack("<Q", klen_bytes)[0]
if klen > 1 << 20:
break
kbytes = f.read(klen)
if len(kbytes) < klen:
break
key = kbytes.decode("utf-8", "replace")
vt_bytes = f.read(4)
if len(vt_bytes) < 4:
break
vtype = struct.unpack("<I", vt_bytes)[0]
if key == wanted_key and vtype == 8:
slen_bytes = f.read(8)
if len(slen_bytes) < 8:
break
slen = struct.unpack("<Q", slen_bytes)[0]
if slen > 1 << 22:
break
sbytes = f.read(slen)
if len(sbytes) < slen:
break
return sbytes.decode("utf-8", "replace")
if not _skip_gguf_value(f, vtype):
break
except (struct.error, UnicodeDecodeError):
break
except OSError as e:
logger.debug(f"_parse_gguf_string: cannot open {path}: {e}")
return None
except Exception as e:
logger.debug(f"_parse_gguf_string: parse failure on {path}: {e}")
return None
return None
def _read_gguf_string(path: str, wanted_key: str) -> Optional[str]:
fkey = _cache_key(path)
if fkey is None:
return None
ckey = (fkey, wanted_key)
with _CACHE_LOCK:
if ckey in _STRING_CACHE:
return _STRING_CACHE[ckey]
result = _parse_gguf_string(path, wanted_key)
with _CACHE_LOCK:
while len(_STRING_CACHE) >= _CACHE_MAX_ENTRIES:
try:
_STRING_CACHE.pop(next(iter(_STRING_CACHE)))
except StopIteration:
break
_STRING_CACHE[ckey] = result
return result
def read_gguf_chat_template(path: str) -> Optional[str]:
template = _read_gguf_string(path, "tokenizer.chat_template")
if isinstance(template, str) and template.strip():
return template
return None
def read_mmproj_audio_capability(path: str) -> Optional[bool]:
"""``clip.has_audio_encoder`` from an mmproj GGUF (e.g. Gemma 4's
gemma4ua): ``True``/``False`` if present, ``None`` if absent/unreadable.

79
studio/backend/utils/models/model_config.py
View file

@ -1249,6 +1249,77 @@ def _iter_gguf_files(directory: Path, recursive: bool = False):
yield f
_GGUF_SPLIT_FILE_RE = re.compile(
r"^(?P<prefix>.+)-(?P<index>\d{5})-of-(?P<total>\d{5})\.gguf$",
re.IGNORECASE,
)
def _colocated_first_split_shard(path: Path) -> tuple[Optional[Path], bool]:
"""Return shard 1 and whether every shard is beside *path*."""
match = _GGUF_SPLIT_FILE_RE.match(path.name)
if match is None:
return None, False
prefix = match.group("prefix").casefold()
total_text = match.group("total")
total = int(total_text)
if total < 1:
return None, False
first: Optional[Path] = None
indices: set[int] = set()
try:
siblings = path.parent.iterdir()
for sibling in siblings:
sibling_match = _GGUF_SPLIT_FILE_RE.match(sibling.name)
if (
sibling_match is None
or sibling_match.group("prefix").casefold() != prefix
or sibling_match.group("total") != total_text
):
continue
try:
if not sibling.is_file():
continue
except OSError:
continue
index = int(sibling_match.group("index"))
if not 1 <= index <= total:
continue
indices.add(index)
if index == 1:
first = sibling
except OSError:
return None, False
return first, first is not None and len(indices) == total
def _local_gguf_load_path(path: Path) -> Path:
"""Choose a loadable local path while preserving complete symlink sets."""
if _GGUF_SPLIT_FILE_RE.match(path.name) is None:
return path.absolute()
first, complete = _colocated_first_split_shard(path)
if complete and first is not None:
return first.absolute()
try:
is_symlink = path.is_symlink()
except OSError:
is_symlink = False
if is_symlink:
try:
target = path.resolve()
except OSError:
return (first or path).absolute()
target_first, _ = _colocated_first_split_shard(target)
return (target_first or target).absolute()
return (first or path).absolute()
def detect_mmproj_file(path: str, search_root: Optional[str] = None) -> Optional[str]:
"""Find the mmproj GGUF for a model.
@ -1434,7 +1505,7 @@ def detect_gguf_model(path: str) -> Optional[str]:
except OSError:
is_dir = False # stat() unavailable in the lock window
if not is_dir:
return str(p.absolute()) # absolute() keeps symlink names readable
return str(_local_gguf_load_path(p))
# Directory named "*.gguf": fall through to the dir scan below.
# Case 2: directory containing .gguf files (skip mmproj / MTP drafter)
@ -1452,7 +1523,7 @@ def detect_gguf_model(path: str) -> Optional[str]:
gguf_files.append(f)
gguf_files.sort(key = lambda f: f.stat().st_size, reverse = True)
if gguf_files:
return str(gguf_files[0].resolve())
return str(_local_gguf_load_path(gguf_files[0]))
return None
@ -1879,7 +1950,7 @@ def _find_local_gguf_by_variant(directory: str, variant: str) -> Optional[str]:
For sharded GGUFs (multiple files sharing a quant label), returns the
first shard (sorted by name), which is what ``llama-server -m`` expects.
Returns the resolved absolute path, or ``None`` if no match.
Returns the absolute path, or ``None`` if no match.
"""
p = _resolve_gguf_dir(Path(directory))
if p is None:
@ -1900,7 +1971,7 @@ def _find_local_gguf_by_variant(directory: str, variant: str) -> Optional[str]:
matches.append(f)
matches.sort()
if matches:
return str(matches[0].resolve())
return str(_local_gguf_load_path(matches[0]))
return None

61
studio/backend/utils/openai_auto_switch_settings.py
View file

@ -30,11 +30,13 @@ from typing import Any, Optional
OPENAI_AUTO_SWITCH_SETTING_KEY = "openai_api_auto_switch_model"
AUTO_UNLOAD_IDLE_SETTING_KEY = "openai_api_auto_unload_idle_seconds"
AUTO_UNLOAD_KEEP_KV_SETTING_KEY = "openai_api_auto_unload_keep_kv"
MODEL_OVERRIDES_SETTING_KEY = "openai_api_auto_switch_overrides"
MODEL_IDLE_TTL_ENV_VAR = "UNSLOTH_MODEL_IDLE_TTL"
DEFAULT_OPENAI_AUTO_SWITCH_ENABLED = False
DEFAULT_AUTO_UNLOAD_IDLE_SECONDS = 0
DEFAULT_AUTO_UNLOAD_KEEP_KV = True
MIN_AUTO_UNLOAD_IDLE_SECONDS = 60
_CACHE_TTL_S = 2.0
@ -158,29 +160,54 @@ def get_auto_unload_idle_seconds() -> int:
return env if env is not None else 0
def set_openai_auto_switch(enabled: Any, idle_seconds: Any) -> tuple[bool, int]:
"""Set both auto-switch flags in one transaction so a settings PUT can't leave
one key updated and the other stale. Both values are coerced before any write,
so an invalid value raises without persisting either."""
def get_auto_unload_keep_kv() -> bool:
"""Whether the idle unload persists slot KV to disk for restore on reload."""
parsed = _coerce_bool(_cached_setting(AUTO_UNLOAD_KEEP_KV_SETTING_KEY, None))
return parsed if parsed is not None else DEFAULT_AUTO_UNLOAD_KEEP_KV
def set_openai_auto_switch(
enabled: Any,
idle_seconds: Any,
keep_kv: Any = None,
) -> tuple[bool, int, bool]:
"""One-transaction write; ``None`` leaves a stored value untouched."""
parsed_enabled = _coerce_bool(enabled)
if parsed_enabled is None:
raise ValueError("OpenAI auto-switch must be true or false.")
parsed_idle = _coerce_int(idle_seconds)
if parsed_idle is None:
raise ValueError("Auto-unload idle seconds must be a non-negative integer.")
if 0 < parsed_idle < MIN_AUTO_UNLOAD_IDLE_SECONDS:
raise ValueError(
f"Auto-unload idle seconds must be 0 (off) or at least "
f"{MIN_AUTO_UNLOAD_IDLE_SECONDS}."
)
parsed_idle = None
if idle_seconds is not None:
parsed_idle = _coerce_int(idle_seconds)
if parsed_idle is None:
raise ValueError("Auto-unload idle seconds must be a non-negative integer.")
if 0 < parsed_idle < MIN_AUTO_UNLOAD_IDLE_SECONDS:
raise ValueError(
f"Auto-unload idle seconds must be 0 (off) or at least "
f"{MIN_AUTO_UNLOAD_IDLE_SECONDS}."
)
parsed_keep_kv = None
if keep_kv is not None:
parsed_keep_kv = _coerce_bool(keep_kv)
if parsed_keep_kv is None:
raise ValueError("Keep KV on idle unload must be true or false.")
from storage.studio_db import upsert_app_settings
upsert_app_settings(
{OPENAI_AUTO_SWITCH_SETTING_KEY: parsed_enabled, AUTO_UNLOAD_IDLE_SETTING_KEY: parsed_idle}
)
updates: dict[str, Any] = {OPENAI_AUTO_SWITCH_SETTING_KEY: parsed_enabled}
if parsed_idle is not None:
updates[AUTO_UNLOAD_IDLE_SETTING_KEY] = parsed_idle
if parsed_keep_kv is not None:
updates[AUTO_UNLOAD_KEEP_KV_SETTING_KEY] = parsed_keep_kv
upsert_app_settings(updates)
_invalidate(OPENAI_AUTO_SWITCH_SETTING_KEY)
_invalidate(AUTO_UNLOAD_IDLE_SETTING_KEY)
return parsed_enabled, parsed_idle
if parsed_idle is not None:
_invalidate(AUTO_UNLOAD_IDLE_SETTING_KEY)
if parsed_keep_kv is not None:
_invalidate(AUTO_UNLOAD_KEEP_KV_SETTING_KEY)
return (
parsed_enabled,
parsed_idle if parsed_idle is not None else get_stored_auto_unload_idle_seconds(),
parsed_keep_kv if parsed_keep_kv is not None else get_auto_unload_keep_kv(),
)
def get_model_overrides() -> dict[str, dict]:

5
studio/backend/utils/paths/storage_roots.py
View file

@ -61,6 +61,11 @@ def cache_root() -> Path:
return studio_root() / "cache"
def llama_slot_cache_root() -> Path:
"""Dir llama-server saves/restores slot KV state in across idle unloads."""
return cache_root() / "llama-slots"
def studio_bin_root() -> Path:
"""Dir for Unsloth-managed executables (the `unsloth` shim, downloaded tools like cloudflared)."""
return studio_root() / "bin"

20
studio/backend/utils/utils.py
View file

@ -123,6 +123,26 @@ def without_hf_auth():
os.environ.pop("HF_HUB_DISABLE_IMPLICIT_TOKEN", None)
def is_hf_authentication_error(error: Exception) -> bool:
"""Return whether an exception chain contains a definitive HF auth failure."""
seen: set[int] = set()
current: BaseException | None = error
while current is not None and id(current) not in seen:
seen.add(id(current))
response = getattr(current, "response", None)
status = getattr(response, "status_code", None)
try:
if status is not None and int(status) == 401:
return True
except (TypeError, ValueError):
pass
message = str(current).lower()
if "invalid user token" in message or "invalid hf token" in message:
return True
current = current.__cause__ or current.__context__
return False
def format_error_message(error: Exception, model_name: str) -> str:
"""
Format a user-friendly error message for common load issues.

61
studio/frontend/src/app/auth-guards.ts
View file

@ -23,19 +23,47 @@ interface AuthStatus {
requires_password_change: boolean;
}
const AUTH_STATUS_TTL_MS = 30_000;
let authStatusCheckedAt = 0;
let authStatusRequest: Promise<AuthStatus> | null = null;
function hasFreshAuthStatus(): boolean {
return (
authStatusCheckedAt !== 0 &&
Date.now() - authStatusCheckedAt < AUTH_STATUS_TTL_MS
);
}
async function fetchAuthStatus(): Promise<AuthStatus> {
try {
const res = await fetch(apiUrl("/api/auth/status"));
if (!res.ok) return { initialized: true, requires_password_change: mustChangePassword() };
const status = (await res.json()) as AuthStatus;
// Server truth wins; keep localStorage in sync both ways.
if (status.requires_password_change !== mustChangePassword()) {
setMustChangePassword(status.requires_password_change);
if (authStatusRequest) return authStatusRequest;
const request = (async () => {
try {
const res = await fetch(apiUrl("/api/auth/status"));
if (!res.ok) {
return {
initialized: true,
requires_password_change: mustChangePassword(),
};
}
const status = (await res.json()) as AuthStatus;
authStatusCheckedAt = Date.now();
// Server truth wins; keep localStorage in sync both ways.
if (status.requires_password_change !== mustChangePassword()) {
setMustChangePassword(status.requires_password_change);
}
return status;
} catch {
return {
initialized: true,
requires_password_change: mustChangePassword(),
};
}
return status;
} catch {
return { initialized: true, requires_password_change: mustChangePassword() };
}
})().finally(() => {
authStatusRequest = null;
});
authStatusRequest = request;
return request;
}
function authRedirect(to: "/login" | "/change-password"): never {
@ -49,12 +77,17 @@ export async function requireAuth(): Promise<void> {
}
if (await hasActiveSession()) {
const { requires_password_change } = await fetchAuthStatus();
if (requires_password_change || mustChangePassword()) {
authRedirect("/change-password");
// Reconcile periodically so local-only routes cannot outlive a server-side
// password-change requirement, while nearby route switches stay local.
if (mustChangePassword() || !hasFreshAuthStatus()) {
const { requires_password_change } = await fetchAuthStatus();
if (requires_password_change || mustChangePassword()) {
authRedirect("/change-password");
}
}
return;
}
const status = await fetchAuthStatus();
if (status.requires_password_change || mustChangePassword()) {
authRedirect("/change-password");

22
studio/frontend/src/app/provider.tsx
View file

@ -51,9 +51,12 @@ async function showSetupWindow(isCurrent: WindowLayoutGuard): Promise<void> {
const { getCurrentWindow, LogicalSize } = await import(
"@tauri-apps/api/window"
);
const { invoke } = await import("@tauri-apps/api/core");
if (!isCurrent()) return;
const win = getCurrentWindow();
await invoke("reset_app_window_layout_initialized");
if (!isCurrent()) return;
await win.setResizable(false);
if (!isCurrent()) return;
await win.setSize(new LogicalSize(SETUP_WINDOW_WIDTH, SETUP_WINDOW_HEIGHT));
@ -98,18 +101,20 @@ async function applyAppWindowLayout(
if (!isCurrent()) return;
const win = getCurrentWindow();
// Decide first-launch vs restore from the on-disk state file BEFORE touching the
// window. Probing the window after restoreStateCurrent is unreliable: on GTK,
// set_size on a hidden window is deferred until show(), so innerSize() reads a
// stale value and a baseline fallback would overwrite the queued restore. On
// macOS the same probe works, hence the inconsistency between prior iterations.
const hasSavedState = await invoke<boolean>("has_saved_window_state");
// Setup-window activity may create plugin state before the full app is ever
// shown, so use a dedicated full-app marker to decide whether restoration is
// appropriate. Keep checking plugin state so a missing/corrupt state file
// falls back to a monitor-safe centered layout.
const [hasInitializedAppLayout, hasSavedState] = await Promise.all([
invoke<boolean>("has_initialized_app_window_layout"),
invoke<boolean>("has_saved_window_state"),
]);
if (!isCurrent()) return;
await win.setResizable(true);
if (!isCurrent()) return;
if (hasSavedState) {
if (hasInitializedAppLayout && hasSavedState) {
// Subsequent launch: plugin restores size/position/maximized, with built-in
// off-screen protection for positions saved on a now-disconnected display.
await restoreStateCurrent(
@ -144,6 +149,9 @@ async function applyAppWindowLayout(
});
if (!isCurrent()) return;
await enforceMinimumWindowSize(win, LogicalSize, isCurrent);
if (!isCurrent()) return;
await invoke("mark_app_window_layout_initialized");
}
async function showWindowFallback(): Promise<void> {

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

@ -16,6 +16,7 @@ import {
type ChatSearch,
} from "@/features/chat";
import { RemoteCodeConsentDialog } from "@/features/security";
import { HfTokenWarningDialog } from "@/features/hf-auth";
import { TransformersUpgradeDialog } from "@/features/transformers-upgrade";
import { useTrainingUnloadGuard } from "@/features/training";
import { useExportRuntimeLifecycle } from "@/features/export";
@ -195,9 +196,6 @@ function RootLayout() {
chatRuntime.setActiveThreadId(null);
chatRuntime.setActiveProjectId(null);
chatRuntime.setIncognito(false);
// Detach the staging UI but keep any in-flight download running, like Hub.
if (chatRuntime.pendingSelection)
chatRuntime.abandonStagedModel({ keepDownload: true });
void navigate({
to: "/chat",
search: { new: crypto.randomUUID() },
@ -220,16 +218,13 @@ function RootLayout() {
chatRuntime.setActiveProjectId(null);
chatRuntime.setActiveThreadId(null);
chatRuntime.setIncognito(false);
// Leaving chat must not kill an in-flight download: detach the staging UI
// but keep the transfer running in the manager, like a Hub download.
if (chatRuntime.pendingSelection)
chatRuntime.abandonStagedModel({ keepDownload: true });
}, [isChatRoute]);
return (
<AppProvider>
<PersonalizationSyncMount />
{!isAuthFlowRoute && <SettingsDialog />}
<HfTokenWarningDialog />
<RemoteCodeConsentDialog />
<TransformersUpgradeDialog />
{hideNavbar ? (
@ -279,7 +274,7 @@ function RootLayout() {
initial={{ opacity: 0 }}
animate={{ opacity: 1 }}
exit={{ opacity: 0 }}
transition={{ duration: 0.15 }}
transition={{ duration: 0.06 }}
className="flex min-h-0 min-w-0 flex-1 basis-0 flex-col overflow-visible"
>
<Suspense fallback={<RouteFallback />}>

10
studio/frontend/src/app/routes/data-recipes.tsx
View file

@ -1,15 +1,13 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import { createRoute } from "@tanstack/react-router";
import { lazy } from "react";
import { createRoute, lazyRouteComponent } from "@tanstack/react-router";
import { requireAuth } from "../auth-guards";
import { Route as rootRoute } from "./__root";
const DataRecipesPage = lazy(() =>
import("@/features/data-recipes").then((m) => ({
default: m.DataRecipesPage,
})),
const DataRecipesPage = lazyRouteComponent(
() => import("@/features/data-recipes"),
"DataRecipesPage",
);
export const Route = createRoute({

10
studio/frontend/src/app/routes/export.tsx
View file

@ -1,15 +1,13 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import { createRoute } from "@tanstack/react-router";
import { lazy } from "react";
import { createRoute, lazyRouteComponent } from "@tanstack/react-router";
import { requireAuth } from "../auth-guards";
import { Route as rootRoute } from "./__root";
const ExportPage = lazy(() =>
import("@/features/export/export-page").then((m) => ({
default: m.ExportPage,
})),
const ExportPage = lazyRouteComponent(
() => import("@/features/export/export-page"),
"ExportPage",
);
export type ExportSearch = {

16
studio/frontend/src/app/routes/hub.tsx
View file

@ -1,15 +1,13 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import { createRoute } from "@tanstack/react-router";
import { lazy } from "react";
import { createRoute, lazyRouteComponent } from "@tanstack/react-router";
import { requireAuth } from "../auth-guards";
import { Route as rootRoute } from "./__root";
const ModelsPage = lazy(() =>
import("@/features/hub/hub-page").then((m) => ({
default: m.ModelsPage,
})),
const ModelsPage = lazyRouteComponent(
() => import("@/features/hub/hub-page"),
"ModelsPage",
);
export interface ModelsSearch {
@ -31,7 +29,11 @@ export const Route = createRoute({
const model = search.model;
if (typeof model === "string" && model.length > 0) next.model = model;
const section = search.section;
if (section === "trending" || section === "latest" || section === "finetune") {
if (
section === "trending" ||
section === "latest" ||
section === "finetune"
) {
next.section = section;
}
const kind = search.kind;

10
studio/frontend/src/app/routes/projects.tsx
View file

@ -1,15 +1,13 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import { createRoute } from "@tanstack/react-router";
import { lazy } from "react";
import { createRoute, lazyRouteComponent } from "@tanstack/react-router";
import { requireAuth } from "../auth-guards";
import { Route as rootRoute } from "./__root";
const ProjectsPage = lazy(() =>
import("@/features/chat/projects-page").then((m) => ({
default: m.ProjectsPage,
})),
const ProjectsPage = lazyRouteComponent(
() => import("@/features/chat/projects-page"),
"ProjectsPage",
);
export const Route = createRoute({

10
studio/frontend/src/app/routes/studio.tsx
View file

@ -1,15 +1,13 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import { createRoute } from "@tanstack/react-router";
import { lazy } from "react";
import { createRoute, lazyRouteComponent } from "@tanstack/react-router";
import { requireAuth } from "../auth-guards";
import { Route as rootRoute } from "./__root";
const StudioPage = lazy(() =>
import("@/features/studio/studio-page").then((m) => ({
default: m.StudioPage,
})),
const StudioPage = lazyRouteComponent(
() => import("@/features/studio/studio-page"),
"StudioPage",
);
export const Route = createRoute({

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

@ -55,6 +55,7 @@ import {
Archive03Icon,
ArrowRight02Icon,
BadgeInfoIcon,
BubbleChatIcon,
ChefHatIcon,
CloudIcon,
CpuIcon,
@ -93,7 +94,12 @@ import {
import { Tooltip as TooltipPrimitive } from "radix-ui";
import { HugeiconsIcon } from "@hugeicons/react";
import { ChevronDown, Moon } from "lucide-react";
import { Link, useNavigate, useRouterState } from "@tanstack/react-router";
import {
Link,
useNavigate,
useRouter,
useRouterState,
} from "@tanstack/react-router";
import {
archiveChatItem,
ChatSearchDialog,
@ -103,6 +109,7 @@ import {
deleteChatItem,
listStoredChatThreads,
moveChatItemToProject,
notifyChatHistoryUpdated,
renameChatItem,
renameChatProject,
useChatRuntimeStore,
@ -110,6 +117,7 @@ import {
useChatSearchStore,
useChatSidebarItems,
usePinnedChatsStore,
usePinnedProjectsStore,
useChatPreferencesStore,
type ProjectRecord,
type SidebarItem,
@ -135,7 +143,15 @@ import {
} from "@/features/training";
import type { TrainingRunSummary } from "@/features/training";
import { useExportRuntimeStore } from "@/features/export";
import { useEffect, useMemo, useRef, useState, type ReactNode } from "react";
import {
Fragment,
useEffect,
useMemo,
useRef,
useState,
type ReactNode,
} from "react";
import { isDownloadCancelled } from "@/lib/native-files";
import { toast } from "@/lib/toast";
import { ShutdownDialog } from "@/components/shutdown-dialog";
import { translate, useT, type TranslationKey } from "@/i18n";
@ -193,6 +209,9 @@ const TestTubeOutlineIcon = TestTube01Icon.slice(
type ConversationExportFormat = "raw-jsonl" | "csv" | "sharegpt-jsonl";
// A pinned project shows this many recent chats before "Show more".
const PINNED_PROJECT_CHAT_LIMIT = 4;
const CHAT_EXPORT_OPTIONS: Array<{
label: string;
format: ConversationExportFormat;
@ -256,6 +275,10 @@ function createNavigationNonce(): string {
return `${Date.now()}-${Math.random().toString(36).slice(2, 10)}`;
}
function preloadSilently(request: Promise<unknown>): void {
void request.catch(() => undefined);
}
function NavItem({
icon,
label,
@ -267,6 +290,7 @@ function NavItem({
className,
spinner,
tooltip,
onIntent,
}: {
icon: typeof ZapIcon;
label: string;
@ -277,6 +301,7 @@ function NavItem({
dataTour?: string;
className?: string;
spinner?: boolean;
onIntent?: () => void;
// Overrides the hover tooltip (defaults to `label`). Used to explain why a
// disabled item (e.g. Train/Export on a chat-only host) is greyed out.
tooltip?: string;
@ -288,6 +313,8 @@ function NavItem({
tooltip={tooltip ?? label}
disabled={disabled}
onClick={onClick}
onPointerEnter={disabled ? undefined : onIntent}
onFocus={disabled ? undefined : onIntent}
isActive={active}
data-tour={dataTour}
className="sidebar-nav-btn h-[33px] rounded-full gap-[8.5px] pl-3 pr-2.5 font-medium group-data-[collapsible=icon]:px-2.5 group-data-[collapsible=icon]:!w-[32px] group-data-[collapsible=icon]:mx-auto"
@ -324,6 +351,7 @@ export function AppSidebar() {
});
const { togglePinned, isMobile, setOpenMobile } = useSidebar();
const navigate = useNavigate();
const router = useRouter();
// Web update detection: `webUpdate` is non-null only when the installed
// (PyPI) version is behind the latest release, so the card is hidden by
@ -430,14 +458,63 @@ export function AppSidebar() {
),
[allChatItems, pinnedIdSet],
);
// Pinned chats, in pin order (most recent first).
const [pinnedOpen, setPinnedOpen] = useState(true);
// "Projects" section: projects the user pinned, in pin order (most recent
// first). The section only appears once at least one project is pinned.
const pinnedProjectIds = usePinnedProjectsStore((s) => s.pinnedIds);
const unpinProject = usePinnedProjectsStore((s) => s.unpin);
const pinnedProjectRecords = useMemo(() => {
const byId = new Map(projects.map((p) => [p.id, p]));
return pinnedProjectIds
.map((id) => byId.get(id))
.filter((p): p is ProjectRecord => Boolean(p));
}, [projects, pinnedProjectIds]);
// Pinned chats, in pin order (most recent first). Includes chats that live
// inside a project: pinning promotes a chat into this list, and it is removed
// from the project's nested list below so it never shows twice.
const pinnedChatItems = useMemo(() => {
const byId = new Map(allChatItems.map((item) => [item.id, item]));
return pinnedIds
.map((id) => byId.get(id))
.filter((item): item is SidebarItem => Boolean(item));
}, [allChatItems, pinnedIds]);
const [pinnedOpen, setPinnedOpen] = useState(true);
// A pinned project reveals its recent chats (most recent first) nested below.
// Pinned chats are excluded here since they render in the pinned-chats list.
const chatsByProjectId = useMemo(() => {
const map = new Map<string, SidebarItem[]>();
for (const item of allChatItems) {
if (!item.projectId) continue;
if (pinnedIdSet.has(item.id)) continue;
const list = map.get(item.projectId);
if (list) list.push(item);
else map.set(item.projectId, [item]);
}
for (const list of map.values())
list.sort((a, b) => b.updatedAt - a.updatedAt);
return map;
}, [allChatItems, pinnedIdSet]);
// Default expanded (not collapsed); the row toggles this. Show-more reveals
// chats past the first PINNED_PROJECT_CHAT_LIMIT.
const [collapsedProjectIds, setCollapsedProjectIds] = useState<Set<string>>(
() => new Set(),
);
const [expandedChatProjectIds, setExpandedChatProjectIds] = useState<
Set<string>
>(() => new Set());
const toggleProjectCollapsed = (id: string) =>
setCollapsedProjectIds((prev) => {
const next = new Set(prev);
if (next.has(id)) next.delete(id);
else next.add(id);
return next;
});
const toggleProjectShowAll = (id: string) =>
setExpandedChatProjectIds((prev) => {
const next = new Set(prev);
if (next.has(id)) next.delete(id);
else next.add(id);
return next;
});
const storeThreadId = useChatRuntimeStore((s) => s.activeThreadId);
const setActiveThreadId = useChatRuntimeStore((s) => s.setActiveThreadId);
const anyChatRunning = useChatRuntimeStore((s) =>
@ -717,6 +794,8 @@ export function AppSidebar() {
const shouldDeleteProjectFiles =
target.kind === "project" && deleteProjectFiles;
setConfirmingDelete(null);
// Reset so the next project delete never inherits this checkbox.
setDeleteProjectFiles(false);
if (target.kind === "chat") {
await deleteChatWithCleanup(target.item);
return;
@ -726,7 +805,20 @@ export function AppSidebar() {
await deleteChatProject(target.project.id, {
deleteFiles: shouldDeleteProjectFiles,
});
if (activeProjectId === target.project.id) {
// Refresh chat history so the project's reparented chats don't linger
// as stale top-level rows.
notifyChatHistoryUpdated();
// activeProjectId is only the ?project= param; on a thread-only URL the
// project is resolved from the thread into the runtime store, so check
// that too or we strand the user on a now-deleted thread. Only redirect
// from a chat route: the runtime store value can be stale elsewhere.
const runtimeProjectId =
useChatRuntimeStore.getState().activeProjectId;
if (
isChatRoute &&
(activeProjectId === target.project.id ||
runtimeProjectId === target.project.id)
) {
useChatRuntimeStore.getState().setActiveProjectId(null);
navigate({ to: "/chat", search: { new: createNavigationNonce() } });
}
@ -813,8 +905,11 @@ export function AppSidebar() {
// pl-3 (12px) over the content's pl-1.5 (6px) = 18px, aligning the
// title with the nav items above.
variant === "project" ? "pl-[39px]" : "pl-3",
// Pinned chats carry a chat icon, so add the nav-item icon gap.
isPinned && variant !== "project" && "gap-[8.5px]",
variant === "project"
? "group-hover/project-chat-item:pr-6 group-has-[.sidebar-row-action[data-state=open]]/project-chat-item:pr-6"
? // Room for the hover pin quick-action plus the kebab.
"group-hover/project-chat-item:pr-14 group-has-[.sidebar-row-action[data-state=open]]/project-chat-item:pr-8"
: isPinned
? // Pinned rows show an extra unpin button on hover, so reserve more room
// (pr-8 when the menu is open keeps the unpin button clear of the title).
@ -874,10 +969,43 @@ export function AppSidebar() {
closeMobileIfOpen();
}}
>
{isPinned && variant !== "project" && (
<HugeiconsIcon icon={BubbleChatIcon} strokeWidth={1.75} className="size-icon! shrink-0" />
)}
<span className="truncate">
{pendingRename?.id === item.id ? pendingRename.title : item.title}
</span>
</SidebarMenuButton>
{variant === "project" && (
<button
type="button"
onClick={(e) => {
e.stopPropagation();
togglePinnedChat(item.id);
}}
aria-label={isPinned ? "Unpin chat" : "Pin chat"}
className="sidebar-row-action is-unpin-action group-hover/project-chat-item:opacity-100 group-hover/project-chat-item:pointer-events-auto focus-visible:opacity-100 focus-visible:pointer-events-auto"
>
<span className="sidebar-row-action-glyph">
<HugeiconsIcon icon={isPinned ? PinOffIcon : PinIcon} strokeWidth={1.75} className="size-icon" />
</span>
</button>
)}
{variant === "recent" && isPinned && (
<button
type="button"
onClick={(e) => {
e.stopPropagation();
togglePinnedChat(item.id);
}}
aria-label="Unpin chat"
className="sidebar-row-action is-unpin-action group-hover/recent-item:opacity-100 group-hover/recent-item:pointer-events-auto focus-visible:opacity-100 focus-visible:pointer-events-auto"
>
<span className="sidebar-row-action-glyph">
<HugeiconsIcon icon={PinOffIcon} strokeWidth={1.75} className="size-icon" />
</span>
</button>
)}
<DropdownMenu>
<DropdownMenuTrigger asChild>
<button
@ -957,11 +1085,13 @@ export function AppSidebar() {
const ids = item.type === "single"
? [item.id]
: (await listStoredChatThreads({ pairId: item.id })).map((t) => t.id);
await Promise.all(
ids.map((id) => exportConversationByFormat(id, format)),
);
} catch {
toast.error("Export failed.");
for (const id of ids) {
await exportConversationByFormat(id, format);
}
} catch (error) {
if (!isDownloadCancelled(error)) {
toast.error("Export failed.");
}
}
}}
>
@ -969,10 +1099,10 @@ export function AppSidebar() {
</DropdownMenuItem>
))}
<DropdownMenuSeparator />
{/* Bulk export and import live in Settings -> Chat -> Data. */}
{/* Bulk export and import live in Settings -> Data. */}
<DropdownMenuItem
onSelect={() =>
useSettingsDialogStore.getState().openDialog("chat")
useSettingsDialogStore.getState().openDialog("data")
}
>
Export all chats…
@ -997,28 +1127,6 @@ export function AppSidebar() {
</DropdownMenuItem>
</DropdownMenuContent>
</DropdownMenu>
{isPinned ? (
<Tooltip>
<TooltipPrimitive.Trigger asChild>
<button
type="button"
onClick={(e) => {
e.stopPropagation();
unpinChat(item.id);
}}
aria-label="Unpin chat"
className={cn(actionClass, "is-unpin-action")}
>
<span className="sidebar-row-action-glyph">
<HugeiconsIcon icon={PinOffIcon} strokeWidth={1.75} className="size-4" />
</span>
</button>
</TooltipPrimitive.Trigger>
<TooltipContent side="bottom" sideOffset={6} className="tooltip-compact">
Unpin
</TooltipContent>
</Tooltip>
) : null}
</SidebarMenuItem>
);
}
@ -1218,6 +1326,9 @@ export function AppSidebar() {
navigate({ to: "/projects" });
closeMobileIfOpen();
}}
onIntent={() => {
preloadSilently(router.preloadRoute({ to: "/projects" }));
}}
className="group/projects-item relative"
>
<button
@ -1248,6 +1359,9 @@ export function AppSidebar() {
navigate({ to: "/hub" });
closeMobileIfOpen();
}}
onIntent={() => {
preloadSilently(router.preloadRoute({ to: "/hub" }));
}}
/>
{/* Train has a labelled section when expanded; plain icon here only when collapsed. */}
<NavItem
@ -1264,6 +1378,9 @@ export function AppSidebar() {
navigate({ to: "/studio" });
closeMobileIfOpen();
}}
onIntent={() => {
preloadSilently(router.preloadRoute({ to: "/studio" }));
}}
className="hidden group-data-[collapsible=icon]:block"
/>
</SidebarMenu>
@ -1293,6 +1410,9 @@ export function AppSidebar() {
navigate({ to: "/studio" });
closeMobileIfOpen();
}}
onIntent={() => {
preloadSilently(router.preloadRoute({ to: "/studio" }));
}}
/>
<NavItem
icon={ChefHatIcon}
@ -1302,6 +1422,16 @@ export function AppSidebar() {
navigate({ to: "/data-recipes" });
closeMobileIfOpen();
}}
onIntent={() => {
preloadSilently(
router.preloadRoute({ to: "/data-recipes" }),
);
preloadSilently(
import("@/features/data-recipes").then((module) =>
module.preloadRecipes(),
),
);
}}
/>
<NavItem
icon={DownloadSquare01Icon}
@ -1312,6 +1442,14 @@ export function AppSidebar() {
navigate({ to: "/export" });
closeMobileIfOpen();
}}
onIntent={() => {
preloadSilently(router.preloadRoute({ to: "/export" }));
preloadSilently(
import(
"@/features/export/export-navigation-cache"
).then((module) => module.preloadExportData()),
);
}}
/>
</SidebarMenu>
</SidebarGroupContent>
@ -1319,28 +1457,156 @@ export function AppSidebar() {
</SidebarGroup>
</Collapsible>
{/* Pinned chats: own section above Recents */}
{!isStudioRoute && !showTrainingRecents && pinnedChatItems.length > 0 && (
<Collapsible open={pinnedOpen} onOpenChange={setPinnedOpen} asChild>
<SidebarGroup className="group-data-[collapsible=icon]:hidden px-0 py-0">
<SidebarGroupLabel className={cn("sidebar-sticky-label sidebar-sticky-label-following", scrolled && "is-scrolled")} asChild>
<CollapsibleTrigger className="cursor-pointer flex w-full items-center gap-1 group/sb-collap">
Pinned
<ChevronDown className="size-3.5 opacity-0 transition-[transform,opacity] duration-200 group-hover/sb-collap:opacity-100 group-focus-visible/sb-collap:opacity-100 data-[state=open]:rotate-0 [[data-state=closed]_&]:rotate-[-90deg] [[data-state=closed]_&]:opacity-100" />
</CollapsibleTrigger>
</SidebarGroupLabel>
<CollapsibleContent>
<SidebarGroupContent className="pl-1.5 pr-2">
<SidebarMenu>
{pinnedChatItems.map((item) =>
renderChatSidebarItem(item, "recent"),
)}
</SidebarMenu>
</SidebarGroupContent>
</CollapsibleContent>
</SidebarGroup>
</Collapsible>
)}
{/* Pinned: pinned projects (with their chats) and pinned chats */}
{!isStudioRoute &&
!showTrainingRecents &&
(pinnedProjectRecords.length > 0 ||
pinnedChatItems.length > 0) && (
<Collapsible open={pinnedOpen} onOpenChange={setPinnedOpen} asChild>
<SidebarGroup className="group-data-[collapsible=icon]:hidden px-0 py-0">
<SidebarGroupLabel className={cn("sidebar-sticky-label sidebar-sticky-label-following", scrolled && "is-scrolled")} asChild>
<CollapsibleTrigger className="cursor-pointer flex w-full items-center gap-1 group/sb-collap">
Pinned
<ChevronDown className="size-3.5 opacity-0 transition-[transform,opacity] duration-200 group-hover/sb-collap:opacity-100 group-focus-visible/sb-collap:opacity-100 data-[state=open]:rotate-0 [[data-state=closed]_&]:rotate-[-90deg] [[data-state=closed]_&]:opacity-100" />
</CollapsibleTrigger>
</SidebarGroupLabel>
<CollapsibleContent>
<SidebarGroupContent className="pl-1.5 pr-2">
<SidebarMenu>
{pinnedProjectRecords.map((project) => {
const projectChats =
chatsByProjectId.get(project.id) ?? [];
const expanded = !collapsedProjectIds.has(project.id);
const showAll = expandedChatProjectIds.has(project.id);
const visibleChats =
expanded && !showAll
? projectChats.slice(0, PINNED_PROJECT_CHAT_LIMIT)
: projectChats;
return (
<Fragment key={project.id}>
<SidebarMenuItem
className="group/recent-item relative"
>
<SidebarMenuButton
// Highlight the folder only on the project home; when
// a chat inside it is open, only that chat row is active.
isActive={activeProjectId === project.id && !activeThreadId}
onClick={() => toggleProjectCollapsed(project.id)}
className="sidebar-nav-btn h-[33px] rounded-full gap-[8.5px] pl-3 pr-2.5 font-medium group-hover/recent-item:pr-16 group-has-[.sidebar-row-action[data-state=open]]/recent-item:pr-8"
>
<HugeiconsIcon icon={Folder01Icon} strokeWidth={1.75} className="size-icon! shrink-0" />
<span className="truncate text-[14.5px] leading-[19px] tracking-nav">{project.name}</span>
</SidebarMenuButton>
{/* New chat in this project */}
<button
type="button"
aria-label="New chat"
onClick={(e) => {
e.stopPropagation();
openNewChat(project.id);
}}
className="sidebar-row-action is-unpin-action group-hover/recent-item:opacity-100 group-hover/recent-item:pointer-events-auto focus-visible:opacity-100 focus-visible:pointer-events-auto"
>
<span className="sidebar-row-action-glyph">
<HugeiconsIcon icon={PencilEdit02Icon} strokeWidth={1.75} className="size-icon" />
</span>
</button>
{/* Project options */}
<DropdownMenu>
<DropdownMenuTrigger asChild>
<button
type="button"
onClick={(e) => e.stopPropagation()}
aria-label="Project options"
className="sidebar-row-action group-hover/recent-item:opacity-100 group-hover/recent-item:pointer-events-auto focus-visible:opacity-100 focus-visible:pointer-events-auto"
>
<span className="sidebar-row-action-glyph">
<HugeiconsIcon icon={MoreVerticalIcon} strokeWidth={1.75} className="size-icon" />
</span>
</button>
</DropdownMenuTrigger>
<DropdownMenuContent
side="bottom"
align="start"
sideOffset={0}
className="unsloth-plus-menu menu-flat-destructive w-56"
>
<DropdownMenuItem onSelect={() => openProject(project.id)}>
<HugeiconsIcon icon={Folder01Icon} strokeWidth={1.75} className="size-icon" />
<span>Project home</span>
</DropdownMenuItem>
<DropdownMenuItem onSelect={() => openNewChat(project.id)}>
<HugeiconsIcon icon={PencilEdit02Icon} strokeWidth={1.75} className="size-icon" />
<span>New chat</span>
</DropdownMenuItem>
<DropdownMenuItem
onSelect={() => {
// Seed the shared draft so the dialog opens
// with the current name, not stale text.
setRenameDraft(project.name);
setRenamingTarget({
kind: "project",
project,
current: project.name,
});
}}
>
<HugeiconsIcon icon={Edit03Icon} strokeWidth={1.75} className="size-icon" />
<span>Rename project</span>
</DropdownMenuItem>
<DropdownMenuItem onSelect={() => unpinProject(project.id)}>
<HugeiconsIcon icon={PinOffIcon} strokeWidth={1.75} className="size-icon" />
<span>Unpin project</span>
</DropdownMenuItem>
<DropdownMenuSeparator />
<DropdownMenuItem
variant="destructive"
onSelect={() => {
// Start each delete with the file toggle off:
// Cancel closes programmatically and skips the
// dialog onOpenChange reset.
setDeleteProjectFiles(false);
setConfirmingDelete({ kind: "project", project });
}}
>
<HugeiconsIcon icon={Delete02Icon} strokeWidth={1.75} className="size-icon" />
<span>Delete project</span>
</DropdownMenuItem>
</DropdownMenuContent>
</DropdownMenu>
</SidebarMenuItem>
{expanded &&
visibleChats.map((chat) =>
renderChatSidebarItem(chat, "project"),
)}
{expanded &&
projectChats.length > PINNED_PROJECT_CHAT_LIMIT && (
<SidebarMenuItem>
<SidebarMenuButton
onClick={() => toggleProjectShowAll(project.id)}
// Force the muted token: .sidebar-nav-btn's own
// color rule outweighs a plain text utility, so
// Show more would otherwise match the chat rows.
className="sidebar-nav-btn h-[30px] rounded-full pl-9 pr-4 font-medium text-nav-fg-muted!"
>
<span className="text-[13px] leading-[18px] tracking-nav">
{showAll ? "Show less" : "Show more"}
</span>
</SidebarMenuButton>
</SidebarMenuItem>
)}
</Fragment>
);
})}
{pinnedChatItems.map((item) =>
renderChatSidebarItem(item, "recent"),
)}
</SidebarMenu>
</SidebarGroupContent>
</CollapsibleContent>
</SidebarGroup>
</Collapsible>
)}
{!isStudioRoute && !showTrainingRecents && (
<Collapsible open={chatOpen} onOpenChange={setChatOpen} asChild>

27
studio/frontend/src/components/assistant-ui/attachment.tsx
View file

@ -7,6 +7,7 @@
import { TooltipIconButton } from "@/components/assistant-ui/tooltip-icon-button";
import {
Dialog,
DialogClose,
DialogContent,
DialogTitle,
DialogTrigger,
@ -27,12 +28,7 @@ import {
import { AudioWave01Icon, File02Icon } from "@hugeicons/core-free-icons";
import { HugeiconsIcon } from "@hugeicons/react";
import { PlusIcon, XIcon } from "lucide-react";
import {
type FC,
type PropsWithChildren,
useEffect,
useState,
} from "react";
import { type FC, type PropsWithChildren, useEffect, useState } from "react";
import { useShallow } from "zustand/shallow";
const useFileSrc = (file: File | undefined): string | undefined => {
@ -83,7 +79,7 @@ const AttachmentPreview: FC<AttachmentPreviewProps> = ({ src }) => {
src={src}
alt="Preview"
className={cn(
"block h-auto max-h-[80vh] w-auto max-w-full object-contain",
"block h-auto max-h-[90dvh] w-auto max-w-[92vw] object-contain",
isLoaded
? "aui-attachment-preview-image-loaded"
: "aui-attachment-preview-image-loading invisible",
@ -108,12 +104,23 @@ const AttachmentPreviewDialog: FC<PropsWithChildren> = ({ children }) => {
>
{children}
</DialogTrigger>
<DialogContent className="aui-attachment-preview-dialog-content p-2 sm:max-w-3xl [&>button]:rounded-full [&>button]:bg-foreground/60 [&>button]:p-1 [&>button]:opacity-100 [&>button]:ring-0! [&_svg]:text-background [&>button]:hover:[&_svg]:text-destructive">
{/* Chrome-free lightbox: the image floats on the dimmed backdrop with
no dialog panel, and the close button sits in the screen corner. */}
<DialogContent
overlayClassName="bg-black/70"
className="aui-attachment-preview-dialog-content top-0 left-0 grid h-dvh w-screen max-w-none translate-x-0 translate-y-0 place-items-center rounded-none border-0 bg-transparent p-0 shadow-none ring-0 sm:max-w-none [&>button]:fixed [&>button]:top-4 [&>button]:right-4 [&>button]:z-20 [&>button]:size-9 [&>button]:rounded-full [&>button]:bg-transparent [&>button]:text-white [&>button]:opacity-100 [&>button]:ring-0! [&>button]:hover:bg-white/25 [&>button]:hover:text-white [&_svg]:text-white"
>
<DialogTitle className="aui-sr-only sr-only">
Image Attachment Preview
</DialogTitle>
<div className="aui-attachment-preview relative mx-auto flex max-h-[80dvh] w-full items-center justify-center overflow-hidden bg-background">
<AttachmentPreview src={src} />
{/* Clicking the backdrop (anywhere off the image) closes the preview. */}
<DialogClose asChild={true}>
<div aria-hidden="true" className="absolute inset-0" />
</DialogClose>
<div className="aui-attachment-preview pointer-events-none relative z-10 flex items-center justify-center">
<span className="pointer-events-auto">
<AttachmentPreview src={src} />
</span>
</div>
</DialogContent>
</Dialog>

69
studio/frontend/src/components/assistant-ui/model-selector/remembered-load-settings.ts
View file

@ -1,69 +0,0 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
// Per-model pre-load inference settings, persisted in localStorage so the load
// dialog can offer "Remember settings for <model>".
const KEY = "unsloth_load_settings";
export interface RememberedLoadSettings {
contextLength: number | null;
kvCacheDtype: string | null;
speculativeType: string | null;
specDraftNMax: number | null;
tensorParallel: boolean;
}
// Storage key for a pick's remembered settings. The remembered knobs are
// VRAM-budget driven (context override, KV-cache dtype, tensor-parallel), so the
// right values differ per quant. An HF repo collapses all its GGUF variants into
// one `id`, so fold the variant in to scope settings per quant. Local .gguf
// paths key by their file path (already file-specific); native drag-drop files
// key by display label, so same-named files in different folders share an entry.
export function rememberedLoadSettingsKey(selection: {
id: string;
ggufVariant?: string | null;
}): string {
return selection.ggufVariant
? `${selection.id}::${selection.ggufVariant}`
: selection.id;
}
function readAll(): Record<string, RememberedLoadSettings> {
try {
return JSON.parse(localStorage.getItem(KEY) ?? "{}");
} catch {
return {};
}
}
function writeAll(all: Record<string, RememberedLoadSettings>) {
try {
localStorage.setItem(KEY, JSON.stringify(all));
} catch {
// Ignore quota / unavailable storage.
}
}
export function loadRememberedLoadSettings(
key: string,
): RememberedLoadSettings | null {
return readAll()[key] ?? null;
}
export function saveRememberedLoadSettings(
key: string,
settings: RememberedLoadSettings,
) {
const all = readAll();
all[key] = settings;
writeAll(all);
}
export function clearRememberedLoadSettings(key: string) {
const all = readAll();
if (key in all) {
delete all[key];
writeAll(all);
}
}

140
studio/frontend/src/components/assistant-ui/thread.tsx
View file

@ -102,6 +102,7 @@ import { applyQwenThinkingParams } from "@/features/chat/utils/qwen-params";
import { isTauri } from "@/lib/api-base";
import { copyToClipboard } from "@/lib/copy-to-clipboard";
import { MicIcon } from "@/lib/mic-icon";
import { downloadFile, isDownloadCancelled } from "@/lib/native-files";
import { toast } from "@/lib/toast";
import { Tick02Icon } from "@/lib/tick-icon";
import { cn } from "@/lib/utils";
@ -175,6 +176,7 @@ import {
} from "react";
import { create } from "zustand";
import { extractTaggedText, updateThreadMessage } from "@/features/chat/utils/update-thread-message";
import { useIsMobile } from "@/hooks/use-mobile";
// True while a file is dragged anywhere over the chat page, so the composer
// can show its "Drop files here" affordance.
@ -1369,7 +1371,13 @@ export const ProjectComposer: FC<{
}> = ({ disabled, placeholder }) => {
return (
<GeneratedImageOverlayProvider>
<ComposerAnimated disabled={disabled} placeholder={placeholder} />
{/* New chat in a project: queuing follow-ups here misbinds the thread,
so the queue only runs once the user is inside a chat session. */}
<ComposerAnimated
disabled={disabled}
placeholder={placeholder}
disableQueue
/>
</GeneratedImageOverlayProvider>
);
};
@ -1379,11 +1387,17 @@ const ComposerAnimated: FC<{
placeholder?: string;
threadId?: string | null;
menuSide?: "top" | "bottom";
}> = ({ disabled, threadId, menuSide }) => {
disableQueue?: boolean;
}> = ({ disabled, threadId, menuSide, disableQueue }) => {
return (
<div className="relative mx-auto min-w-0 w-full max-w-[46rem]">
<div className="relative z-10 w-full">
<Composer disabled={disabled} threadId={threadId} menuSide={menuSide} />
<Composer
disabled={disabled}
threadId={threadId}
menuSide={menuSide}
disableQueue={disableQueue}
/>
</div>
</div>
);
@ -1418,7 +1432,8 @@ const Composer: FC<{
placeholder?: string;
threadId?: string | null;
menuSide?: "top" | "bottom";
}> = ({ disabled, threadId, menuSide }) => {
disableQueue?: boolean;
}> = ({ disabled, threadId, menuSide, disableQueue }) => {
const aui = useAui();
const pageDragging = useContext(PageDragContext);
const { overlay, closeOverlay } = useGeneratedImageOverlay();
@ -1434,18 +1449,17 @@ const Composer: FC<{
const artifactsEnabled = useChatRuntimeStore((s) => s.artifactsEnabled);
const mcpEnabledForChat = useChatRuntimeStore((s) => s.mcpEnabledForChat);
const ragEnabled = useChatRuntimeStore((s) => s.ragEnabled);
const permissionMode = useChatRuntimeStore((s) => s.permissionMode);
// More than 4 pills: collapse to icons only. Search and Code always show; the
// permission pill shows in every mode except "off" (it renders null there);
// Images, RAG, Canvas and MCP are conditional.
const pillsCompact =
2 +
(permissionMode !== "off" ? 1 : 0) +
(ragEnabled ? 1 : 0) +
(supportsBuiltinImageGeneration ? 1 : 0) +
(artifactsEnabled ? 1 : 0) +
(mcpEnabledForChat ? 1 : 0) >
4;
// More than 4 pills: collapse to icons only. Search, Code, and permissions
// always show; Images, RAG, Canvas and MCP are conditional. Narrow viewports
// collapse too: the labelled row is wider than a phone-width composer.
const isMobile = useIsMobile();
const pillCount =
3 +
(ragEnabled ? 1 : 0) +
(supportsBuiltinImageGeneration ? 1 : 0) +
(artifactsEnabled ? 1 : 0) +
(mcpEnabledForChat ? 1 : 0);
const pillsCompact = isMobile || pillCount > 4;
const activeThreadId = useChatRuntimeStore((s) => s.activeThreadId);
const setPendingImageEditReference = useChatRuntimeStore(
(s) => s.setPendingImageEditReference,
@ -1556,20 +1570,6 @@ const Composer: FC<{
const t = setTimeout(() => writeComposerDraft(draftKey, composerText), 300);
return () => clearTimeout(t);
}, [composerText, draftKey]);
// Two-row layout shows once the input wraps or a tool is on. Tools can
// pre-select before a model loads, so an active toggle expands it either way.
// Keep the composer expanded whenever the permission pill is visible.
const composerExpanded =
isMultiline ||
hasAttachments ||
hasPendingAudio ||
toolsEnabled ||
codeToolsEnabled ||
imageToolsEnabled ||
ragEnabled ||
artifactsEnabled ||
mcpEnabledForChat ||
permissionMode !== "off";
// react-textarea-autosize re-measures only on value change or window resize,
// not on the width swap from expanding, so it keeps the taller height and
// leaves a stray blank row. Nudge a resize whenever input width changes.
@ -1742,6 +1742,11 @@ const Composer: FC<{
if (threadIsRunning || promptQueueActive) {
event.preventDefault();
// Project new-chat composer: never queue, just ask the user to wait.
if (disableQueue) {
toast.error("Wait for the current response to finish");
return;
}
if (!canQueueCurrentPrompt) {
if (overlay || hasAttachments || hasPendingAudio) {
toast.error(
@ -1819,6 +1824,7 @@ const Composer: FC<{
composerText,
createPromptQueueTarget,
disabled,
disableQueue,
hasAttachments,
hasPendingAudio,
interceptSend,
@ -1838,9 +1844,12 @@ const Composer: FC<{
const startQueue = useCallback(
(items: string[], waitForCurrentRun = threadIsRunning) => {
// Saved-prompt Run-list calls this directly, so honour disableQueue here
// too: queuing from the project new-chat composer misbinds the thread.
if (disableQueue) return;
startPromptQueue(items, createPromptQueueTarget(), waitForCurrentRun);
},
[createPromptQueueTarget, threadIsRunning],
[createPromptQueueTarget, threadIsRunning, disableQueue],
);
const queueContextValue: PromptQueueCallbacks = { startQueue, stopQueue };
@ -1856,27 +1865,25 @@ const Composer: FC<{
<ToolStatusDisplay />
<div
className="unsloth-composer-line"
data-expanded={composerExpanded ? "true" : "false"}
// The permission pill is always visible, so keep the two-row layout
// expanded and leave the primary tool toggles accessible in every mode.
data-expanded="true"
>
<div
className="unsloth-composer-left"
data-pill-compact={pillsCompact ? "true" : undefined}
>
<ComposerToolsMenu side={effectiveMenuSide} />
{/* Permission-level pill: always visible, even while the pill row
is collapsed; opens the permission level dropdown. */}
{/* Permission-level pill: always visible and opens the permission
level dropdown. */}
<PermissionModeComposerPill side={effectiveMenuSide} />
{composerExpanded ? (
<>
<WebSearchToggle />
<CodeToolsToggle />
<ImagesToggle />
<KnowledgeBaseComposerButton side={effectiveMenuSide} />
{artifactsEnabled ? <ArtifactsToggle /> : null}
{mcpEnabledForChat ? (
<McpComposerButton side={effectiveMenuSide} />
) : null}
</>
<WebSearchToggle />
<CodeToolsToggle />
<ImagesToggle />
<KnowledgeBaseComposerButton side={effectiveMenuSide} />
{artifactsEnabled ? <ArtifactsToggle /> : null}
{mcpEnabledForChat ? (
<McpComposerButton side={effectiveMenuSide} />
) : null}
</div>
<ComposerPrimitive.Input
@ -1902,8 +1909,11 @@ const Composer: FC<{
isComposing ||
hasPendingAttachments
}
queueDisabled={!canQueueCurrentPrompt}
// disableQueue (project new-chat composer) also blocks the queue
// button, so a running thread shows Stop instead of Queue.
queueDisabled={disableQueue || !canQueueCurrentPrompt}
onQueueClick={() => {
if (disableQueue) return;
const queuedPrompt = composerText.trim();
if (queuedPrompt.length === 0) {
return;
@ -2930,9 +2940,9 @@ const ComposerToolsMenu: FC<{ side?: "top" | "bottom" }> = ({
<DropdownMenuItem
onSelect={() => {
if (!activeThreadId) return;
exportConversationRawJsonl(activeThreadId).catch(() =>
toast.error("Export failed."),
);
exportConversationRawJsonl(activeThreadId).catch((error) => {
if (!isDownloadCancelled(error)) toast.error("Export failed.");
});
}}
>
Raw JSONL
@ -2940,9 +2950,9 @@ const ComposerToolsMenu: FC<{ side?: "top" | "bottom" }> = ({
<DropdownMenuItem
onSelect={() => {
if (!activeThreadId) return;
exportConversationCsv(activeThreadId).catch(() =>
toast.error("Export failed."),
);
exportConversationCsv(activeThreadId).catch((error) => {
if (!isDownloadCancelled(error)) toast.error("Export failed.");
});
}}
>
CSV
@ -2950,9 +2960,9 @@ const ComposerToolsMenu: FC<{ side?: "top" | "bottom" }> = ({
<DropdownMenuItem
onSelect={() => {
if (!activeThreadId) return;
exportConversationShareGPT(activeThreadId).catch(() =>
toast.error("Export failed."),
);
exportConversationShareGPT(activeThreadId).catch((error) => {
if (!isDownloadCancelled(error)) toast.error("Export failed.");
});
}}
>
ShareGPT JSONL
@ -3915,6 +3925,21 @@ const EditAssistantMessageButton: FC = () => {
);
};
async function exportMessageMarkdown(content: string): Promise<void> {
try {
await downloadFile(
content,
`message-${Date.now()}.md`,
"text/markdown",
);
} catch (error) {
if (!isDownloadCancelled(error)) {
toast.error("Could not save Markdown export.", {
description: error instanceof Error ? error.message : String(error),
});
}
}
}
const AssistantActionBar: FC = () => {
const { forkMessage, forkDisabled } = useForkMessageAction();
const [detailsOpen, setDetailsOpen] = useState(false);
@ -3983,7 +4008,10 @@ const AssistantActionBar: FC = () => {
<GitBranchIcon strokeWidth={1.75} className="size-icon" />
Fork in new chat
</ActionBarMorePrimitive.Item>
<ActionBarPrimitive.ExportMarkdown asChild={true}>
<ActionBarPrimitive.ExportMarkdown
asChild={true}
onExport={exportMessageMarkdown}
>
<ActionBarMorePrimitive.Item className="aui-action-bar-more-item flex cursor-pointer select-none items-center gap-2 rounded-[12px] px-3 py-2 text-sm outline-none hover:bg-accent hover:text-accent-foreground focus:bg-accent focus:text-accent-foreground">
<HugeiconsIcon
icon={Download01Icon}

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

@ -70,13 +70,18 @@ export function FloatingMonitor() {
(sum, device) => sum + (device.memory_total_gb ?? 0),
0,
);
const vramUsed = devices.reduce(
(sum, device) => sum + (device.vram_used_gb ?? 0),
0,
);
// null usage = unknown (e.g. Windows ROCm perf counter): treating it as 0
// fabricates a 0-used readout, so the aggregate is unknown if any device is.
const vramUsageKnown =
devices.length > 0 &&
devices.every((device) => Number.isFinite(device.vram_used_gb));
const vramUsed = vramUsageKnown
? devices.reduce((sum, device) => sum + (device.vram_used_gb ?? 0), 0)
: 0;
const vramPercent = clampPercent(
vramTotal > 0 ? (vramUsed / vramTotal) * 100 : 0,
vramUsageKnown && vramTotal > 0 ? (vramUsed / vramTotal) * 100 : 0,
);
const unknownLabel = t("settings.resources.environment.unknown");
const hasGpu = (systemInfo.gpu?.available ?? false) && devices.length > 0;
@ -164,17 +169,20 @@ export function FloatingMonitor() {
<span
className={cn(
"shrink-0 tabular-nums",
usageTextClass(vramPercent),
vramUsageKnown
? usageTextClass(vramPercent)
: "text-muted-foreground",
)}
>
{Math.round(vramPercent)}%
{vramUsageKnown ? `${Math.round(vramPercent)}%` : "--"}
</span>
</div>
<div className="text-xs text-muted-foreground font-mono tabular-nums">
{formatGiB(vramUsed)} / {formatGiB(vramTotal)}
{vramUsageKnown ? formatGiB(vramUsed) : unknownLabel} /{" "}
{formatGiB(vramTotal)}
</div>
<Progress
value={vramPercent}
value={vramUsageKnown ? vramPercent : 0}
className="mt-1 h-1.5 rounded-full bg-muted"
indicatorClassName={usageIndicatorClass(vramPercent)}
/>

2
studio/frontend/src/components/navbar.tsx
View file

@ -22,7 +22,7 @@ export function Navbar() {
);
}
return (
<header className="absolute top-0 inset-x-0 z-40 h-[48px] pointer-events-none">
<header className="absolute top-0 inset-x-0 z-[45] h-[48px] pointer-events-none">
<div className="flex h-full items-start pt-[11px] pl-2">
<SidebarTrigger className="pointer-events-auto !size-[34px]" />
</div>

10
studio/frontend/src/components/tauri/window-titlebar.tsx
View file

@ -115,7 +115,7 @@ export function WindowTitlebar({
? "var(--studio-sidebar-expanded-width,17.5rem)"
: "var(--studio-sidebar-collapsed-width,3rem)"
: "0px";
const contentBorderLeft = `calc(${sidebarWidth} + 12px)`;
const contentBorderLeft = pinned ? `calc(${sidebarWidth} + 12px)` : "0px";
const refreshMaximized = useCallback(async () => {
if (!enabled) {
@ -232,10 +232,10 @@ export function WindowTitlebar({
)}
aria-label="Window titlebar"
>
{showSidebarSurface && (
{showSidebarSurface && pinned && (
<div
aria-hidden="true"
className="pointer-events-none absolute top-full h-3 w-px -translate-x-px bg-sidebar"
className="pointer-events-none absolute top-full size-3 -translate-x-px bg-sidebar"
style={{ left: sidebarWidth }}
/>
)}
@ -246,7 +246,7 @@ export function WindowTitlebar({
style={{ left: contentBorderLeft, right: 0 }}
/>
)}
{showSidebarSurface && (
{showSidebarSurface && pinned && (
<div
aria-hidden="true"
className="pointer-events-none absolute top-full size-3 -translate-x-px rounded-tl-[12px] border-l border-t border-sidebar-border bg-background"
@ -257,7 +257,7 @@ export function WindowTitlebar({
<div
className={cn(
"pointer-events-auto absolute left-0 top-0 flex h-full min-w-0 items-center",
pinned ? "gap-2 px-3" : "justify-center",
pinned ? "gap-2 pl-3" : "justify-center",
)}
style={{ width: sidebarWidth }}
onMouseDown={handleDragMouseDown}

2
studio/frontend/src/components/ui/input-group.tsx
View file

@ -122,7 +122,7 @@ function InputGroupInput({
<Input
data-slot="input-group-control"
className={cn(
"rounded-none border-0 bg-transparent shadow-none ring-0 focus-visible:ring-0 aria-invalid:ring-0 dark:bg-transparent flex-1",
"rounded-none border-0 bg-transparent shadow-none ring-0 focus-visible:ring-0 aria-invalid:ring-0 dark:bg-transparent dark:focus-visible:bg-transparent flex-1",
className,
)}
{...props}

7
studio/frontend/src/components/ui/tooltip.tsx
View file

@ -67,9 +67,14 @@ function TooltipTrigger({
const handleClick = useCallback(
(e: React.MouseEvent<HTMLButtonElement>) => {
// Run the composed handler first: when this trigger wraps another Radix
// trigger (e.g. DialogTrigger around an attachment tile), that trigger's
// action is skipped if the event is already default-prevented.
onClick?.(e);
// preventDefault keeps Radix Tooltip's internal close-on-click from
// undoing the tap-toggle below (its composed handler checks it).
e.preventDefault();
toggle?.();
onClick?.(e);
},
[toggle, onClick],
);

157
studio/frontend/src/features/chat/api/chat-adapter.ts
View file

@ -2,10 +2,7 @@
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import { getAuthToken } from "@/features/auth";
import {
loadRememberedLoadSettings,
rememberedLoadSettingsKey,
} from "@/components/assistant-ui/model-selector/remembered-load-settings";
import { resolveInitialConfig } from "@/features/model-picker";
import { projectHasSources } from "@/features/rag/api/rag-api";
import { apiUrl } from "@/lib/api-base";
import { parseParamCountB } from "@/lib/model-size";
@ -45,12 +42,18 @@ import {
import {
type PendingImageEditReference,
type RagAutoInject,
GPU_LAYERS_AUTO,
loadedGpuMemoryFields,
reconcilePersistedGpuIds,
resolveLoadedSpeculativeSettings,
resolveSpeculativeSettingsForLoad,
persistGpuMemoryModeOnLoad,
resolveToolsEnabledOnLoad,
saveSpeculativeType,
useChatRuntimeStore,
} from "../stores/chat-runtime-store";
import { resolveFitMaxSeqLength, resolveManualAutoCtxPin } from "../presets/preset-policy";
import { ensureGpuDeviceCache } from "@/hooks/use-gpu-info";
import { useExternalProvidersStore } from "../stores/external-providers-store";
import {
shouldPreserveFullOutput,
@ -1489,6 +1492,13 @@ async function autoLoadSmallestModel(): Promise<{
max_seq_length: number;
is_lora: boolean;
gguf_variant?: string | null;
// GGUF-only: scopes the training guard to the same placement policy /load
// will use. Manual mode must match because it makes placement user-owned.
// The layer/MoE/split/KV/spec knobs are deliberately not sent: Auto mode's
// guard sizes conservatively, while Manual mode bypasses that estimate.
// The safetensors fallback omits both fields and uses HF auto-placement.
gpu_ids?: number[];
gpu_memory_mode?: "auto" | "manual";
}): Promise<boolean> {
const validation = await validateModel({
...payload,
@ -1520,33 +1530,69 @@ async function autoLoadSmallestModel(): Promise<{
return false;
}
const currentStore = useChatRuntimeStore.getState();
const remembered = loadRememberedLoadSettings(
rememberedLoadSettingsKey({
id: candidate.id,
ggufVariant: candidate.ggufVariant,
}),
);
const { config } = resolveInitialConfig(candidate.id, candidate.ggufVariant);
const effectiveMaxSeqLength = resolveLoadMaxSeqLength({
modelId: candidate.id,
ggufVariant: candidate.ggufVariant,
isGguf: candidate.kind === "gguf",
customContextLength: remembered?.contextLength ?? null,
customContextLength: config.customContextLength,
ggufContextLength: null,
currentCheckpoint: currentStore.params.checkpoint,
activeGgufVariant: currentStore.activeGgufVariant,
maxSeqLength: candidate.maxSeqLength,
maxSeqLength: config.maxSeqLength ?? candidate.maxSeqLength,
presetSource: currentStore.activePresetSource,
});
// The GPU knobs are per-model, so read them from the same per-model config
// that fed effectiveMaxSeqLength -- on a background auto-load the live store
// holds session defaults, not the saved Manual mode / layer pin / GPU pick.
// Absent fields fall back like the interactive restore: the mode to the store
// (a persisted standing preference), the per-model knobs to their defaults.
// The saved GPU pick is reconciled against the GPUs present now.
const effectiveGpuMemoryMode =
config.gpuMemoryMode ?? currentStore.gpuMemoryMode;
const effectiveGpuLayers = config.gpuLayers ?? GPU_LAYERS_AUTO;
const effectiveNCpuMoe = config.nCpuMoe ?? 0;
if (config.selectedGpuIds != null) {
// Warm the device cache first: on a cold cache the reconcile passes the
// saved pick through unvalidated, and a stale cross-host pick then fails
// the load with the picker hidden.
await ensureGpuDeviceCache();
}
const effectiveGpuIds =
config.selectedGpuIds !== undefined
? reconcilePersistedGpuIds(config.selectedGpuIds)
: null;
// Under Manual GPU memory + Auto layers, llama.cpp's --fit owns context
// sizing, so send 0 (or the pinned length). GGUF-only; a no-op otherwise.
// The context pin is per-model too, so it comes from the saved config, not
// the live store.
const fitMaxSeqLength = resolveFitMaxSeqLength(
candidate.kind === "gguf",
effectiveGpuMemoryMode,
effectiveGpuLayers,
config.customContextLength ?? null,
effectiveMaxSeqLength,
);
const effectiveSpeculativeType =
remembered?.speculativeType ?? specSettings.speculativeType;
config.speculativeType ?? specSettings.speculativeType;
const effectiveSpecDraftNMax =
remembered?.specDraftNMax ?? specSettings.specDraftNMax;
config.specDraftNMax ?? specSettings.specDraftNMax;
const effectiveChatTemplateOverride = config.chatTemplateOverride?.trim()
? config.chatTemplateOverride
: null;
if (
!(await canAutoLoad({
model_path: candidate.id,
max_seq_length: effectiveMaxSeqLength,
max_seq_length: fitMaxSeqLength,
is_lora: false,
gguf_variant: candidate.ggufVariant,
// The same remembered-derived GPU pick the load below sends.
...(candidate.kind === "gguf"
? {
gpu_ids: effectiveGpuIds ?? undefined,
gpu_memory_mode: effectiveGpuMemoryMode,
}
: {}),
}))
) {
skippedAutoLoadCandidates.add(
@ -1558,17 +1604,37 @@ async function autoLoadSmallestModel(): Promise<{
const loadResp = await loadModel({
model_path: candidate.id,
hf_token: hfToken,
max_seq_length: effectiveMaxSeqLength,
max_seq_length: fitMaxSeqLength,
load_in_4bit: true,
is_lora: false,
gguf_variant: candidate.ggufVariant,
trust_remote_code: trustRemoteCode,
cache_type_kv: remembered?.kvCacheDtype ?? null,
chat_template_override: effectiveChatTemplateOverride,
cache_type_kv: config.kvCacheDtype,
speculative_type: effectiveSpeculativeType,
spec_draft_n_max: effectiveSpecDraftNMax,
tensor_parallel: remembered?.tensorParallel ?? false,
tensor_parallel: config.tensorParallel,
// GGUF-only: the safetensors fallback loads via HF auto-placement (no
// explicit pins). The split ratio is deliberately never remembered
// (positionally bound to an exact GPU set), so auto-load leaves llama.cpp's
// free-VRAM default in charge rather than sending a stale store value.
...(candidate.kind === "gguf"
? {
gpu_memory_mode: effectiveGpuMemoryMode,
gpu_layers: effectiveGpuLayers,
n_cpu_moe: effectiveNCpuMoe,
gpu_ids: effectiveGpuIds ?? undefined,
}
: {}),
});
saveSpeculativeType(effectiveSpeculativeType);
// Only persist the global preference when the value came from the global
// settings. A per-model config's choice must stay load-local, or autoloading
// a remembered model on startup would rewrite the global default.
if (config.speculativeType == null) {
saveSpeculativeType(effectiveSpeculativeType);
}
// Self-gates on is_gguf (skips diffusion), so persists only for a real GGUF load.
persistGpuMemoryModeOnLoad(loadResp, effectiveGpuMemoryMode);
useChatRuntimeStore
.getState()
.setCheckpoint(candidate.id, candidate.ggufVariant ?? undefined);
@ -1578,6 +1644,9 @@ async function autoLoadSmallestModel(): Promise<{
);
store.setParams({
...store.params,
...(candidate.kind === "gguf"
? {}
: { maxSeqLength: effectiveMaxSeqLength }),
maxTokens:
candidate.kind === "gguf"
? loadResp.context_length ?? 131072
@ -1597,6 +1666,15 @@ async function autoLoadSmallestModel(): Promise<{
store.setModels([...store.models, autoModel]);
}
if (candidate.kind === "gguf") {
// Keep an explicit Manual+Auto context pin the load just applied (so a
// later Apply doesn't silently revert it to auto-fit sizing), mirroring
// the interactive path's keepCustomCtx; other cases baseline on
// ggufContextLength.
const keepCustomCtx = resolveManualAutoCtxPin(
effectiveGpuMemoryMode,
effectiveGpuLayers,
config.customContextLength ?? null,
);
useChatRuntimeStore.setState({
ggufContextLength: loadResp.context_length ?? 131072,
ggufMaxContextLength:
@ -1613,9 +1691,14 @@ async function autoLoadSmallestModel(): Promise<{
loadedKvCacheDtype: loadResp.cache_type_kv ?? null,
tensorParallel: loadResp.tensor_parallel ?? false,
loadedTensorParallel: loadResp.tensor_parallel ?? false,
...loadedGpuMemoryFields(loadResp),
loadedCustomContextLength: keepCustomCtx,
defaultChatTemplate: loadResp.chat_template ?? null,
chatTemplateOverride: null,
loadedChatTemplateOverride: null,
chatTemplateOverride: effectiveChatTemplateOverride,
loadedChatTemplateOverride: effectiveChatTemplateOverride,
// Retain the saved requested context so re-saving the config keeps the
// override; null stays null (auto/VRAM-fit).
customContextLength: config.customContextLength,
loadedIsMultimodal: isMultimodalResponse(loadResp),
loadedIsDiffusion: loadResp.is_diffusion ?? false,
...resolveLoadedSpeculativeSettings(loadResp),
@ -1633,9 +1716,13 @@ async function autoLoadSmallestModel(): Promise<{
loadedKvCacheDtype: loadResp.cache_type_kv ?? null,
tensorParallel: loadResp.tensor_parallel ?? false,
loadedTensorParallel: loadResp.tensor_parallel ?? false,
// Non-GGUF response: clears any stale GPU baseline a prior manual-GPU
// GGUF load left, matching the interactive/status sibling load paths.
...loadedGpuMemoryFields(loadResp),
defaultChatTemplate: loadResp.chat_template ?? null,
chatTemplateOverride: null,
loadedChatTemplateOverride: null,
chatTemplateOverride: effectiveChatTemplateOverride,
loadedChatTemplateOverride: effectiveChatTemplateOverride,
customContextLength: null,
...resolveLoadedSpeculativeSettings(loadResp),
loadedIsMultimodal: isMultimodalResponse(loadResp),
loadedIsDiffusion: loadResp.is_diffusion ?? false,
@ -1820,12 +1907,17 @@ async function autoLoadSmallestModel(): Promise<{
duration: 30000,
});
try {
const rt = useChatRuntimeStore.getState();
if (
!(await canAutoLoad({
model_path: "unsloth/Qwen3.5-4B-MTP-GGUF",
max_seq_length: 0,
is_lora: false,
gguf_variant: "UD-Q4_K_XL",
// The same live-store GPU pick the load below sends (a fresh default
// model has no remembered settings to prefer).
gpu_ids: rt.selectedGpuIds ?? undefined,
gpu_memory_mode: rt.gpuMemoryMode,
}))
) {
toast.dismiss(toastId);
@ -1835,6 +1927,9 @@ async function autoLoadSmallestModel(): Promise<{
const loadResp = await loadModel({
model_path: "unsloth/Qwen3.5-4B-MTP-GGUF",
hf_token: hfToken,
// Model default under both modes: Auto layers + no pin means
// resolveFitMaxSeqLength returns 0 for every mode (the canAutoLoad
// preflight above sends the same).
max_seq_length: 0,
load_in_4bit: true,
is_lora: false,
@ -1842,8 +1937,20 @@ async function autoLoadSmallestModel(): Promise<{
trust_remote_code: trustRemoteCode,
speculative_type: specSettings.speculativeType,
spec_draft_n_max: specSettings.specDraftNMax,
// GPU Memory mode is a standing preference, so honor it on auto-load.
// The layer/MoE/split knobs and the context pin are per-model: the live
// store may hold edits drafted for a staged pick, and a fresh default
// model has no remembered settings, so those stay at their defaults like
// the cached-candidate path. The GPU pick deliberately differs (it's the
// picker's current on-screen selection, which the canAutoLoad preflight
// above already committed to).
gpu_memory_mode: rt.gpuMemoryMode,
gpu_layers: GPU_LAYERS_AUTO,
n_cpu_moe: 0,
gpu_ids: rt.selectedGpuIds ?? undefined,
});
saveSpeculativeType(specSettings.speculativeType);
persistGpuMemoryModeOnLoad(loadResp, rt.gpuMemoryMode);
useChatRuntimeStore
.getState()
.setCheckpoint("unsloth/Qwen3.5-4B-MTP-GGUF", "UD-Q4_K_XL");
@ -1880,6 +1987,10 @@ async function autoLoadSmallestModel(): Promise<{
loadedKvCacheDtype: loadResp.cache_type_kv ?? null,
tensorParallel: loadResp.tensor_parallel ?? false,
loadedTensorParallel: loadResp.tensor_parallel ?? false,
...loadedGpuMemoryFields(loadResp),
// Drives the GPU Memory controls' diffusion gate; set alongside the
// GPU fields on every load path so the gate can't read stale.
loadedIsDiffusion: loadResp.is_diffusion ?? false,
defaultChatTemplate: loadResp.chat_template ?? null,
chatTemplateOverride: null,
loadedIsMultimodal: isMultimodalResponse(loadResp),

173
studio/frontend/src/features/chat/api/chat-api.ts
View file

@ -2,7 +2,12 @@
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import { authFetch } from "@/features/auth";
import { prepareHfTokenForUse } from "@/features/hf-auth";
// These helpers are deliberately API-layer-only and are not part of their
// features' React-facing public barrels.
// eslint-disable-next-line no-restricted-imports
import { hubTokenHeader } from "@/features/hub/lib/hub-token-header";
// eslint-disable-next-line no-restricted-imports
import { consumeNativePathToken } from "@/features/native-intents/api";
import { formatFastApiDetail } from "@/lib/format-fastapi-error";
import type {
@ -28,6 +33,7 @@ import type {
} from "../types/api";
export const CHAT_HISTORY_UPDATED_EVENT = "unsloth-chat-history-updated";
export const CHAT_PROJECTS_UPDATED_EVENT = "unsloth-chat-projects-updated";
/**
* Thrown when the chat SSE stream ends without a terminal signal (`[DONE]` or a
@ -50,6 +56,13 @@ export function notifyChatHistoryUpdated(): void {
}
}
function notifyChatProjectsUpdated(): void {
notifyChatHistoryUpdated();
if (typeof window !== "undefined") {
window.dispatchEvent(new Event(CHAT_PROJECTS_UPDATED_EVENT));
}
}
function parseErrorText(status: number, body: unknown): string {
if (body && typeof body === "object") {
const detail = (body as { detail?: unknown }).detail;
@ -104,11 +117,14 @@ export async function getApiMonitorEntry(id: string): Promise<ApiMonitorEntry> {
export async function loadModel(
payload: LoadModelRequest,
): Promise<LoadModelResponse> {
const preparedToken = await prepareHfTokenForUse(payload.hf_token);
if (!preparedToken.proceed) throw new Error("Model load cancelled.");
const response = await authFetch("/api/inference/load", {
method: "POST",
headers: { "Content-Type": "application/json" },
body: JSON.stringify({
...payload,
hf_token: preparedToken.token,
native_path_lease: payload.nativePathLease ?? null,
nativePathLease: undefined,
}),
@ -119,36 +135,48 @@ export async function loadModel(
export async function validateModel(
payload: LoadModelRequest,
): Promise<ValidateModelResponse> {
const preparedToken = await prepareHfTokenForUse(payload.hf_token);
if (!preparedToken.proceed) throw new Error("Model load cancelled.");
const response = await authFetch("/api/inference/validate", {
method: "POST",
headers: { "Content-Type": "application/json" },
body: JSON.stringify({
model_path: payload.model_path,
native_path_lease: payload.nativePathLease ?? null,
hf_token: payload.hf_token,
hf_token: preparedToken.token,
gguf_variant: payload.gguf_variant ?? null,
// Send the intended load settings so validate's VRAM check matches the
// follow-up /load and doesn't unload for a load /load would then reject.
// Intended load settings so validate's preflight matches the follow-up
// /load. Default placement is sized against the selected GPUs.
max_seq_length: payload.max_seq_length,
load_in_4bit: payload.load_in_4bit,
gpu_ids: payload.gpu_ids,
// Manual placement is an explicit override: Auto layers use llama.cpp
// --fit, while a pinned layer count is owned by the user. Tell validate
// so it applies the same training-guard policy as /load.
gpu_memory_mode: payload.gpu_memory_mode,
}),
});
return parseJsonOrThrow<ValidateModelResponse>(response);
}
/**
* Read a GGUF's native context length from its local header (no GPU load, no
* download). Returns null when the file isn't downloaded yet, the model isn't a
* GGUF, or it's gated. For a native (drag-drop / picked) file, pass
* `nativePathToken` so the backend reads the granted local path. Used by the
* deferred-load staging flow to fill the context slider before the single load.
* Read a GGUF's header dims (native context length, total layer count, MoE
* expert-layer count) from its local file (no GPU load, no download). All are
* null when the file isn't downloaded yet, the model isn't a GGUF, or it's
* gated. For a native (drag-drop / picked) file, pass `nativePathToken` so the
* backend reads the granted local path. Used by the deferred-load staging flow
* to size the context, GPU-layers and MoE sliders before the single load.
*/
export async function fetchGgufContextLength(payload: {
export async function fetchGgufStagedMetadata(payload: {
model_path: string;
gguf_variant?: string | null;
hf_token?: string | null;
nativePathToken?: string | null;
}): Promise<number | null> {
}): Promise<{
contextLength: number | null;
layerCount: number | null;
moeLayerCount: number | null;
}> {
let nativePathLease: string | null = null;
if (payload.nativePathToken) {
try {
@ -156,8 +184,8 @@ export async function fetchGgufContextLength(payload: {
await consumeNativePathToken(payload.nativePathToken, "validate-model")
).nativePathLease;
} catch {
// Lease expired / revoked: degrade to no context (the load can re-mint).
return null;
// Lease expired / revoked: degrade to no metadata (the load can re-mint).
return { contextLength: null, layerCount: null, moeLayerCount: null };
}
}
const response = await authFetch("/api/inference/validate", {
@ -172,7 +200,11 @@ export async function fetchGgufContextLength(payload: {
}),
});
const res = await parseJsonOrThrow<ValidateModelResponse>(response);
return res.context_length ?? null;
return {
contextLength: res.context_length ?? null,
layerCount: res.layer_count ?? null,
moeLayerCount: res.moe_layer_count ?? null,
};
}
export async function unloadModel(payload: UnloadModelRequest): Promise<void> {
@ -310,13 +342,17 @@ interface LocalModelListResponse {
models: LocalModelInfo[];
}
export async function listLocalModels(): Promise<LocalModelListResponse> {
const response = await authFetch("/api/models/local");
export async function listLocalModels(
signal?: AbortSignal,
): Promise<LocalModelListResponse> {
const response = await authFetch("/api/models/local", { signal });
return parseJsonOrThrow<LocalModelListResponse>(response);
}
export async function listCachedGguf(): Promise<CachedGgufRepo[]> {
const response = await authFetch("/api/models/cached-gguf");
export async function listCachedGguf(
signal?: AbortSignal,
): Promise<CachedGgufRepo[]> {
const response = await authFetch("/api/models/cached-gguf", { signal });
const data = await parseJsonOrThrow<{ cached: CachedGgufRepo[] }>(response);
return data.cached;
}
@ -331,22 +367,43 @@ export interface CachedModelRepo {
export async function listCachedModels(
hfToken?: string | null,
signal?: AbortSignal,
): Promise<CachedModelRepo[]> {
const response = await authFetch("/api/models/cached-models", {
headers: hubTokenHeader(hfToken),
signal,
});
const data = await parseJsonOrThrow<{ cached: CachedModelRepo[] }>(response);
return data.cached;
}
export async function deleteCachedModel(
export interface CachedModelPath {
path: string;
is_dir: boolean;
}
/** Absolute on-disk path of a cached repo or one of its GGUF variants. */
export async function getCachedModelPath(
repoId: string,
variant?: string,
): Promise<CachedModelPath> {
const params = new URLSearchParams({ repo_id: repoId });
if (variant) params.set("variant", variant);
const response = await authFetch(
`/api/models/cached-model-path?${params.toString()}`,
);
return parseJsonOrThrow<CachedModelPath>(response);
}
/** Reveal a cached repo (or one GGUF variant's file) in the OS file manager. */
export async function revealCachedModel(
repoId: string,
variant?: string,
): Promise<void> {
const payload: Record<string, string> = { repo_id: repoId };
if (variant) payload.variant = variant;
const response = await authFetch("/api/models/delete-cached", {
method: "DELETE",
const response = await authFetch("/api/models/reveal-cached-model", {
method: "POST",
headers: { "Content-Type": "application/json" },
body: JSON.stringify(payload),
});
@ -423,6 +480,73 @@ export async function listChatThreads(
return Array.isArray(data.threads) ? data.threads : [];
}
/** One chat message attachment, as listed for the settings uploaded-files view. */
export interface ChatAttachmentRecord {
id: string;
messageId: string;
threadId: string;
pairId?: string | null;
threadTitle?: string | null;
name: string;
type?: string | null;
contentType?: string | null;
sizeBytes?: number | null;
createdAt?: number | null;
}
export interface ChatAttachmentPage {
attachments: ChatAttachmentRecord[];
nextOffset: number | null;
}
export async function listChatAttachments(
offset = 0,
limit = 50,
): Promise<ChatAttachmentPage> {
const params = new URLSearchParams({
limit: String(limit),
offset: String(offset),
});
const response = await authFetch(`/api/chat/attachments?${params}`);
const data = await parseJsonOrThrow<{
attachments: ChatAttachmentRecord[];
nextOffset: number | null;
}>(response);
return {
attachments: Array.isArray(data.attachments) ? data.attachments : [],
nextOffset:
typeof data.nextOffset === "number" && Number.isFinite(data.nextOffset)
? data.nextOffset
: null,
};
}
/** Stored attachment content (image bytes or extracted text) as a Blob. */
export async function fetchChatAttachmentBlob(
messageId: string,
attachmentId: string,
): Promise<Blob> {
const response = await authFetch(
`/api/chat/attachments/${encodeURIComponent(messageId)}/${encodeURIComponent(attachmentId)}/file`,
);
if (!response.ok) {
const body = await response.json().catch(() => null);
throw new Error(parseErrorText(response.status, body));
}
return response.blob();
}
export async function deleteChatAttachment(
messageId: string,
attachmentId: string,
): Promise<void> {
const response = await authFetch(
`/api/chat/attachments/${encodeURIComponent(messageId)}/${encodeURIComponent(attachmentId)}`,
{ method: "DELETE" },
);
await parseJsonOrThrow<{ ok: boolean }>(response);
}
export async function getChatThread(
threadId: string,
): Promise<ThreadRecord | null> {
@ -547,7 +671,7 @@ export async function saveChatProject(
body: JSON.stringify(project),
});
const saved = await parseJsonOrThrow<ProjectRecord>(response);
notifyChatHistoryUpdated();
notifyChatProjectsUpdated();
return saved;
}
@ -564,7 +688,7 @@ export async function updateChatProject(
},
);
const project = await parseJsonOrThrow<ProjectRecord>(response);
notifyChatHistoryUpdated();
notifyChatProjectsUpdated();
return project;
}
@ -580,7 +704,7 @@ export async function deleteChatProject(
{ method: "DELETE" },
);
await parseJsonOrThrow<ProjectRecord>(response);
notifyChatHistoryUpdated();
notifyChatProjectsUpdated();
}
export async function listChatMessages(
@ -946,7 +1070,8 @@ export async function* streamChatCompletions(
parsed.type === "reasoning_summary"
) {
yield {
_reasoningDurationMs: (parsed as { duration_ms?: number }).duration_ms,
_reasoningDurationMs: (parsed as { duration_ms?: number })
.duration_ms,
} as unknown as OpenAIChatChunk;
separatorIndex = buffer.search(/\r?\n\r?\n/);
continue;

20
studio/frontend/src/features/chat/bypass-permissions-menu-item.tsx
View file

@ -20,16 +20,12 @@ import {
DropdownMenuSubTrigger,
} from "@/components/ui/dropdown-menu";
import { useChatRuntimeStore } from "@/features/chat/stores/chat-runtime-store";
import { PermissionModeMenuItems } from "./permission-mode-select";
import {
FULL_ACCESS_WARNING,
PermissionModeMenuItems,
} from "./permission-mode-select";
// "Bypass permissions" entry for the composer "+" -> More menu. Like the MCP
// pill, it opens a submenu where the user picks the permission level (Ask for
// approval / Approve for me / Full access). Picking Full access demands the
// danger warning; the other levels apply immediately. The menu closes normally
// on select (no preventDefault) -- the warning dialog lives outside the menu
// (BypassPermissionsConfirmDialog, mounted once at the chat-page root and
// driven by the store), so it survives the menu unmounting and the "+"/More
// popovers don't stay frozen.
// Tool permissions entry for the composer "+" menu.
export function BypassPermissionsMenuItem() {
const permissionMode = useChatRuntimeStore((s) => s.permissionMode);
const setBypassConfirmOpen = useChatRuntimeStore(
@ -44,7 +40,7 @@ export function BypassPermissionsMenuItem() {
}
>
<HugeiconsIcon icon={ShieldBanIcon} strokeWidth={2} />
Bypass permissions
Tool permissions
</DropdownMenuSubTrigger>
<DropdownMenuSubContent className="unsloth-plus-menu w-[300px]">
<PermissionModeMenuItems
@ -75,9 +71,7 @@ export function BypassPermissionsConfirmDialog() {
<AlertDialogHeader>
<AlertDialogTitle>Enable Full access?</AlertDialogTitle>
<AlertDialogDescription>
Full access (Bypass permissions) is dangerous since the AI model
might delete, corrupt your machine, and or cause real world damage
to you or the world - only accept if you are certain
{FULL_ACCESS_WARNING}
</AlertDialogDescription>
</AlertDialogHeader>
<AlertDialogFooter>

1005
studio/frontend/src/features/chat/chat-page.tsx
View file

File diff suppressed because it is too large Load diff

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