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

Merge commit '3d379cdb81' into r7248-string

Browse source
This commit is contained in:
danielhanchen 2026-07-21 02:25:52 +00:00
commit f10e785a53
120 changed files with 14553 additions and 1440 deletions
Download patch file Download diff file Expand all files Collapse all files

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

16
.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
@ -297,6 +301,8 @@ jobs:
logs/install.log
logs/server-logs/
logs/playwright
logs/playwright-permissions-*
logs/playwright_extra
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)

163
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,7 +2119,9 @@ 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
@ -2102,6 +2171,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 +2259,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 +2280,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 +2305,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 +2335,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 +2354,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 +2366,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 +2394,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)
@ -2335,8 +2438,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 +2450,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)

353
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"
@ -1575,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
@ -1636,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")"
@ -2001,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)"
@ -2069,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.
@ -2197,6 +2280,45 @@ _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
@ -2277,14 +2399,14 @@ _install_torch_default_index() {
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 $TORCH_INDEX_URL -- installing the newest supported release instead" "$C_WARN"
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 torchaudio \
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 torchaudio \
run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" "$TORCH_CONSTRAINT" "$TORCHVISION_CONSTRAINT" "$TORCHAUDIO_CONSTRAINT" \
--default-index "$TORCH_INDEX_URL" "$@"
fi
}
@ -2292,13 +2414,21 @@ _install_torch_default_index() {
# 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
}
@ -2308,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/,
@ -2561,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)
@ -2572,29 +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 and the CUDA cu12x/cu13x indexes now ship torch 2.11.x, so widen the
# ceiling to <2.12.0 (matches the base image and _CUDA_TORCH_PKG_SPEC in
# studio/install_python_stack.py). Keep the >=2.4 floor so an older CUDA index
# (e.g. cu118) still resolves. Match on _torch_index_leaf, not the full URL, so
# a mirror whose base path contains cu*/rocm7.2 but resolves to a cpu/older-rocm
# leaf keeps the default <2.11.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) TORCH_CONSTRAINT="torch>=2.11.0,<2.12.0" ;;
cu[0-9]*) TORCH_CONSTRAINT="torch>=2.4,<2.12.0" ;;
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 && \
@ -2671,10 +2886,14 @@ 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
@ -2821,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
@ -2867,8 +3086,8 @@ for _p in ('torch', 'torchvision', 'torchaudio'):
}
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
@ -2877,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.
@ -2894,7 +3113,7 @@ if [ "$_MIGRATED" = true ]; then
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
@ -2909,18 +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)..."
_install_torch_default_index --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)
@ -3074,7 +3289,7 @@ 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"
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})..."
@ -3095,7 +3310,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
fi
fi
else
substep "[WARN] Radeon repo unavailable; falling back to ROCm index ($TORCH_INDEX_URL)" "$C_WARN"
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
@ -3103,19 +3318,15 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
_install_torch_default_index
fi
else
substep "installing PyTorch ($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 the torch Step 1 installed
tauri_log "STEP" "Installing Unsloth"
@ -3126,7 +3337,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
# runtime deps (typer, safetensors, transformers, etc.) with --no-deps.
run_install_cmd_retry "install unsloth (no-torch)" uv pip install --python "$_VENV_PY" --no-deps \
--upgrade-package unsloth --upgrade-package unsloth-zoo \
"unsloth>=2026.7.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
@ -3145,7 +3356,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
elif [ "$STUDIO_LOCAL_INSTALL" = true ]; then
run_install_cmd_retry "install unsloth (local)" uv pip install --python "$_VENV_PY" \
${_UNSLOTH_TORCH_OVERRIDES:+--overrides "$_UNSLOTH_TORCH_OVERRIDES"} \
--upgrade-package unsloth "unsloth>=2026.7.3" "unsloth-zoo>=2026.7.3"
--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..."
@ -3161,23 +3372,19 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
_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)..."
_install_torch_default_index --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..."
@ -3217,7 +3424,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",

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

@ -23,6 +23,7 @@ import subprocess
import sys
import threading
import time
import uuid
from pathlib import Path
from typing import (
Callable,
@ -453,6 +454,23 @@ def _hf_offline_if_dns_dead():
os.environ.pop("TRANSFORMERS_OFFLINE", None)
try:
_SLOT_SAVE_MAX_BYTES = int(os.environ.get("UNSLOTH_SLOT_SAVE_MAX_BYTES") or (10 << 30))
except ValueError:
_SLOT_SAVE_MAX_BYTES = 10 << 30
# The idle loop holds the lifecycle gate across a slot save, so a newly arriving
# request waits on the in-flight save's HTTP call. Bound it (was 120s) so a slow
# or stuck save can't stall the next request for minutes; best-effort save just
# falls back to a plain unload. Override with UNSLOTH_SLOT_SAVE_TIMEOUT (seconds).
try:
_SLOT_SAVE_HTTP_TIMEOUT = float(os.environ.get("UNSLOTH_SLOT_SAVE_TIMEOUT") or 30.0)
except ValueError:
_SLOT_SAVE_HTTP_TIMEOUT = 30.0
if _SLOT_SAVE_HTTP_TIMEOUT <= 0:
_SLOT_SAVE_HTTP_TIMEOUT = 30.0
def _swa_cache_path() -> Path:
home = os.environ.get("UNSLOTH_STUDIO_HOME") or os.environ.get("STUDIO_HOME")
base = Path(home) if home else Path.home() / ".unsloth" / "studio"
@ -2000,6 +2018,12 @@ class LlamaCppBackend:
self._llama_log_path: Optional[Path] = None
self._cancel_event = threading.Event()
self._api_key: Optional[str] = None
self._slot_save_dir: Optional[str] = None
self._slot_save_binary: Optional[tuple[str, int]] = None
# (gguf_identity, launch_fingerprint) snapshotted at load, so a later slot
# save can tell whether the model files were swapped on disk since load.
self._slot_loaded_identity: Optional[tuple] = None
self._prompt_cache_disabled: bool = False
# True once a probe has completed; cleared on transient failure.
self._is_audio: bool = False
self._audio_type: Optional[str] = None
@ -2638,6 +2662,7 @@ class LlamaCppBackend:
"supports_ctx_checkpoints": False,
"supports_no_cache_prompt": False,
"supports_metrics": False,
"supports_slot_save": False,
}
try:
mtime = int(Path(bin_path).stat().st_mtime)
@ -2658,6 +2683,7 @@ class LlamaCppBackend:
supports_ctx_checkpoints = False
supports_no_cache_prompt = False
supports_metrics = False
supports_slot_save = False
try:
probe_env = cls._llama_server_env_for_binary(bin_path)
result = subprocess.run(
@ -2756,6 +2782,7 @@ class LlamaCppBackend:
supports_ctx_checkpoints = _is_real("--ctx-checkpoints")
supports_no_cache_prompt = _is_real("--no-cache-prompt")
supports_metrics = _is_real("--metrics")
supports_slot_save = _is_real("--slot-save-path")
except (OSError, subprocess.SubprocessError) as exc:
logger.debug(f"llama-server --help probe failed: {exc}")
@ -2773,6 +2800,7 @@ class LlamaCppBackend:
"supports_ctx_checkpoints": supports_ctx_checkpoints,
"supports_no_cache_prompt": supports_no_cache_prompt,
"supports_metrics": supports_metrics,
"supports_slot_save": supports_slot_save,
}
cls._capability_cache[cache_key] = info
return info
@ -7332,6 +7360,26 @@ class LlamaCppBackend:
# when the binary advertises it (older/custom binaries may not).
if server_caps.get("supports_metrics"):
cmd.append("--metrics")
self._slot_save_dir = None
self._slot_save_binary = None
self._prompt_cache_disabled = False
if server_caps.get("supports_slot_save"):
try:
from utils.paths.storage_roots import ( # noqa: WPS433
llama_slot_cache_root,
)
slot_dir = llama_slot_cache_root()
slot_dir.mkdir(parents = True, exist_ok = True)
# Saved KV encodes chat content; keep it from other local users.
with contextlib.suppress(OSError):
os.chmod(slot_dir, 0o700)
cmd.extend(["--slot-save-path", str(slot_dir)])
self._slot_save_dir = str(slot_dir)
self._slot_save_binary = (binary, Path(binary).stat().st_mtime_ns)
except OSError:
self._slot_save_dir = None
self._slot_save_binary = None
cmd.extend(
self._ctx_integrity_flags(
n_parallel,
@ -7515,8 +7563,9 @@ class LlamaCppBackend:
else:
self._api_key = None
# Windows + full offload: disable KV checkpoints (WDDM/PCI-E
# overhead). CPU/partial offload keeps prompt caching. #5692.
# Windows + full offload: drop the host-RAM KV checkpoints that cause
# WDDM/PCI-E overhead, but keep prompt caching (in-VRAM prefix reuse) so
# a repeated prompt is not re-prefilled on every request. #5692.
if sys.platform == "win32" and full_offload_tuning_active:
unsupported_cache_flags: list[str] = []
if server_caps.get("supports_cache_ram"):
@ -7527,10 +7576,6 @@ class LlamaCppBackend:
cmd.extend(["--ctx-checkpoints", "0"])
else:
unsupported_cache_flags.append("--ctx-checkpoints")
if server_caps.get("supports_no_cache_prompt"):
cmd.append("--no-cache-prompt")
else:
unsupported_cache_flags.append("--no-cache-prompt")
if unsupported_cache_flags:
logger.info(
"Skipping unsupported Windows cache flags for llama-server: %s",
@ -8105,6 +8150,15 @@ class LlamaCppBackend:
if not self._healthy:
return False
# Snapshot the files the server actually loaded. If a GGUF shard or a
# LoRA/control-vector sidecar is swapped on disk afterwards while the
# old weights stay mapped, save_slots_for_resume() compares against
# this and refuses to persist KV that a reload could misapply.
if self._slot_save_dir:
self._slot_loaded_identity = (
self._gguf_file_identity(self._gguf_path),
self._slot_launch_fingerprint(),
)
return True
def _build_speculative_flags(
@ -8690,6 +8744,10 @@ class LlamaCppBackend:
self._effective_context_length = None
self._max_context_length = None
self._reset_effective_parallel_slots()
self._slot_save_dir = None
self._slot_save_binary = None
self._slot_loaded_identity = None
self._prompt_cache_disabled = False
self._chat_template = None
self._chat_template_override = None
self._supports_reasoning = False
@ -9216,6 +9274,237 @@ class LlamaCppBackend:
return False
return True
def _slot_launch_fingerprint(self) -> tuple:
# KV validity keys on extra args, stat'd sidecar weights, effective ctx.
sidecars = []
for path in self._sidecar_weight_files():
try:
st = os.stat(path)
sidecars.append((path, st.st_size, st.st_mtime_ns))
except OSError:
sidecars.append((path, None, None))
return (
tuple(self._extra_args or ()),
tuple(sidecars),
self._requested_n_ctx,
self._effective_context_length,
getattr(self, "_cache_type_kv", None),
self.effective_parallel_slots,
)
def _gguf_file_identity(self, path) -> Optional[tuple]:
# (size, mtime_ns) per shard: a split GGUF keys KV validity on every sibling.
p = Path(path)
paths = [p]
m = _SHARD_FULL_RE.match(p.name)
if m:
prefix, _first, total = m.groups()
paths = [
p.with_name(f"{prefix}-{i:05d}-of-{total}{p.suffix}")
for i in range(1, int(total) + 1)
]
try:
return tuple((sp.stat().st_size, sp.stat().st_mtime_ns) for sp in paths)
except OSError:
return None
_SIDECAR_WEIGHT_FLAGS = (
"--lora",
"--lora-scaled",
"--control-vector",
"--control-vector-scaled",
)
def _sidecar_weight_files(self) -> list[str]:
# llama.cpp: comma-separated paths, FNAME:SCALE on -scaled (older builds: FNAME SCALE).
args = [str(a).strip() for a in (self._extra_args or ())]
files: list[str] = []
for i, arg in enumerate(args):
flag, sep, inline = arg.partition("=")
if flag not in self._SIDECAR_WEIGHT_FLAGS:
continue
operand = inline if sep else (args[i + 1] if i + 1 < len(args) else "")
if not operand:
continue
candidates = [operand]
pieces = [p for p in operand.split(",") if p]
if len(pieces) > 1:
candidates.extend(pieces)
if flag.endswith("-scaled"):
for item in list(candidates):
# ":<number>" tail is a scale; rpartition spares drive letters.
head, colon, tail = item.rpartition(":")
if not (colon and head):
continue
try:
float(tail)
except ValueError:
continue
candidates.append(head)
for cand in candidates:
if cand not in files:
files.append(cand)
return files
def _prompt_cache_off(self) -> bool:
# Caching off makes restores useless; last prompt-cache flag wins, env only when unset.
last = None
for arg in self._extra_args or ():
flag = arg.strip().split("=", 1)[0]
if flag in ("--cache-prompt", "--no-cache-prompt"):
last = flag
if last is not None:
return last == "--no-cache-prompt"
if self._prompt_cache_disabled:
return True
if os.environ.get("LLAMA_ARG_NO_CACHE_PROMPT") is not None:
return True
env = (os.environ.get("LLAMA_ARG_CACHE_PROMPT") or "").strip().lower()
return env in {"off", "disabled", "false", "0"}
def save_slots_for_resume(
self, should_abort: Optional[Callable[[], bool]] = None
) -> Optional[dict]:
if (
not self.is_loaded
or not self._slot_save_dir
or not self._gguf_path
or self._prompt_cache_off()
):
return None
save_dir = Path(self._slot_save_dir)
gguf_stat = self._gguf_file_identity(self._gguf_path)
if gguf_stat is None:
return None
launch = self._slot_launch_fingerprint()
# If the GGUF or a sidecar was swapped on disk while the original weights
# stayed mapped, the live KV belongs to the old weights but a reload would
# load the new file. Persisting it would let restore misapply stale KV.
if self._slot_loaded_identity is not None and self._slot_loaded_identity != (
gguf_stat,
launch,
):
logger.debug("Skipping slot save: model files changed on disk since load")
return None
try:
estimate = self._estimate_kv_cache_bytes(
self._effective_context_length or self._context_length or 0,
self._cache_type_kv,
n_parallel = self.effective_parallel_slots,
)
# Skip before writing anything when the estimate alone blows the cap,
# rather than fully writing a slot and discarding it afterwards.
if estimate > _SLOT_SAVE_MAX_BYTES:
logger.debug(
"Skipping slot save: estimated %d bytes exceeds cap %d",
estimate,
_SLOT_SAVE_MAX_BYTES,
)
return None
# A 0 estimate means metadata was insufficient, not a zero-byte cache:
# a slot can still be many GiB, so demand room for the whole cap before
# trusting the post-write check.
required = (estimate if estimate > 0 else _SLOT_SAVE_MAX_BYTES) + (1 << 30)
if shutil.disk_usage(save_dir).free < required:
logger.debug("Skipping slot save: insufficient free disk")
return None
except Exception:
pass
token = uuid.uuid4().hex[:8]
entries: list[dict] = []
total_bytes = 0
for slot in range(self.effective_parallel_slots):
# A request pending mid-save waits on the gate; stop wasting its time.
if should_abort is not None and should_abort():
break
filename = f"resume-{token}-slot{slot}.bin"
path = save_dir / filename
try:
resp = httpx.post(
f"{self.base_url}/slots/{slot}",
params = {"action": "save"},
json = {"filename": filename},
headers = self._auth_headers,
timeout = _SLOT_SAVE_HTTP_TIMEOUT,
trust_env = False,
)
except Exception as e:
logger.debug(f"slot {slot} save failed: {e}")
with contextlib.suppress(OSError):
path.unlink()
break
if resp.status_code != 200:
logger.debug(f"slot {slot} save returned HTTP {resp.status_code}")
with contextlib.suppress(OSError):
path.unlink()
continue
try:
body = resp.json()
if not isinstance(body, dict):
raise ValueError("slot save response was not a JSON object")
n_saved = int(body.get("n_saved") or 0)
except Exception as e:
# A 200 that still wrote a file but returns a malformed body must
# clean up like the transport/HTTP error paths above, or the file
# (which holds chat KV) is orphaned until the next startup sweep.
logger.debug(f"slot {slot} save returned an invalid response: {e}")
with contextlib.suppress(OSError):
path.unlink()
continue
if n_saved <= 0:
with contextlib.suppress(OSError):
path.unlink()
continue
# Account by the bytes actually on disk, not the server-reported
# count, so the cap holds even if a custom binary under-reports.
try:
n_written = path.stat().st_size
except OSError:
n_written = 0
total_bytes += n_written
entries.append({"id": slot, "filename": filename, "n_saved": n_saved})
if total_bytes > _SLOT_SAVE_MAX_BYTES:
break # already over the cap; the discard below cleans up
if not entries:
return None
if total_bytes > _SLOT_SAVE_MAX_BYTES:
logger.debug(
"Discarding slot save: %d bytes exceeds cap %d",
total_bytes,
_SLOT_SAVE_MAX_BYTES,
)
for entry in entries:
with contextlib.suppress(OSError):
(save_dir / entry["filename"]).unlink()
return None
return {
"dir": self._slot_save_dir,
"binary": self._slot_save_binary,
"gguf": str(self._gguf_path),
"gguf_stat": gguf_stat,
"launch": launch,
"slots": entries,
}
def restore_slots_for_resume(self, manifest: dict) -> None:
if not self.is_loaded or not self._slot_save_dir:
return
for entry in manifest.get("slots") or []:
try:
resp = httpx.post(
f"{self.base_url}/slots/{int(entry['id'])}",
params = {"action": "restore"},
json = {"filename": str(entry["filename"])},
headers = self._auth_headers,
timeout = _SLOT_SAVE_HTTP_TIMEOUT,
trust_env = False,
)
except Exception as e:
logger.debug(f"slot restore failed: {e}")
break
if resp.status_code != 200:
logger.debug(f"slot {entry.get('id')} restore returned HTTP {resp.status_code}")
def _maybe_recover_from_mtp_crash(self, exc: Optional[BaseException] = None) -> bool:
"""Schedule one background reload without MTP after a mid-generation death.
@ -9614,6 +9903,75 @@ class LlamaCppBackend:
except Exception:
logger.debug("Could not close httpx client", exc_info = True)
@staticmethod
def _install_cancel_aware_read(
client: "httpx.Client",
cancel_event: threading.Event,
response: Optional["httpx.Response"] = None,
poll_s: float = 0.2,
) -> None:
"""Wrap the httpcore stream so the reader interrupts its own blocked recv() on cancel.
A cross-thread socket shutdown wakes a parked recv() on POSIX but not on
Windows (Winsock), so read in short slices and poll cancel_event between them
(plain or TLS); slice timeouts are swallowed so a slow-but-alive stream survives.
httpcore snapshots request.extensions["timeout"]["read"] once at body start, so
given ``response`` we re-read the live value per call to honor the post-first-token
stall timeout instead of the long prefill timeout."""
import httpcore
def _live_read_timeout() -> Optional[float]:
if response is None:
return None
try:
ext = response.request.extensions.get("timeout")
if isinstance(ext, dict):
value = ext.get("read")
if isinstance(value, (int, float)):
return float(value)
except Exception:
pass
return None
try:
pool = getattr(getattr(client, "_transport", None), "_pool", None)
for connection in list(getattr(pool, "_connections", []) or []):
inner = getattr(connection, "_connection", None)
stream = getattr(inner, "_network_stream", None)
if stream is None or getattr(stream, "_unsloth_cancel_wrapped", False):
continue
orig_read = stream.read
def read(
max_bytes,
timeout = None,
_orig = orig_read,
):
live = _live_read_timeout()
effective = live if live is not None else timeout
deadline = None if effective is None else time.monotonic() + effective
while True:
if cancel_event.is_set():
raise httpcore.ReadError("stream cancelled by user")
if deadline is None:
step = poll_s
else:
remaining = deadline - time.monotonic()
if remaining <= 0:
raise httpcore.ReadTimeout("read operation timed out")
step = min(poll_s, remaining)
try:
return _orig(max_bytes, timeout = step)
except httpcore.ReadTimeout:
if deadline is not None and time.monotonic() >= deadline:
raise
continue # slow but alive: keep reading
stream.read = read
stream._unsloth_cancel_wrapped = True
except Exception:
logger.debug("Could not install cancel-aware read", exc_info = True)
@staticmethod
@contextlib.contextmanager
def _stream_with_retry(
@ -9671,6 +10029,11 @@ class LlamaCppBackend:
headers = headers,
) as response:
_response_ref[0] = response
if cancel_event is not None:
# Portable mid-stream cancel: the reader polls cancel itself, so
# Stop interrupts a stalled read where the watcher's Windows socket
# shutdown does not. Pass response to honor the live stall timeout.
LlamaCppBackend._install_cancel_aware_read(client, cancel_event, response)
if cancel_event is not None and cancel_event.is_set():
raise _LlamaStreamCancelled
yield response

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:

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

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

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

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

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

13
studio/backend/main.py
View file

@ -312,6 +312,7 @@ 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 hub.schemas.downloads import TransportCapabilities
from hub.utils.download_registry import (
@ -547,8 +548,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).
@ -992,6 +994,7 @@ 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(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,11 +1151,15 @@ 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)

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

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

@ -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
@ -4710,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
@ -5244,6 +5244,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'",

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

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

@ -7,6 +7,7 @@ Like auth/storage.py (module-level functions, raw sqlite3, per-function
connections) plus WAL mode and PRAGMA foreign_keys = ON for CASCADE deletes.
"""
import hashlib
import json
import logging
import os
@ -100,6 +101,7 @@ _schema_lock = threading.Lock()
_schema_ready = False
_SQLITE_IN_CHUNK_SIZE = 900
_PROJECT_WORKSPACE_SUBDIRS = ("sandbox",)
_CHAT_ATTACHMENT_INVENTORY_VERSION = 1
def _project_slug(name: str) -> str:
@ -313,6 +315,141 @@ def _ensure_schema(conn: sqlite3.Connection) -> None:
)
"""
)
tombstone_schema = """
CREATE TABLE chat_attachment_tombstones (
thread_id TEXT NOT NULL REFERENCES chat_threads(id) ON DELETE CASCADE,
message_id TEXT NOT NULL,
attachment_id TEXT NOT NULL,
deleted_at INTEGER NOT NULL,
PRIMARY KEY(thread_id, message_id, attachment_id)
) WITHOUT ROWID
"""
tombstone_table = conn.execute(
"""
SELECT 1 FROM sqlite_master
WHERE type = 'table' AND name = 'chat_attachment_tombstones'
"""
).fetchone()
if tombstone_table is None:
conn.execute(tombstone_schema)
else:
tombstone_columns = {
row[1] for row in conn.execute("PRAGMA table_info(chat_attachment_tombstones)")
}
tombstone_fk_targets = {
row[2] for row in conn.execute("PRAGMA foreign_key_list(chat_attachment_tombstones)")
}
if "thread_id" not in tombstone_columns or "chat_threads" not in tombstone_fk_targets:
# The first implementation cascaded through chat_messages, which
# erased deletion knowledge during pruneMissing. Rebuild once,
# retaining every tombstone whose owning thread still exists.
conn.execute("SAVEPOINT migrate_chat_attachment_tombstones")
try:
conn.execute(
"ALTER TABLE chat_attachment_tombstones "
"RENAME TO chat_attachment_tombstones_legacy"
)
conn.execute(tombstone_schema)
if "thread_id" in tombstone_columns:
conn.execute(
"""
INSERT OR IGNORE INTO chat_attachment_tombstones
(thread_id, message_id, attachment_id, deleted_at)
SELECT legacy.thread_id, legacy.message_id,
legacy.attachment_id, legacy.deleted_at
FROM chat_attachment_tombstones_legacy legacy
JOIN chat_threads thread ON thread.id = legacy.thread_id
"""
)
else:
conn.execute(
"""
INSERT OR IGNORE INTO chat_attachment_tombstones
(thread_id, message_id, attachment_id, deleted_at)
SELECT message.thread_id, legacy.message_id,
legacy.attachment_id, legacy.deleted_at
FROM chat_attachment_tombstones_legacy legacy
JOIN chat_messages message ON message.id = legacy.message_id
"""
)
conn.execute("DROP TABLE chat_attachment_tombstones_legacy")
conn.execute("RELEASE SAVEPOINT migrate_chat_attachment_tombstones")
except Exception:
conn.execute("ROLLBACK TO SAVEPOINT migrate_chat_attachment_tombstones")
conn.execute("RELEASE SAVEPOINT migrate_chat_attachment_tombstones")
raise
conn.execute(
"""
CREATE TABLE IF NOT EXISTS chat_attachment_inventory (
message_id TEXT NOT NULL REFERENCES chat_messages(id) ON DELETE CASCADE,
attachment_id TEXT NOT NULL,
name TEXT NOT NULL,
type TEXT,
content_type TEXT,
size_bytes INTEGER,
PRIMARY KEY(message_id, attachment_id)
) WITHOUT ROWID
"""
)
conn.execute(
"""
CREATE TABLE IF NOT EXISTS chat_attachment_inventory_state (
singleton INTEGER NOT NULL PRIMARY KEY CHECK(singleton = 1),
inventory_version INTEGER NOT NULL DEFAULT 0,
dirty INTEGER NOT NULL DEFAULT 1,
backfilled_at INTEGER NOT NULL
)
"""
)
inventory_state_columns = {
row[1] for row in conn.execute("PRAGMA table_info(chat_attachment_inventory_state)")
}
if "inventory_version" not in inventory_state_columns:
conn.execute(
"ALTER TABLE chat_attachment_inventory_state "
"ADD COLUMN inventory_version INTEGER NOT NULL DEFAULT 0"
)
if "dirty" not in inventory_state_columns:
conn.execute(
"ALTER TABLE chat_attachment_inventory_state "
"ADD COLUMN dirty INTEGER NOT NULL DEFAULT 1"
)
conn.execute(
"""
CREATE TRIGGER IF NOT EXISTS chat_attachment_inventory_dirty_insert
AFTER INSERT ON chat_messages
BEGIN
INSERT INTO chat_attachment_inventory_state
(singleton, inventory_version, dirty, backfilled_at)
VALUES (1, 0, 1, 0)
ON CONFLICT(singleton) DO UPDATE SET dirty = 1;
END
"""
)
conn.execute(
"""
CREATE TRIGGER IF NOT EXISTS chat_attachment_inventory_dirty_update
AFTER UPDATE ON chat_messages
BEGIN
INSERT INTO chat_attachment_inventory_state
(singleton, inventory_version, dirty, backfilled_at)
VALUES (1, 0, 1, 0)
ON CONFLICT(singleton) DO UPDATE SET dirty = 1;
END
"""
)
conn.execute(
"""
CREATE TRIGGER IF NOT EXISTS chat_attachment_inventory_dirty_delete
AFTER DELETE ON chat_messages
BEGIN
INSERT INTO chat_attachment_inventory_state
(singleton, inventory_version, dirty, backfilled_at)
VALUES (1, 0, 1, 0)
ON CONFLICT(singleton) DO UPDATE SET dirty = 1;
END
"""
)
conn.execute(
"CREATE INDEX IF NOT EXISTS idx_chat_threads_model_type_created_at ON chat_threads(model_type, created_at)"
)
@ -391,6 +528,21 @@ def _ensure_schema(conn: sqlite3.Connection) -> None:
conn.execute(
"CREATE INDEX IF NOT EXISTS idx_prompt_lists_created_at ON prompt_lists(created_at)"
)
inventory_state = conn.execute(
"""
SELECT inventory_version, dirty
FROM chat_attachment_inventory_state
WHERE singleton = 1
"""
).fetchone()
if (
inventory_state is None
or inventory_state["inventory_version"] != _CHAT_ATTACHMENT_INVENTORY_VERSION
or inventory_state["dirty"]
):
_rebuild_chat_attachment_inventory(conn)
_mark_chat_attachment_inventory_clean(conn)
conn.commit()
def _prompt_entry_from_row(row: sqlite3.Row) -> dict:
@ -1219,7 +1371,14 @@ def delete_chat_threads(ids: list[str]) -> None:
return
conn = get_connection()
try:
conn.execute("BEGIN IMMEDIATE")
_ensure_chat_attachment_inventory_current(conn)
conn.executemany(
"DELETE FROM chat_attachment_tombstones WHERE thread_id = ?",
[(id,) for id in ids],
)
conn.executemany("DELETE FROM chat_threads WHERE id = ?", [(id,) for id in ids])
_mark_chat_attachment_inventory_clean(conn)
conn.commit()
finally:
conn.close()
@ -1228,7 +1387,11 @@ def delete_chat_threads(ids: list[str]) -> None:
def clear_chat_history() -> None:
conn = get_connection()
try:
conn.execute("BEGIN IMMEDIATE")
_ensure_chat_attachment_inventory_current(conn)
conn.execute("DELETE FROM chat_attachment_tombstones")
conn.execute("DELETE FROM chat_threads")
_mark_chat_attachment_inventory_clean(conn)
conn.commit()
finally:
conn.close()
@ -1354,6 +1517,7 @@ def delete_chat_project(id: str, delete_files: bool = False) -> Optional[dict]:
conn = get_connection()
try:
conn.execute("BEGIN IMMEDIATE")
_ensure_chat_attachment_inventory_current(conn)
row = conn.execute("SELECT * FROM chat_projects WHERE id = ?", (id,)).fetchone()
if row is None:
conn.rollback()
@ -1361,6 +1525,7 @@ def delete_chat_project(id: str, delete_files: bool = False) -> Optional[dict]:
project = _chat_project_from_row(row)
conn.execute("DELETE FROM chat_threads WHERE project_id = ?", (id,))
conn.execute("DELETE FROM chat_projects WHERE id = ?", (id,))
_mark_chat_attachment_inventory_clean(conn)
conn.commit()
if delete_files:
_delete_project_workspace(project)
@ -1483,15 +1648,285 @@ def _recompute_chat_thread_updated_at(conn: sqlite3.Connection, thread_id: str)
)
_CONTENT_PART_ID_PREFIX = "content-part-sha256-"
_URI_SCHEME_RE = re.compile(r"^[A-Za-z][A-Za-z0-9+.-]*:")
def _is_locally_stored_blob(value: str) -> bool:
"""True for data URIs or bare base64, never external/blob URI references."""
candidate = value.lstrip()
if not candidate:
return False
if candidate[:5].lower() == "data:":
return True
if candidate.startswith(("//", "\\\\")):
return False
return _URI_SCHEME_RE.match(candidate) is None
def _managed_content_part_payload(part: dict) -> Optional[tuple[str, Any]]:
"""Return the locally stored blob payload used to identify a content part."""
image = part.get("image")
if isinstance(image, str) and image[:5].lower() == "data:":
return "image", image
audio = part.get("audio")
if isinstance(audio, str) and _is_locally_stored_blob(audio):
return "audio", audio
if isinstance(audio, dict):
data = audio.get("data")
if isinstance(data, str) and _is_locally_stored_blob(data):
return "audio", audio
return None
def _content_part_id(part: dict) -> Optional[str]:
"""Stable managed id derived from blob data, without mutating inference content."""
payload = _managed_content_part_payload(part)
if payload is None:
return None
canonical = json.dumps(
payload,
ensure_ascii = False,
separators = (",", ":"),
sort_keys = True,
).encode("utf-8")
return f"{_CONTENT_PART_ID_PREFIX}{hashlib.sha256(canonical).hexdigest()}"
def _chat_attachment_tombstones_for_messages(
conn: sqlite3.Connection, thread_id: str, message_ids: list[str]
) -> dict[str, set[str]]:
tombstones = {message_id: set() for message_id in message_ids}
unique_ids = list(dict.fromkeys(message_ids))
for start in range(0, len(unique_ids), _SQLITE_IN_CHUNK_SIZE):
chunk = unique_ids[start : start + _SQLITE_IN_CHUNK_SIZE]
placeholders = ",".join("?" for _ in chunk)
rows = conn.execute(
f"""
SELECT message_id, attachment_id
FROM chat_attachment_tombstones
WHERE thread_id = ? AND message_id IN ({placeholders})
""",
(thread_id, *chunk),
).fetchall()
for row in rows:
tombstones[row["message_id"]].add(row["attachment_id"])
return tombstones
def _reconcile_chat_message_uploads(message: dict, tombstones: set[str]) -> dict:
"""Strip uploads previously deleted through the Data tab from a stale write."""
if not tombstones:
return message
reconciled = dict(message)
attachments = message.get("attachments")
if isinstance(attachments, list):
reconciled["attachments"] = [
attachment
for attachment in attachments
if not (isinstance(attachment, dict) and str(attachment.get("id") or "") in tombstones)
]
content = message.get("content")
if isinstance(content, list):
reconciled["content"] = [
part
for part in content
if not (isinstance(part, dict) and (_content_part_id(part) or "") in tombstones)
]
return reconciled
def _chat_attachment_metadata_text(value, fallback: Optional[str] = None) -> Optional[str]:
"""Keep untyped legacy/import metadata safe for SQLite binding."""
if value is None:
return fallback
if isinstance(value, str):
return value or fallback
if isinstance(value, (bool, int, float)):
return str(value)
# Objects and arrays are not useful display metadata and sqlite3 rejects
# binding them directly.
return fallback
def _chat_attachment_inventory_entries(
attachments_json: Optional[str],
content_json: Optional[str],
tombstones: Optional[set[str]] = None,
) -> list[dict]:
tombstones = tombstones or set()
attachments = _json_loads(attachments_json, None)
if not isinstance(attachments, list):
attachments = []
attachments = [
attachment
for attachment in attachments
if isinstance(attachment, dict) and attachment.get("id")
]
attachments.extend(_content_part_attachments(content_json))
entries: list[dict] = []
seen: set[str] = set()
for attachment in attachments:
attachment_id = str(attachment["id"])
if attachment_id in seen or attachment_id in tombstones:
continue
seen.add(attachment_id)
entries.append(
{
"id": attachment_id,
"name": _chat_attachment_metadata_text(attachment.get("name"), "attachment"),
"type": _chat_attachment_metadata_text(attachment.get("type")),
"contentType": _chat_attachment_metadata_text(attachment.get("contentType")),
"sizeBytes": _chat_attachment_size_bytes(attachment),
}
)
return entries
def _replace_chat_attachment_inventory(
conn: sqlite3.Connection,
message_id: str,
attachments_json: Optional[str],
content_json: Optional[str],
tombstones: Optional[set[str]] = None,
) -> None:
conn.execute("DELETE FROM chat_attachment_inventory WHERE message_id = ?", (message_id,))
entries = _chat_attachment_inventory_entries(
attachments_json,
content_json,
tombstones,
)
conn.executemany(
"""
INSERT INTO chat_attachment_inventory
(message_id, attachment_id, name, type, content_type, size_bytes)
VALUES (?, ?, ?, ?, ?, ?)
""",
[
(
message_id,
entry["id"],
entry["name"],
entry["type"],
entry["contentType"],
entry["sizeBytes"],
)
for entry in entries
],
)
def _mark_chat_attachment_inventory_clean(conn: sqlite3.Connection) -> None:
conn.execute(
"""
INSERT INTO chat_attachment_inventory_state
(singleton, inventory_version, dirty, backfilled_at)
VALUES (1, ?, 0, ?)
ON CONFLICT(singleton) DO UPDATE SET
inventory_version = excluded.inventory_version,
dirty = 0,
backfilled_at = excluded.backfilled_at
""",
(
_CHAT_ATTACHMENT_INVENTORY_VERSION,
int(datetime.now(timezone.utc).timestamp() * 1000),
),
)
def _rebuild_chat_attachment_inventory(conn: sqlite3.Connection) -> None:
"""Rebuild after schema upgrade or a write from an older Studio build."""
conn.execute("DELETE FROM chat_attachment_inventory")
tombstones: dict[tuple[str, str], set[str]] = {}
for row in conn.execute(
"SELECT thread_id, message_id, attachment_id FROM chat_attachment_tombstones"
).fetchall():
tombstones.setdefault((row["thread_id"], row["message_id"]), set()).add(
row["attachment_id"]
)
rows = conn.execute(
"SELECT id, thread_id, attachments_json, content_json FROM chat_messages"
).fetchall()
for row in rows:
_replace_chat_attachment_inventory(
conn,
row["id"],
row["attachments_json"],
row["content_json"],
tombstones.get((row["thread_id"], row["id"]), set()),
)
def _ensure_chat_attachment_inventory_current(conn: sqlite3.Connection) -> None:
state = conn.execute(
"""
SELECT inventory_version, dirty
FROM chat_attachment_inventory_state
WHERE singleton = 1
"""
).fetchone()
if (
state is not None
and state["inventory_version"] == _CHAT_ATTACHMENT_INVENTORY_VERSION
and not state["dirty"]
):
return
owns_transaction = not conn.in_transaction
if owns_transaction:
conn.execute("BEGIN IMMEDIATE")
try:
state = conn.execute(
"""
SELECT inventory_version, dirty
FROM chat_attachment_inventory_state
WHERE singleton = 1
"""
).fetchone()
if (
state is None
or state["inventory_version"] != _CHAT_ATTACHMENT_INVENTORY_VERSION
or state["dirty"]
):
_rebuild_chat_attachment_inventory(conn)
_mark_chat_attachment_inventory_clean(conn)
if owns_transaction:
conn.commit()
except Exception:
if owns_transaction:
conn.rollback()
raise
def upsert_chat_message(message: dict) -> dict:
conn = get_connection()
try:
conn.execute("BEGIN IMMEDIATE")
_ensure_chat_attachment_inventory_current(conn)
_raise_if_chat_message_thread_conflicts(
conn,
message["threadId"],
[message["id"]],
)
tombstones = _chat_attachment_tombstones_for_messages(
conn,
message["threadId"],
[message["id"]],
)
reconciled = _reconcile_chat_message_uploads(
message,
tombstones.get(message["id"], set()),
)
content_json = json.dumps(reconciled.get("content", []))
attachments_json = (
json.dumps(reconciled.get("attachments"))
if reconciled.get("attachments") is not None
else None
)
conn.execute(
"""
INSERT INTO chat_messages
@ -1507,23 +1942,32 @@ def upsert_chat_message(message: dict) -> dict:
WHERE excluded.thread_id = chat_messages.thread_id
""",
(
message["id"],
message["threadId"],
message.get("parentId"),
message["role"],
json.dumps(message.get("content", [])),
json.dumps(message.get("attachments"))
if message.get("attachments") is not None
reconciled["id"],
reconciled["threadId"],
reconciled.get("parentId"),
reconciled["role"],
content_json,
attachments_json,
json.dumps(reconciled.get("metadata"))
if reconciled.get("metadata") is not None
else None,
json.dumps(message.get("metadata"))
if message.get("metadata") is not None
else None,
int(message["createdAt"]),
int(reconciled["createdAt"]),
),
)
_bump_chat_thread_updated_at(conn, message["threadId"], int(message["createdAt"]))
_replace_chat_attachment_inventory(
conn,
reconciled["id"],
attachments_json,
content_json,
)
_bump_chat_thread_updated_at(
conn,
reconciled["threadId"],
int(reconciled["createdAt"]),
)
_mark_chat_attachment_inventory_clean(conn)
conn.commit()
return message
return reconciled
except Exception:
conn.rollback()
raise
@ -1539,13 +1983,28 @@ def sync_chat_messages(
conn = get_connection()
try:
conn.execute("BEGIN IMMEDIATE")
_ensure_chat_attachment_inventory_current(conn)
_raise_if_chat_message_thread_conflicts(
conn,
thread_id,
[m["id"] for m in messages],
)
if prune_missing:
conn.execute("DELETE FROM chat_messages WHERE thread_id = ?", (thread_id,))
tombstones = _chat_attachment_tombstones_for_messages(
conn,
thread_id,
[m["id"] for m in messages],
)
reconciled_messages = [
_reconcile_chat_message_uploads(m, tombstones.get(m["id"], set())) for m in messages
]
serialized_messages = [
(
m,
json.dumps(m.get("content", [])),
json.dumps(m.get("attachments")) if m.get("attachments") is not None else None,
)
for m in reconciled_messages
]
conn.executemany(
"""
INSERT INTO chat_messages
@ -1566,20 +2025,46 @@ def sync_chat_messages(
thread_id,
m.get("parentId"),
m["role"],
json.dumps(m.get("content", [])),
json.dumps(m.get("attachments")) if m.get("attachments") is not None else None,
content_json,
attachments_json,
json.dumps(m.get("metadata")) if m.get("metadata") is not None else None,
int(m["createdAt"]),
)
for m in messages
for m, content_json, attachments_json in serialized_messages
],
)
if prune_missing:
_recompute_chat_thread_updated_at(conn, thread_id)
elif messages:
_bump_chat_thread_updated_at(
conn, thread_id, max(int(m["createdAt"]) for m in messages)
for m, content_json, attachments_json in serialized_messages:
_replace_chat_attachment_inventory(
conn,
m["id"],
attachments_json,
content_json,
)
if prune_missing:
retained_ids = {m["id"] for m in reconciled_messages}
existing_ids = {
row["id"]
for row in conn.execute(
"SELECT id FROM chat_messages WHERE thread_id = ?",
(thread_id,),
).fetchall()
}
missing_ids = sorted(existing_ids - retained_ids)
for start in range(0, len(missing_ids), _SQLITE_IN_CHUNK_SIZE):
chunk = missing_ids[start : start + _SQLITE_IN_CHUNK_SIZE]
placeholders = ",".join("?" for _ in chunk)
conn.execute(
f"DELETE FROM chat_messages WHERE thread_id = ? AND id IN ({placeholders})",
(thread_id, *chunk),
)
_recompute_chat_thread_updated_at(conn, thread_id)
elif reconciled_messages:
_bump_chat_thread_updated_at(
conn,
thread_id,
max(int(m["createdAt"]) for m in reconciled_messages),
)
_mark_chat_attachment_inventory_clean(conn)
conn.commit()
return list_chat_messages(thread_id)
except ChatMessageConflictError:
@ -1613,6 +2098,7 @@ def fork_chat_thread(
conn = get_connection()
try:
conn.execute("BEGIN IMMEDIATE")
_ensure_chat_attachment_inventory_current(conn)
src = conn.execute(
"SELECT * FROM chat_threads WHERE id = ?", (source_thread_id,)
).fetchone()
@ -1686,6 +2172,14 @@ def fork_chat_thread(
for row in ancestry
],
)
for row in ancestry:
_replace_chat_attachment_inventory(
conn,
id_map[row["id"]],
row["attachments_json"],
row["content_json"],
)
_mark_chat_attachment_inventory_clean(conn)
conn.commit()
thread_row = conn.execute(
"SELECT * FROM chat_threads WHERE id = ?", (new_thread_id,)
@ -1744,6 +2238,279 @@ def get_chat_message(thread_id: str, message_id: str) -> Optional[dict]:
conn.close()
def _blob_part_base64_len(part: dict) -> int:
"""Base64 payload length of an image or audio content part, or 0."""
image = part.get("image")
if isinstance(image, str) and image[:5].lower() == "data:":
return len(image.rsplit(",", 1)[-1])
audio = part.get("audio")
if isinstance(audio, str) and _is_locally_stored_blob(audio):
return len(audio.rsplit(",", 1)[-1])
if isinstance(audio, dict):
data = audio.get("data")
if isinstance(data, str) and _is_locally_stored_blob(data):
return len(data)
return 0
def _chat_attachment_size_bytes(attachment: dict) -> Optional[int]:
"""Approximate stored size of one attachment's content parts.
Image and audio parts hold base64 payloads (decoded bytes ~= 3/4 of the
encoded length); text parts count their character length. None when there
is no sizable content (e.g. a stripped/legacy attachment).
"""
total = 0
found = False
for part in attachment.get("content") or []:
if not isinstance(part, dict):
continue
blob_len = _blob_part_base64_len(part)
if blob_len > 0:
total += (blob_len * 3) // 4
found = True
continue
text = part.get("text")
if isinstance(text, str) and text:
total += len(text.encode("utf-8", errors = "ignore"))
found = True
return total if found else None
def _content_part_attachments(content_json: Optional[str]) -> list[dict]:
"""Managed local blobs stored in content_json, with stable payload ids.
Exact duplicate blobs intentionally share one inventory id. Deleting that
id removes every identical copy, avoiding ambiguous index-based addressing.
"""
content = _json_loads(content_json, None)
if not isinstance(content, list):
return []
out: list[dict] = []
seen: set[str] = set()
for part in content:
if not isinstance(part, dict):
continue
attachment_id = _content_part_id(part)
payload = _managed_content_part_payload(part)
if attachment_id is None or payload is None or attachment_id in seen:
continue
seen.add(attachment_id)
kind, value = payload
content_type = None
if kind == "image" and isinstance(value, str):
content_type = value[5:].split(";", 1)[0].split(",", 1)[0] or None
out.append(
{
"id": attachment_id,
"type": kind,
"name": "Chat image" if kind == "image" else "Chat audio",
"contentType": content_type,
"content": [part],
}
)
return out
def list_chat_attachments_page(
limit: int = 50, offset: int = 0
) -> tuple[list[dict], Optional[int]]:
"""One bounded page from the normalized attachment inventory."""
if not 1 <= limit <= 100:
raise ValueError("limit must be between 1 and 100")
if offset < 0:
raise ValueError("offset must be non-negative")
conn = get_connection()
try:
_ensure_chat_attachment_inventory_current(conn)
rows = conn.execute(
"""
SELECT i.attachment_id, i.name, i.type, i.content_type,
i.size_bytes, m.id AS message_id, m.thread_id,
m.created_at, t.title AS thread_title, t.pair_id
FROM chat_attachment_inventory i
JOIN chat_messages m ON m.id = i.message_id
LEFT JOIN chat_threads t ON t.id = m.thread_id
ORDER BY m.created_at DESC, m.id ASC, i.attachment_id ASC
LIMIT ? OFFSET ?
""",
(limit + 1, offset),
).fetchall()
finally:
conn.close()
has_more = len(rows) > limit
page_rows = rows[:limit]
attachments = [
{
"id": row["attachment_id"],
"messageId": row["message_id"],
"threadId": row["thread_id"],
"pairId": row["pair_id"],
"threadTitle": row["thread_title"],
"name": row["name"],
"type": row["type"],
"contentType": row["content_type"],
"sizeBytes": row["size_bytes"],
"createdAt": row["created_at"],
}
for row in page_rows
]
return attachments, offset + limit if has_more else None
def list_chat_attachments() -> list[dict]:
"""Compatibility helper returning the full normalized inventory."""
attachments: list[dict] = []
offset = 0
while True:
page, next_offset = list_chat_attachments_page(limit = 100, offset = offset)
attachments.extend(page)
if next_offset is None:
return attachments
offset = next_offset
def get_chat_attachment(message_id: str, attachment_id: str) -> Optional[dict]:
"""One attachment record (full content) from a message, or None."""
conn = get_connection()
try:
row = conn.execute(
"""
SELECT message.attachments_json, message.content_json,
EXISTS(
SELECT 1 FROM chat_attachment_tombstones tombstone
WHERE tombstone.thread_id = message.thread_id
AND tombstone.message_id = message.id
AND tombstone.attachment_id = ?
) AS tombstoned
FROM chat_messages message
WHERE message.id = ?
""",
(attachment_id, message_id),
).fetchone()
finally:
conn.close()
if row is None or row["tombstoned"]:
return None
attachments = _json_loads(row["attachments_json"], None)
if isinstance(attachments, list):
for attachment in attachments:
if isinstance(attachment, dict) and str(attachment.get("id") or "") == attachment_id:
return attachment
if attachment_id.startswith(_CONTENT_PART_ID_PREFIX):
for attachment in _content_part_attachments(row["content_json"]):
if attachment["id"] == attachment_id:
return attachment
return None
def _record_chat_attachment_tombstone(
conn: sqlite3.Connection, thread_id: str, message_id: str, attachment_id: str
) -> None:
conn.execute(
"""
INSERT INTO chat_attachment_tombstones
(thread_id, message_id, attachment_id, deleted_at)
VALUES (?, ?, ?, ?)
ON CONFLICT(thread_id, message_id, attachment_id) DO UPDATE SET
deleted_at = excluded.deleted_at
""",
(
thread_id,
message_id,
attachment_id,
int(datetime.now(timezone.utc).timestamp() * 1000),
),
)
def delete_chat_attachment(message_id: str, attachment_id: str) -> bool:
"""Remove one stored upload from a message.
The tombstone is retained while the thread exists, so pruning and later
recreating the same message id cannot restore the deleted upload. If an
ordinary attachment id collides with a content-blob id, both are deleted as
one managed item.
"""
conn = get_connection()
try:
conn.execute("BEGIN IMMEDIATE")
_ensure_chat_attachment_inventory_current(conn)
row = conn.execute(
"""
SELECT thread_id, attachments_json, content_json
FROM chat_messages WHERE id = ?
""",
(message_id,),
).fetchone()
if row is None:
conn.rollback()
return False
attachments = _json_loads(row["attachments_json"], None)
updated_attachments_json = row["attachments_json"]
deleted_attachment = False
if isinstance(attachments, list):
remaining_attachments = [
attachment
for attachment in attachments
if not (
isinstance(attachment, dict)
and str(attachment.get("id") or "") == attachment_id
)
]
deleted_attachment = len(remaining_attachments) != len(attachments)
if deleted_attachment:
updated_attachments_json = json.dumps(remaining_attachments)
content = _json_loads(row["content_json"], None)
updated_content_json = row["content_json"]
deleted_content = False
if attachment_id.startswith(_CONTENT_PART_ID_PREFIX) and isinstance(content, list):
remaining_content = [
part
for part in content
if not (isinstance(part, dict) and _content_part_id(part) == attachment_id)
]
deleted_content = len(remaining_content) != len(content)
if deleted_content:
updated_content_json = json.dumps(remaining_content)
if not deleted_attachment and not deleted_content:
conn.rollback()
return False
conn.execute(
"""
UPDATE chat_messages
SET attachments_json = ?, content_json = ?
WHERE id = ?
""",
(updated_attachments_json, updated_content_json, message_id),
)
_record_chat_attachment_tombstone(
conn,
row["thread_id"],
message_id,
attachment_id,
)
_replace_chat_attachment_inventory(
conn,
message_id,
updated_attachments_json,
updated_content_json,
)
_mark_chat_attachment_inventory_clean(conn)
conn.commit()
return True
except Exception:
conn.rollback()
raise
finally:
conn.close()
def list_chat_messages_for_threads(thread_ids: list[str]) -> list[dict]:
if not thread_ids:
return []

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"

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 == [

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

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

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

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

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

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.

2
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";
@ -230,6 +231,7 @@ function RootLayout() {
<AppProvider>
<PersonalizationSyncMount />
{!isAuthFlowRoute && <SettingsDialog />}
<HfTokenWarningDialog />
<RemoteCodeConsentDialog />
<TransformersUpgradeDialog />
{hideNavbar ? (

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

@ -969,10 +969,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…

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>

809
studio/frontend/src/components/assistant-ui/model-selector/pickers.tsx
View file

File diff suppressed because it is too large Load diff

72
studio/frontend/src/components/assistant-ui/model-selector/pinned-models.ts Normal file
View file

@ -0,0 +1,72 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
// Pinned models for the model selector's On Device list, persisted in
// localStorage so pins survive reloads. GGUF quants pin individually
// (repoId + quant); non-GGUF repos pin as a whole. Pinned entries surface
// in a "Pinned" section above the Unsloth/Downloaded group.
import { create } from "zustand";
const KEY = "unsloth_pinned_models";
// Entries are stored as strings: "repoId" pins a whole (non-GGUF) repo,
// "repoId::quant" pins one GGUF quant. Neither part contains "::".
export function pinKey(repoId: string, quant?: string): string {
return quant ? `${repoId}::${quant}` : repoId;
}
export interface PinnedQuantEntry {
repoId: string;
quant: string;
}
/** The pinned GGUF quants, in pin order. Plain repo pins are excluded. */
export function pinnedQuantEntries(pinned: string[]): PinnedQuantEntry[] {
const out: PinnedQuantEntry[] = [];
for (const key of pinned) {
const sep = key.indexOf("::");
if (sep <= 0) continue;
const repoId = key.slice(0, sep);
const quant = key.slice(sep + 2);
if (repoId && quant) out.push({ repoId, quant });
}
return out;
}
function readPinned(): string[] {
try {
const raw = JSON.parse(localStorage.getItem(KEY) ?? "[]");
return Array.isArray(raw)
? raw.filter((v): v is string => typeof v === "string")
: [];
} catch {
return [];
}
}
function writePinned(pinned: string[]): void {
try {
localStorage.setItem(KEY, JSON.stringify(pinned));
} catch {
// Ignore unavailable storage; pins stay session-only.
}
}
interface PinnedModelsState {
pinned: string[];
togglePinned: (repoId: string, quant?: string) => void;
}
export const usePinnedModelsStore = create<PinnedModelsState>((set) => ({
pinned: readPinned(),
togglePinned: (repoId, quant) =>
set((state) => {
const key = pinKey(repoId, quant);
const next = state.pinned.includes(key)
? state.pinned.filter((id) => id !== key)
: [...state.pinned, key];
writePinned(next);
return { pinned: next };
}),
}));

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

@ -1434,13 +1434,10 @@ 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.
// More than 4 pills: collapse to icons only. Search, Code, and permissions
// always show; Images, RAG, Canvas and MCP are conditional.
const pillsCompact =
2 +
(permissionMode !== "off" ? 1 : 0) +
3 +
(ragEnabled ? 1 : 0) +
(supportsBuiltinImageGeneration ? 1 : 0) +
(artifactsEnabled ? 1 : 0) +
@ -1556,20 +1553,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.
@ -1856,27 +1839,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

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)}
/>

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],
);

96
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 {
@ -104,11 +109,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,13 +127,15 @@ 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,
// Intended load settings so validate's preflight matches the follow-up
// /load. Default placement is sized against the selected GPUs.
@ -324,13 +334,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;
}
@ -345,9 +359,11 @@ 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;
@ -437,6 +453,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> {
@ -960,7 +1043,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>

28
studio/frontend/src/features/chat/chat-settings-sheet.tsx
View file

@ -723,11 +723,11 @@ export function ChatSettingsPanel({
const autoLayers = isManual && gpuLayers < 0;
// GPUs actually in use: the picked subset, or all visible when none picked.
const gpusInUse = selectedGpuIds ?? gpuDevices.map((d) => d.index);
// TP is off with fewer than 2 GPUs in use (single GPU, or the picker narrowed
// to one): tensor split is a no-op there and aborts on some archs. Mirrors the
// multi-GPU gate on the GPU picker / Split ratio. (Under Auto layers the whole
// TP control is hidden -- llama.cpp's --fit aborts under --split-mode tensor.)
const tpDisabled = gpusInUse.length <= 1;
// The picker must keep one GPU selected.
const singleGpuInUse = gpusInUse.length <= 1;
// TP needs at least two GPUs because tensor split is a no-op on one and may
// abort. Auto layers hides TP because --fit aborts under --split-mode tensor.
const tpDisabled = singleGpuInUse;
// Manual gpu-layers ceiling = model layer count + 1 (else a safe fallback):
// llama.cpp counts the output layer as one more offloadable layer past the
// repeating blocks ("offloaded 33/33" needs -ngl 33 on a 32-block model), so
@ -1537,7 +1537,7 @@ export function ChatSettingsPanel({
Which GPUs this model may use. Unchecked GPUs are hidden
from llama.cpp (CUDA_VISIBLE_DEVICES, or
HIP_VISIBLE_DEVICES on ROCm). Leave all checked to use
every GPU.
every GPU. At least one GPU must stay selected.
</InfoHint>
</div>
<div className="flex flex-col gap-2">
@ -1557,7 +1557,10 @@ export function ChatSettingsPanel({
checked={isGpuChecked(d.index)}
onCheckedChange={() => toggleGpu(d.index)}
data-test-id={`gpu-pick-${d.index}`}
disabled={modelControlsDisabled}
disabled={
modelControlsDisabled ||
(isGpuChecked(d.index) && singleGpuInUse)
}
/>
</div>
))}
@ -2432,13 +2435,14 @@ function ConfirmToolCallsToggle() {
<InfoHint>
When on, every local Unsloth tool call pauses for your approval
before it runs (the "Ask for approval" level). When off, tool calls
run without prompts inside the sandbox (the "Off" level).
run without prompts inside the sandbox (the "Run automatically"
level).
Provider-hosted tools are not gated here.
</InfoHint>
</div>
{permissionMode === "full" ? (
<span className="text-[11px] text-muted-foreground">
Overridden by Full access (Bypass permissions)
Overridden by Full access
</span>
) : null}
</div>
@ -2459,11 +2463,11 @@ function BypassPermissionsToggle() {
<div className="flex flex-col gap-2">
<div className="flex min-w-0 items-center gap-1.5">
<span className="whitespace-nowrap text-[13px] font-medium leading-[1.25] tracking-nav text-nav-fg">
Bypass permissions
Tool permissions
</span>
<InfoHint>
How Unsloth approves tool calls before they run. Full access is
dangerous: it disables confirmations and the code sandbox.
Choose how Unsloth approves tool calls before they run. Full access
disables confirmations and the code sandbox.
</InfoHint>
</div>
{/* Full width, styled like the panel selects/preset input. */}

34
studio/frontend/src/features/chat/hooks/use-chat-sidebar-items.ts
View file

@ -217,6 +217,40 @@ export async function archiveChatItem(
notifyChatHistoryUpdated();
}
export async function archiveAllChatItems(
activeId?: string,
onSelect?: (view: { mode: "single"; newThreadNonce: string }) => void,
): Promise<number> {
const threads = await listStoredChatThreads({ includeArchived: true });
// Boolean() mirrors groupThreads: legacy records may have archived
// undefined/null, which must count as "not archived".
const toArchive = threads.filter((t) => !t.archived);
if (toArchive.length === 0) return 0;
for (const t of toArchive) cancelIfRunning(t.id);
await Promise.all(
toArchive.map((t) => updateStoredChatThread(t.id, { archived: true })),
);
// Reset only when this action archived the active single thread or compare
// pair. An already-archived chat opened from the archive is not in
// toArchive and must stay open.
const archivedActive =
activeId !== undefined &&
toArchive.some(
(thread) => thread.id === activeId || thread.pairId === activeId,
);
if (archivedActive) {
useChatRuntimeStore.getState().setActiveThreadId(null);
onSelect?.({ mode: "single", newThreadNonce: crypto.randomUUID() });
}
notifyChatHistoryUpdated();
// Report sidebar items, not raw threads: a compare pair reads as one chat.
return groupThreads(toArchive).length;
}
export async function unarchiveChatItem(item: SidebarItem): Promise<void> {
const threadIds: string[] =
item.type === "single"

19
studio/frontend/src/features/chat/index.ts
View file

@ -3,10 +3,15 @@
export { ChatPage, validateChatSearch, type ChatSearch } from "./chat-page";
export {
deleteChatAttachment,
fetchChatAttachmentBlob,
getInferenceStatus,
listChatAttachments,
listGgufVariants,
listLocalModels,
loadModel,
type ChatAttachmentPage,
type ChatAttachmentRecord,
type LocalModelInfo,
} from "./api/chat-api";
export type { GgufVariantDetail } from "./types/api";
@ -17,6 +22,10 @@ export {
type Preset,
} from "./chat-settings-sheet";
export { useChatRuntimeStore } from "./stores/chat-runtime-store";
export {
CHAT_RAG_CAPTION_KEY,
CHAT_RAG_OCR_KEY,
} from "./stores/chat-runtime-store";
export {
preferFullToolOutput,
toolOutputKey,
@ -46,12 +55,16 @@ export { setTrainingCompareHandoff } from "./lib/training-compare-handoff";
export type { ProjectRecord } from "./types";
export { clearAllChats, countAllChats } from "./utils/clear-all-chats";
export { listStoredChatThreads } from "./utils/chat-history-storage";
export { emitChatAttachmentDeleted } from "./utils/chat-attachment-events";
export { ArtifactCard } from "./artifacts/artifact-card";
export {
useChatArtifactsStore,
useSelectedChatArtifact,
} from "./artifacts/store";
export { downloadChatExport } from "./utils/export-chat-history";
export {
downloadChatExport,
downloadArchivedChatExport,
} from "./utils/export-chat-history";
export {
clearNewChatDraft,
composerDraftKey,
@ -60,10 +73,14 @@ export {
} from "./utils/composer-draft";
export {
EXPORT_FORMATS_LIST,
buildFineTuneJsonl,
bulkExportConversationsByScope,
exportFineTuneJsonl,
importConversationsFromFile,
type FineTuneFormat,
} from "./prompt-storage/prompt-storage-dialog";
export {
archiveAllChatItems,
archiveChatItem,
deleteChatItem,
renameChatItem,

57
studio/frontend/src/features/chat/permission-mode-select.tsx
View file

@ -7,7 +7,6 @@ import {
CircleOff,
Hand,
ShieldCheck,
XIcon,
} from "lucide-react";
import { useState } from "react";
@ -39,9 +38,8 @@ import {
} from "./stores/chat-runtime-store";
/**
* Permission levels for the Bypass permissions dropdowns (General settings,
* chat settings sheet, composer "+" menu). Off sits last as the toggle that
* turns the feature off entirely.
* Permission levels for tool calls. Full access stays last because it disables
* both approval prompts and the code sandbox.
*/
export const PERMISSION_MODE_OPTIONS: readonly {
value: PermissionMode;
@ -61,6 +59,12 @@ export const PERMISSION_MODE_OPTIONS: readonly {
description: "Only ask for actions detected as potentially unsafe",
icon: ShieldCheck,
},
{
value: "off",
label: "Run automatically",
description: "Run tool calls without approval prompts inside the sandbox",
icon: CircleOff,
},
{
value: "full",
label: "Full access",
@ -68,14 +72,11 @@ export const PERMISSION_MODE_OPTIONS: readonly {
"Unrestricted: no approval prompts and the code sandbox is disabled",
icon: CircleAlert,
},
{
value: "off",
label: "Off",
description: "Turn off bypass permissions",
icon: CircleOff,
},
] as const;
export const FULL_ACCESS_WARNING =
"Full access lets tool calls run without approval prompts or the code sandbox. They can modify or delete files, run commands, and make network requests. Enable it only when you trust the current task.";
export function permissionModeOption(mode: PermissionMode) {
return (
PERMISSION_MODE_OPTIONS.find((option) => option.value === mode) ??
@ -100,10 +101,10 @@ export function PermissionModeMenuItems({
<DropdownMenuItem
key={option.value}
onSelect={() => {
// Reselecting the active level toggles the feature off.
if (option.value === permissionMode) {
setPermissionMode("off");
} else if (option.value === "full") {
return;
}
if (option.value === "full") {
onRequestFullAccess();
} else {
setPermissionMode(option.value);
@ -154,9 +155,7 @@ export function FullAccessConfirmDialog({
<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>
@ -260,15 +259,10 @@ export function PermissionModeComposerPill({
const setBypassConfirmOpen = useChatRuntimeStore(
(s) => s.setBypassConfirmOpen,
);
const setPermissionMode = useChatRuntimeStore((s) => s.setPermissionMode);
const active = permissionModeOption(permissionMode);
const ActiveIcon = active.icon;
const fullAccess = permissionMode === "full";
// Off means the feature is off: no pill (re-enable via the "+" menu or
// settings, like the pre-levels bypass badge).
if (permissionMode === "off") return null;
return (
<DropdownMenu>
<DropdownMenuTrigger asChild={true}>
@ -278,30 +272,11 @@ export function PermissionModeComposerPill({
data-pill-label={active.label}
data-active={fullAccess ? "true" : "false"}
data-variant={fullAccess ? "danger" : undefined}
data-keep-label="true"
aria-label="Permission level for tool calls"
title={`${active.label}: ${active.description}`}
>
{/* The icon doubles as an off switch (mirrors the MCP pill): hover
swaps it to an X; clicking it turns bypass permissions Off (no
prompts, sandbox on) without opening the menu. data-keep-label
exempts this pill from compact icon-only mode, so the off switch
stays clickable even while the other pills are collapsed. */}
<span
role="button"
aria-label="Turn off bypass permissions"
tabIndex={-1}
onPointerDown={(e) => {
e.stopPropagation();
}}
onClick={(e) => {
e.stopPropagation();
setPermissionMode("off");
}}
className="composer-pill-glyph cursor-pointer"
>
<span className="composer-pill-glyph">
<ActiveIcon className="size-[15px]" strokeWidth={2} />
<XIcon className="composer-pill-x" />
</span>
<span>{active.label}</span>
<HugeiconsIcon

223
studio/frontend/src/features/chat/prompt-storage/prompt-storage-dialog.tsx
View file

@ -186,7 +186,16 @@ function contentBlocksToText(content: unknown): string {
// predate the user's next message); the parent chain is timestamp-independent.
type _Msg = { id: string; parentId?: string | null; createdAt?: number };
function orderByParentChain<T extends _Msg>(messages: T[]): T[] {
function orderByParentChain<T extends _Msg>(
messages: T[],
options: {
/** Append messages off the selected chain (abandoned branches) at the
* end. Full exports keep everything; fine-tune conversion must not,
* since alternate replies would merge into one conversation. */
includeSiblings?: boolean;
} = {},
): T[] {
const { includeSiblings = true } = options;
const byId = new Map<string, T>(messages.map((m) => [m.id, m]));
const childrenOf = new Map<string | null, T[]>();
for (const m of messages) {
@ -207,7 +216,9 @@ function orderByParentChain<T extends _Msg>(messages: T[]): T[] {
byId.delete(next.id);
}
for (const [, m] of byId) result.push(m);
if (includeSiblings) {
for (const [, m] of byId) result.push(m);
}
return result;
}
@ -543,6 +554,214 @@ export async function exportProjectConversations(
);
}
// ── Fine-tuning export ─────────────────────────────────────────────────────
// One JSONL line per conversation: {"messages": [{"role", "content"}]} with
// string-only content in system/user/assistant turns. Unsloth's training tab
// detects this as ChatML natively (no column mapping, no standardization) and
// it works with train-on-completions masking, which only trains on assistant
// turns. Reasoning, tool calls, and images are dropped: clean SFT targets.
export type FineTuneMessage = {
role: "system" | "user" | "assistant";
content: string;
};
const FINE_TUNE_ROLES = new Set(["system", "user", "assistant"]);
/** Plain text of a message: text blocks plus text-type attachment parts. */
function messageToPlainText(msg: {
content: unknown;
attachments?: unknown;
}): string {
const parts: string[] = [];
const collect = (blocks: unknown) => {
// Legacy and imported histories can store content as a plain string.
if (typeof blocks === "string") {
if (blocks.trim()) parts.push(blocks);
return;
}
if (!Array.isArray(blocks)) return;
for (const b of blocks) {
if (!b || typeof b !== "object") {
continue;
}
const block = b as Record<string, unknown>;
if (block.type === "text" && typeof block.text === "string" && block.text) {
parts.push(block.text);
}
}
};
collect(msg.content);
if (Array.isArray(msg.attachments)) {
for (const attachment of msg.attachments as Array<{ content?: unknown }>) {
collect(attachment?.content);
}
}
return parts.join("\n\n").trim();
}
/** Merge consecutive same-role turns so chat templates format cleanly. */
function mergeSameRoleTurns(turns: FineTuneMessage[]): FineTuneMessage[] {
const merged: FineTuneMessage[] = [];
for (const turn of turns) {
const last = merged[merged.length - 1];
if (last && last.role === turn.role) {
last.content += `\n\n${turn.content}`;
} else {
merged.push({ ...turn });
}
}
return merged;
}
/** Conversation turns for fine-tuning, or null when the thread has no
* usable user + assistant exchange. Consecutive same-role turns merge,
* assistant turns before the first user turn drop (an assistant target
* with no prompt teaches nothing), and trailing non-assistant turns drop
* so chat templates format cleanly. */
function messagesToFineTuneTurns(
messages: Array<{ role: unknown; content: unknown; attachments?: unknown }>,
): FineTuneMessage[] | null {
const raw: FineTuneMessage[] = [];
for (const msg of messages) {
const role = msg.role as FineTuneMessage["role"];
if (!FINE_TUNE_ROLES.has(role)) continue;
const content = messageToPlainText(msg);
if (!content) continue;
raw.push({ role, content });
}
const firstUser = raw.findIndex((t) => t.role === "user");
if (firstUser === -1) return null;
const turns = mergeSameRoleTurns(
raw.filter((t, i) => i >= firstUser || t.role === "system"),
);
while (turns.length > 0 && turns[turns.length - 1].role !== "assistant") {
turns.pop();
}
const hasUser = turns.some((t) => t.role === "user");
const hasAssistant = turns.some((t) => t.role === "assistant");
return hasUser && hasAssistant ? turns : null;
}
export type FineTuneExportResult = {
lines: string[];
conversations: number;
skipped: number;
};
/** Dataset shapes the Train tab detects without column mapping. */
export type FineTuneFormat = "openai" | "sharegpt" | "alpaca";
const SHAREGPT_FROM: Record<FineTuneMessage["role"], string> = {
system: "system",
user: "human",
assistant: "gpt",
};
/** JSONL lines for one conversation in the chosen format. Alpaca is
* single-turn, so each user to assistant pair becomes its own record with
* the system prompt and earlier exchange carried in the input field. */
function turnsToFineTuneLines(
turns: FineTuneMessage[],
format: FineTuneFormat,
): string[] {
if (format === "sharegpt") {
return [
JSON.stringify({
conversations: turns.map((t) => ({
from: SHAREGPT_FROM[t.role],
value: t.content,
})),
}),
];
}
if (format === "alpaca") {
const lines: string[] = [];
const context: string[] = [];
let system = "";
let pendingUser: string | null = null;
for (const t of turns) {
if (t.role === "system") {
system = system ? `${system}\n\n${t.content}` : t.content;
continue;
}
if (t.role === "user") {
pendingUser = t.content;
continue;
}
if (pendingUser === null) continue;
const inputParts = [];
if (system) inputParts.push(system);
if (context.length > 0) inputParts.push(context.join("\n"));
lines.push(
JSON.stringify({
instruction: pendingUser,
input: inputParts.join("\n\n"),
output: t.content,
}),
);
context.push(`User: ${pendingUser}`, `Assistant: ${t.content}`);
pendingUser = null;
}
return lines;
}
return [JSON.stringify({ messages: turns })];
}
/** Every non-archived chat (Recents and Projects) as training-ready JSONL. */
export async function buildFineTuneJsonl(
format: FineTuneFormat = "openai",
): Promise<FineTuneExportResult> {
const threads = await listStoredChatThreads({ includeArchived: false });
const ids = [...new Set(threads.map((t) => t.id))];
const lines: string[] = [];
let conversations = 0;
let skipped = 0;
for (const id of ids) {
const raw = await listStoredChatMessages(id);
const hasParentIds = raw.some(
(m) => (m as { parentId?: unknown }).parentId != null,
);
// Chain only: retries/regenerations leave sibling branches, and mixing
// alternate replies into one conversation corrupts the training targets.
const ordered = hasParentIds
? (orderByParentChain(raw, { includeSiblings: false }) as typeof raw)
: raw;
const turns = messagesToFineTuneTurns(ordered);
const converted = turns ? turnsToFineTuneLines(turns, format) : [];
if (converted.length === 0) {
skipped += 1;
continue;
}
conversations += 1;
lines.push(...converted);
}
return { lines, conversations, skipped };
}
/** Download the fine-tuning JSONL; returns the conversation count. */
export async function exportFineTuneJsonl(
format: FineTuneFormat = "openai",
): Promise<number> {
const { lines, conversations, skipped } = await buildFineTuneJsonl(format);
if (conversations === 0) {
toast.info("No chats with a user and assistant exchange to export.");
return 0;
}
const suffix = format === "openai" ? "" : `-${format}`;
downloadBlob(
lines.join("\n"),
`chat-finetune${suffix}-${exportTs()}.jsonl`,
"application/x-ndjson",
);
if (skipped > 0) {
toast.success(
`Exported ${conversations} conversation${conversations === 1 ? "" : "s"} (${skipped} without a full exchange skipped).`,
);
}
return conversations;
}
// role:"tool" results are absorbed into the preceding assistant tool-call
// part's `result` field rather than becoming separate records.
function oaiMessagesToRecords(

167
studio/frontend/src/features/chat/runtime-provider.tsx
View file

@ -56,6 +56,11 @@ import { AudioAttachmentAdapter } from "./audio-attachment-adapter";
import { useChatRuntimeStore } from "./stores/chat-runtime-store";
import { ToolPaneScopeContext, toolPaneScope } from "./tool-output-scope";
import type { MessageRecord, ModelType, ThreadRecord } from "./types";
import {
chatContentPartAttachmentIdFromSignature,
chatContentPartAttachmentSignature,
onChatAttachmentDeleted,
} from "./utils/chat-attachment-events";
import {
deleteStoredChatThreads,
ensureStoredChatThread,
@ -890,6 +895,168 @@ function useStudioRuntimeAdapters(
): StudioRuntimeAdapters {
const aui = useAui();
// Mirror Data-tab attachment deletions into the loaded thread. The in-memory
// repository otherwise keeps the attachment, and a later repo-to-storage sync
// (e.g. deleting a message in the thread) would write it back.
useEffect(() => {
let active = true;
let pendingDeletion = Promise.resolve();
const unsubscribe = onChatAttachmentDeleted((event) => {
pendingDeletion = pendingDeletion.then(async () => {
if (!active) return;
const { messageId, attachmentId } = event;
try {
const thread = aui.thread();
if (attachmentId.startsWith("content-part-sha256-")) {
for (let attempt = 0; attempt < 3 && active; attempt += 1) {
const exported = thread.export();
const target = exported.messages.find(
(item) => item.message.id === messageId,
);
if (!target || !Array.isArray(target.message.content)) return;
const content = target.message.content;
const signatures = content.map((part) =>
chatContentPartAttachmentSignature(part),
);
const ids = await Promise.all(
signatures.map((signature) =>
signature === null
? null
: chatContentPartAttachmentIdFromSignature(signature),
),
);
const targetAttachments = (
target.message as {
attachments?: readonly { id: string }[];
}
).attachments;
const hasTargetAttachment =
Array.isArray(targetAttachments) &&
targetAttachments.some(
(attachment) => attachment.id === attachmentId,
);
if (
(!ids.includes(attachmentId) && !hasTargetAttachment) ||
!active
) {
return;
}
// Preserve any messages added or streamed while WebCrypto ran.
// Retry if the target's managed content itself changed.
const latest = thread.export();
const latestTarget = latest.messages.find(
(item) => item.message.id === messageId,
);
const latestContent = latestTarget?.message.content;
if (!Array.isArray(latestContent)) return;
const latestSignatures = latestContent.map((part) =>
chatContentPartAttachmentSignature(part),
);
if (
signatures.length !== latestSignatures.length ||
signatures.some(
(signature, index) => signature !== latestSignatures[index],
)
) {
continue;
}
const messages = latest.messages.map((item) => {
if (item.message.id !== messageId) return item;
const attachments = (
item.message as {
attachments?: readonly { id: string }[];
}
).attachments;
return {
...item,
message: {
...item.message,
content: latestContent.filter(
(_, index) => ids[index] !== attachmentId,
),
...(Array.isArray(attachments)
? {
attachments: attachments.filter(
(attachment) =>
attachment.id !== attachmentId,
),
}
: {}),
} as typeof item.message,
};
});
if (active) thread.import({ ...latest, messages });
return;
}
return;
}
const exported = thread.export();
let changed = false;
const messages = exported.messages.map((item) => {
if (item.message.id !== messageId) return item;
const message = item.message;
const attachments = (
message as { attachments?: readonly { id: string }[] }
).attachments;
if (
Array.isArray(attachments) &&
attachments.some(
(attachment) => attachment.id === attachmentId,
)
) {
changed = true;
return {
...item,
message: {
...message,
attachments: attachments.filter(
(attachment) => attachment.id !== attachmentId,
),
} as typeof message,
};
}
if (/^content-part-[0-9]+$/.test(attachmentId)) {
// Legacy synthetic id for a blob stored as a message content part.
const idx = Number(attachmentId.slice("content-part-".length));
const content = message.content;
if (
!Array.isArray(content) ||
!Number.isInteger(idx) ||
idx < 0 ||
idx >= content.length
) {
return item;
}
const part = content[idx] as { type?: string };
if (part?.type !== "image" && part?.type !== "audio") return item;
changed = true;
return {
...item,
message: {
...message,
content: content.filter((_, i) => i !== idx),
} as typeof message,
};
}
return item;
});
if (changed && active) thread.import({ ...exported, messages });
} catch {
// No active thread mounted: storage already holds the truth.
}
});
return pendingDeletion;
});
return () => {
active = false;
unsubscribe();
};
}, [aui]);
const history = useMemo<ThreadHistoryAdapter>(
() => ({
async load() {

12
studio/frontend/src/features/chat/shared-composer.tsx
View file

@ -619,7 +619,6 @@ export function SharedComposer({
const artifactsEnabled = useChatRuntimeStore((s) => s.artifactsEnabled);
const setArtifactsEnabled = useChatRuntimeStore((s) => s.setArtifactsEnabled);
const showCanvasMenuItem = useChatRuntimeStore((s) => s.showCanvasMenuItem);
const permissionMode = useChatRuntimeStore((s) => s.permissionMode);
const mcpEnabledForChat = useChatRuntimeStore((s) => s.mcpEnabledForChat);
const setMcpEnabledForChat = useChatRuntimeStore(
(s) => s.setMcpEnabledForChat,
@ -790,15 +789,12 @@ export function SharedComposer({
// can still be pre-selected, matching Web search/Code/MCP.
const ragDisabled = modelLoaded && (isExternalModel || !supportsTools);
const showRagPill = !isExternalModel;
// Above 4 pills, collapse to icons only to cut clutter. Compare, Search and
// Code always show; the permission pill shows in every mode except "off"
// (it renders null there); the rest are conditional.
const permissionPillVisible = permissionMode !== "off";
// Above 4 pills, collapse to icons only. Compare, Search, Code, and
// permissions always show; the rest are conditional.
const pillsCompact =
3 +
(permissionPillVisible ? 1 : 0) +
4 +
(showImagePill ? 1 : 0) +
(showRagPill && ragEnabled && !ragDisabled ? 1 : 0) +
(showRagPill && ragEnabled ? 1 : 0) +
(showWebFetchPill ? 1 : 0) +
(artifactsEnabled ? 1 : 0) +
(mcpEnabledForChat ? 1 : 0) >

62
studio/frontend/src/features/chat/utils/archived-chat-export.ts Normal file
View file

@ -0,0 +1,62 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
// Minimal views over the `unknown[]` export fields we filter on.
type ExportThreadView = {
id?: string;
archived?: boolean;
projectId?: string | null;
};
type ExportMessageView = { threadId?: string };
type ExportProjectView = { id?: string };
// Full chat-export backup shape, kept structural so the pure filter below
// stays decoupled from the storage layer that produces it.
export interface ChatExportData {
exportedAt?: string;
version?: number;
threadCount: number;
projects?: unknown[];
threads: unknown[];
messages: unknown[];
}
// Restrict a full chat export to archived threads, their messages and the
// projects those threads belong to. Pure: never mutates the input, and keeps
// the original thread/message objects so the backup re-imports unchanged.
export function filterArchivedChatExport<T extends ChatExportData>(
full: T,
): { data: T; archivedCount: number } {
const archivedThreads = (full.threads as ExportThreadView[]).filter(
(thread) => thread.archived === true,
);
const archivedThreadIds = new Set(
archivedThreads
.map((thread) => thread.id)
.filter((id): id is string => typeof id === "string"),
);
const messages = (full.messages as ExportMessageView[]).filter(
(message) =>
typeof message.threadId === "string" &&
archivedThreadIds.has(message.threadId),
);
const referencedProjectIds = new Set(
archivedThreads
.map((thread) => thread.projectId)
.filter((id): id is string => typeof id === "string"),
);
const projects = (full.projects as ExportProjectView[] | undefined)?.filter(
(project) =>
typeof project.id === "string" && referencedProjectIds.has(project.id),
);
return {
data: {
...full,
threadCount: archivedThreads.length,
projects: projects ?? [],
threads: archivedThreads as unknown[],
messages: messages as unknown[],
},
archivedCount: archivedThreads.length,
};
}

123
studio/frontend/src/features/chat/utils/chat-attachment-events.ts Normal file
View file

@ -0,0 +1,123 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
/**
* Notifies loaded chat runtimes when the Data tab deletes a stored attachment.
* Without this, the active thread's in-memory repository still holds the
* attachment, and any later repo-to-storage sync (e.g. deleting a message in
* that thread) writes it back, undoing the deletion.
*/
import forge from "node-forge";
export type ChatAttachmentDeletedEvent = {
messageId: string;
attachmentId: string;
};
const CONTENT_PART_ID_PREFIX = "content-part-sha256-";
const URI_SCHEME_RE = /^[A-Za-z][A-Za-z0-9+.-]*:/;
function isLocallyStoredBlob(value: string): boolean {
const candidate = value.trimStart();
if (!candidate) return false;
if (candidate.slice(0, 5).toLowerCase() === "data:") return true;
if (candidate.startsWith("//") || candidate.startsWith("\\\\")) {
return false;
}
return !URI_SCHEME_RE.test(candidate);
}
function stableJson(value: unknown): string {
if (Array.isArray(value)) {
return `[${value
.map((item) => (item === undefined ? "null" : stableJson(item)))
.join(",")}]`;
}
if (value && typeof value === "object") {
const record = value as Record<string, unknown>;
return `{${Object.keys(record)
.filter((key) => record[key] !== undefined)
.sort()
.map((key) => `${JSON.stringify(key)}:${stableJson(record[key])}`)
.join(",")}}`;
}
return JSON.stringify(value) ?? "null";
}
/** Canonical payload used to detect whether an async hash still describes the
* current message content. */
export function chatContentPartAttachmentSignature(
part: unknown,
): string | null {
if (!part || typeof part !== "object") return null;
const record = part as Record<string, unknown>;
let payload: ["image" | "audio", unknown] | null = null;
if (
typeof record.image === "string" &&
record.image.slice(0, 5).toLowerCase() === "data:"
) {
payload = ["image", record.image];
} else if (
typeof record.audio === "string" &&
isLocallyStoredBlob(record.audio)
) {
payload = ["audio", record.audio];
} else if (record.audio && typeof record.audio === "object") {
const data = (record.audio as Record<string, unknown>).data;
if (typeof data === "string" && isLocallyStoredBlob(data)) {
payload = ["audio", record.audio];
}
}
if (!payload) return null;
return stableJson(payload);
}
/** Mirrors the backend's stable content-part identity without adding private
* metadata to the message payload sent to inference. */
export async function chatContentPartAttachmentIdFromSignature(
signature: string,
): Promise<string> {
let hex: string | null = null;
const subtle = globalThis.crypto?.subtle;
if (subtle) {
try {
const digest = await subtle.digest(
"SHA-256",
new TextEncoder().encode(signature),
);
hex = Array.from(new Uint8Array(digest), (byte) =>
byte.toString(16).padStart(2, "0"),
).join("");
} catch {
// Fall through to the pure-JS implementation below. Some embedded
// browsers expose crypto.subtle but reject it outside a secure context.
}
}
if (hex === null) {
const digest = forge.md.sha256.create();
digest.update(signature, "utf8");
hex = digest.digest().toHex();
}
return `${CONTENT_PART_ID_PREFIX}${hex}`;
}
type Listener = (event: ChatAttachmentDeletedEvent) => void | Promise<void>;
const listeners = new Set<Listener>();
export function onChatAttachmentDeleted(listener: Listener): () => void {
listeners.add(listener);
return () => {
listeners.delete(listener);
};
}
export function emitChatAttachmentDeleted(
event: ChatAttachmentDeletedEvent,
): void {
for (const listener of [...listeners]) {
void listener(event);
}
}

18
studio/frontend/src/features/chat/utils/download-json.ts Normal file
View file

@ -0,0 +1,18 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
// Save `data` as a pretty-printed JSON file via a temporary object URL. Uses
// only the standard Blob/anchor download path so it works in every browser.
export function triggerJsonDownload(data: unknown, filename: string): void {
const blob = new Blob([JSON.stringify(data, null, 2)], {
type: "application/json",
});
const url = URL.createObjectURL(blob);
const a = document.createElement("a");
a.href = url;
a.download = filename;
document.body.appendChild(a);
a.click();
a.remove();
URL.revokeObjectURL(url);
}

35
studio/frontend/src/features/chat/utils/export-chat-history.ts
View file

@ -1,21 +1,34 @@
// 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 { filterArchivedChatExport } from "./archived-chat-export";
import { buildStoredChatExport } from "./chat-history-storage";
import { triggerJsonDownload } from "./download-json";
export const buildChatExport = buildStoredChatExport;
function dateStamp(): string {
// Date only (no colons) so the filename is valid on every OS.
return new Date().toISOString().slice(0, 10);
}
export async function downloadChatExport(): Promise<void> {
const data = await buildChatExport();
const blob = new Blob([JSON.stringify(data, null, 2)], {
type: "application/json",
});
const url = URL.createObjectURL(blob);
const a = document.createElement("a");
a.href = url;
a.download = `unsloth-chats-${new Date().toISOString().slice(0, 10)}.json`;
document.body.appendChild(a);
a.click();
a.remove();
URL.revokeObjectURL(url);
triggerJsonDownload(data, `unsloth-chats-${dateStamp()}.json`);
}
// Full backup restricted to archived chats. Returns the archived thread count.
export async function buildArchivedChatExport() {
return filterArchivedChatExport(await buildChatExport());
}
// Download only the archived chats. Returns how many were exported; skips the
// download entirely when there are none.
export async function downloadArchivedChatExport(): Promise<number> {
const { data, archivedCount } = await buildArchivedChatExport();
if (archivedCount === 0) {
return 0;
}
triggerJsonDownload(data, `unsloth-archived-chats-${dateStamp()}.json`);
return archivedCount;
}

13
studio/frontend/src/features/export/export-page.tsx
View file

@ -44,6 +44,7 @@ import {
import { usePlatformStore } from "@/config/env";
import { useHubModelSearch } from "@/features/hub/hooks/use-hub-model-search";
import { confirmRemoteCodeIfNeeded } from "@/features/security";
import { prepareHfTokenForUse } from "@/features/hf-auth";
import { GuidedTour, useGuidedTourController } from "@/features/tour";
import {
type LocalModelInfo,
@ -724,11 +725,17 @@ export function ExportPage() {
const checkpointPath = selectedCp?.path ?? null;
const pushToHub = destination === "hub";
const preparedToken = await prepareHfTokenForUse(hfToken, {
allowAnonymous: !pushToHub,
});
if (!preparedToken.proceed) return;
const actionHfToken = preparedToken.token ?? "";
const repoId =
pushToHub && hfUsername && modelName
? `${hfUsername}/${modelName}`
: undefined;
const token = pushToHub && hfToken ? hfToken : undefined;
const token = pushToHub && actionHfToken ? actionHfToken : undefined;
// The GGUF method with the LoRA target reuses the LoRA-adapter export path.
const effectiveMethod: ExportMethod = ggufAsLora ? "lora" : exportMethod;
const emitLoraGguf =
@ -747,7 +754,7 @@ export function ExportPage() {
if (sourceMode !== "checkpoint") {
const remoteCodeOk = await confirmRemoteCodeIfNeeded({
modelName: source,
hfToken: hfToken || null,
hfToken: actionHfToken || null,
// An HF source can need trust_remote_code via its YAML default with no
// auto_map to review; signal it so a YAML-only model does not export
// with it false.
@ -767,7 +774,7 @@ export function ExportPage() {
modelSource,
trustRemoteCode,
approvedRemoteCodeFingerprint,
loadToken: hfToken || null,
loadToken: actionHfToken || null,
exportMethod: effectiveMethod,
isAdapter: adapterExport,
quantLevels,

44
studio/frontend/src/features/hf-auth/api.ts 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
import { authFetch } from "@/features/auth";
// This header helper is API-layer-only and is not part of the feature's
// React-facing public barrel.
// eslint-disable-next-line no-restricted-imports
import { hubTokenHeader } from "@/features/hub/lib/hub-token-header";
export type HfTokenValidationStatus =
| "missing"
| "valid"
| "invalid"
| "rate_limited"
| "unavailable";
export interface HfTokenValidationResult {
status: HfTokenValidationStatus;
retryAfterSeconds: number | null;
}
export async function validateHfToken(
token: string | null | undefined,
): Promise<HfTokenValidationResult> {
const normalized = token?.trim() ?? "";
if (!normalized) {
return { status: "missing", retryAfterSeconds: null };
}
const response = await authFetch("/api/hub/token/validate", {
method: "POST",
headers: hubTokenHeader(normalized),
});
if (!response.ok) {
return { status: "unavailable", retryAfterSeconds: null };
}
const body = (await response.json()) as {
status?: HfTokenValidationStatus;
retry_after_seconds?: number | null;
};
return {
status: body.status ?? "unavailable",
retryAfterSeconds: body.retry_after_seconds ?? null,
};
}

63
studio/frontend/src/features/hf-auth/confirm-token.ts Normal file
View file

@ -0,0 +1,63 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
// These stores are used outside React and are not part of their features'
// React-facing public barrels.
// eslint-disable-next-line no-restricted-imports
import { useHfTokenStore } from "@/features/hub/stores/hf-token-store";
// eslint-disable-next-line no-restricted-imports
import { useSettingsDialogStore } from "@/features/settings/stores/settings-dialog-store";
import { validateHfToken } from "./api";
import { useHfTokenWarningStore } from "./store";
export interface PreparedHfToken {
proceed: boolean;
token: string | null;
}
interface PrepareHfTokenOptions {
allowAnonymous?: boolean;
}
// A caller can retain the pre-dialog payload while the shared store is cleared.
// Remember that one-session choice so a follow-up /load does not prompt again
// after its preceding /validate already continued anonymously.
const anonymousForSession = new Set<string>();
export async function prepareHfTokenForUse(
token: string | null | undefined,
options: PrepareHfTokenOptions = {},
): Promise<PreparedHfToken> {
const normalized = token?.trim() ?? "";
if (!normalized) return { proceed: true, token: null };
const allowAnonymous = options.allowAnonymous ?? true;
if (allowAnonymous && anonymousForSession.has(normalized)) {
return { proceed: true, token: null };
}
let validation;
try {
validation = await validateHfToken(normalized);
} catch {
// Validation is advisory. Let the real operation retain its own error.
return { proceed: true, token: normalized };
}
if (validation.status !== "invalid") {
// A connectivity failure or rate limit cannot prove that a token is bad.
// Let the real operation proceed and retain its repository-specific error.
return { proceed: true, token: normalized };
}
const decision = await useHfTokenWarningStore
.getState()
.requestDecision(allowAnonymous);
if (decision === "anonymous") {
anonymousForSession.add(normalized);
useHfTokenStore.getState().clearToken();
return { proceed: true, token: null };
}
if (decision === "replace") {
useSettingsDialogStore.getState().openDialog("general");
}
return { proceed: false, token: normalized };
}

66
studio/frontend/src/features/hf-auth/hf-token-warning-dialog.tsx Normal file
View file

@ -0,0 +1,66 @@
// 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 {
AlertDialog,
AlertDialogAction,
AlertDialogCancel,
AlertDialogContent,
AlertDialogDescription,
AlertDialogFooter,
AlertDialogHeader,
AlertDialogTitle,
} from "@/components/ui/alert-dialog";
import { Button } from "@/components/ui/button";
import { AlertTriangle } from "lucide-react";
import { useHfTokenWarningStore } from "./store";
export function HfTokenWarningDialog() {
const open = useHfTokenWarningStore((state) => state.open);
const allowAnonymous = useHfTokenWarningStore(
(state) => state.allowAnonymous,
);
const resolve = useHfTokenWarningStore((state) => state.resolve);
return (
<AlertDialog
open={open}
onOpenChange={(next) => {
if (!next) resolve("cancel");
}}
>
<AlertDialogContent className="max-w-md">
<AlertDialogHeader>
<div className="flex items-start gap-3">
<div className="flex size-9 shrink-0 items-center justify-center rounded-full bg-amber-500/10 text-amber-600 dark:text-amber-400">
<AlertTriangle className="size-5" />
</div>
<div className="space-y-1 text-left">
<AlertDialogTitle>Hugging Face token is invalid</AlertDialogTitle>
<AlertDialogDescription>
{allowAnonymous
? "Hugging Face rejected the saved token. Replace it to access private or gated repositories, or continue without it for public and fully downloaded models."
: "Hugging Face rejected the saved token. Replace it before uploading to the Hub."}
</AlertDialogDescription>
</div>
</div>
</AlertDialogHeader>
<AlertDialogFooter className="sm:justify-between">
<AlertDialogCancel onClick={() => resolve("cancel")}>
Cancel
</AlertDialogCancel>
<div className="flex flex-col-reverse gap-2 sm:flex-row">
{allowAnonymous ? (
<Button variant="outline" onClick={() => resolve("anonymous")}>
Continue without token
</Button>
) : null}
<AlertDialogAction onClick={() => resolve("replace")}>
Replace token
</AlertDialogAction>
</div>
</AlertDialogFooter>
</AlertDialogContent>
</AlertDialog>
);
}

10
studio/frontend/src/features/hf-auth/index.ts Normal file
View file

@ -0,0 +1,10 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
export { validateHfToken } from "./api";
export type {
HfTokenValidationResult,
HfTokenValidationStatus,
} from "./api";
export { prepareHfTokenForUse } from "./confirm-token";
export { HfTokenWarningDialog } from "./hf-token-warning-dialog";

33
studio/frontend/src/features/hf-auth/store.ts Normal file
View file

@ -0,0 +1,33 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import { create } from "zustand";
export type HfTokenWarningDecision = "anonymous" | "replace" | "cancel";
type Resolver = (decision: HfTokenWarningDecision) => void;
let pendingResolver: Resolver | null = null;
interface HfTokenWarningStore {
open: boolean;
allowAnonymous: boolean;
requestDecision: (allowAnonymous: boolean) => Promise<HfTokenWarningDecision>;
resolve: (decision: HfTokenWarningDecision) => void;
}
export const useHfTokenWarningStore = create<HfTokenWarningStore>((set) => ({
open: false,
allowAnonymous: true,
requestDecision: (allowAnonymous) =>
new Promise<HfTokenWarningDecision>((resolve) => {
pendingResolver?.("cancel");
pendingResolver = resolve;
set({ open: true, allowAnonymous });
}),
resolve: (decision) => {
const resolver = pendingResolver;
pendingResolver = null;
set({ open: false, allowAnonymous: true });
resolver?.(decision);
},
}));

23
studio/frontend/src/features/hub/catalog/model-inspector.tsx
View file

@ -59,6 +59,13 @@ import { ModelReadme } from "./model-readme";
import { OwnerAvatar } from "./owner-avatar";
import { AccessChip, CapabilityPill } from "./shared";
// HF pipeline_tag values authoritative for embedding-only repos; capability
// labels (code/vision/audio) can leak onto them via name or tags.
const EMBEDDING_PIPELINE_TAGS: ReadonlySet<string> = new Set([
"feature-extraction",
"sentence-similarity",
]);
function ViewRepositoryButton({
repoId,
isDataset,
@ -531,11 +538,27 @@ export const ModelInspector = memo(function ModelInspector({
? formatCompact(model.totalParams)
: "N/A";
const unslothSupported = unslothSupport.status !== "unsupported";
// Embedding-only non-GGUF repos have no generative head, so keep them out of
// the Run gate. Prefer the pipeline tag, else the capability heuristic.
const isEmbeddingOnly =
!model.isGguf &&
model.capabilities.some((c) => c.key === "embedding") &&
(EMBEDDING_PIPELINE_TAGS.has(model.pipelineTag?.toLowerCase() ?? "") ||
!model.capabilities.some(
(c) =>
c.key === "conversational" ||
c.key === "tools" ||
c.key === "reasoning" ||
c.key === "code" ||
c.key === "vision" ||
c.key === "audio",
));
// Chat-only hosts (no supported GPU / usable MLX) run inference only through
// llama.cpp, so only GGUF is loadable.
const canRunModel =
!isDataset &&
(model.runtimeCapabilities?.canChat ?? true) &&
!isEmbeddingOnly &&
(model.isGguf || (!chatOnly && unslothSupported));
const canTrainModel =
!isDataset &&

27
studio/frontend/src/features/rag/api/rag-api.ts
View file

@ -10,6 +10,7 @@ import type {
KnowledgeBase,
PreviewTarget,
RagDocument,
UploadedDocument,
} from "../types/rag";
const RAG_BASE = "/api/rag";
@ -194,10 +195,25 @@ export function invalidateProjectSources(projectId: string): void {
projectSourcesCache.delete(projectId);
}
export function deleteDocument(documentId: string): Promise<{ ok: boolean }> {
return ragRequest(`/documents/${encodeURIComponent(documentId)}`, {
method: "DELETE",
});
export async function listAllDocuments(): Promise<UploadedDocument[]> {
const data = await ragRequest<{ documents: UploadedDocument[] }>(
"/documents",
);
return data.documents ?? [];
}
export async function deleteDocument(
documentId: string,
projectId?: string | null,
): Promise<{ ok: boolean }> {
const result = await ragRequest<{ ok: boolean }>(
`/documents/${encodeURIComponent(documentId)}`,
{
method: "DELETE",
},
);
if (projectId) invalidateProjectSources(projectId);
return result;
}
export function getJob(jobId: string): Promise<IndexJob> {
@ -237,7 +253,8 @@ export async function* streamJobEvents(
const dataLines: string[] = [];
for (const line of rawEvent.split(/\r?\n/)) {
if (line.startsWith("data:")) dataLines.push(line.slice(5).trimStart());
if (line.startsWith("data:"))
dataLines.push(line.slice(5).trimStart());
}
if (dataLines.length > 0) {
const dataText = dataLines.join("\n");

136
studio/frontend/src/features/rag/components/use-rag-documents.ts
View file

@ -2,12 +2,11 @@
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import { useCallback, useEffect, useRef, useState } from "react";
import { useChatRuntimeStore } from "@/features/chat";
import {
CHAT_RAG_CAPTION_KEY,
CHAT_RAG_OCR_KEY,
} from "@/features/chat/stores/chat-runtime-store";
useChatRuntimeStore,
} from "@/features/chat";
import { toast } from "@/lib/toast";
import {
deleteDocument,
@ -60,9 +59,7 @@ export function useRagDocuments(
if (ids.size === 0) return false;
const docs = documentsRef.current.filter((d) => ids.has(d.id));
if (docs.length === 0) return false; // sig tracked but doc gone -> allow re-upload
return docs.some(
(d) => d.status !== "completed" || (d.numChunks ?? 0) > 0,
);
return docs.some((d) => d.status !== "completed" || (d.numChunks ?? 0) > 0);
}, []);
// True while upload() runs, so the scope-change effect can tell a real switch
// from lazy thread materialization mid-upload (which must not reset).
@ -80,9 +77,7 @@ export function useRagDocuments(
const patchDoc = useCallback(
(documentId: string, patch: Partial<TrackedDocument>) => {
setDocuments((rows) =>
rows.map((row) =>
row.id === documentId ? { ...row, ...patch } : row,
),
rows.map((row) => (row.id === documentId ? { ...row, ...patch } : row)),
);
},
[],
@ -176,49 +171,61 @@ export function useRagDocuments(
[patchDoc],
);
const refresh = useCallback(async (opts?: { quiet?: boolean }) => {
if (!scope) return;
if (!opts?.quiet) setLoading(true);
try {
// Merge server truth with local progress so a refresh mid-index keeps a
// live "running %" chip. Failed docs hidden (toast warned at upload).
const rows = (await lister()).filter((row) => row.status !== "failed");
setDocuments((prev) => {
const merged = rows.map((row) => {
const tracked = prev.find((p) => p.id === row.id);
return tracked && tracked.progress != null && row.status !== "completed"
? { ...row, progress: tracked.progress }
: row;
const refresh = useCallback(
async (opts?: { quiet?: boolean }) => {
if (!scope) return;
if (!opts?.quiet) setLoading(true);
try {
// Merge server truth with local progress so a refresh mid-index keeps a
// live "running %" chip. Failed docs hidden (toast warned at upload).
const rows = (await lister()).filter((row) => row.status !== "failed");
setDocuments((prev) => {
const merged = rows.map((row) => {
const tracked = prev.find((p) => p.id === row.id);
return tracked &&
tracked.progress != null &&
row.status !== "completed"
? { ...row, progress: tracked.progress }
: row;
});
// Keep optimistic chips (not yet listed) so a refresh racing an upload
// can't make them vanish.
const serverIds = new Set(rows.map((row) => row.id));
const pendingLocal = prev.filter(
(row) => row.id.startsWith("pending_") && !serverIds.has(row.id),
);
return [...merged, ...pendingLocal];
});
// Keep optimistic chips (not yet listed) so a refresh racing an upload
// can't make them vanish.
const serverIds = new Set(rows.map((row) => row.id));
const pendingLocal = prev.filter(
(row) => row.id.startsWith("pending_") && !serverIds.has(row.id),
);
return [...merged, ...pendingLocal];
});
} catch (err) {
toast.error("Failed to load documents", {
description: err instanceof Error ? err.message : String(err),
});
} finally {
if (!opts?.quiet) setLoading(false);
}
}, [scope, lister]);
} catch (err) {
toast.error("Failed to load documents", {
description: err instanceof Error ? err.message : String(err),
});
} finally {
if (!opts?.quiet) setLoading(false);
}
},
[scope, lister],
);
// A real switch (thread/KB swap) resets + reloads; first acquiring a scope just
// loads. Skip both during materialization mid-upload (scope null -> new thread
// while upload() runs) so we don't abort tracking or wipe optimistic chips.
useEffect(() => {
const jobs = trackedJobs.current;
const prev = prevScopeKeyRef.current;
prevScopeKeyRef.current = scopeKey;
if (prev !== null && prev !== scopeKey) {
for (const controller of trackedJobs.current.values()) controller.abort();
trackedJobs.current.clear();
for (const controller of jobs.values()) controller.abort();
jobs.clear();
sigByDocId.current.clear();
// Scope changes intentionally clear the old scope before fetching the new
// one. Keep this synchronous so React StrictMode's setup/cleanup replay
// cannot cancel the only refresh after prevScopeKeyRef has advanced.
setDocuments([]);
if (scope) void refresh();
if (scope) {
// eslint-disable-next-line react-hooks/set-state-in-effect
void refresh();
}
} else if (prev === null && scope && !uploadInFlightRef.current) {
void refresh();
}
@ -226,8 +233,8 @@ export function useRagDocuments(
// Preserve in-flight tracking when cleanup is the materialization flip,
// not a real switch/unmount.
if (uploadInFlightRef.current) return;
for (const controller of trackedJobs.current.values()) controller.abort();
trackedJobs.current.clear();
for (const controller of jobs.values()) controller.abort();
jobs.clear();
};
// eslint-disable-next-line react-hooks/exhaustive-deps
}, [scopeKey]);
@ -260,21 +267,41 @@ export function useRagDocuments(
// otherwise backend env defaults own the ingest policy.
const state = useChatRuntimeStore.getState();
const hasLocal = (key: string) =>
typeof window !== "undefined" && window.localStorage.getItem(key) !== null;
const ocr = hasLocal(CHAT_RAG_OCR_KEY) ? state.ragOcrScanned : undefined;
typeof window !== "undefined" &&
window.localStorage.getItem(key) !== null;
const ocr = hasLocal(CHAT_RAG_OCR_KEY)
? state.ragOcrScanned
: undefined;
const caption = hasLocal(CHAT_RAG_CAPTION_KEY)
? state.ragCaptionFigures
: undefined;
const result =
activeScope.type === "kb"
? await uploadKnowledgeBaseDocument(activeScope.kbId, file, ocr, caption)
? await uploadKnowledgeBaseDocument(
activeScope.kbId,
file,
ocr,
caption,
)
: activeScope.type === "project"
? await uploadProjectDocument(activeScope.projectId, file, ocr, caption)
: await uploadThreadDocument(activeScope.threadId, file, ocr, caption);
? await uploadProjectDocument(
activeScope.projectId,
file,
ocr,
caption,
)
: await uploadThreadDocument(
activeScope.threadId,
file,
ocr,
caption,
);
sigByDocId.current.set(result.documentId, fileSignature(file));
if (seenIds.has(result.documentId)) {
setDocuments((rows) => rows.filter((row) => row.id !== tempId));
toast.info(`${result.filename || file.name} is already indexed - skipping`);
toast.info(
`${result.filename || file.name} is already indexed - skipping`,
);
return;
}
seenIds.add(result.documentId);
@ -341,7 +368,9 @@ export function useRagDocuments(
]);
const resolved =
overrideScope instanceof Promise ? await overrideScope : overrideScope;
overrideScope instanceof Promise
? await overrideScope
: overrideScope;
const activeScope = resolved ?? scope;
if (!activeScope) {
// Materialization failed: drop the chips so they don't hang "pending".
@ -377,7 +406,10 @@ export function useRagDocuments(
const prevSig = sigByDocId.current.get(documentId);
sigByDocId.current.delete(documentId);
try {
await deleteDocument(documentId);
await deleteDocument(
documentId,
scope?.type === "project" ? scope.projectId : undefined,
);
} catch (err) {
setDocuments(prev);
if (prevSig !== undefined) sigByDocId.current.set(documentId, prevSig);
@ -386,7 +418,7 @@ export function useRagDocuments(
});
}
},
[documents],
[documents, scope],
);
return { documents, loading, uploading, refresh, upload, remove };

7
studio/frontend/src/features/rag/index.ts
View file

@ -5,4 +5,9 @@ export { KnowledgeBaseComposerButton } from "./components/knowledge-base-compose
export { KnowledgeBaseDialog } from "./components/knowledge-base-dialog";
export { RetrievalSettingsSection } from "./components/retrieval-settings-section";
export { ThreadDocumentsBar } from "./components/thread-documents-bar";
export type { KnowledgeBase, RagDocument } from "./types/rag";
export {
deleteDocument,
getDocumentFileUrl,
listAllDocuments,
} from "./api/rag-api";
export type { KnowledgeBase, RagDocument, UploadedDocument } from "./types/rag";

7
studio/frontend/src/features/rag/types/rag.ts
View file

@ -24,6 +24,13 @@ export interface RagDocument {
createdAt?: string | null;
}
/** RagDocument enriched for the global uploaded-files list (settings Data tab). */
export interface UploadedDocument extends RagDocument {
sizeBytes?: number | null;
kbName?: string | null;
projectName?: string | null;
}
export interface DocumentUploadResult {
documentId: string;
jobId: string;

19
studio/frontend/src/features/settings/api/openai-auto-switch.ts
View file

@ -11,6 +11,8 @@ export type OpenAIAutoSwitchSettings = {
// True when the idle-unload loop will actually unload (e.g. enabled via the
// UNSLOTH_MODEL_IDLE_TTL env var even while the toggle is off).
idleUnloadActive: boolean;
// Persist the KV cache to disk on idle unload and restore it on reload.
autoUnloadKeepKv: boolean;
};
type ApiOpenAIAutoSwitchSettings = {
@ -21,6 +23,8 @@ type ApiOpenAIAutoSwitchSettings = {
default_enabled: boolean;
// biome-ignore lint/style/useNamingConvention: API schema
idle_unload_active?: boolean;
// biome-ignore lint/style/useNamingConvention: API schema
auto_unload_keep_kv?: boolean;
};
let cachedSettings: OpenAIAutoSwitchSettings | null = null;
@ -34,6 +38,7 @@ function fromApi(
autoUnloadIdleSeconds: settings.auto_unload_idle_seconds,
defaultEnabled: settings.default_enabled,
idleUnloadActive: settings.idle_unload_active ?? false,
autoUnloadKeepKv: settings.auto_unload_keep_kv ?? true,
};
}
@ -66,15 +71,23 @@ export async function loadOpenAIAutoSwitchSettings() {
export async function updateOpenAIAutoSwitchSettings(
enabled: boolean,
autoUnloadIdleSeconds: number,
autoUnloadIdleSeconds?: number,
autoUnloadKeepKv?: boolean,
): Promise<OpenAIAutoSwitchSettings> {
const res = await authFetch("/api/settings/openai-auto-switch", {
method: "PUT",
headers: { "Content-Type": "application/json" },
body: JSON.stringify({
enabled,
// biome-ignore lint/style/useNamingConvention: API schema
auto_unload_idle_seconds: autoUnloadIdleSeconds,
// Omitted fields keep their stored value.
...(autoUnloadIdleSeconds === undefined
? {}
: // biome-ignore lint/style/useNamingConvention: API schema
{ auto_unload_idle_seconds: autoUnloadIdleSeconds }),
...(autoUnloadKeepKv === undefined
? {}
: // biome-ignore lint/style/useNamingConvention: API schema
{ auto_unload_keep_kv: autoUnloadKeepKv }),
}),
});
if (!res.ok) {

135
studio/frontend/src/features/settings/components/archived-chats-dialog.tsx
View file

@ -12,18 +12,12 @@ import {
AlertDialogTitle,
} from "@/components/ui/alert-dialog";
import {
Dialog,
DialogContent,
DialogHeader,
DialogTitle,
} from "@/components/ui/dialog";
import {
type SidebarItem,
deleteChatItem,
unarchiveChatItem,
useChatPreferencesStore,
useChatRuntimeStore,
useChatSidebarItems,
type SidebarItem,
} from "@/features/chat";
import { toast } from "@/lib/toast";
import { ArchiveRestoreIcon, Delete02Icon } from "@hugeicons/core-free-icons";
@ -40,13 +34,7 @@ function formatCreatedAt(ms: number): string {
});
}
export function ArchivedChatsDialog({
open,
onOpenChange,
}: {
open: boolean;
onOpenChange: (open: boolean) => void;
}) {
export function ArchivedChatsView() {
const { archivedItems } = useChatSidebarItems({ requireMessages: false });
const navigate = useNavigate();
const closeSettings = useSettingsDialogStore((s) => s.closeDialog);
@ -74,7 +62,6 @@ export function ArchivedChatsDialog({
search:
item.type === "single" ? { thread: item.id } : { compare: item.id },
});
onOpenChange(false);
closeSettings();
}
@ -114,72 +101,66 @@ export function ArchivedChatsDialog({
}
return (
<Dialog open={open} onOpenChange={onOpenChange}>
<DialogContent className="max-w-2xl">
<DialogHeader>
<DialogTitle>Archived chats</DialogTitle>
</DialogHeader>
{archivedItems.length === 0 ? (
<p className="py-8 text-center text-sm text-muted-foreground">
No archived chats.
</p>
) : (
<div className="max-h-[60vh] overflow-y-auto">
<div className="flex items-center gap-4 border-b border-border/60 px-1 pb-2 text-xs font-semibold text-foreground">
<span className="flex-1">Name</span>
<span className="w-32 shrink-0">Date created</span>
<span className="w-16 shrink-0" />
</div>
{archivedItems.map((item) => (
<div
key={item.id}
className="group flex items-center gap-4 border-b border-border/40 px-1 py-2.5 text-sm last:border-0"
<div className="flex flex-col gap-4">
{archivedItems.length === 0 ? (
<p className="py-8 text-center text-sm text-muted-foreground">
No archived chats.
</p>
) : (
<div>
<div className="flex items-center gap-4 border-b border-border/60 px-1 pb-2 text-xs font-semibold text-foreground">
<span className="flex-1">Name</span>
<span className="w-32 shrink-0">Date created</span>
<span className="w-16 shrink-0" />
</div>
{archivedItems.map((item) => (
<div
key={item.id}
className="group flex items-center gap-4 border-b border-border/40 px-1 py-2.5 text-sm last:border-0"
>
<button
type="button"
onClick={() => openChat(item)}
className="min-w-0 flex-1 truncate text-left text-primary hover:underline"
title={item.title}
>
{item.title}
</button>
<span className="w-32 shrink-0 text-muted-foreground tabular-nums">
{formatCreatedAt(item.createdAt)}
</span>
<span className="flex w-16 shrink-0 items-center justify-end gap-1">
<button
type="button"
onClick={() => openChat(item)}
className="min-w-0 flex-1 truncate text-left text-primary hover:underline"
title={item.title}
onClick={() => void handleUnarchive(item)}
aria-label="Unarchive chat"
title="Unarchive"
className="inline-flex size-7 items-center justify-center rounded-full text-muted-foreground transition-colors hover:bg-muted hover:text-foreground"
>
{item.title}
<HugeiconsIcon
icon={ArchiveRestoreIcon}
strokeWidth={1.75}
className="size-4"
/>
</button>
<span className="w-32 shrink-0 text-muted-foreground tabular-nums">
{formatCreatedAt(item.createdAt)}
</span>
<span className="flex w-16 shrink-0 items-center justify-end gap-1">
<button
type="button"
onClick={() => void handleUnarchive(item)}
aria-label="Unarchive chat"
title="Unarchive"
className="inline-flex size-7 items-center justify-center rounded-full text-muted-foreground transition-colors hover:bg-muted hover:text-foreground"
>
<HugeiconsIcon
icon={ArchiveRestoreIcon}
strokeWidth={1.75}
className="size-4"
/>
</button>
<button
type="button"
onClick={() => requestDelete(item)}
aria-label="Delete chat"
title="Delete"
className="inline-flex size-7 items-center justify-center rounded-full text-muted-foreground transition-colors hover:bg-destructive/10 hover:text-destructive"
>
<HugeiconsIcon
icon={Delete02Icon}
strokeWidth={1.75}
className="size-4"
/>
</button>
</span>
</div>
))}
</div>
)}
</DialogContent>
<button
type="button"
onClick={() => requestDelete(item)}
aria-label="Delete chat"
title="Delete"
className="inline-flex size-7 items-center justify-center rounded-full text-muted-foreground transition-colors hover:bg-destructive/10 hover:text-destructive"
>
<HugeiconsIcon
icon={Delete02Icon}
strokeWidth={1.75}
className="size-4"
/>
</button>
</span>
</div>
))}
</div>
)}
<AlertDialog
open={confirmingDelete !== null}
@ -213,6 +194,6 @@ export function ArchivedChatsDialog({
</AlertDialogFooter>
</AlertDialogContent>
</AlertDialog>
</Dialog>
</div>
);
}

98
studio/frontend/src/features/settings/components/finetune-recipe.ts Normal file
View file

@ -0,0 +1,98 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
// Settings Data tab glue: turn chat history into a fine-tuning JSONL, stage
// it as a Data Recipe seed upload, and open a new recipe on that file.
import { type FineTuneFormat, buildFineTuneJsonl } from "@/features/chat";
import { saveRecipe } from "@/features/data-recipes/data/recipes-db";
import { createEmptyRecipePayload } from "@/features/recipe-studio";
import { inspectSeedUpload } from "@/features/recipe-studio/api";
import { uploadTrainingDataset } from "@/features/training/api/datasets-api";
import { useTrainingConfigStore } from "@/features/training/stores/training-config-store";
import { toast } from "@/lib/toast";
/** btoa cannot handle code points above latin-1, so encode UTF-8 bytes. */
function base64FromString(value: string): string {
const bytes = new TextEncoder().encode(value);
let binary = "";
const CHUNK = 0x8000;
for (let i = 0; i < bytes.length; i += CHUNK) {
binary += String.fromCharCode(...bytes.subarray(i, i + CHUNK));
}
return btoa(binary);
}
/** Builds the JSONL, uploads it as a local recipe seed, and saves a new
* recipe whose seed block points at the file. Returns the recipe id, or
* null when there is nothing to export. */
export async function createFineTuneRecipeFromChats(
format: FineTuneFormat = "openai",
): Promise<string | null> {
const { lines, conversations } = await buildFineTuneJsonl(format);
if (conversations === 0) {
toast.info("No chats with a user and assistant exchange to export.");
return null;
}
const dateLabel = new Date().toISOString().slice(0, 10);
const suffix = format === "openai" ? "" : `-${format}`;
const filename = `chat-finetune${suffix}-${dateLabel}.jsonl`;
const inspected = await inspectSeedUpload({
filename,
// biome-ignore lint/style/useNamingConvention: api schema
content_base64: base64FromString(lines.join("\n")),
// biome-ignore lint/style/useNamingConvention: api schema
preview_size: 10,
});
const payload = createEmptyRecipePayload();
payload.recipe.seed_config = {
source: {
// biome-ignore lint/style/useNamingConvention: api schema
seed_type: "local",
path: inspected.resolved_path,
},
// biome-ignore lint/style/useNamingConvention: api schema
sampling_strategy: "ordered",
// biome-ignore lint/style/useNamingConvention: api schema
selection_strategy: null,
};
payload.ui.nodes = [{ id: "seed", x: 0, y: 0, width: 400 }];
payload.ui.seed_source_type = "local";
payload.ui.seed_columns = inspected.columns;
payload.ui.seed_preview_rows = inspected.preview_rows ?? [];
payload.ui.local_file_name = filename;
const record = await saveRecipe({
name: `Chat fine-tuning ${dateLabel}`,
payload,
});
return record.id;
}
/** Builds the JSONL, uploads it as a training dataset, and selects it in the
* Train tab's config store so the Train page opens with it loaded. Returns
* false when there is nothing to export. */
export async function loadFineTuneDatasetInTrainTab(
format: FineTuneFormat = "openai",
): Promise<boolean> {
const { lines, conversations } = await buildFineTuneJsonl(format);
if (conversations === 0) {
toast.info("No chats with a user and assistant exchange to export.");
return false;
}
const dateLabel = new Date().toISOString().slice(0, 10);
const suffix = format === "openai" ? "" : `-${format}`;
const file = new File(
[lines.join("\n")],
`chat-finetune${suffix}-${dateLabel}.jsonl`,
{ type: "application/x-ndjson" },
);
const uploaded = await uploadTrainingDataset(file);
// Selecting also kicks off the dataset format check, so the Train tab
// shows the detected format as soon as it mounts.
useTrainingConfigStore.getState().selectLocalDataset(uploaded.stored_path);
return true;
}

26
studio/frontend/src/features/settings/components/model-auto-switch-section.tsx
View file

@ -62,13 +62,18 @@ export function ModelAutoSwitchSection() {
const persist = async (
enabled: boolean,
idleSeconds: number,
idleSeconds: number | undefined,
syncDraft = true,
keepKv?: boolean,
) => {
setIsSaving(true);
setError(null);
try {
const saved = await updateOpenAIAutoSwitchSettings(enabled, idleSeconds);
const saved = await updateOpenAIAutoSwitchSettings(
enabled,
idleSeconds,
keepKv,
);
setSettings(saved);
if (syncDraft) {
setDraftIdleSeconds(String(saved.autoUnloadIdleSeconds));
@ -107,6 +112,11 @@ export function ModelAutoSwitchSection() {
void persist(true, idleSeconds);
};
const handleKeepKvToggle = (keepKv: boolean) => {
if (!settings) return;
void persist(settings.enabled, undefined, false, keepKv);
};
return (
<SettingsSection title={t("settings.general.modelAutoSwitch.sectionTitle")}>
<SettingsRow
@ -166,6 +176,18 @@ export function ModelAutoSwitchSection() {
) : null}
</div>
</SettingsRow>
{settings?.idleUnloadActive ? (
<SettingsRow
label={t("settings.general.modelAutoSwitch.keepKv")}
description={t("settings.general.modelAutoSwitch.keepKvDescription")}
>
<Switch
checked={settings.autoUnloadKeepKv}
disabled={isSaving}
onCheckedChange={handleKeepKvToggle}
/>
</SettingsRow>
) : null}
</SettingsSection>
);
}

644
studio/frontend/src/features/settings/components/uploaded-files-dialog.tsx Normal file
View file

@ -0,0 +1,644 @@
// 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 {
AlertDialog,
AlertDialogAction,
AlertDialogCancel,
AlertDialogContent,
AlertDialogDescription,
AlertDialogFooter,
AlertDialogHeader,
AlertDialogTitle,
} from "@/components/ui/alert-dialog";
import { Spinner } from "@/components/ui/spinner";
import {
type ChatAttachmentRecord,
deleteChatAttachment,
emitChatAttachmentDeleted,
fetchChatAttachmentBlob,
listChatAttachments,
} from "@/features/chat";
import {
deleteDocument,
getDocumentFileUrl,
listAllDocuments,
type UploadedDocument,
} from "@/features/rag";
import { toast } from "@/lib/toast";
import {
ArrowUpRight01Icon,
Delete02Icon,
File02Icon,
} from "@hugeicons/core-free-icons";
import { HugeiconsIcon } from "@hugeicons/react";
import { useNavigate } from "@tanstack/react-router";
import { type ReactNode, useEffect, useRef, useState } from "react";
import { useSettingsDialogStore } from "../stores/settings-dialog-store";
function formatUploadedAt(value: string | number | null | undefined): string {
if (value === null || value === undefined || value === "") return "-";
// Chat attachments carry ms epoch numbers; RAG documents carry SQLite
// ISO-ish strings (no timezone). Unparseable strings fall through raw.
const parsed = new Date(value);
if (Number.isNaN(parsed.getTime())) return String(value);
return parsed.toLocaleDateString(undefined, {
year: "numeric",
month: "long",
day: "numeric",
});
}
function formatSize(bytes: number | null | undefined): string {
if (bytes === null || bytes === undefined) return "-";
if (bytes < 1024) return `${bytes} B`;
const units = ["KB", "MB", "GB"];
let value = bytes;
let unit = "B";
for (const next of units) {
if (value < 1024) break;
value /= 1024;
unit = next;
}
return `${value >= 10 ? Math.round(value) : value.toFixed(1)} ${unit}`;
}
function ragLocationLabel(doc: UploadedDocument): string {
if (doc.kbId) return doc.kbName ? `KB · ${doc.kbName}` : "Knowledge base";
if (doc.projectId) {
return doc.projectName ? `Project · ${doc.projectName}` : "Project";
}
if (doc.threadId) return "Chat files (RAG)";
return "-";
}
/** Short uppercase file-type label from the filename extension, falling back
* to the content-type subtype (e.g. "image/webp" gives WEBP). */
function fileTypeLabel(
name: string,
contentType?: string | null,
): string | null {
const dot = name.lastIndexOf(".");
const ext = dot > 0 ? name.slice(dot + 1).trim() : "";
if (ext && ext.length <= 5) return ext.toUpperCase();
const subtype = contentType?.split("/")[1]?.split("+")[0]?.trim();
return subtype && subtype.length <= 10 ? subtype.toUpperCase() : null;
}
/** Lazy image thumbnail for a chat attachment; a file icon until it loads.
* The stored blob only downloads once the row scrolls into view, so a long
* history of screenshots does not fetch every image on open. */
function ChatImageThumb({
messageId,
attachmentId,
}: {
messageId: string;
attachmentId: string;
}) {
const [src, setSrc] = useState<string | null>(null);
const [visible, setVisible] = useState(false);
const holderRef = useRef<HTMLSpanElement>(null);
useEffect(() => {
const el = holderRef.current;
if (!el) return;
if (typeof IntersectionObserver === "undefined") {
return;
}
const observer = new IntersectionObserver((entries) => {
if (entries.some((entry) => entry.isIntersecting)) {
setVisible(true);
observer.disconnect();
}
});
observer.observe(el);
return () => observer.disconnect();
}, []);
useEffect(() => {
if (!visible) return;
let cancelled = false;
let url: string | null = null;
fetchChatAttachmentBlob(messageId, attachmentId)
.then((blob) => {
if (cancelled) return;
url = URL.createObjectURL(blob);
setSrc(url);
})
.catch(() => {
// Keep the file icon on failure.
});
return () => {
cancelled = true;
if (url) URL.revokeObjectURL(url);
};
}, [visible, messageId, attachmentId]);
if (!src) {
return (
<span
ref={holderRef}
className="flex h-full w-full items-center justify-center"
>
<FileIconThumb />
</span>
);
}
return <img src={src} alt="" className="h-full w-full object-cover" />;
}
function FileIconThumb() {
return (
<HugeiconsIcon
icon={File02Icon}
strokeWidth={1.75}
className="size-4 text-muted-foreground"
/>
);
}
/** One display row: a RAG document or a chat message attachment. */
interface UploadedFileRow {
key: string;
source: "rag" | "chat";
name: string;
location: string;
sizeBytes?: number | null;
createdAt?: string | number | null;
failed?: boolean;
/** Epoch ms for sorting; rows with unknown dates sort last. */
sortTime: number;
typeLabel: string | null;
/** Image rows render a thumbnail; others show a file icon. */
thumb: ReactNode;
/** Chat rows link back to their thread. */
threadId?: string | null;
/** Compare-chat rows navigate by pair id instead of opening one pane alone. */
pairId?: string | null;
open: () => Promise<void>;
remove: () => Promise<void>;
deleteDescription: string;
}
function toSortTime(value: string | number | null | undefined): number {
if (value === null || value === undefined || value === "") return 0;
const parsed = new Date(value).getTime();
return Number.isNaN(parsed) ? 0 : parsed;
}
// Safari and Firefox block window.open after an await (the user gesture is
// gone), so open a blank tab synchronously and point it at the URL once
// resolved. A blocked synchronous open is surfaced instead of silently losing
// the file after the asynchronous URL lookup.
async function openResolvedUrl(resolve: () => Promise<string>): Promise<void> {
const win = window.open("", "_blank");
if (!win) {
throw new Error(
"Your browser blocked the new tab. Allow popups and retry.",
);
}
win.opener = null;
let url: string;
try {
url = await resolve();
} catch (err) {
win.close();
throw err;
}
win.location.replace(url);
}
function ragRow(doc: UploadedDocument): UploadedFileRow {
return {
key: `rag-${doc.id}`,
source: "rag",
name: doc.filename,
location: ragLocationLabel(doc),
sizeBytes: doc.sizeBytes,
createdAt: doc.createdAt,
failed: doc.status === "failed",
sortTime: toSortTime(doc.createdAt),
typeLabel: fileTypeLabel(doc.filename),
// RAG uploads are documents (pdf, txt, md, docx, html), not images.
thumb: <FileIconThumb />,
open: () => openResolvedUrl(() => getDocumentFileUrl(doc.id)),
remove: async () => {
await deleteDocument(doc.id, doc.projectId);
},
deleteDescription:
"The file and its indexed content are removed. This cannot be undone.",
};
}
function chatAttachmentRow(att: ChatAttachmentRecord): UploadedFileRow {
const isImage =
att.type === "image" || Boolean(att.contentType?.startsWith("image/"));
return {
key: `chat-${att.messageId}-${att.id}`,
source: "chat",
name: att.name,
location: att.threadTitle ? `Chat · ${att.threadTitle}` : "Chat",
sizeBytes: att.sizeBytes,
createdAt: att.createdAt,
sortTime: toSortTime(att.createdAt),
typeLabel: fileTypeLabel(att.name, att.contentType),
threadId: att.threadId,
pairId: att.pairId,
thumb: isImage ? (
<ChatImageThumb messageId={att.messageId} attachmentId={att.id} />
) : (
<FileIconThumb />
),
open: () =>
openResolvedUrl(async () => {
const blob = await fetchChatAttachmentBlob(att.messageId, att.id);
const url = URL.createObjectURL(blob);
// Give the new tab time to load the blob before revoking.
setTimeout(() => URL.revokeObjectURL(url), 60_000);
return url;
}),
remove: async () => {
await deleteChatAttachment(att.messageId, att.id);
// Patch any loaded runtime copy so a later repo sync cannot write the
// deleted attachment back to storage.
emitChatAttachmentDeleted({
messageId: att.messageId,
attachmentId: att.id,
});
},
deleteDescription:
"The attachment is removed from its chat message; the message text is kept. This cannot be undone.",
};
}
type SourceLoad<T> = {
status: "loading" | "ready" | "error";
data: T;
error: string | null;
};
function errorMessage(error: unknown, fallback: string): string {
return error instanceof Error ? error.message : fallback;
}
/** Inline settings page listing uploaded files from each available source. */
export function UploadedFilesView() {
const [ragFiles, setRagFiles] = useState<SourceLoad<UploadedDocument[]>>({
status: "loading",
data: [],
error: null,
});
const [chatFiles, setChatFiles] = useState<
SourceLoad<ChatAttachmentRecord[]>
>({ status: "loading", data: [], error: null });
const [chatNextOffset, setChatNextOffset] = useState<number | null>(null);
const [loadingMore, setLoadingMore] = useState(false);
const [confirmingDelete, setConfirmingDelete] =
useState<UploadedFileRow | null>(null);
const navigate = useNavigate();
const rows = [
...ragFiles.data.map(ragRow),
...chatFiles.data.map(chatAttachmentRow),
].sort((a, b) => b.sortTime - a.sortTime);
// Jump to the chat thread the attachment lives in, closing the settings
// dialog so the thread is actually visible.
function goToChat(row: UploadedFileRow) {
if (!row.threadId) return;
useSettingsDialogStore.getState().closeDialog();
if (row.pairId) {
void navigate({ to: "/chat", search: { compare: row.pairId } });
} else {
void navigate({ to: "/chat", search: { thread: row.threadId } });
}
}
useEffect(() => {
let cancelled = false;
void listAllDocuments().then(
(data) => {
if (!cancelled) setRagFiles({ status: "ready", data, error: null });
},
(error: unknown) => {
if (!cancelled) {
setRagFiles({
status: "error",
data: [],
error: errorMessage(error, "Failed to load RAG documents"),
});
}
},
);
void listChatAttachments().then(
(page) => {
if (!cancelled) {
setChatFiles({
status: "ready",
data: page.attachments,
error: null,
});
setChatNextOffset(page.nextOffset);
}
},
(error: unknown) => {
if (!cancelled) {
setChatFiles({
status: "error",
data: [],
error: errorMessage(error, "Failed to load chat attachments"),
});
}
},
);
return () => {
cancelled = true;
};
}, []);
function retryRagFiles() {
setRagFiles((current) => ({ ...current, status: "loading", error: null }));
void listAllDocuments().then(
(data) => setRagFiles({ status: "ready", data, error: null }),
(error: unknown) =>
setRagFiles((current) => ({
...current,
status: "error",
error: errorMessage(error, "Failed to load RAG documents"),
})),
);
}
async function loadChatPage(offset: number, append: boolean) {
setLoadingMore(true);
setChatFiles((current) => ({ ...current, status: "loading", error: null }));
try {
const page = await listChatAttachments(offset);
setChatFiles((current) => ({
status: "ready",
data: append
? [
...current.data,
...page.attachments.filter(
(incoming) =>
!current.data.some(
(existing) =>
existing.id === incoming.id &&
existing.messageId === incoming.messageId,
),
),
]
: page.attachments,
error: null,
}));
setChatNextOffset(page.nextOffset);
} catch (error) {
setChatFiles((current) => ({
...current,
status: "error",
error: errorMessage(error, "Failed to load chat attachments"),
}));
} finally {
setLoadingMore(false);
}
}
function retryChatFiles() {
const append = chatFiles.data.length > 0 && chatNextOffset !== null;
void loadChatPage(append ? chatNextOffset : 0, append);
}
async function handleOpen(row: UploadedFileRow) {
try {
await row.open();
} catch (err) {
toast.error("Failed to open file", {
description: err instanceof Error ? err.message : undefined,
});
}
}
async function handleDelete(row: UploadedFileRow) {
// Offset pages and destructive mutations must not race: a deletion shifts
// the boundary used by an in-flight page request.
if (loadingMore) return;
try {
await row.remove();
if (row.source === "rag") {
setRagFiles((current) => ({
...current,
data: current.data.filter((doc) => `rag-${doc.id}` !== row.key),
}));
} else {
setChatFiles((current) => ({
...current,
data: current.data.filter(
(attachment) =>
`chat-${attachment.messageId}-${attachment.id}` !== row.key,
),
}));
// Offset pagination is relative to the current server inventory. A
// deletion before the next page shifts every later row back by one.
setChatNextOffset((current) =>
current === null ? null : Math.max(0, current - 1),
);
}
toast.success("File deleted");
} catch (err) {
toast.error("Failed to delete file", {
description: err instanceof Error ? err.message : undefined,
});
}
}
return (
<div className="flex flex-col gap-4">
{ragFiles.status === "error" ? (
<div className="flex flex-wrap items-center justify-between gap-2 rounded-md border border-destructive/30 bg-destructive/5 px-3 py-2 text-sm">
<span>RAG documents unavailable: {ragFiles.error}</span>
<button
type="button"
onClick={retryRagFiles}
className="font-medium underline underline-offset-2"
>
Retry
</button>
</div>
) : null}
{chatFiles.status === "error" ? (
<div className="flex flex-wrap items-center justify-between gap-2 rounded-md border border-destructive/30 bg-destructive/5 px-3 py-2 text-sm">
<span>Chat attachments unavailable: {chatFiles.error}</span>
<button
type="button"
onClick={retryChatFiles}
className="font-medium underline underline-offset-2"
>
Retry
</button>
</div>
) : null}
{rows.length === 0 &&
(ragFiles.status === "loading" || chatFiles.status === "loading") ? (
<div className="flex justify-center py-8">
<Spinner className="size-5 text-muted-foreground" />
</div>
) : rows.length === 0 &&
ragFiles.status !== "error" &&
chatFiles.status !== "error" ? (
<p className="py-8 text-center text-sm text-muted-foreground">
No uploaded files.
</p>
) : rows.length > 0 ? (
<div>
<div className="hidden items-center gap-3 border-b border-border/60 px-1 pb-2 text-xs font-semibold text-foreground sm:flex">
<span className="flex-1">Name</span>
<span className="w-36 shrink-0">Location</span>
<span className="w-24 shrink-0">Uploaded</span>
<span className="w-16 shrink-0" />
</div>
{rows.map((row) => (
<div
key={row.key}
className="group flex flex-wrap items-center gap-x-3 gap-y-1 border-b border-border/40 px-1 py-2.5 text-sm last:border-0 sm:flex-nowrap"
>
{/* Clicking the file jumps to its chat; files without one
open directly. The theme scales rounded-md up to a near
circle at this size, so the thumb pins a small radius. */}
<button
type="button"
onClick={() =>
row.threadId ? goToChat(row) : void handleOpen(row)
}
title={
row.threadId ? `Go to ${row.location}` : `Open ${row.name}`
}
className="group/name flex min-w-0 flex-1 basis-[calc(100%-5rem)] items-center gap-2.5 overflow-hidden text-left sm:basis-auto"
>
<span className="flex size-8 shrink-0 items-center justify-center overflow-hidden rounded-[7px] border border-border/50 bg-muted/40">
{row.thumb}
</span>
<span className="flex min-w-0 flex-1 flex-col">
<span className="flex min-w-0 items-center gap-2">
{/* Floor keeps the name visible when the chip and fixed
columns squeeze the cell at narrow widths. */}
<span className="min-w-[3.5rem] truncate underline-offset-2 group-hover/name:underline">
{row.name}
</span>
{row.typeLabel ? (
<span className="shrink-0 rounded-md bg-black/[0.06] px-1.5 py-px text-[9px] font-medium uppercase tracking-wide text-muted-foreground dark:bg-white/[0.1]">
{row.typeLabel}
</span>
) : null}
{row.failed ? (
<span className="shrink-0 text-xs text-destructive">
failed
</span>
) : null}
</span>
<span className="text-xs text-muted-foreground tabular-nums">
{formatSize(row.sizeBytes)}
</span>
</span>
</button>
{row.threadId ? (
<button
type="button"
onClick={() => goToChat(row)}
title={`Go to ${row.location}`}
className="order-3 w-full truncate pl-10 text-left text-muted-foreground underline-offset-2 transition-colors hover:text-foreground hover:underline sm:order-none sm:w-36 sm:pl-0"
>
{row.location}
</button>
) : (
<span
className="order-3 w-full truncate pl-10 text-muted-foreground sm:order-none sm:w-36 sm:pl-0"
title={row.location}
>
{row.location}
</span>
)}
<span className="order-4 w-full pl-10 text-muted-foreground tabular-nums sm:order-none sm:w-24 sm:pl-0">
{formatUploadedAt(row.createdAt)}
</span>
<span className="flex w-16 shrink-0 items-center justify-end gap-1">
<button
type="button"
onClick={() => void handleOpen(row)}
aria-label={`Open ${row.name}`}
title="Open"
className="inline-flex size-7 items-center justify-center rounded-full text-muted-foreground transition-colors hover:bg-muted hover:text-foreground"
>
<HugeiconsIcon
icon={ArrowUpRight01Icon}
strokeWidth={1.75}
className="size-4"
/>
</button>
<button
type="button"
disabled={loadingMore}
onClick={() => setConfirmingDelete(row)}
aria-label={`Delete ${row.name}`}
title="Delete"
className="inline-flex size-7 items-center justify-center rounded-full text-muted-foreground transition-colors hover:bg-destructive/10 hover:text-destructive disabled:cursor-wait disabled:opacity-50"
>
<HugeiconsIcon
icon={Delete02Icon}
strokeWidth={1.75}
className="size-4"
/>
</button>
</span>
</div>
))}
{chatNextOffset !== null ? (
<div className="flex justify-center pt-3">
<button
type="button"
disabled={loadingMore}
onClick={() => void loadChatPage(chatNextOffset, true)}
className="rounded-md border border-border px-3 py-1.5 text-sm font-medium hover:bg-muted disabled:cursor-wait disabled:opacity-60"
>
{loadingMore ? "Loading..." : "Load more chat attachments"}
</button>
</div>
) : null}
</div>
) : null}
<AlertDialog
open={confirmingDelete !== null}
onOpenChange={(o) => {
if (!o) setConfirmingDelete(null);
}}
>
<AlertDialogContent>
<AlertDialogHeader>
<AlertDialogTitle>Delete file</AlertDialogTitle>
<AlertDialogDescription>
Delete{" "}
<span className="font-medium text-foreground">
&quot;{confirmingDelete?.name}&quot;
</span>
? {confirmingDelete?.deleteDescription}
</AlertDialogDescription>
</AlertDialogHeader>
<AlertDialogFooter>
<AlertDialogCancel>Cancel</AlertDialogCancel>
<AlertDialogAction
variant="destructive"
onClick={() => {
const row = confirmingDelete;
setConfirmingDelete(null);
if (row) void handleDelete(row);
}}
>
Delete
</AlertDialogAction>
</AlertDialogFooter>
</AlertDialogContent>
</AlertDialog>
</div>
);
}

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

@ -15,6 +15,7 @@ import {
Cancel01Icon,
CloudIcon,
CpuIcon,
DatabaseSettingIcon,
Globe02Icon,
HelpCircleIcon,
Message01Icon,
@ -43,6 +44,7 @@ import { ApiKeysTab } from "./tabs/api-keys-tab";
import { AppearanceTab } from "./tabs/appearance-tab";
import { ChatTab } from "./tabs/chat-tab";
import { ConnectionsTab } from "./tabs/connections-tab";
import { DataTab } from "./tabs/data-tab";
import { GeneralTab } from "./tabs/general-tab";
import { ProfileTab } from "./tabs/profile-tab";
import { ResourcesTab } from "./tabs/resources-tab";
@ -93,6 +95,12 @@ const TABS: TabDef[] = [
iconComponent: MicIcon,
badgeKey: "common.new",
},
{
id: "data",
labelKey: "settings.tabs.data",
icon: DatabaseSettingIcon,
badgeKey: "common.new",
},
{ id: "about", labelKey: "settings.tabs.about", icon: HelpCircleIcon },
];
@ -112,6 +120,8 @@ function renderTab(tab: SettingsTab) {
return <VoiceTab />;
case "connections":
return <ConnectionsTab />;
case "data":
return <DataTab />;
case "api-keys":
return <ApiKeysTab />;
case "about":
@ -210,6 +220,7 @@ export function SettingsDialog() {
chat: null,
voice: null,
connections: null,
data: null,
"api-keys": null,
about: null,
});

13
studio/frontend/src/features/settings/settings-search.ts
View file

@ -85,12 +85,19 @@ export const SETTINGS_SEARCH_INDEX: Record<SettingsTab, TranslationKey[]> = {
"settings.chat.artifacts.title",
"settings.chat.artifacts.collapseHtmlBlocks",
"settings.chat.artifacts.allowNetworkAccess",
"settings.chat.data",
"settings.chat.modelDisclaimer",
],
// Chat data management moved to the Data tab; keep these rows findable there.
data: [
"settings.data.fineTuneExport",
"settings.data.archivedChats",
"settings.data.archiveAllChats",
"settings.data.confirmBeforeDeleting",
"settings.data.uploadedFiles",
"settings.chat.exportHistory",
"settings.chat.exportConversations",
"settings.chat.importChats",
"settings.chat.clearAllChats",
"settings.chat.exportHistory",
"settings.chat.modelDisclaimer",
],
"api-keys": [
"settings.apiKeys.title",

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

@ -11,6 +11,7 @@ export type SettingsTab =
| "chat"
| "voice"
| "connections"
| "data"
| "api-keys"
| "about";
@ -30,7 +31,7 @@ interface SettingsDialogState {
// explicitly via onCloseAutoFocus.
opener: HTMLElement | null;
// Set when something asks to jump straight to the archived chats list (the
// archive toast). ChatTab consumes it to open the dialog, then clears it.
// archive toast). DataTab uses it as its initial subpage, then clears it.
archivedChatsRequested: boolean;
openDialog: (tab?: SettingsTab, options?: OpenDialogOptions) => void;
openArchivedChats: () => void;
@ -66,6 +67,7 @@ function loadInitialTab(): SettingsTab {
"chat",
"voice",
"connections",
"data",
"api-keys",
"about",
];
@ -90,7 +92,7 @@ export const useSettingsDialogStore = create<SettingsDialogState>((set) => ({
openArchivedChats: () =>
set({
open: true,
activeTab: "chat",
activeTab: "data",
scrollTarget: null,
archivedChatsRequested: true,
opener: captureOpener(),

331
studio/frontend/src/features/settings/tabs/chat-tab.tsx
View file

@ -1,43 +1,16 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import { Button } from "@/components/ui/button";
import {
Dialog,
DialogContent,
DialogDescription,
DialogFooter,
DialogHeader,
DialogTitle,
} from "@/components/ui/dialog";
import {
DropdownMenu,
DropdownMenuContent,
DropdownMenuItem,
DropdownMenuSeparator,
DropdownMenuSub,
DropdownMenuSubContent,
DropdownMenuSubTrigger,
DropdownMenuTrigger,
} from "@/components/ui/dropdown-menu";
import { Switch } from "@/components/ui/switch";
import {
EXPORT_FORMATS_LIST,
type PlusMenuItemId,
bulkExportConversationsByScope,
clearAllChats,
countAllChats,
downloadChatExport,
importConversationsFromFile,
useChatPreferencesStore,
useChatRuntimeStore,
usePlusMenuPrefsStore,
} from "@/features/chat";
import { useT } from "@/i18n";
import { toast } from "@/lib/toast";
import {
Bookmark02Icon,
Delete02Icon,
Download01Icon,
FileDatabaseIcon,
Folder01Icon,
@ -45,19 +18,16 @@ import {
PencilRulerIcon,
Settings02Icon,
ShieldBanIcon,
Upload01Icon,
} from "@hugeicons/core-free-icons";
import { HugeiconsIcon } from "@hugeicons/react";
import { Columns2Icon, PlusIcon } from "lucide-react";
import { useEffect, useRef, useState } from "react";
import { useEffect } from "react";
import type { ReactNode } from "react";
import { ArchivedChatsDialog } from "../components/archived-chats-dialog";
import { SettingsRow } from "../components/settings-row";
import {
SettingsGroupDivider,
SettingsSection,
} from "../components/settings-section";
import { useSettingsDialogStore } from "../stores/settings-dialog-store";
// Adjustable "+" menu items shown in settings, in display order. Icons mirror
// the ones used in the composer + menu itself.
@ -140,7 +110,7 @@ const PLUS_MENU_SETTINGS: {
},
{
id: "bypassPermissions",
label: "Bypass permissions",
label: "Tool permissions",
icon: (
<HugeiconsIcon
icon={ShieldBanIcon}
@ -155,24 +125,6 @@ export function ChatTab() {
const t = useT();
const plusPins = usePlusMenuPrefsStore((state) => state.pins);
const togglePlusPin = usePlusMenuPrefsStore((state) => state.togglePin);
const [confirmOpen, setConfirmOpen] = useState(false);
const [archivedOpen, setArchivedOpen] = useState(false);
const [count, setCount] = useState<number | null>(null);
const archivedChatsRequested = useSettingsDialogStore(
(s) => s.archivedChatsRequested,
);
const consumeArchivedChatsRequest = useSettingsDialogStore(
(s) => s.consumeArchivedChatsRequest,
);
// Open the archived list when the archive toast asked to jump here.
useEffect(() => {
if (!archivedChatsRequested) return;
setArchivedOpen(true);
consumeArchivedChatsRequest();
}, [archivedChatsRequested, consumeArchivedChatsRequest]);
const [exporting, setExporting] = useState(false);
const [clearing, setClearing] = useState(false);
const autoTitle = useChatRuntimeStore((state) => state.autoTitle);
const setAutoTitle = useChatRuntimeStore((state) => state.setAutoTitle);
const showCanvasMenuItem = useChatRuntimeStore(
@ -212,12 +164,6 @@ export function ChatTab() {
const setShowAllQuantizations = useChatRuntimeStore(
(state) => state.setShowAllQuantizations,
);
const confirmDeleteChats = useChatPreferencesStore(
(state) => state.confirmDeleteChats,
);
const setConfirmDeleteChats = useChatPreferencesStore(
(state) => state.setConfirmDeleteChats,
);
const showModelDisclaimer = useChatPreferencesStore(
(state) => state.showModelDisclaimer,
);
@ -232,95 +178,9 @@ export function ChatTab() {
);
useEffect(() => {
void countAllChats().then(setCount);
void hydratePersistedSettings();
}, [hydratePersistedSettings]);
const handleExport = async () => {
setExporting(true);
try {
await downloadChatExport();
} finally {
setExporting(false);
}
};
const importInputRef = useRef<HTMLInputElement>(null);
const handleImport = async (file: File) => {
try {
const imported = await importConversationsFromFile(file, null);
if (imported === 0) {
toast.info(t("settings.chat.importNoConversations"));
} else {
toast.success(
imported === 1
? t("settings.chat.importedOneChat")
: t("settings.chat.importedChatCount", { count: imported }),
);
setCount(await countAllChats().catch(() => count));
}
} catch {
toast.error(t("settings.chat.importFailed"));
}
};
const handleClear = async () => {
setClearing(true);
try {
const result = await clearAllChats();
const clearedCount = result.deletedThreadIds.length;
const hasFailedStore =
result.backend === "failed" || result.legacy === "failed";
if (!hasFailedStore && result.failedThreadIds.length === 0) {
setCount(0);
setConfirmOpen(false);
toast.success(
clearedCount === 0
? t("settings.chat.clearedAllChats")
: clearedCount === 1
? t("settings.chat.clearedOneChat")
: t("settings.chat.clearedChatCount", { count: clearedCount }),
);
return;
}
const fallbackRemaining =
result.failedThreadIds.length > 0
? result.failedThreadIds.length
: (count ?? 0);
const remaining = await countAllChats().catch(() => fallbackRemaining);
setCount(remaining);
setConfirmOpen(false);
toast.warning(t("settings.chat.someChatsCouldNotBeCleared"), {
description:
result.failedThreadIds.length > 0
? clearedCount === 1 && result.failedThreadIds.length === 1
? t("settings.chat.oneChatClearedRemainOne")
: clearedCount === 1
? t("settings.chat.oneChatClearedRemain", {
remainingCount: result.failedThreadIds.length,
})
: result.failedThreadIds.length === 1
? t("settings.chat.chatsClearedRemainOne", { clearedCount })
: t("settings.chat.chatsClearedRemain", {
clearedCount,
remainingCount: result.failedThreadIds.length,
})
: remaining === 1
? t("settings.chat.storageClearFailedOne")
: t("settings.chat.storageClearFailed", { count: remaining }),
});
} catch (error) {
const remaining = await countAllChats().catch(() => count);
setCount(remaining);
toast.error(t("settings.chat.failedToClearChats"), {
description: error instanceof Error ? error.message : undefined,
});
} finally {
setClearing(false);
}
};
return (
<div className="flex flex-col gap-6">
<header className="flex flex-col gap-1">
@ -347,7 +207,7 @@ export function ChatTab() {
<span className="font-mono text-xs text-foreground">
Q4_K_M
</span>
<span className="text-[9px] font-medium text-green-400">
<span className="text-[9px] font-medium text-green-600/90 dark:text-green-400/80">
downloaded
</span>
<span className="text-[10px] text-muted-foreground">16 GB</span>
@ -481,191 +341,6 @@ export function ChatTab() {
/>
</SettingsRow>
</SettingsSection>
<SettingsSection title={t("settings.chat.data")}>
<SettingsRow
label="Archived chats"
description="View and manage chats you have archived."
>
<Button
variant="outline"
size="sm"
onClick={() => setArchivedOpen(true)}
>
Manage
</Button>
</SettingsRow>
<SettingsRow
label="Confirm before deleting"
description="Ask for confirmation before a chat is deleted. Turn off to delete instantly."
>
<Switch
checked={confirmDeleteChats}
onCheckedChange={setConfirmDeleteChats}
/>
</SettingsRow>
<SettingsRow
label={t("settings.chat.exportHistory")}
description={t("settings.chat.exportHistoryDescription")}
>
<Button
variant="outline"
size="sm"
onClick={handleExport}
disabled={exporting || count === 0}
>
<HugeiconsIcon icon={Download01Icon} className="size-3.5 mr-1.5" />
{exporting
? t("settings.chat.exportingAction")
: t("settings.chat.exportAction")}
</Button>
</SettingsRow>
<SettingsRow
label={t("settings.chat.exportConversations")}
description={t("settings.chat.exportConversationsDescription")}
>
<DropdownMenu>
<DropdownMenuTrigger asChild={true}>
<Button variant="outline" size="sm" disabled={count === 0}>
<HugeiconsIcon
icon={Download01Icon}
className="size-3.5 mr-1.5"
/>
{t("settings.chat.exportConversationsAction")}
</Button>
</DropdownMenuTrigger>
<DropdownMenuContent align="end" className="w-56">
{(
[
{ scope: "recents", label: "exportScopeRecents" },
{ scope: "all", label: "exportScopeAll" },
] as const
).map(({ scope, label }) => (
<DropdownMenuSub key={scope}>
<DropdownMenuSubTrigger>
<HugeiconsIcon
icon={Download01Icon}
className="size-3.5 mr-1"
/>
{t(`settings.chat.${label}`)}
</DropdownMenuSubTrigger>
<DropdownMenuSubContent className="w-56">
{EXPORT_FORMATS_LIST.map(({ fmt, label: fmtLabel }) => (
<DropdownMenuItem
key={`${scope}-m-${fmt}`}
onSelect={() =>
void bulkExportConversationsByScope(scope, fmt, true)
}
>
{fmtLabel} {t("settings.chat.exportCombinedSuffix")}
</DropdownMenuItem>
))}
<DropdownMenuSeparator />
{EXPORT_FORMATS_LIST.map(({ fmt, label: fmtLabel }) => (
<DropdownMenuItem
key={`${scope}-s-${fmt}`}
onSelect={() =>
void bulkExportConversationsByScope(scope, fmt, false)
}
>
{fmtLabel} {t("settings.chat.exportPerChatSuffix")}
</DropdownMenuItem>
))}
</DropdownMenuSubContent>
</DropdownMenuSub>
))}
</DropdownMenuContent>
</DropdownMenu>
</SettingsRow>
<SettingsRow
label={t("settings.chat.importChats")}
description={t("settings.chat.importChatsDescription")}
>
<Button
variant="outline"
size="sm"
onClick={() => importInputRef.current?.click()}
>
<HugeiconsIcon icon={Upload01Icon} className="size-3.5 mr-1.5" />
{t("settings.chat.importChatsAction")}
</Button>
<input
ref={importInputRef}
type="file"
accept=".jsonl,.ndjson,.csv"
className="hidden"
onChange={(e) => {
const file = e.target.files?.[0];
e.target.value = "";
if (file) void handleImport(file);
}}
/>
</SettingsRow>
<SettingsRow
destructive={true}
label={t("settings.chat.clearAllChats")}
description={
count === null
? t("settings.chat.clearAllChatsDescription")
: count === 0
? t("settings.chat.noChatsToClear")
: count === 1
? t("settings.chat.clearOneChatDescription")
: t("settings.chat.clearChatCountDescription", { count })
}
>
<Button
variant="outline"
size="sm"
onClick={() => setConfirmOpen(true)}
disabled={count === 0}
className="text-destructive hover:text-destructive hover:border-destructive/60"
>
<HugeiconsIcon icon={Delete02Icon} className="size-3.5 mr-1.5" />
{t("settings.chat.clearChatsAction")}
</Button>
</SettingsRow>
</SettingsSection>
<ArchivedChatsDialog open={archivedOpen} onOpenChange={setArchivedOpen} />
<Dialog open={confirmOpen} onOpenChange={setConfirmOpen}>
<DialogContent className="max-w-md">
<DialogHeader>
<DialogTitle>
{count === 1
? t("settings.chat.clearOneChatTitle")
: t("settings.chat.clearChatsTitle", { count: count ?? 0 })}
</DialogTitle>
<DialogDescription>
{t("settings.chat.clearChatsConfirmDescription")}
</DialogDescription>
</DialogHeader>
<DialogFooter>
<Button variant="outline" onClick={() => setConfirmOpen(false)}>
{t("common.cancel")}
</Button>
<Button
onClick={handleClear}
disabled={clearing}
className="bg-destructive hover:bg-destructive/90 text-destructive-foreground"
>
{clearing
? t("settings.chat.clearingAction")
: count === 1
? t("settings.chat.clearOneChatAction")
: t("settings.chat.clearChatCountAction", {
count: count ?? 0,
})}
</Button>
</DialogFooter>
</DialogContent>
</Dialog>
</div>
);
}

727
studio/frontend/src/features/settings/tabs/data-tab.tsx Normal file
View file

@ -0,0 +1,727 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import { Button } from "@/components/ui/button";
import {
Dialog,
DialogContent,
DialogDescription,
DialogFooter,
DialogHeader,
DialogTitle,
} from "@/components/ui/dialog";
import {
DropdownMenu,
DropdownMenuContent,
DropdownMenuItem,
DropdownMenuSeparator,
DropdownMenuSub,
DropdownMenuSubContent,
DropdownMenuSubTrigger,
DropdownMenuTrigger,
} from "@/components/ui/dropdown-menu";
import { Spinner } from "@/components/ui/spinner";
import { Switch } from "@/components/ui/switch";
import { usePlatformStore } from "@/config/env";
import {
EXPORT_FORMATS_LIST,
type FineTuneFormat,
archiveAllChatItems,
bulkExportConversationsByScope,
clearAllChats,
countAllChats,
downloadArchivedChatExport,
downloadChatExport,
exportFineTuneJsonl,
importConversationsFromFile,
useChatPreferencesStore,
useChatRuntimeStore,
useChatSidebarItems,
} from "@/features/chat";
import { useT } from "@/i18n";
import {
ChevronDownStandardIcon,
ChevronRightStandardIcon,
} from "@/lib/chevron-icons";
import { toast } from "@/lib/toast";
import {
Archive02Icon,
ArrowLeft01Icon,
Delete02Icon,
Download01Icon,
Tick02Icon,
Upload01Icon,
} from "@hugeicons/core-free-icons";
import { HugeiconsIcon } from "@hugeicons/react";
import { useNavigate, useRouterState } from "@tanstack/react-router";
import { useEffect, useRef, useState } from "react";
import { ArchivedChatsView } from "../components/archived-chats-dialog";
import {
createFineTuneRecipeFromChats,
loadFineTuneDatasetInTrainTab,
} from "../components/finetune-recipe";
import { SettingsRow } from "../components/settings-row";
import { SettingsSection } from "../components/settings-section";
import { UploadedFilesView } from "../components/uploaded-files-dialog";
import { useSettingsDialogStore } from "../stores/settings-dialog-store";
export function DataTab() {
const t = useT();
const navigate = useNavigate();
const archivedChatsRequested = useSettingsDialogStore(
(s) => s.archivedChatsRequested,
);
const consumeArchivedChatsRequest = useSettingsDialogStore(
(s) => s.consumeArchivedChatsRequest,
);
const [confirmOpen, setConfirmOpen] = useState(false);
const [archiveConfirmOpen, setArchiveConfirmOpen] = useState(false);
// Subpages swap the Data tab body instead of opening nested dialogs.
const [subpage, setSubpage] = useState<"main" | "archived" | "files">(
archivedChatsRequested ? "archived" : "main",
);
const [count, setCount] = useState<number | null>(null);
const [exporting, setExporting] = useState(false);
const [archivedExporting, setArchivedExporting] = useState(false);
// Gates the archived subpage Export button.
const { archivedItems } = useChatSidebarItems({ requireMessages: false });
const [clearing, setClearing] = useState(false);
const [archiving, setArchiving] = useState(false);
const [fineTuneExporting, setFineTuneExporting] = useState(false);
const [openingRecipe, setOpeningRecipe] = useState(false);
const [loadingTraining, setLoadingTraining] = useState(false);
// Chat-only hosts redirect /studio back to /chat, so loading a dataset in
// the Train tab would upload it and then strand the user; gate the action
// the same way the sidebar gates Train.
const chatOnly = usePlatformStore((s) => s.isChatOnly());
const [fineTuneAction, setFineTuneAction] = useState<
"train" | "recipes" | "export"
>(chatOnly ? "export" : "train");
// Chat Completions (OpenAI messages) is the only export format we ship.
const fineTuneFormat: FineTuneFormat = "openai";
// The MLX self-heal can flip chat-only while the dialog is open.
useEffect(() => {
if (chatOnly) {
setFineTuneAction((a) => (a === "train" ? "export" : a));
}
}, [chatOnly]);
// Requests can arrive after Data is already mounted (for example from the
// archive-all toast), so always switch before consuming the flag.
useEffect(() => {
if (!archivedChatsRequested) return;
let cancelled = false;
queueMicrotask(() => {
if (cancelled) return;
setSubpage("archived");
consumeArchivedChatsRequest();
});
return () => {
cancelled = true;
};
}, [archivedChatsRequested, consumeArchivedChatsRequest]);
const confirmDeleteChats = useChatPreferencesStore(
(state) => state.confirmDeleteChats,
);
const setConfirmDeleteChats = useChatPreferencesStore(
(state) => state.setConfirmDeleteChats,
);
const storeThreadId = useChatRuntimeStore((s) => s.activeThreadId);
// Open chat id from the route (single thread or compare pair), mirroring
// ArchivedChatsView: compare panes only live in the search params.
const openChatId = useRouterState({
select: (s) => {
if (!s.location.pathname.startsWith("/chat")) return undefined;
const search = s.location.search as Record<string, string | undefined>;
return search.thread ?? search.compare ?? storeThreadId ?? undefined;
},
});
useEffect(() => {
void countAllChats().then(setCount);
}, []);
const handleExport = async () => {
setExporting(true);
try {
await downloadChatExport();
} finally {
setExporting(false);
}
};
const handleExportArchived = async () => {
setArchivedExporting(true);
try {
const exported = await downloadArchivedChatExport();
toast.success(
exported === 0
? t("settings.data.noArchivedChatsToExport")
: exported === 1
? t("settings.data.exportedOneArchivedChat")
: t("settings.data.exportedArchivedChatCount", { count: exported }),
);
} catch (error) {
toast.error(t("settings.data.failedToExportArchivedChats"), {
description: error instanceof Error ? error.message : undefined,
});
} finally {
setArchivedExporting(false);
}
};
const importInputRef = useRef<HTMLInputElement>(null);
const handleImport = async (file: File) => {
try {
const imported = await importConversationsFromFile(file, null);
if (imported === 0) {
toast.info(t("settings.chat.importNoConversations"));
} else {
toast.success(
imported === 1
? t("settings.chat.importedOneChat")
: t("settings.chat.importedChatCount", { count: imported }),
);
setCount(await countAllChats().catch(() => count));
}
} catch {
toast.error(t("settings.chat.importFailed"));
}
};
const handleArchiveAll = async () => {
setArchiving(true);
try {
const archived = await archiveAllChatItems(openChatId, (view) => {
navigate({ to: "/chat", search: { new: view.newThreadNonce } });
});
setArchiveConfirmOpen(false);
toast.success(
archived === 0
? t("settings.data.noChatsToArchive")
: archived === 1
? t("settings.data.archivedOneChat")
: t("settings.data.archivedChatCount", { count: archived }),
);
} catch (error) {
toast.error(t("settings.data.failedToArchiveChats"), {
description: error instanceof Error ? error.message : undefined,
});
} finally {
setArchiving(false);
}
};
const handleFineTuneExport = async () => {
setFineTuneExporting(true);
try {
await exportFineTuneJsonl(fineTuneFormat);
} catch (error) {
toast.error(t("settings.data.fineTuneExportFailed"), {
description: error instanceof Error ? error.message : undefined,
});
} finally {
setFineTuneExporting(false);
}
};
const handleOpenInRecipes = async () => {
setOpeningRecipe(true);
try {
const recipeId = await createFineTuneRecipeFromChats(fineTuneFormat);
if (!recipeId) return;
useSettingsDialogStore.getState().closeDialog();
void navigate({ to: "/data-recipes/$recipeId", params: { recipeId } });
} catch (error) {
toast.error(t("settings.data.fineTuneRecipeFailed"), {
description: error instanceof Error ? error.message : undefined,
});
} finally {
setOpeningRecipe(false);
}
};
const handleUseInTraining = async () => {
setLoadingTraining(true);
try {
const loaded = await loadFineTuneDatasetInTrainTab(fineTuneFormat);
if (!loaded) return;
useSettingsDialogStore.getState().closeDialog();
void navigate({ to: "/studio" });
} catch (error) {
toast.error(t("settings.data.fineTuneTrainFailed"), {
description: error instanceof Error ? error.message : undefined,
});
} finally {
setLoadingTraining(false);
}
};
const fineTuneActionLabels = {
train: t("settings.data.fineTuneTrainAction"),
recipes: t("settings.data.fineTuneOpenRecipesAction"),
export: t("settings.data.fineTuneExportAction"),
} as const;
const fineTuneBusy = loadingTraining || openingRecipe || fineTuneExporting;
const runFineTuneAction = () => {
if (fineTuneAction === "train") {
if (chatOnly) return;
void handleUseInTraining();
} else if (fineTuneAction === "recipes") void handleOpenInRecipes();
else void handleFineTuneExport();
};
const handleClear = async () => {
setClearing(true);
try {
const result = await clearAllChats();
const clearedCount = result.deletedThreadIds.length;
const hasFailedStore =
result.backend === "failed" || result.legacy === "failed";
if (!hasFailedStore && result.failedThreadIds.length === 0) {
setCount(0);
setConfirmOpen(false);
toast.success(
clearedCount === 0
? t("settings.chat.clearedAllChats")
: clearedCount === 1
? t("settings.chat.clearedOneChat")
: t("settings.chat.clearedChatCount", { count: clearedCount }),
);
return;
}
const fallbackRemaining =
result.failedThreadIds.length > 0
? result.failedThreadIds.length
: (count ?? 0);
const remaining = await countAllChats().catch(() => fallbackRemaining);
setCount(remaining);
setConfirmOpen(false);
toast.warning(t("settings.chat.someChatsCouldNotBeCleared"), {
description:
result.failedThreadIds.length > 0
? clearedCount === 1 && result.failedThreadIds.length === 1
? t("settings.chat.oneChatClearedRemainOne")
: clearedCount === 1
? t("settings.chat.oneChatClearedRemain", {
remainingCount: result.failedThreadIds.length,
})
: result.failedThreadIds.length === 1
? t("settings.chat.chatsClearedRemainOne", { clearedCount })
: t("settings.chat.chatsClearedRemain", {
clearedCount,
remainingCount: result.failedThreadIds.length,
})
: remaining === 1
? t("settings.chat.storageClearFailedOne")
: t("settings.chat.storageClearFailed", { count: remaining }),
});
} catch (error) {
const remaining = await countAllChats().catch(() => count);
setCount(remaining);
toast.error(t("settings.chat.failedToClearChats"), {
description: error instanceof Error ? error.message : undefined,
});
} finally {
setClearing(false);
}
};
if (subpage === "archived") {
return (
<div className="flex flex-col gap-6">
<header className="flex items-center gap-2">
<button
type="button"
onClick={() => setSubpage("main")}
aria-label={`Back to ${t("settings.data.title")}`}
className="inline-flex size-7 items-center justify-center rounded-full text-muted-foreground transition-colors hover:bg-muted hover:text-foreground"
>
<HugeiconsIcon icon={ArrowLeft01Icon} className="size-4" />
</button>
<h1 className="text-xl font-semibold font-heading">
{t("settings.data.title")}
</h1>
</header>
<div className="flex items-start justify-between gap-4">
<div className="flex flex-col gap-1">
<h2 className="text-sm font-semibold">
{t("settings.data.archivedChats")}
</h2>
<p className="text-xs text-muted-foreground">
{t("settings.data.archivedChatsDescription")}
</p>
</div>
{archivedItems.length > 0 && (
<Button
variant="outline"
size="sm"
className="shrink-0"
onClick={handleExportArchived}
disabled={archivedExporting}
>
{archivedExporting ? (
<Spinner className="size-4" />
) : (
<HugeiconsIcon
icon={Download01Icon}
strokeWidth={1.75}
className="size-4"
/>
)}
{archivedExporting
? t("settings.data.exportingArchivedChats")
: t("settings.data.exportArchivedChats")}
</Button>
)}
</div>
<ArchivedChatsView />
</div>
);
}
if (subpage === "files") {
return (
<div className="flex flex-col gap-6">
<header className="flex items-center gap-2">
<button
type="button"
onClick={() => setSubpage("main")}
aria-label={`Back to ${t("settings.data.title")}`}
className="inline-flex size-7 items-center justify-center rounded-full text-muted-foreground transition-colors hover:bg-muted hover:text-foreground"
>
<HugeiconsIcon icon={ArrowLeft01Icon} className="size-4" />
</button>
<h1 className="text-xl font-semibold font-heading">
{t("settings.data.title")}
</h1>
</header>
<div className="flex flex-col gap-1">
<h2 className="text-sm font-semibold">
{t("settings.data.uploadedFiles")}
</h2>
<p className="text-xs text-muted-foreground">
{t("settings.data.uploadedFilesDescription")}
</p>
</div>
<UploadedFilesView />
</div>
);
}
return (
<div className="flex flex-col gap-6">
<header className="flex flex-col gap-1">
<h1 className="text-xl font-semibold font-heading">
{t("settings.data.title")}
</h1>
<p className="text-xs text-muted-foreground">
{t("settings.data.description")}
</p>
</header>
<div className="flex flex-col divide-y divide-border/60">
<SettingsRow
alignTop={true}
label={t("settings.data.fineTuneExport")}
description={t("settings.data.fineTuneExportDescription")}
>
<div className="flex items-center gap-2">
<DropdownMenu>
<DropdownMenuTrigger asChild={true}>
{/* Fixed width so switching actions never resizes the row. */}
<Button
variant="outline"
size="sm"
disabled={count === 0}
className="w-44 justify-between"
>
<span className="truncate">
{fineTuneActionLabels[fineTuneAction]}
</span>
<HugeiconsIcon
icon={ChevronDownStandardIcon}
className="size-3.5 shrink-0"
/>
</Button>
</DropdownMenuTrigger>
<DropdownMenuContent align="end" className="w-56">
{(["export", "train", "recipes"] as const).map((action) => (
<DropdownMenuItem
key={action}
disabled={action === "train" && chatOnly}
onSelect={() => setFineTuneAction(action)}
>
<span className="flex-1">
{fineTuneActionLabels[action]}
</span>
{fineTuneAction === action ? (
<HugeiconsIcon icon={Tick02Icon} className="size-4" />
) : null}
</DropdownMenuItem>
))}
</DropdownMenuContent>
</DropdownMenu>
<Button
size="icon-sm"
onClick={runFineTuneAction}
disabled={fineTuneBusy || count === 0}
aria-label={t("settings.data.fineTuneRunAction")}
title={`${t("settings.data.fineTuneRunAction")}: ${fineTuneActionLabels[fineTuneAction]}`}
className="shrink-0 rounded-full"
>
{fineTuneBusy ? (
<Spinner className="size-4" />
) : (
<HugeiconsIcon
icon={ChevronRightStandardIcon}
strokeWidth={2.5}
className="size-4"
/>
)}
</Button>
</div>
</SettingsRow>
<SettingsRow
label={t("settings.data.archivedChats")}
description={t("settings.data.archivedChatsDescription")}
>
<Button
variant="outline"
size="sm"
onClick={() => setSubpage("archived")}
>
{t("settings.data.manageAction")}
</Button>
</SettingsRow>
<SettingsRow
label={t("settings.data.archiveAllChats")}
description={t("settings.data.archiveAllChatsDescription")}
>
<Button
variant="outline"
size="sm"
onClick={() => setArchiveConfirmOpen(true)}
>
<HugeiconsIcon icon={Archive02Icon} className="size-3.5 mr-1.5" />
{t("settings.data.archiveAllAction")}
</Button>
</SettingsRow>
<SettingsRow
label={t("settings.data.confirmBeforeDeleting")}
description={t("settings.data.confirmBeforeDeletingDescription")}
>
<Switch
checked={confirmDeleteChats}
onCheckedChange={setConfirmDeleteChats}
/>
</SettingsRow>
<SettingsRow
label={t("settings.chat.exportHistory")}
description={t("settings.chat.exportHistoryDescription")}
>
<Button
variant="outline"
size="sm"
onClick={handleExport}
disabled={exporting || count === 0}
>
<HugeiconsIcon icon={Download01Icon} className="size-3.5 mr-1.5" />
{exporting
? t("settings.chat.exportingAction")
: t("settings.chat.exportAction")}
</Button>
</SettingsRow>
<SettingsRow
label={t("settings.chat.exportConversations")}
description={t("settings.chat.exportConversationsDescription")}
>
<DropdownMenu>
<DropdownMenuTrigger asChild={true}>
<Button variant="outline" size="sm" disabled={count === 0}>
<HugeiconsIcon
icon={Download01Icon}
className="size-3.5 mr-1.5"
/>
{t("settings.chat.exportConversationsAction")}
</Button>
</DropdownMenuTrigger>
<DropdownMenuContent align="end" className="w-56">
{(
[
{ scope: "recents", label: "exportScopeRecents" },
{ scope: "all", label: "exportScopeAll" },
] as const
).map(({ scope, label }) => (
<DropdownMenuSub key={scope}>
<DropdownMenuSubTrigger>
<HugeiconsIcon
icon={Download01Icon}
className="size-3.5 mr-1"
/>
{t(`settings.chat.${label}`)}
</DropdownMenuSubTrigger>
<DropdownMenuSubContent className="w-56">
{EXPORT_FORMATS_LIST.map(({ fmt, label: fmtLabel }) => (
<DropdownMenuItem
key={`${scope}-m-${fmt}`}
onSelect={() =>
void bulkExportConversationsByScope(scope, fmt, true)
}
>
{fmtLabel} {t("settings.chat.exportCombinedSuffix")}
</DropdownMenuItem>
))}
<DropdownMenuSeparator />
{EXPORT_FORMATS_LIST.map(({ fmt, label: fmtLabel }) => (
<DropdownMenuItem
key={`${scope}-s-${fmt}`}
onSelect={() =>
void bulkExportConversationsByScope(scope, fmt, false)
}
>
{fmtLabel} {t("settings.chat.exportPerChatSuffix")}
</DropdownMenuItem>
))}
</DropdownMenuSubContent>
</DropdownMenuSub>
))}
</DropdownMenuContent>
</DropdownMenu>
</SettingsRow>
<SettingsRow
destructive={true}
// divide-y already draws the row separator; drop the extra border.
className="border-t-0 mt-0 pt-3"
label={t("settings.chat.clearAllChats")}
description={
count === null
? t("settings.chat.clearAllChatsDescription")
: count === 0
? t("settings.chat.noChatsToClear")
: count === 1
? t("settings.chat.clearOneChatDescription")
: t("settings.chat.clearChatCountDescription", { count })
}
>
<Button
variant="outline"
size="sm"
onClick={() => setConfirmOpen(true)}
disabled={count === 0}
className="text-destructive hover:text-destructive hover:border-destructive/60"
>
<HugeiconsIcon icon={Delete02Icon} className="size-3.5 mr-1.5" />
{t("settings.chat.clearChatsAction")}
</Button>
</SettingsRow>
<SettingsRow
label={t("settings.chat.importChats")}
description={t("settings.chat.importChatsDescription")}
>
<Button
variant="outline"
size="sm"
onClick={() => importInputRef.current?.click()}
>
<HugeiconsIcon icon={Upload01Icon} className="size-3.5 mr-1.5" />
{t("settings.chat.importChatsAction")}
</Button>
<input
ref={importInputRef}
type="file"
accept=".jsonl,.ndjson,.csv"
className="hidden"
onChange={(e) => {
const file = e.target.files?.[0];
e.target.value = "";
if (file) void handleImport(file);
}}
/>
</SettingsRow>
</div>
<SettingsSection title={t("settings.data.filesSection")}>
<SettingsRow
label={t("settings.data.uploadedFiles")}
description={t("settings.data.uploadedFilesDescription")}
>
<Button
variant="outline"
size="sm"
onClick={() => setSubpage("files")}
>
{t("settings.data.manageAction")}
</Button>
</SettingsRow>
</SettingsSection>
<Dialog open={archiveConfirmOpen} onOpenChange={setArchiveConfirmOpen}>
<DialogContent className="max-w-md">
<DialogHeader>
<DialogTitle>{t("settings.data.archiveAllChatsTitle")}</DialogTitle>
<DialogDescription>
{t("settings.data.archiveAllChatsConfirmDescription")}
</DialogDescription>
</DialogHeader>
<DialogFooter>
<Button
variant="outline"
onClick={() => setArchiveConfirmOpen(false)}
>
{t("common.cancel")}
</Button>
<Button onClick={handleArchiveAll} disabled={archiving}>
{archiving
? t("settings.data.archivingAction")
: t("settings.data.archiveAllAction")}
</Button>
</DialogFooter>
</DialogContent>
</Dialog>
<Dialog open={confirmOpen} onOpenChange={setConfirmOpen}>
<DialogContent className="max-w-md">
<DialogHeader>
<DialogTitle>
{count === 1
? t("settings.chat.clearOneChatTitle")
: t("settings.chat.clearChatsTitle", { count: count ?? 0 })}
</DialogTitle>
<DialogDescription>
{t("settings.chat.clearChatsConfirmDescription")}
</DialogDescription>
</DialogHeader>
<DialogFooter>
<Button variant="outline" onClick={() => setConfirmOpen(false)}>
{t("common.cancel")}
</Button>
<Button
onClick={handleClear}
disabled={clearing}
className="bg-destructive hover:bg-destructive/90 text-destructive-foreground"
>
{clearing
? t("settings.chat.clearingAction")
: count === 1
? t("settings.chat.clearOneChatAction")
: t("settings.chat.clearChatCountAction", {
count: count ?? 0,
})}
</Button>
</DialogFooter>
</DialogContent>
</Dialog>
</div>
);
}

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

@ -21,6 +21,7 @@ import {
setShowLlamaUpdateBanner,
useShowLlamaUpdateBanner,
} from "@/hooks/use-llama-update-pref";
import { useHfTokenValidation } from "@/hooks";
import { LOCALE_STORAGE_KEY, useT } from "@/i18n";
import { isTauri } from "@/lib/api-base";
import { copyToClipboard } from "@/lib/copy-to-clipboard";
@ -216,11 +217,18 @@ export function GeneralTab() {
if (trimmed !== hfToken) setHfToken(trimmed);
};
const clearHfToken = () => {
draftRef.current = "";
setDraftToken("");
setHfToken("");
};
// Show an "accepted" tick once a non-empty token has been committed to the
// store and the field still matches it (i.e. not mid-edit). Gives the user
// feedback that a pasted token was saved.
const tokenSaved =
draftToken.trim().length > 0 && draftToken.trim() === (hfToken ?? "");
const tokenValidation = useHfTokenValidation(hfToken ?? "");
useEffect(() => {
let cancelled = false;
@ -498,46 +506,70 @@ export function GeneralTab() {
label={t("settings.general.huggingFaceToken")}
description={t("settings.general.huggingFaceTokenDescription")}
>
<div className="relative w-[260px]">
<Input
type={showToken ? "text" : "password"}
placeholder="hf_…"
value={draftToken}
onChange={(e) => setDraftToken(e.target.value)}
onBlur={commitToken}
className={cn(
"h-8 w-full font-mono text-xs",
tokenSaved ? "pr-14" : "pr-8",
)}
/>
{tokenSaved ? (
// Decorative: pointer-events-none lets clicks reach the input
// underneath so the field still focuses anywhere.
<span
className="pointer-events-none absolute right-7 top-1/2 flex size-5 -translate-y-1/2 items-center justify-center text-emerald-600 duration-150 animate-in fade-in zoom-in dark:text-emerald-500"
role="img"
aria-label={t("settings.general.tokenSaved")}
<div className="flex flex-col items-end gap-1.5">
<div className="flex items-center gap-2">
<div className="relative w-[260px]">
<Input
type={showToken ? "text" : "password"}
name="hf-token"
autoComplete="new-password"
spellCheck={false}
placeholder="hf_…"
value={draftToken}
onChange={(e) => setDraftToken(e.target.value)}
onBlur={commitToken}
className={cn(
"h-8 w-full font-mono text-xs",
tokenSaved ? "pr-14" : "pr-8",
)}
/>
{tokenSaved ? (
// Decorative: pointer-events-none lets clicks reach the input
// underneath so the field still focuses anywhere.
<span
className="pointer-events-none absolute right-7 top-1/2 flex size-5 -translate-y-1/2 items-center justify-center text-emerald-600 duration-150 animate-in fade-in zoom-in dark:text-emerald-500"
role="img"
aria-label={t("settings.general.tokenSaved")}
>
<Check className="size-4" strokeWidth={2.5} />
</span>
) : null}
<button
type="button"
onClick={() => setShowToken((s) => !s)}
className="absolute right-1.5 top-1/2 flex size-5 -translate-y-1/2 items-center justify-center rounded text-muted-foreground transition-colors hover:text-foreground"
aria-label={
showToken
? t("settings.general.hideToken")
: t("settings.general.showToken")
}
tabIndex={-1}
>
{showToken ? (
<EyeOff className="size-3.5" />
) : (
<Eye className="size-3.5" />
)}
</button>
</div>
<Button
type="button"
variant="outline"
size="sm"
className="h-8"
disabled={!draftToken && !hfToken}
onClick={clearHfToken}
>
<Check className="size-4" strokeWidth={2.5} />
</span>
Clear
</Button>
</div>
{tokenValidation.isChecking ? (
<p className="text-xs text-muted-foreground">Checking token…</p>
) : tokenValidation.error ? (
<p className="max-w-[330px] text-right text-xs text-destructive">
{tokenValidation.error}
</p>
) : null}
<button
type="button"
onClick={() => setShowToken((s) => !s)}
className="absolute right-1.5 top-1/2 flex size-5 -translate-y-1/2 items-center justify-center rounded text-muted-foreground transition-colors hover:text-foreground"
aria-label={
showToken
? t("settings.general.hideToken")
: t("settings.general.showToken")
}
tabIndex={-1}
>
{showToken ? (
<EyeOff className="size-3.5" />
) : (
<Eye className="size-3.5" />
)}
</button>
</div>
</SettingsRow>
{/* The desktop app authenticates via desktop auto-auth with a generated

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

@ -90,8 +90,11 @@ function MetricTile({
label: string;
value: string;
detail: string;
percent: number;
// null = usage unknown (e.g. Windows ROCm perf counter): show a dash and
// empty bar rather than a fabricated 0%.
percent: number | null;
}) {
const percentKnown = isFiniteNumber(percent);
const safePercent = clampPercent(percent);
return (
<div className="flex min-w-0 flex-col gap-2 rounded-md border border-border/60 bg-muted/20 p-3">
@ -102,10 +105,10 @@ function MetricTile({
<span
className={cn(
"shrink-0 font-mono text-xs tabular-nums",
usageTextClass(safePercent),
percentKnown ? usageTextClass(safePercent) : "text-muted-foreground",
)}
>
{formatPercent(safePercent)}
{percentKnown ? formatPercent(safePercent) : "--"}
</span>
</div>
<div className="min-w-0">
@ -117,7 +120,7 @@ function MetricTile({
</div>
</div>
<Progress
value={safePercent}
value={percentKnown ? safePercent : 0}
aria-label={label}
className="h-1.5 rounded-full bg-muted"
indicatorClassName={usageIndicatorClass(safePercent)}
@ -194,18 +197,30 @@ export function ResourcesTab() {
(sum, device) => sum + (device.memory_total_gb ?? 0),
0,
);
const vramUsed = devices.reduce(
(sum, device) => sum + (device.vram_used_gb ?? 0),
0,
);
const vramFree = devices.reduce(
(sum, device) =>
sum +
(device.vram_free_gb ??
Math.max(0, (device.memory_total_gb ?? 0) - (device.vram_used_gb ?? 0))),
0,
);
const vramPercent = vramTotal > 0 ? (vramUsed / vramTotal) * 100 : 0;
// null usage = unknown (e.g. Windows ROCm perf counter): treating it as 0
// fabricates a 0-used total, so the aggregate is unknown if any device is.
const vramUsageKnown =
devices.length > 0 &&
devices.every((device) => isFiniteNumber(device.vram_used_gb));
const vramUsed = vramUsageKnown
? devices.reduce((sum, device) => sum + (device.vram_used_gb ?? 0), 0)
: null;
const vramFree = vramUsageKnown
? devices.reduce(
(sum, device) =>
sum +
(device.vram_free_gb ??
Math.max(
0,
(device.memory_total_gb ?? 0) - (device.vram_used_gb ?? 0),
)),
0,
)
: null;
const vramPercent =
vramUsageKnown && isFiniteNumber(vramUsed) && vramTotal > 0
? (vramUsed / vramTotal) * 100
: 0;
return {
devices,
@ -218,6 +233,7 @@ export function ResourcesTab() {
vramUsed,
vramFree,
vramPercent,
vramUsageKnown,
};
}, [systemInfo]);
@ -259,6 +275,7 @@ export function ResourcesTab() {
: modelsFolderLoaded
? t("settings.resources.environment.unknown")
: t("common.loading");
const unknownLabel = t("settings.resources.environment.unknown");
return (
<div className="flex flex-col gap-6">
@ -327,17 +344,21 @@ export function ResourcesTab() {
label={t("settings.resources.liveMonitor.vram")}
value={
hasGpu
? `${formatGiB(metrics.vramUsed)} / ${formatGiB(metrics.vramTotal)}`
? metrics.vramUsageKnown
? `${formatGiB(metrics.vramUsed)} / ${formatGiB(metrics.vramTotal)}`
: `${unknownLabel} / ${formatGiB(metrics.vramTotal)}`
: t("settings.resources.liveMonitor.noGpu")
}
detail={
hasGpu
? t("settings.resources.liveMonitor.free", {
value: formatGiB(metrics.vramFree),
})
? metrics.vramUsageKnown
? t("settings.resources.liveMonitor.free", {
value: formatGiB(metrics.vramFree),
})
: unknownLabel
: backendLabel
}
percent={metrics.vramPercent}
percent={metrics.vramUsageKnown ? metrics.vramPercent : null}
/>
</div>
</SettingsSection>
@ -346,13 +367,33 @@ export function ResourcesTab() {
{hasGpu ? (
metrics.devices.map((device, index) => {
const ordinal = deviceOrdinal(device);
const total = device.memory_total_gb ?? 0;
const used = device.vram_used_gb ?? 0;
const free = device.vram_free_gb ?? Math.max(0, total - used);
// Preserve null (unknown, e.g. Windows ROCm perf counter); coercing
// to 0 would render a fabricated 0 used / full free.
const total = device.memory_total_gb ?? null;
const used = device.vram_used_gb ?? null;
const free =
device.vram_free_gb ??
(isFiniteNumber(total) && isFiniteNumber(used)
? Math.max(0, total - used)
: null);
const percent =
device.vram_utilization_pct ??
(total > 0 ? (used / total) * 100 : null);
(isFiniteNumber(total) && total > 0 && isFiniteNumber(used)
? (used / total) * 100
: null);
const safePercent = clampPercent(percent);
const usedText = isFiniteNumber(used)
? formatGiB(used)
: unknownLabel;
const freeText = isFiniteNumber(free)
? formatGiB(free)
: unknownLabel;
const totalText = isFiniteNumber(total)
? formatGiB(total)
: unknownLabel;
const percentText = isFiniteNumber(percent)
? formatPercent(safePercent)
: unknownLabel;
return (
<div
key={`${device.index ?? index}-${device.name ?? "gpu"}`}
@ -374,7 +415,7 @@ export function ResourcesTab() {
</div>
<div className="shrink-0 font-mono text-xs tabular-nums text-muted-foreground">
<span>
{formatPercent(safePercent)}{" "}
{percentText}{" "}
{t("settings.resources.gpu.vramUtilization")}
</span>
</div>
@ -382,17 +423,17 @@ export function ResourcesTab() {
<div className="grid gap-1 text-xs text-muted-foreground sm:grid-cols-3 sm:gap-2">
<span className="min-w-0 truncate font-mono tabular-nums">
{t("settings.resources.gpu.used", {
value: formatGiB(used),
value: usedText,
})}
</span>
<span className="min-w-0 truncate font-mono tabular-nums sm:text-center">
{t("settings.resources.gpu.free", {
value: formatGiB(free),
value: freeText,
})}
</span>
<span className="min-w-0 truncate font-mono tabular-nums sm:text-right">
{t("settings.resources.gpu.total", {
value: formatGiB(total),
value: totalText,
})}
</span>
</div>

5
studio/frontend/src/features/training/api/train-api.ts
View file

@ -2,6 +2,7 @@
// 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";
import { readFastApiError } from "@/lib/format-fastapi-error";
import type {
TrainingStartRequest,
@ -30,10 +31,12 @@ async function parseJson<T>(response: Response): Promise<T> {
export async function startTraining(
payload: TrainingStartRequest,
): Promise<TrainingStartResponse> {
const preparedToken = await prepareHfTokenForUse(payload.hf_token);
if (!preparedToken.proceed) throw new Error("Training start cancelled.");
const response = await authFetch("/api/train/start", {
method: "POST",
headers: { "Content-Type": "application/json" },
body: JSON.stringify(payload),
body: JSON.stringify({ ...payload, hf_token: preparedToken.token }),
});
return parseJson<TrainingStartResponse>(response);
}

17
studio/frontend/src/features/training/hooks/use-training-actions.ts
View file

@ -2,6 +2,7 @@
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import { primeNativeNotificationPermission } from "@/lib/native-notifications";
import { prepareHfTokenForUse } from "@/features/hf-auth";
import { confirmRemoteCodeIfNeeded } from "@/features/security";
import { useCallback } from "react";
import { toast } from "@/lib/toast";
@ -43,7 +44,7 @@ export function useTrainingActions() {
const startError = useTrainingRuntimeStore((state) => state.startError);
const startTrainingRun = useCallback(async (): Promise<boolean> => {
const config = useTrainingConfigStore.getState();
let config = useTrainingConfigStore.getState();
const runtimeStore = useTrainingRuntimeStore.getState();
const dialogStore = useDatasetPreviewDialogStore.getState();
@ -54,6 +55,13 @@ export function useTrainingActions() {
return false;
}
const preparedToken = await prepareHfTokenForUse(config.hfToken);
if (!preparedToken.proceed) return false;
if ((preparedToken.token ?? "") !== config.hfToken) {
config.setHfToken(preparedToken.token ?? "");
config = useTrainingConfigStore.getState();
}
primeNativeNotificationPermission().catch(() => undefined);
runtimeStore.setStartResources(
@ -226,6 +234,13 @@ export function useTrainingActions() {
resume_from_checkpoint: outputDir,
} as TrainingStartRequest;
const preparedToken = await prepareHfTokenForUse(payload.hf_token);
if (!preparedToken.proceed) {
runtimeStore.setStarting(false);
return false;
}
payload.hf_token = preparedToken.token;
runtimeStore.setStartResources(
payload.model_name,
payload.hf_dataset,

110
studio/frontend/src/hooks/use-hf-token-validation.ts
View file

@ -1,8 +1,8 @@
// 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 { whoAmI } from "@huggingface/hub";
import { useCallback, useEffect, useRef, useState } from "react";
import { validateHfToken } from "@/features/hf-auth";
import { useEffect, useRef, useState } from "react";
import { useDebouncedValue } from "./use-debounced-value";
export interface HfTokenValidationState {
@ -17,6 +17,20 @@ const INITIAL: HfTokenValidationState = {
isChecking: false,
};
interface CompletedValidation extends HfTokenValidationState {
token: string;
}
const NO_COMPLETED_VALIDATION: CompletedValidation = {
...INITIAL,
token: "",
};
// Current user access tokens contain 34 characters after the hf_ prefix.
// Action-time validation still accepts legacy shapes without spending quota
// on every intermediate value typed into a live form field.
const COMPLETE_HF_TOKEN = /^hf_[A-Za-z0-9]{34}$/;
/**
* Validates the HF token via the whoami-v2 API, debounced to avoid excessive
* requests while typing. isValid is null until checked.
@ -26,39 +40,73 @@ export function useHfTokenValidation(token: string): HfTokenValidationState {
token.trim().replace(/^["']+|["']+$/g, ""),
500,
);
const [state, setState] = useState<HfTokenValidationState>(INITIAL);
const [completed, setCompleted] = useState<CompletedValidation>(
NO_COMPLETED_VALIDATION,
);
const versionRef = useRef(0);
const runCheck = useCallback(async (t: string) => {
if (!t) {
setState({ isValid: null, error: null, isChecking: false });
return;
}
const v = ++versionRef.current;
setState((prev) => ({ ...prev, isChecking: true, error: null }));
try {
await whoAmI({ accessToken: t });
if (versionRef.current !== v) return;
setState({ isValid: true, error: null, isChecking: false });
} catch {
if (versionRef.current !== v) return;
setState({
isValid: false,
error: "invalid or expired token",
isChecking: false,
});
}
}, []);
const shouldValidate = COMPLETE_HF_TOKEN.test(debouncedToken);
useEffect(() => {
if (!debouncedToken) {
setState(INITIAL);
if (!shouldValidate) {
versionRef.current += 1;
return;
}
runCheck(debouncedToken);
}, [debouncedToken, runCheck]);
const version = ++versionRef.current;
void validateHfToken(debouncedToken).then(
(result) => {
if (versionRef.current !== version) return;
if (result.status === "valid") {
setCompleted({
token: debouncedToken,
isValid: true,
error: null,
isChecking: false,
});
} else if (result.status === "invalid") {
setCompleted({
token: debouncedToken,
isValid: false,
error: "invalid or expired token",
isChecking: false,
});
} else if (result.status === "rate_limited") {
const wait = result.retryAfterSeconds
? ` Try again in about ${Math.ceil(result.retryAfterSeconds / 60)} minute(s).`
: " Try again later.";
setCompleted({
token: debouncedToken,
isValid: null,
error: `Token verification is rate limited.${wait}`,
isChecking: false,
});
} else {
setCompleted({
token: debouncedToken,
isValid: null,
error: "Could not verify the token. Check your connection and try again.",
isChecking: false,
});
}
},
() => {
if (versionRef.current !== version) return;
setCompleted({
token: debouncedToken,
isValid: null,
error: "Could not verify the token. Check your connection and try again.",
isChecking: false,
});
},
);
}, [debouncedToken, shouldValidate]);
return state;
if (!shouldValidate) return INITIAL;
if (completed.token !== debouncedToken) {
return { isValid: null, error: null, isChecking: true };
}
return {
isValid: completed.isValid,
error: completed.error,
isChecking: false,
};
}

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

@ -99,6 +99,7 @@ export const en = {
chat: "Chat",
voice: "Voice",
connections: "Connections",
data: "Data",
apiKeys: "API",
about: "About",
},
@ -232,6 +233,9 @@ export const en = {
loadError: "Failed to load model auto-switch settings.",
saveError: "Failed to save model auto-switch settings.",
idleError: "Enter 0 to keep the model loaded, or at least 60 seconds.",
keepKv: "Keep chat context across idle unload",
keepKvDescription:
"Save the model's KV cache to disk before an idle unload and restore it on reload, so resumed chats skip re-reading their history. Chat context is written to disk (up to 10 GB) until it is restored or cleaned up.",
},
previewSharing: {
sectionTitle: "Preview sharing",
@ -254,7 +258,7 @@ export const en = {
},
permissions: {
sectionTitle: "Permissions",
bypassLabel: "Bypass permissions",
bypassLabel: "Tool permissions",
bypassDescription:
"How Unsloth approves chat tool calls (terminal, python, web, MCP) before they run. Full access disables approvals and the code sandbox.",
},
@ -508,7 +512,7 @@ export const en = {
},
chat: {
title: "Chat",
description: "Manage chat history stored on this device.",
description: "Customize how chat behaves on this device.",
modelDisclaimer: "Show model disclaimer",
modelDisclaimerDescription:
'Show "LLMs can make mistakes" under the chat box.',
@ -577,6 +581,53 @@ export const en = {
"A storage clear failed; {count} chats may remain. Please retry.",
failedToClearChats: "Failed to clear chats",
},
data: {
title: "Data",
description:
"Manage chat history and uploaded files stored on this device.",
archivedChats: "Archived chats",
archivedChatsDescription: "View and manage chats you have archived.",
manageAction: "Manage",
exportArchivedChats: "Export",
exportingArchivedChats: "Exporting...",
exportedOneArchivedChat: "Exported 1 archived chat",
exportedArchivedChatCount: "Exported {count} archived chats",
noArchivedChatsToExport: "No archived chats to export.",
failedToExportArchivedChats: "Failed to export archived chats",
archiveAllChats: "Archive all chats",
archiveAllChatsDescription:
"Move every chat in Recents and Projects to the archive.",
noChatsToArchive: "No chats to archive.",
archiveAllAction: "Archive all",
archivingAction: "Archiving...",
archiveAllChatsTitle: "Archive all chats?",
archiveAllChatsConfirmDescription:
"Moves every chat on this device to the archive. Archived chats stay available and can be unarchived at any time.",
archivedAllChats: "Archived all chats",
archivedOneChat: "Archived 1 chat",
archivedChatCount: "Archived {count} chats",
failedToArchiveChats: "Failed to archive chats",
confirmBeforeDeleting: "Confirm before deleting",
confirmBeforeDeletingDescription:
"Ask for confirmation before a chat is deleted. Turn off to delete instantly.",
filesSection: "Files",
uploadedFiles: "Uploaded files",
uploadedFilesDescription:
"View and manage files uploaded to chats, projects, and knowledge bases.",
fineTuneExport: "Use chats as training data",
fineTuneExportDescription:
"Create a fine-tuning JSONL dataset from your chats. Load it in Train, refine in Recipes, or export it.",
fineTuneExportAction: "Export JSONL",
fineTuneRunAction: "Run",
fineTuneExportingAction: "Exporting...",
fineTuneOpenRecipesAction: "Open in Recipes",
fineTuneOpeningRecipesAction: "Opening...",
fineTuneTrainAction: "Load in Train tab",
fineTuneTrainingAction: "Loading...",
fineTuneExportFailed: "Failed to export training data",
fineTuneRecipeFailed: "Failed to open chats in Recipes",
fineTuneTrainFailed: "Failed to load dataset in the Train tab",
},
connections: {
title: "Connections",
description: "Manage providers and external connections.",

4
studio/frontend/src/index.css
View file

@ -1532,12 +1532,12 @@ html[data-chat-font] .aui-root {
}
/* Hovering an active pill swaps the icon for an X (click to turn off). */
.composer-pill-btn[data-active="true"]:hover .composer-pill-glyph > :not(.composer-pill-x),
.composer-pill-btn[data-active="true"]:has(.composer-pill-x):hover .composer-pill-glyph > :not(.composer-pill-x),
.unsloth-thinking-pill[data-active="true"]:hover .composer-pill-glyph > :not(.composer-pill-x) {
@apply opacity-0;
}
.composer-pill-btn[data-active="true"]:hover .composer-pill-x,
.composer-pill-btn[data-active="true"]:has(.composer-pill-x):hover .composer-pill-x,
.unsloth-thinking-pill[data-active="true"]:hover .composer-pill-x {
@apply opacity-100;
}

69
studio/install_llama_prebuilt.py
View file

@ -6450,6 +6450,7 @@ def write_prebuilt_metadata(
choice: AssetChoice,
approved_checksums: ApprovedReleaseChecksums,
prebuilt_fallback_used: bool,
force_cpu: bool = False,
) -> None:
source_asset_name, source_sha256 = selected_source_archive_metadata(
approved_checksums,
@ -6474,6 +6475,10 @@ def write_prebuilt_metadata(
"release_tag": release_tag,
"published_repo": approved_checksums.repo,
"asset": choice.name,
# True only for a deliberate CPU choice (--force-cpu). The updater re-asserts it
# so a forced CPU install is not re-routed to a GPU bundle (#7213). An automatic
# --cpu-fallback (e.g. arm64 GPU-build recovery) stays False so it can heal to GPU.
"force_cpu": force_cpu,
"asset_sha256": choice.expected_sha256,
"source": choice.source_label,
# Binary-side repo/tag for non-fork sources (e.g. the ggml-org upstream
@ -6501,6 +6506,24 @@ def write_prebuilt_metadata(
(install_dir / "UNSLOTH_PREBUILT_INFO.json").write_text(json.dumps(metadata, indent = 2) + "\n")
def sync_marker_force_cpu(install_dir: Path, persist_force_cpu: bool) -> None:
"""Sync only the force_cpu flag of an existing marker when the resolved bundle is
unchanged, so the install is skipped without a full metadata rewrite. A deliberate
--force-cpu on top of a naturally installed CPU bundle (same asset) must still be
recorded, else the updater will not re-assert it and can re-route the install to a
GPU/Vulkan bundle that revives the crash (#7213)."""
marker_path = install_dir / "UNSLOTH_PREBUILT_INFO.json"
try:
marker = json.loads(marker_path.read_text())
except (OSError, ValueError):
return
if not isinstance(marker, dict) or bool(marker.get("force_cpu")) == persist_force_cpu:
return
marker["force_cpu"] = persist_force_cpu
marker_path.write_text(json.dumps(marker, indent = 2) + "\n")
log(f"existing install reused; recorded force_cpu={persist_force_cpu} from this run")
def expected_install_fingerprint(
*,
llama_tag: str,
@ -6746,6 +6769,7 @@ def validate_prebuilt_choice(
approved_checksums: ApprovedReleaseChecksums,
prebuilt_fallback_used: bool,
quantized_path: Path,
force_cpu: bool = False,
) -> tuple[Path, Path]:
source_repo, source_ref, source_archive, exact_source = preferred_source_archive(
approved_checksums, llama_tag
@ -6786,6 +6810,7 @@ def validate_prebuilt_choice(
choice = choice,
approved_checksums = approved_checksums,
prebuilt_fallback_used = prebuilt_fallback_used,
force_cpu = force_cpu,
)
# Hashless external prebuilts are not in the approved-sha256
# manifest and rely on the functional smoke test as their only integrity gate,
@ -6828,6 +6853,7 @@ def validate_prebuilt_attempts(
approved_checksums: ApprovedReleaseChecksums,
initial_fallback_used: bool = False,
existing_install_dir: Path | None = None,
force_cpu: bool = False,
) -> tuple[AssetChoice, Path, bool]:
attempt_list = list(attempts)
if not attempt_list:
@ -6880,6 +6906,7 @@ def validate_prebuilt_attempts(
approved_checksums = approved_checksums,
prebuilt_fallback_used = tried_fallback,
quantized_path = quantized_path,
force_cpu = force_cpu,
)
except Exception as exc:
remove_tree(staging_dir)
@ -6939,8 +6966,8 @@ def _route_to_vulkan_prebuilt(
"""Point a Vulkan-capable host at the upstream ggml-org Vulkan prebuilt.
The unsloth published repo ships only CUDA/ROCm/CPU assets, so Vulkan comes
from UPSTREAM_REPO. Two triggers route here, both suppressed under
--cpu-fallback (the explicit "give me CPU" last resort wins):
from UPSTREAM_REPO. Two triggers route here, both suppressed when a CPU flag
(--cpu-fallback or --force-cpu, folded into force_cpu) wins:
* UNSLOTH_FORCE_VULKAN forces Vulkan over the detected CUDA/ROCm backend;
* an auto-detected Intel GPU with NO physical NVIDIA/ROCm -- the purpose
of the has_intel_gpu probe, since the fork manifest ships no Vulkan asset.
@ -7020,8 +7047,11 @@ def install_prebuilt(
override_has_rocm: bool = False,
override_rocm_gfx: str | None = None,
force_cpu: bool = False,
persist_force_cpu: bool = False,
instruction_cleanup_root: Path | None = None,
) -> None:
# force_cpu drops GPU detection (mechanism, both --cpu-fallback and --force-cpu);
# persist_force_cpu records the deliberate choice so the updater re-asserts it.
host = detect_host()
host = _apply_host_overrides(
host,
@ -7072,6 +7102,9 @@ def install_prebuilt(
"existing llama.cpp install already matches selected release "
f"{current.release_tag} upstream_tag={current.llama_tag}; skipping download and install"
)
# Reused bundle is unchanged, but a fresh --force-cpu still must be
# recorded so the updater re-asserts it (#7213).
sync_marker_force_cpu(install_dir, persist_force_cpu)
return
with tempfile.TemporaryDirectory(prefix = "unsloth-llama-prebuilt-") as tmp:
work_dir = Path(tmp)
@ -7092,6 +7125,7 @@ def install_prebuilt(
"existing llama.cpp install already matches fallback release "
f"{plan.release_tag} upstream_tag={plan.llama_tag}; skipping reinstall"
)
sync_marker_force_cpu(install_dir, persist_force_cpu)
return
log(
"selected "
@ -7112,6 +7146,8 @@ def install_prebuilt(
initial_fallback_used = release_index > 0,
# Skip is gated per-attempt inside, so pass the dir always.
existing_install_dir = install_dir,
# Persist only the deliberate choice, not a transient fallback.
force_cpu = persist_force_cpu,
)
except ExistingInstallSatisfied:
return
@ -7209,8 +7245,21 @@ def parse_args() -> argparse.Namespace:
default = False,
help = (
"Select the CPU prebuilt for this OS/arch even when a GPU is present. "
"setup.sh uses this as a last resort for arm64 Linux GPU hosts whose "
"source build failed (no arm64 CUDA prebuilt exists anywhere)."
"Automatic/transient: setup.sh uses this as a last resort for arm64 Linux "
"GPU hosts whose source build failed. Does NOT persist, so a later update "
"heals back to a GPU bundle once one is available (#6097). Use --force-cpu "
"for a deliberate CPU-only choice that survives updates."
),
)
parser.add_argument(
"--force-cpu",
action = "store_true",
default = False,
help = (
"Deliberate CPU-only install (UNSLOTH_LLAMA_CPP_BACKEND=cpu). Drops GPU "
"detection like --cpu-fallback but also records force_cpu in the marker, so "
"the in-app updater re-asserts CPU and never re-routes to a GPU/Vulkan "
"bundle that would revive the Intel iGPU crash (#7213)."
),
)
resolve_group = parser.add_mutually_exclusive_group()
@ -7333,16 +7382,19 @@ def main() -> int:
# Host-aware "is a prebuilt available" probe, no download. Every host now
# plans against the fork (args.published_repo defaults to it); an explicit
# --published-repo overrides. PrebuiltFallback == source build.
# Both flags drop GPU detection; --force-cpu additionally persists (install
# path only). The probe only needs the mechanism, so OR them.
_cpu_mechanism = args.cpu_fallback or args.force_cpu
host = _apply_host_overrides(
detect_host(),
override_has_rocm = args.has_rocm,
override_rocm_gfx = args.rocm_gfx,
force_cpu = args.cpu_fallback,
force_cpu = _cpu_mechanism,
)
# Same Vulkan routing the install path applies, so the probe's answer
# matches what would install (an Intel/forced-Vulkan host -> upstream).
host, repo, release_tag = _route_to_vulkan_prebuilt(
host, args.published_repo, args.published_release_tag or "", force_cpu = args.cpu_fallback
host, args.published_repo, args.published_release_tag or "", force_cpu = _cpu_mechanism
)
try:
_requested, plans = resolve_simple_install_release_plans(
@ -7380,7 +7432,10 @@ def main() -> int:
published_release_tag = args.published_release_tag or "",
override_has_rocm = args.has_rocm,
override_rocm_gfx = args.rocm_gfx,
force_cpu = args.cpu_fallback,
# Both drop GPU detection; only --force-cpu (deliberate) is recorded so the
# updater re-asserts it. --cpu-fallback stays transient and heals to GPU.
force_cpu = args.cpu_fallback or args.force_cpu,
persist_force_cpu = args.force_cpu,
instruction_cleanup_root = install_arg.absolute(),
)
return EXIT_SUCCESS

872
studio/install_python_stack.py
View file

File diff suppressed because it is too large Load diff

415
studio/setup.ps1
View file

@ -32,6 +32,10 @@ $PackageDir = Split-Path -Parent $ScriptDir
# (no matching GitHub release), forces a source build, and causes HTTP 422
# errors. Only use "master" temporarily when the latest release is missing
# support for a new model architecture.
#
# UNSLOTH_LLAMA_CPP_BACKEND : "auto" (default) or "cpu". When "cpu", forces
# the CPU-only prebuilt bundle on GPU hosts. Fixes Intel iGPU Vulkan
# crashes (#7213).
$DefaultLlamaPrForce = ""
$DefaultLlamaSource = "https://github.com/ggml-org/llama.cpp"
$DefaultLlamaTag = "latest"
@ -427,6 +431,167 @@ function Get-PytorchCudaTag {
return "cu126"
}
# 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)
}
# Explicit torch-index pin (UNSLOTH_TORCH_INDEX_URL / _FAMILY), shared by the stale-venv check
# and install selection so a pinned index wins over GPU probing (parity with the other
# installers). URL is verbatim; _FAMILY is the leaf joined to the mirror base.
function Get-PinnedTorchIndexUrl {
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)) {
$base = if ($env:UNSLOTH_PYTORCH_MIRROR) { $env:UNSLOTH_PYTORCH_MIRROR.TrimEnd('/') } else { "https://download.pytorch.org/whl" }
return "$base/$($env:UNSLOTH_TORCH_INDEX_FAMILY.Trim().Trim('/'))"
}
return $null
}
# Last path segment of a wheel index URL, query/fragment dropped first so a token-authenticated
# pin (.../cu128?token=x) classifies as cu128 (else it reinstalls every update). Classification
# only. Shared with the py / install.sh leaf extractors.
function Get-TorchIndexLeaf {
param([string]$Url)
if ([string]::IsNullOrWhiteSpace($Url)) { return $null }
$path = ($Url -split '[?#]', 2)[0]
if ([string]::IsNullOrWhiteSpace($path)) { return $null }
return ($path.TrimEnd('/') -split '/')[-1].ToLowerInvariant()
}
# 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>'
}
# AMD per-arch leaves needing the torch 2.11 floor (the _grouped_mm <2.11 bug). MUST match
# the install-spec path below and the other installers; other leaves ship <2.11 and stay default.
function Test-RocmGfx211Leaf {
param([string]$Leaf)
return @('gfx120x-all', 'gfx1151', 'gfx1150') -contains $Leaf
}
# rocmX.Y versions KNOWN to ship torch 2.11: rocm7.2 only today. Do NOT floor an unknown newer
# rocm speculatively. MUST match _ROCM_KNOWN_TORCH211_VERSIONS and the rocm7.2 leaf elsewhere.
function Test-RocmKnown211Version {
param([int]$Major, [int]$Minor)
return ($Major -eq 7 -and $Minor -eq 2)
}
# True only for a real CUDA family leaf: "cu" + digits (cu118, cu128, ...). A bare -like 'cu*'
# would match "custom"/"current" and rebuild the venv every run. Mirrors _is_cuda_family_leaf.
function Test-CudaFamilyLeaf {
param([string]$Leaf)
if ([string]::IsNullOrWhiteSpace($Leaf)) { return $false }
# EXACT cu+digits: cu128-private routes through the unknown-leaf path instead.
return $Leaf -match '^cu[0-9]+$'
}
# True only for a real pip ROCm family leaf: EXACT rocm<digits>[.<digits>] or a gfx leaf. A leaf
# that merely STARTS with rocm (rocm-rel-7.2.1, rocm7.2-private) is a custom pin the verbatim
# path owns, so anchor the match. Mirrors _is_pip_rocm_family_leaf / install.sh.
function Test-PipRocmFamilyLeaf {
param([string]$Leaf)
if ([string]::IsNullOrWhiteSpace($Leaf)) { return $false }
# gfx must be followed by a digit (an architecture leaf); gfx-private is custom.
return ($Leaf -match '^gfx[0-9]') -or ($Leaf -match '^rocm[0-9]+(\.[0-9]+)?$')
}
# Stale-venv ROCm comparison for a pinned gfx*/rocm* index. Returns @{ Expected; Installed } so
# the caller rebuilds when they differ. Mirrors _rocm_pin_family_mismatch (same rocmX.Y / gfx
# cases). An untagged (no +rocm) wheel never satisfies a ROCm pin -> stale.
function Get-RocmPinStaleTags {
param([string]$PinLeaf, [string]$TorchVersion)
$_pinRocm = [regex]::Match($PinLeaf, '^rocm(\d+)\.(\d+)')
$_pinVer = if ($_pinRocm.Success) { "$($_pinRocm.Groups[1].Value).$($_pinRocm.Groups[2].Value)" } else { $null }
# The family classifier accepts a major-only rocm<d> leaf too (rocm7).
$_pinMajorOnly = [regex]::Match($PinLeaf, '^rocm(\d+)$')
# Installed rocm version and whether the wheel is a per-arch (three-part) build.
$_instRocm = [regex]::Match($TorchVersion, '\+rocm(\d+)\.(\d+)')
$_instVer = if ($_instRocm.Success) { "$($_instRocm.Groups[1].Value).$($_instRocm.Groups[2].Value)" } else { $null }
$_instPerArch = [regex]::IsMatch($TorchVersion, '\+rocm\d+\.\d+\.\d+')
# A ROCm build MUST carry a +rocm tag; an untagged wheel can't satisfy any ROCm pin.
$_instHasRocm = [regex]::IsMatch($TorchVersion, '\+rocm')
$_instRel = [regex]::Match($TorchVersion, '^(\d+)\.(\d+)')
$_instIs211 = $false
if ($_instRel.Success) {
$_instIs211 = ([int]$_instRel.Groups[1].Value -gt 2) -or ([int]$_instRel.Groups[1].Value -eq 2 -and [int]$_instRel.Groups[2].Value -ge 11)
}
if ($PinLeaf -like 'gfx*') {
if (Test-RocmGfx211Leaf $PinLeaf) {
# Expect the AMD per-arch (three-part) 2.11 wheel: satisfied only when BOTH
# a 2.11 release AND a three-part rocm tag are installed.
$installed = if ($_instIs211 -and $_instPerArch) { "rocm-perarch(torch>=2.11)" } else { "rocm-generic-or-old" }
return @{ Expected = "rocm-perarch(torch>=2.11)"; Installed = $installed }
}
# Non-2.11 gfx leaf (<2.11 spec): stale on an untagged wheel or a 2.11+ build.
$installed = if (-not $_instHasRocm) { "not-rocm" } elseif ($_instIs211) { "rocm(torch>=2.11)" } else { "rocm(torch<2.11)" }
return @{
Expected = "rocm(torch<2.11)"
Installed = $installed
}
}
# Major-only rocm pin (rocm7): compare majors only -- a +rocm6.4 wheel under a rocm7
# pin is stale, any +rocm7.x wheel satisfies it (no pinned minor to compare, and the
# 2.11-line fallback below would invert both verdicts). Mirrors _rocm_pin_family_mismatch.
if ($_pinMajorOnly.Success) {
$_pinMaj = [int]$_pinMajorOnly.Groups[1].Value
if ($_instVer) {
$_instMaj = [int]$_instRocm.Groups[1].Value
$expected = if ($_instMaj -eq $_pinMaj) { "rocm$_instVer" } else { "rocm$_pinMaj.x" }
return @{ Expected = $expected; Installed = "rocm$_instVer" }
}
# Untagged wheel never satisfies a ROCm pin; a +rocm tag with an unreadable
# version is accepted (matches the lenient unreadable fallback below).
$installed = if ($_instHasRocm) { "rocm" } else { "not-rocm" }
return @{ Expected = "rocm"; Installed = $installed }
}
# rocmX.Y pin.
if ($_pinVer -and $_instVer) {
# Both readable: exact compare. When they match AND the pin is KNOWN-2.11, the
# installed release must also be 2.11 (a +rocm7.2 wheel drifted to 2.12 shares the
# tag but violates the spec), so fold the release into the tag. Mirrors _rocm_pin_family_mismatch.
$_pinKnown211 = Test-RocmKnown211Version -Major ([int]$_pinRocm.Groups[1].Value) -Minor ([int]$_pinRocm.Groups[2].Value)
$_instOn211 = $_instRel.Success -and [int]$_instRel.Groups[1].Value -eq 2 -and [int]$_instRel.Groups[2].Value -eq 11
if ($_pinKnown211 -and -not $_instOn211) {
return @{ Expected = "rocm$_pinVer(torch2.11)"; Installed = "rocm$_instVer(torch-off-2.11)" }
}
return @{ Expected = "rocm$_pinVer"; Installed = "rocm$_instVer" }
}
$_pinNeeds211 = $false
if ($_pinRocm.Success) {
# Only KNOWN-2.11 rocm (rocm7.2) is on the 2.11 line (no speculative floor).
# Matches _ROCM_KNOWN_TORCH211_VERSIONS.
$_pinNeeds211 = Test-RocmKnown211Version -Major ([int]$_pinRocm.Groups[1].Value) -Minor ([int]$_pinRocm.Groups[2].Value)
}
# Fallback (installed rocm version unreadable): compare on the 2.11 line; an untagged
# wheel never satisfies a rocmX.Y pin -> stale.
$installed = if (-not $_instHasRocm) { "not-rocm" } elseif ($_instIs211) { "rocm(torch>=2.11)" } else { "rocm(torch<2.11)" }
return @{
Expected = if ($_pinNeeds211) { "rocm(torch>=2.11)" } else { "rocm(torch<2.11)" }
Installed = $installed
}
}
# VS generator -> MSBuild BuildCustomizations dir; toolset tracks the VS major
# (18->v180, 17->v170), defaulting to v170 when unparseable.
function Get-VcBuildCustomizationsDir {
@ -809,11 +974,14 @@ function Invoke-SetupCommand {
# 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/pip echo index URLs (credentials and all) in
# their 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
@ -2531,6 +2699,8 @@ if ((Test-Path -LiteralPath $VenvDir -PathType Container) -and -not $NoTorchMode
$VenvPyExe = Join-Path $VenvDir "Scripts\python.exe"
$installedTorchTag = $null
$shouldRebuild = $false
# Set when a stale venv under a pin is repaired in place (force-reinstall) not wiped.
$script:PinChangedForceReinstall = $false
if (Test-Path -LiteralPath $VenvPyExe) {
try {
@ -2547,10 +2717,14 @@ if ((Test-Path -LiteralPath $VenvDir -PathType Container) -and -not $NoTorchMode
if ($finished -and $proc.ExitCode -eq 0 -and $torchVer) {
if ($torchVer -match '\+(cu\d+)') {
$installedTorchTag = $Matches[1]
} elseif ($torchVer -match '\+rocm') {
# Any +rocm / gfx wheel -> generic "rocm" flavor (the exact version is
# repaired later by install_python_stack.py; here we only need the flavor).
$installedTorchTag = "rocm"
} elseif ($torchVer -match '\+cpu') {
$installedTorchTag = "cpu"
} else {
# Untagged wheel (plain "2.x.y" from PyPI) -- treat as cpu
# Untagged wheel (plain "2.x.y" from PyPI) -> cpu.
$installedTorchTag = "cpu"
}
} else {
@ -2566,12 +2740,71 @@ if ((Test-Path -LiteralPath $VenvDir -PathType Container) -and -not $NoTorchMode
}
if (-not $shouldRebuild) {
$expectedTorchTag = if ($HasNvidiaSmi) { Get-PytorchCudaTag } else { "cpu" }
if ($installedTorchTag -and $installedTorchTag -ne $expectedTorchTag) {
$_pinnedIdx = Get-PinnedTorchIndexUrl
$_expectedKnown = $true
if ($_pinnedIdx) {
$_pinLeaf = Get-TorchIndexLeaf $_pinnedIdx
# Digit-gated like the install selection: a custom rocm-* leaf (rocm-current /
# rocm-rel-7.2.1) is NOT a ROCm family and must not be stale-compared.
if (Test-PipRocmFamilyLeaf $_pinLeaf) {
# Don't collapse a pinned ROCm/gfx leaf to a generic "rocm" (would mask a family
# change, rocm6.4 -> gfx1151). Get-RocmPinStaleTags uses the SAME 2.11 allowlist
# as the install path, so a gfx110X-all/gfx90a/gfx908 pin on a <2.11 wheel is NOT stale.
$_rocmTags = Get-RocmPinStaleTags -PinLeaf $_pinLeaf -TorchVersion $torchVer
$expectedTorchTag = $_rocmTags.Expected
$installedTorchTag = $_rocmTags.Installed
} elseif ((Test-CudaFamilyLeaf $_pinLeaf) -or $_pinLeaf -eq 'cpu') {
# cu*/cpu leaves stay specific so a cu126-vs-cu128 mismatch rebuilds;
# /custom and /current fall through to the unknown-index branch below.
$expectedTorchTag = $_pinLeaf
} else {
# Custom index whose leaf is not a torch flavor (a /simple mirror): the
# flavor can't be inferred, so never treat the venv as stale over it.
$_expectedKnown = $false
$expectedTorchTag = $installedTorchTag
}
} elseif ($HasNvidiaSmi) {
$expectedTorchTag = Get-PytorchCudaTag
} elseif ($HasROCm -or $script:ROCmGfxArch) {
# AMD/ROCm host with no explicit pin: an existing +rocm wheel is correct (gfx arch
# counts even when $HasROCm is false). But only the arches the install path maps to a
# repo.amd.com index get ROCm torch; an unmapped arch installs CPU, so expect "cpu"
# for those or a correct CPU venv rebuilds every update.
$_rocmWheelArches = @(
"gfx1201", "gfx1200", # RDNA 4
"gfx1151", "gfx1150", # RDNA 3.5 (Strix Halo/Point)
"gfx1103", "gfx1102", "gfx1101", "gfx1100", # RDNA 3
"gfx90a", "gfx908" # MI200 / MI100
)
if ($script:ROCmGfxArch -and ($_rocmWheelArches -contains $script:ROCmGfxArch)) {
# A correct +rocm wheel is not stale. A CPU wheel on a supported AMD arch is
# NOT wiped either (the AMD Windows ROCm override below upgrades it in place);
# expect "cpu" for that case. A wrong CUDA wheel still rebuilds.
if ($installedTorchTag -eq "cpu") {
$expectedTorchTag = "cpu"
} else {
$expectedTorchTag = "rocm"
}
} else {
$expectedTorchTag = "cpu"
}
} else {
$expectedTorchTag = "cpu"
}
if ($_expectedKnown -and $installedTorchTag -and $installedTorchTag -ne $expectedTorchTag) {
$shouldRebuild = $true
}
}
# A stale venv under a pin whose torch still imports is repaired IN PLACE (the dependency
# pass force-reinstalls from the pin). The rebuild path wipes the venv and would strand a
# direct `studio update`; only a broken venv or an unpinned drift wipes.
if ($shouldRebuild -and $_pinnedIdx -and $installedTorchTag) {
substep "Torch-index pin changed ($installedTorchTag) -- reinstalling torch from the pin in place." "Cyan"
$script:PinChangedForceReinstall = $true
$shouldRebuild = $false
}
if ($shouldRebuild) {
$reason = if ($installedTorchTag) { "torch $installedTorchTag != required $expectedTorchTag" } else { "torch could not be imported" }
if ($InstallerManagedSetup) {
@ -2649,23 +2882,41 @@ if (Get-Command uv -ErrorAction SilentlyContinue) {
# Helper: install a package, preferring uv with pip fallback
function Fast-Install {
param([Parameter(ValueFromRemainingArguments=$true)]$Args_)
if ($UseUv) {
$VenvPy = (Get-Command python).Source
# An explicit --index-url must win. Inherited uv index env vars otherwise
# override it and pull CPU torch over the CUDA/ROCm build (#6898), so drop
# them only for index-pinned installs; mirrors still apply elsewhere.
$saved = @{}
if (@($Args_) -contains '--index-url') {
foreach ($n in 'UV_DEFAULT_INDEX', 'UV_INDEX_URL', 'UV_INDEX', 'UV_EXTRA_INDEX_URL') {
$saved[$n] = [Environment]::GetEnvironmentVariable($n)
Remove-Item "Env:$n" -ErrorAction SilentlyContinue
}
# An explicit --index-url must win: inherited uv index vars otherwise pull CPU torch over
# the CUDA/ROCm build (#6898), so drop them for pinned installs (scrub covers the whole
# function since the pip fallback honours PIP_* too). UV_TORCH_BACKEND / UV_FIND_LINKS also
# reroute; UV_NO_CONFIG=1 (+ dropping UV_CONFIG_FILE) stops a uv.toml index outranking the
# pin (uv 0.10); PIP_NO_INDEX / PIP_INDEX_URL would defeat the pinned --index-url in pip.
$saved = @{}
$pinned = @($Args_) -contains '--index-url'
if ($pinned) {
foreach ($n in 'UV_DEFAULT_INDEX', 'UV_INDEX_URL', 'UV_INDEX', 'UV_EXTRA_INDEX_URL',
'UV_TORCH_BACKEND', 'UV_FIND_LINKS', 'PIP_EXTRA_INDEX_URL', 'PIP_FIND_LINKS',
'PIP_NO_INDEX', 'PIP_INDEX_URL',
'UV_CONFIG_FILE', 'UV_NO_CONFIG', 'PIP_CONFIG_FILE') {
$saved[$n] = [Environment]::GetEnvironmentVariable($n)
Remove-Item "Env:$n" -ErrorAction SilentlyContinue
}
try { $result = & uv pip install --python $VenvPy @Args_ 2>&1 }
finally { foreach ($n in $saved.Keys) { if ($null -ne $saved[$n]) { Set-Item "Env:$n" $saved[$n] } } }
if ($LASTEXITCODE -eq 0) { return }
$env:UV_NO_CONFIG = '1'
# A `pip config` global.extra-index-url still adds indexes to the pip FALLBACK;
# PIP_CONFIG_FILE = 'nul' (Windows devnull) loads NO config (uv ignores pip config).
$env:PIP_CONFIG_FILE = 'nul'
}
try {
if ($UseUv) {
$VenvPy = (Get-Command python).Source
$result = & uv pip install --python $VenvPy @Args_ 2>&1
if ($LASTEXITCODE -eq 0) { return }
}
& python -m pip install @Args_ 2>&1
}
finally {
if ($pinned) {
Remove-Item "Env:UV_NO_CONFIG" -ErrorAction SilentlyContinue
Remove-Item "Env:PIP_CONFIG_FILE" -ErrorAction SilentlyContinue
}
foreach ($n in $saved.Keys) { if ($null -ne $saved[$n]) { Set-Item "Env:$n" $saved[$n] } }
}
& python -m pip install @Args_ 2>&1
}
# ── Check if Python deps need updating ──
@ -2748,6 +2999,10 @@ sys.exit(0 if (major, minor) >= (4, 14) else 1)
# pip install unsloth 2>&1 | Out-Null
# }
# A torch-index pin change repairs in place: force the dependency pass so the torch install
# below force-reinstalls from the new pin (else the fast path keeps the old wheel).
if ($script:PinChangedForceReinstall) { $SkipPythonDeps = $false }
if (-not $SkipPythonDeps) {
if ($script:UnslothVerbose) {
@ -2775,7 +3030,13 @@ $env:TORCHINDUCTOR_CACHE_DIR = $TorchCacheDir
[Environment]::SetEnvironmentVariable('TORCHINDUCTOR_CACHE_DIR', $TorchCacheDir, 'User')
substep "TORCHINDUCTOR_CACHE_DIR set to $TorchCacheDir (avoids MAX_PATH issues)"
if ($HasNvidiaSmi) {
# Explicit pin (URL or family) wins over GPU probing and suppresses the AMD reroute below;
# matches install.sh / install.ps1 / install_python_stack.py.
$PinnedTorchIndexUrl = Get-PinnedTorchIndexUrl
$TorchIndexPinned = [bool]$PinnedTorchIndexUrl
if ($PinnedTorchIndexUrl) {
$CuTag = Get-TorchIndexLeaf $PinnedTorchIndexUrl
} elseif ($HasNvidiaSmi) {
$CuTag = Get-PytorchCudaTag
} else {
$CuTag = "cpu"
@ -2796,7 +3057,7 @@ $ROCmIndexUrl = $null
# SDK -- which flips Unsloth out of chat-only (CHAT_ONLY) and enables Train/Export.
# Gating on $HasROCm alone left Strix Halo / Radeon 8060S on CPU torch; a failed
# ROCm install still falls back to CPU below, so this is safe.
if (($HasROCm -or $ROCmGfxArch) -and $CuTag -eq "cpu") {
if (-not $TorchIndexPinned -and ($HasROCm -or $ROCmGfxArch) -and $CuTag -eq "cpu") {
$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
@ -2846,8 +3107,45 @@ if (($HasROCm -or $ROCmGfxArch) -and $CuTag -eq "cpu") {
}
}
# A pinned gfx*/rocm index skips the auto-reroute above; route it through the ROCm install path
# with the same floor/companions the unpinned AMD path uses (mirrors install.ps1), else the CUDA
# branch installs bare torch and resolves a known-bad wheel for gfx115x/gfx120x/rocm>=7.2.
if ($TorchIndexPinned -and -not $ROCmIndexUrl -and $PinnedTorchIndexUrl) {
$_pinLeaf = Get-TorchIndexLeaf $PinnedTorchIndexUrl
$_pinRocm211 = $false
# Anchor the match ($) so a suffixed custom leaf (rocm7.2-private) falls through to the
# verbatim install instead of being floored by its rocm7.2 prefix.
if ($_pinLeaf -match '^rocm(\d+)\.(\d+)$') {
# Only KNOWN-2.11 rocm (rocm7.2) gets the floor (no speculative floor). Matches
# Test-RocmKnown211Version / _ROCM_KNOWN_TORCH211_VERSIONS.
$_pinRocm211 = Test-RocmKnown211Version -Major ([int]$Matches[1]) -Minor ([int]$Matches[2])
}
# Only the 2.11 gfx arches need the floor; others publish <2.11 and stay bare. Reuse
# Test-RocmGfx211Leaf so this allowlist and the stale-venv check never diverge.
$_pinGfx211 = Test-RocmGfx211Leaf $_pinLeaf
if ($_pinGfx211 -or $_pinRocm211) {
$ROCmIndexUrl = $PinnedTorchIndexUrl
$ROCmTorchSpec = "torch>=2.11.0,<2.12.0"
$ROCmVisionSpec = "torchvision>=0.26.0,<0.27.0"
$ROCmAudioSpec = "torchaudio>=2.11.0,<2.12.0"
substep "pinned ROCm index ($_pinLeaf) -- enforcing $ROCmTorchSpec" "Cyan"
} elseif (Test-PipRocmFamilyLeaf $_pinLeaf) {
# Other gfx / older rocm (<=7.1) ship torch <2.11; route via the ROCm path with
# bare specs. Only EXACT rocm<digits> and gfx* are --index-url families; a suffixed
# leaf stays on the verbatim path. Mirrors install.ps1 / _is_pip_rocm_family_leaf.
$ROCmIndexUrl = $PinnedTorchIndexUrl
$ROCmTorchSpec = "torch"
$ROCmVisionSpec = "torchvision"
$ROCmAudioSpec = "torchaudio"
}
}
$PyTorchWhlBase = if ($env:UNSLOTH_PYTORCH_MIRROR) { $env:UNSLOTH_PYTORCH_MIRROR.TrimEnd('/') } else { "https://download.pytorch.org/whl" }
# A full URL pin is used verbatim; a family pin already set $CuTag. A pinned ROCm install
# goes through $ROCmIndexUrl; on failure the fallback uses the CPU index, not the ROCm pin.
$TorchInstallIndexUrl = if ($ROCmIndexUrl) { "$PyTorchWhlBase/cpu" } elseif ($PinnedTorchIndexUrl) { $PinnedTorchIndexUrl } else { "$PyTorchWhlBase/$CuTag" }
$ROCmCpuFallback = $false
if ($ROCmIndexUrl) {
substep "installing PyTorch (AMD ROCm, $ROCmGfxArch)..."
@ -2855,7 +3153,7 @@ if ($ROCmIndexUrl) {
substep " enforcing $ROCmTorchSpec $ROCmVisionSpec $ROCmAudioSpec (known _grouped_mm bug in older wheels)" "Cyan"
}
if ($script:UnslothVerbose) {
Fast-Install $ROCmTorchSpec $ROCmVisionSpec $ROCmAudioSpec --force-reinstall --index-url $ROCmIndexUrl
Fast-Install $ROCmTorchSpec $ROCmVisionSpec $ROCmAudioSpec --force-reinstall --index-url $ROCmIndexUrl | ForEach-Object { Redact-InstallOutput "$_" } | Out-Host
$torchInstallExit = $LASTEXITCODE
$output = ""
} else {
@ -2864,7 +3162,7 @@ if ($ROCmIndexUrl) {
}
if ($torchInstallExit -ne 0) {
Write-Host "[WARN] AMD ROCm PyTorch install failed -- falling back to CPU" -ForegroundColor Yellow
Write-Host $output -ForegroundColor Yellow
Write-Host (Redact-InstallOutput $output) -ForegroundColor Yellow
$ROCmIndexUrl = $null
$ROCmCpuFallback = $true
} else {
@ -2874,42 +3172,70 @@ if ($ROCmIndexUrl) {
}
}
if (-not $ROCmIndexUrl -and $CuTag -eq "cpu") {
if (-not $ROCmIndexUrl -and ($CuTag -eq "cpu" -or $ROCmCpuFallback)) {
substep "installing PyTorch (CPU-only)..."
# After an AMD ROCm fallback, force-reinstall so a partially-installed ROCm torch
# (which still satisfies the CPU torch>= range) is replaced by the CPU build. Skip
# the forced reinstall on a genuine CPU-only host so the common path stays fast.
# Build the array directly: an if-expression collapses @("x") to a scalar string,
# which @splat would then enumerate char-by-char into broken single-letter args.
# After an AMD ROCm fallback, force-reinstall so a partial ROCm torch (which satisfies the
# CPU torch>= range) is replaced by the CPU build; skip on a genuine CPU host to stay fast.
# $ROCmCpuFallback matters when a PINNED ROCm index failed ($CuTag is still the rocm leaf).
# Build the array directly: an if-expression collapses @("x") to a scalar @splat would
# enumerate char-by-char.
$cpuForce = @()
if ($ROCmCpuFallback) { $cpuForce = @("--force-reinstall") }
# --force-reinstall on a pin change: a stale +cu / +rocm wheel still satisfies the CPU
# torch>= range, so uv would keep it and only swap companions.
if ($script:PinChangedForceReinstall) { $cpuForce = @("--force-reinstall") }
# A PINNED cpu index installs the bounded trio (parity with _CPU_TORCH_PKG_SPEC): the /cpu
# index serves newer torch, and _ensure_cpu_torch keeps any CPU build, so a bare trio could
# land an unsupported version. Unpinned CPU hosts keep the bare trio (pre-pin behavior).
$cpuTorchSpec = "torch"; $cpuVisionSpec = "torchvision"; $cpuAudioSpec = "torchaudio"
if ($TorchIndexPinned) {
$cpuTorchSpec = "torch>=2.4,<2.12.0"
$cpuVisionSpec = "torchvision>=0.19,<0.27.0"
$cpuAudioSpec = "torchaudio>=2.4,<2.12.0"
}
if ($script:UnslothVerbose) {
Fast-Install torch torchvision torchaudio @cpuForce --index-url "$PyTorchWhlBase/cpu"
Fast-Install $cpuTorchSpec $cpuVisionSpec $cpuAudioSpec @cpuForce --index-url $TorchInstallIndexUrl | ForEach-Object { Redact-InstallOutput "$_" } | Out-Host
$torchInstallExit = $LASTEXITCODE
$output = ""
} else {
$output = Fast-Install torch torchvision torchaudio @cpuForce --index-url "$PyTorchWhlBase/cpu" | Out-String
$output = Fast-Install $cpuTorchSpec $cpuVisionSpec $cpuAudioSpec @cpuForce --index-url $TorchInstallIndexUrl | Out-String
$torchInstallExit = $LASTEXITCODE
}
if ($torchInstallExit -ne 0) {
Write-Host "[FAILED] PyTorch install failed (exit code $torchInstallExit)" -ForegroundColor Red
Write-Host $output -ForegroundColor Red
Write-Host (Redact-InstallOutput $output) -ForegroundColor Red
exit 1
}
} elseif (-not $ROCmIndexUrl) {
substep "installing PyTorch with CUDA support ($CuTag)..."
substep "(This download is ~2.8 GB -- may take a few minutes)"
# --force-reinstall on a pin change: an installed cuXXX wheel satisfies the bare torch
# requirement (PEP 440 ignores the +cuXXX tag), so without it a changed CUDA pin (cu126
# -> cu128) never applies.
$cudaForce = @()
if ($script:PinChangedForceReinstall) { $cudaForce = @("--force-reinstall") }
# An unknown-leaf custom pin (/simple, /current) routes here with $CuTag as that leaf. Bound
# the trio like the fresh custom-pin paths so a mirror can't pull an ABI-newer companion
# against the capped torch. Known cu* leaves keep bare specs.
$cudaTorchSpec = "torch"
$cudaVisionSpec = "torchvision"
$cudaAudioSpec = "torchaudio"
if ($TorchIndexPinned -and -not (Test-CudaFamilyLeaf $CuTag)) {
$cudaTorchSpec = "torch>=2.4,<2.11.0"
$cudaVisionSpec = "torchvision>=0.19,<0.26.0"
$cudaAudioSpec = "torchaudio>=2.4,<2.11.0"
}
if ($script:UnslothVerbose) {
Fast-Install torch torchvision torchaudio --index-url "$PyTorchWhlBase/$CuTag"
Fast-Install $cudaTorchSpec $cudaVisionSpec $cudaAudioSpec @cudaForce --index-url $TorchInstallIndexUrl | ForEach-Object { Redact-InstallOutput "$_" } | Out-Host
$torchInstallExit = $LASTEXITCODE
$output = ""
} else {
$output = Fast-Install torch torchvision torchaudio --index-url "$PyTorchWhlBase/$CuTag" | Out-String
$output = Fast-Install $cudaTorchSpec $cudaVisionSpec $cudaAudioSpec @cudaForce --index-url $TorchInstallIndexUrl | Out-String
$torchInstallExit = $LASTEXITCODE
}
if ($torchInstallExit -ne 0) {
Write-Host "[FAILED] PyTorch CUDA install failed (exit code $torchInstallExit)" -ForegroundColor Red
Write-Host $output -ForegroundColor Red
Write-Host (Redact-InstallOutput $output) -ForegroundColor Red
exit 1
}
@ -2925,7 +3251,7 @@ if (-not $ROCmIndexUrl -and $CuTag -eq "cpu") {
}
if ($tritonInstallExit -ne 0) {
substep "Triton install failed -- torch.compile may not work" "Yellow"
Write-Host $output -ForegroundColor Yellow
Write-Host (Redact-InstallOutput $output) -ForegroundColor Yellow
} else {
substep "Triton for Windows installed (enables torch.compile)"
}
@ -3022,7 +3348,7 @@ foreach ($pkg in @("transformers==5.3.0", "huggingface_hub==1.8.0", "hf_xet==1.4
}
if ($t5PkgExit -ne 0) {
Write-Host "[FAIL] Could not install $pkg into .venv_t5_530/" -ForegroundColor Red
Write-Host $output -ForegroundColor Red
Write-Host (Redact-InstallOutput $output) -ForegroundColor Red
$ErrorActionPreference = $prevEAP_t5
exit 1
}
@ -3057,7 +3383,7 @@ foreach ($pkg in @("transformers==5.5.0", "huggingface_hub==1.8.0", "hf_xet==1.4
}
if ($t5PkgExit -ne 0) {
Write-Host "[FAIL] Could not install $pkg into .venv_t5_550/" -ForegroundColor Red
Write-Host $output -ForegroundColor Red
Write-Host (Redact-InstallOutput $output) -ForegroundColor Red
$ErrorActionPreference = $prevEAP_t5
exit 1
}
@ -3092,7 +3418,7 @@ foreach ($pkg in @("transformers==5.10.2", "huggingface_hub==1.8.0", "hf_xet==1.
}
if ($t5PkgExit -ne 0) {
Write-Host "[FAIL] Could not install $pkg into .venv_t5_510/" -ForegroundColor Red
Write-Host $output -ForegroundColor Red
Write-Host (Redact-InstallOutput $output) -ForegroundColor Red
$ErrorActionPreference = $prevEAP_t5
exit 1
}
@ -3367,6 +3693,15 @@ if ($LocalLlamaCppLinked) {
if ($env:UNSLOTH_LLAMA_RELEASE_TAG) {
$prebuiltArgs += @("--published-release-tag", $env:UNSLOTH_LLAMA_RELEASE_TAG)
}
# UNSLOTH_LLAMA_CPP_BACKEND=cpu (case-insensitive, whitespace-trimmed) forces the
# CPU-only prebuilt via --force-cpu (persisted so updates keep it). Fixes Intel
# iGPU Vulkan crash (#7213).
$llamaBackend = "$($env:UNSLOTH_LLAMA_CPP_BACKEND)".Trim().ToLowerInvariant()
if ($llamaBackend -eq "cpu") {
$prebuiltArgs += "--force-cpu"
} elseif ($llamaBackend -and $llamaBackend -ne "auto") {
Write-Host "[WARN] Ignoring UNSLOTH_LLAMA_CPP_BACKEND='$($env:UNSLOTH_LLAMA_CPP_BACKEND)' (expected 'auto' or 'cpu')" -ForegroundColor Yellow
}
$prevEAPPrebuilt = $ErrorActionPreference
$ErrorActionPreference = "Continue"
$previousNativeErrorPreference = $null

20
studio/setup.sh
View file

@ -36,6 +36,10 @@ fi
# forces a source build, and causes HTTP 422 errors.
# Only use "master" temporarily when the latest release
# is missing support for a new model architecture.
#
# UNSLOTH_LLAMA_CPP_BACKEND : "auto" (default) or "cpu". When "cpu", forces
# the CPU-only prebuilt bundle on GPU hosts.
# Fixes Intel iGPU Vulkan crashes (#7213).
# ──────────────────────────────────────────────────────────────────────────
_DEFAULT_LLAMA_PR_FORCE=""
_DEFAULT_LLAMA_SOURCE="https://github.com/ggml-org/llama.cpp"
@ -1359,6 +1363,22 @@ else
# present so it can still attempt a prebuilt. Mirrors setup.ps1 behaviour.
_PREBUILT_CMD+=(--has-rocm)
fi
# UNSLOTH_LLAMA_CPP_BACKEND=cpu (case-insensitive, trimmed) forces the CPU-only
# prebuilt via --force-cpu, bypassing Vulkan/CUDA/ROCm. Fixes Intel iGPU crash (#7213).
# No effect on macOS: the universal bundle already runs on CPU (Metal is a runtime
# -ngl choice), so warn instead of writing a misleading forced-CPU marker.
_llama_backend="$(printf '%s' "${UNSLOTH_LLAMA_CPP_BACKEND:-auto}" | awk '{$1=$1; print tolower($0)}')"
case "$_llama_backend" in
cpu)
if [ "$_HOST_SYSTEM" = "Darwin" ]; then
step "llama.cpp" "UNSLOTH_LLAMA_CPP_BACKEND=cpu has no effect on macOS (universal build; use -ngl 0 at runtime for CPU-only)" "$C_WARN" >&2
else
_PREBUILT_CMD+=(--force-cpu)
fi
;;
""|auto) ;;
*) step "llama.cpp" "Ignoring UNSLOTH_LLAMA_CPP_BACKEND='$UNSLOTH_LLAMA_CPP_BACKEND' (expected 'auto' or 'cpu')" "$C_WARN" >&2 ;;
esac
_PREBUILT_LOG="$(mktemp)"
set +e
if _is_verbose; then

606
tests/python/test_cross_platform_parity.py
View file

@ -10,6 +10,8 @@ import pytest
REPO_ROOT = Path(__file__).resolve().parents[2]
INSTALL_SH = REPO_ROOT / "install.sh"
INSTALL_PS1 = REPO_ROOT / "install.ps1"
SETUP_PS1 = REPO_ROOT / "studio" / "setup.ps1"
STACK_PY = REPO_ROOT / "studio" / "install_python_stack.py"
class TestNoTorchBackendAutoInInstallSh:
@ -180,3 +182,607 @@ class TestUvBytecodeCompileTimeout:
assert (
'$env:UV_COMPILE_BYTECODE_TIMEOUT = "180"' in text
), "install.ps1 should default UV_COMPILE_BYTECODE_TIMEOUT"
class TestTorchIndexOverrideParity:
"""Every installer must honor UNSLOTH_TORCH_INDEX_URL / _FAMILY so a pinned wheel
index wins over GPU probing on all platforms (no asymmetric, per-OS coverage)."""
@pytest.mark.parametrize(
"path",
[INSTALL_SH, INSTALL_PS1, SETUP_PS1, STACK_PY],
ids = ["install.sh", "install.ps1", "setup.ps1", "install_python_stack.py"],
)
def test_installer_reads_override_env(self, path):
text = path.read_text(encoding = "utf-8")
for var in ("UNSLOTH_TORCH_INDEX_URL", "UNSLOTH_TORCH_INDEX_FAMILY"):
assert var in text, f"{path.name} does not honor {var}"
@pytest.mark.parametrize(
"path",
[INSTALL_PS1, SETUP_PS1],
ids = ["install.ps1", "setup.ps1"],
)
def test_amd_reroute_guarded_when_pinned(self, path):
# The AMD ROCm reroute must be skipped when the index is explicitly pinned,
# so an explicit cpu / cu* / rocm pin on an AMD host is not overwritten.
text = path.read_text(encoding = "utf-8")
assert (
"TorchIndexPinned" in text
), f"{path.name} should gate the AMD ROCm reroute on a pinned-index flag"
def test_cuda_pin_overrides_cvd_hide_gate(self):
# A pinned cu* index skips ALL host-GPU probing, so the CUDA repair must clear the
# CUDA_VISIBLE_DEVICES hide gate too (else the GPU-less CI case bails).
text = STACK_PY.read_text(encoding = "utf-8")
m = re.search(r"def _ensure_cuda_torch\(\).*?(?=\ndef )", text, re.DOTALL)
assert m, "could not locate _ensure_cuda_torch"
body = m.group(0)
assert "_cuda_pinned" in body, (
"_ensure_cuda_torch should compute a CUDA-pin flag so the pin can "
"override the CVD hide gate"
)
assert re.search(
r"if not _cuda_pinned and _cvd is not None", body
), "the CVD hide gate must be bypassed when a CUDA index is pinned"
def test_cpu_repair_pins_supported_torch_range(self):
# The explicit-CPU repair must use the bounded CPU/CUDA spec, not a bare trio (the
# /cpu index serves torch 2.11+, so a bare install could resolve out of range).
text = STACK_PY.read_text(encoding = "utf-8")
m = re.search(r"def _ensure_cpu_torch\(\).*?(?=\ndef )", text, re.DOTALL)
assert m, "could not locate _ensure_cpu_torch"
body = m.group(0)
assert "_CPU_TORCH_PKG_SPEC" in body, (
"_ensure_cpu_torch should install the bounded _CPU_TORCH_PKG_SPEC, "
"not a bare torch/torchvision/torchaudio trio"
)
def test_setup_ps1_stale_check_gates_rocm_on_supported_arch(self):
# The stale check must expect ROCm torch only for arches the install path maps to a
# repo.amd.com index; expecting "rocm" for an unmapped arch marks a good CPU venv stale.
text = SETUP_PS1.read_text(encoding = "utf-8")
assert "_rocmWheelArches" in text, (
"setup.ps1 stale check should restrict the ROCm expected-tag to the "
"supported gfx wheel arches"
)
class TestGfx211AllowlistParity:
"""The gfx per-arch 2.11-floor leaves (gfx120X-all / gfx1151 / gfx1150) must be the
SAME set in every installer and its stale/mismatch check. When they diverged, a
pinned gfx110X-all / gfx90a / gfx908 wheel (<2.11) was force-reinstalled every update."""
EXPECTED = {"gfx120x-all", "gfx1151", "gfx1150"}
def test_install_sh_allowlist(self):
text = INSTALL_SH.read_text(encoding = "utf-8").lower()
# install.sh: the TORCH_CONSTRAINT case (rocm7.2|gfx120x-all|gfx1151|gfx1150).
m = re.search(r"rocm7\.2\|gfx120x-all\|gfx1151\|gfx1150", text)
assert m, "install.sh gfx-2.11 allowlist case not found / changed"
def test_install_ps1_allowlist(self):
text = INSTALL_PS1.read_text(encoding = "utf-8").lower()
m = re.search(r"@\('gfx120x-all',\s*'gfx1151',\s*'gfx1150'\)", text)
assert m, "install.ps1 $_pinGfx211 allowlist not found / changed"
def test_setup_ps1_defines_single_allowlist_helper(self):
# setup.ps1 must define the allowlist once (Test-RocmGfx211Leaf) and reuse it, so
# the stale check and install spec can't disagree.
text = SETUP_PS1.read_text(encoding = "utf-8")
assert (
"function Test-RocmGfx211Leaf" in text
), "setup.ps1 should define a single Test-RocmGfx211Leaf allowlist helper"
assert re.search(
r"@\('gfx120x-all',\s*'gfx1151',\s*'gfx1150'\)", text.lower()
), "Test-RocmGfx211Leaf should hold the gfx-2.11 allowlist"
assert "$_pinGfx211 = Test-RocmGfx211Leaf" in text, (
"setup.ps1 install-spec path should reuse Test-RocmGfx211Leaf, not "
"re-hardcode the allowlist (they must not diverge)"
)
def test_stack_py_allowlist(self):
text = STACK_PY.read_text(encoding = "utf-8").lower()
assert (
'"gfx120x-all", "gfx1151", "gfx1150"' in text
), "install_python_stack.py _ROCM_GFX_TORCH211_LEAVES not found / changed"
class TestCudaLeafDigitParity:
"""A wheel-family leaf is CUDA only when it is "cu" + digits (cu118/cu128/...).
A bare cu* glob wrongly catches mirror leaves like /custom or /current; when
that happened the venv was marked stale and rebuilt on every run. Every
installer must require a digit after "cu" in its family/CUDA classification."""
def test_stack_py_requires_cu_digit(self):
text = STACK_PY.read_text(encoding = "utf-8")
# EXACT cu+digits: a custom leaf like cu128-private must route to the
# verbatim/unknown path, not be compared against the installed +cu128 tag.
assert re.search(
r'r"cu\[0-9\]\+"', text
), "install_python_stack.py _is_cuda_family_leaf must fullmatch cu[0-9]+"
def test_setup_ps1_requires_cu_digit(self):
text = SETUP_PS1.read_text(encoding = "utf-8")
# EXACT cu+digits: cu128-private must not classify as CUDA (it would become
# the expected tag and rebuild the venv on every update).
assert re.search(
r"'\^cu\[0-9\]\+\$'", text
), "setup.ps1 Test-CudaFamilyLeaf must match ^cu[0-9]+$, not a cu* prefix"
# The stale-venv branch must go through the digit-guarded helper.
assert (
"Test-CudaFamilyLeaf $_pinLeaf" in text
), "setup.ps1 stale check should classify CUDA via Test-CudaFamilyLeaf"
def test_install_ps1_requires_cu_digit_in_gpu_branch(self):
text = INSTALL_PS1.read_text(encoding = "utf-8")
assert re.search(
r"'\^cu\[0-9\]'", text
), "install.ps1 Get-TauriGpuBranch must require a digit after cu"
def test_install_sh_requires_cu_digit_in_gpu_branch(self):
text = INSTALL_SH.read_text(encoding = "utf-8")
# The _tauri_gpu_branch cuda case must be cu[0-9]*, not a bare cu*.
assert re.search(
r"cu\[0-9\]\*\)\s*echo \"cuda\"", text
), "install.sh _tauri_gpu_branch cuda case must be cu[0-9]*, not cu*"
def test_install_sh_backend_export_requires_cu_digit(self):
text = INSTALL_SH.read_text(encoding = "utf-8")
# Brand CUDA only on cu[0-9]*; a bare catch-all *) -> cuda would mis-brand
# /current, /custom pins and skip ROCm repair on AMD hosts.
assert re.search(
r'cu\[0-9\]\*\)\s*export UNSLOTH_TORCH_BACKEND="cuda"', text
), "install.sh backend export must brand cuda only on cu[0-9]*"
# An unknown leaf must NOT commit a cuda backend (it unsets instead).
assert re.search(
r"\*\)\s*unset UNSLOTH_TORCH_BACKEND", text
), "install.sh backend export must unset (not force cuda) on an unknown leaf"
def test_install_sh_lowercases_backend_leaf(self):
text = INSTALL_SH.read_text(encoding = "utf-8")
# The leaf feeding both the backend case and the 2.11 floor case must be
# lowercased so the canonical gfx120X-all (capital X) matches.
assert re.search(
r"_torch_index_leaf=\$\(printf '%s' \"\$_torch_index_leaf\" \| tr '\[:upper:\]' '\[:lower:\]'\)",
text,
), "install.sh must lowercase _torch_index_leaf before the gfx/rocm/cu case matches"
class TestKnown211SetParity:
"""The KNOWN-2.11 rocm/gfx set must be identical across all four installers:
exactly {rocm7.2} plus the gfx allowlist {gfx120x-all, gfx1151, gfx1150}.
rocm7.3 / torch 2.12 do not exist, so no side may floor them speculatively."""
def test_install_sh_known_211_leaf_is_rocm72_and_gfx_allowlist(self):
text = INSTALL_SH.read_text(encoding = "utf-8")
# The 2.11 floor case matches exactly rocm7.2 + the three gfx leaves.
assert re.search(
r"rocm7\.2\|gfx120x-all\|gfx1151\|gfx1150\)", text
), "install.sh 2.11 floor must be exactly rocm7.2|gfx120x-all|gfx1151|gfx1150"
# No speculative rocm7.3 anywhere.
assert "rocm7.3" not in text, "install.sh must not reference a non-existent rocm7.3"
def test_python_known_211_versions_is_only_rocm72(self):
text = STACK_PY.read_text(encoding = "utf-8")
assert "_ROCM_KNOWN_TORCH211_VERSIONS" in text
# The frozenset literal is exactly {(7, 2)}.
m = re.search(r"_ROCM_KNOWN_TORCH211_VERSIONS[^=]*=\s*frozenset\(\{([^}]*)\}\)", text)
assert m is not None, "install_python_stack.py must define _ROCM_KNOWN_TORCH211_VERSIONS"
assert "(7, 2)" in m.group(1)
assert "7, 3" not in m.group(1) and "7, 1" not in m.group(1)
def test_setup_ps1_known_211_helper_is_only_rocm72(self):
text = SETUP_PS1.read_text(encoding = "utf-8")
assert "Test-RocmKnown211Version" in text
# The predicate is Major -eq 7 -and Minor -eq 2 (only rocm7.2).
assert re.search(
r"Test-RocmKnown211Version[\s\S]{0,400}\$Major -eq 7 -and \$Minor -eq 2", text
), "setup.ps1 Test-RocmKnown211Version must accept only rocm7.2"
def test_install_ps1_pin_floor_is_only_rocm72(self):
text = INSTALL_PS1.read_text(encoding = "utf-8")
# The pinned-ROCm install-spec floor must be Major -eq 7 -and Minor -eq 2,
# not the speculative >= 2 that would floor a non-existent rocm7.3.
assert re.search(
r"\$_pinRocm211 = \(\[int\]\$Matches\[1\] -eq 7 -and \[int\]\$Matches\[2\] -eq 2\)",
text,
), "install.ps1 pinned-ROCm floor must be rocm7.2 only (no speculative >= 2)"
def test_ps1_pin_floor_gate_is_anchored(self):
"""The floor-selection gate that reads $_pinRocm211 from the raw leaf must anchor
the rocm match ($), or a suffixed custom leaf (rocm7.2-private) matches the rocm7.2
prefix, takes the 2.11-floor branch, and is force-routed through the ROCm path
before the exact-match elseif can send it to the verbatim install (Codex P2)."""
for path, label in ((INSTALL_PS1, "install.ps1"), (SETUP_PS1, "setup.ps1")):
text = path.read_text(encoding = "utf-8")
assert "-match '^rocm(\\d+)\\.(\\d+)$'" in text, (
f"{label} floor gate must anchor the rocm match (^rocm(\\d+)\\.(\\d+)$) so a "
"suffixed custom leaf is not floored/routed as rocm7.2"
)
assert (
"-match '^rocm(\\d+)\\.(\\d+)'\n" not in text
), f"{label} floor gate must not use the unanchored ^rocm(\\d+)\\.(\\d+) prefix"
def test_install_ps1_bounds_unknown_leaf_pinned_torch(self):
"""install.ps1's pinned-torch install must bound BOTH companions on EVERY
index, cu<digits> families included: torchaudio 2.11 dropped its exact torch
pin from the wheel metadata, so a bare companion beside torch<2.11 can
resolve a mismatched 2.11.0 build (Codex P2, then unconditional per the
torchaudio 2.11 unpinning)."""
text = INSTALL_PS1.read_text(encoding = "utf-8")
assert (
'$_pinVisionSpec = "torchvision>=0.19,<0.26.0"' in text
), "install.ps1 custom-pin install must bound torchvision (>=0.19,<0.26.0)"
assert (
'$_pinAudioSpec = "torchaudio>=2.4,<2.11.0"' in text
), "install.ps1 custom-pin install must bound torchaudio (>=2.4,<2.11.0)"
# No cu-family exemption: the bounds apply unconditionally.
assert (
"$_pinCuLeaf" not in text
), "install.ps1 must bound companions on every index (no cu-family exemption)"
# The bounded companions must actually be passed to the install command.
assert re.search(
r'"torch>=2\.4,<2\.11\.0" \$_pinVisionSpec \$_pinAudioSpec --default-index \$TorchIndexUrl',
text,
), "install.ps1 custom-pin install must pass the bounded companion specs to uv"
def test_gfx_allowlist_matches_across_installers(self):
# The gfx 2.11 allowlist {gfx120x-all, gfx1151, gfx1150} must appear in each.
gfx = ("gfx120x-all", "gfx1151", "gfx1150")
for path, label in (
(INSTALL_SH, "install.sh"),
(INSTALL_PS1, "install.ps1"),
(SETUP_PS1, "setup.ps1"),
(STACK_PY, "install_python_stack.py"),
):
low = path.read_text(encoding = "utf-8").lower()
for g in gfx:
assert g in low, f"{label} missing gfx 2.11 allowlist member {g}"
class TestPinnedRocmLeafDigitParity:
"""A pinned index is a pip ROCm --default-index family only when its leaf is an
EXACT rocm+digits (rocm7 / rocm7.2) or gfx*. A ^rocm[0-9] PREFIX (or a bare rocm*
glob) wrongly catches a custom mirror / find-links leaf (rocm-current /
rocm-rel-7.2.1) AND a suffixed private-mirror leaf (rocm7.2-private / rocm7-current),
routing it through the ROCm install path (which silently falls back to CPU on
failure) or skipping the custom-index companion bounds, instead of the verbatim
--default-index install. All installers must match the family EXACTLY: Python and
install.sh via a shared _is_pip_rocm_family_leaf, setup.ps1 via Test-PipRocmFamilyLeaf,
install.ps1 via an anchored ^rocm[0-9]+(\\.[0-9]+)?$ reroute."""
def test_install_ps1_pinned_reroute_requires_rocm_digit(self):
text = INSTALL_PS1.read_text(encoding = "utf-8")
# The pinned gfx*/rocm reroute must match rocm EXACTLY (anchored), so a suffixed
# rocm7.2-private / rocm-current falls through to the verbatim --default-index path.
assert "-match '^rocm[0-9]+(\\.[0-9]+)?$'" in text, (
"install.ps1 pinned-index reroute must anchor the rocm match "
"(^rocm[0-9]+(\\.[0-9]+)?$), not a bare -like 'rocm*' or an unanchored ^rocm\\d"
)
# Neither the broad glob nor the unanchored prefix may drive that reroute.
assert (
"-like 'rocm*'" not in text
), "install.ps1 must not route a pinned index on a bare -like 'rocm*' glob"
assert (
"-match '^rocm\\d'" not in text
), "install.ps1 must not route a pinned index on an unanchored -match '^rocm\\d'"
def test_setup_ps1_pinned_reroute_requires_rocm_digit(self):
text = SETUP_PS1.read_text(encoding = "utf-8")
# setup.ps1 routes every family decision through Test-PipRocmFamilyLeaf, which
# anchors the rocm match so a suffixed custom leaf stays on the verbatim path.
assert (
"function Test-PipRocmFamilyLeaf" in text
), "setup.ps1 must define Test-PipRocmFamilyLeaf (the exact rocm/gfx family gate)"
assert "'^rocm[0-9]+(\\.[0-9]+)?$'" in text, (
"setup.ps1 Test-PipRocmFamilyLeaf must anchor the rocm match "
"(^rocm[0-9]+(\\.[0-9]+)?$) so rocm7.2-private / rocm-current stay verbatim"
)
pinned_block = text[text.find("$_pinGfx211 = Test-RocmGfx211Leaf") :][:2000]
assert (
"-like 'rocm*'" not in pinned_block
), "setup.ps1 pinned reroute must not route on a bare -like 'rocm*' glob"
def test_install_sh_repairable_requires_rocm_digit(self):
text = INSTALL_SH.read_text(encoding = "utf-8")
# _torch_index_repairable routes rocm/gfx through the exact-match helper.
assert (
"_is_pip_rocm_family_leaf" in text
), "install.sh must define/use _is_pip_rocm_family_leaf for the exact rocm gate"
# gfx needs a following digit: gfx-private / gfxfoo are custom verbatim pins.
assert re.search(
r'case "\$1" in\n\s*gfx\[0-9\]\*\) return 0', text
), "install.sh _is_pip_rocm_family_leaf must treat only gfx<digit>* as a family"
assert not re.search(
r'case "\$1" in\n\s*gfx\*\) return 0', text
), "install.sh _is_pip_rocm_family_leaf must not family-match a bare gfx* glob"
def test_stack_py_pip_rocm_family_requires_digit(self):
text = STACK_PY.read_text(encoding = "utf-8")
assert re.search(
r'fullmatch\(r"rocm\\d\+\(\?:\\\.\\d\+\)\?", leaf\)', text
), "install_python_stack.py _is_pip_rocm_family_leaf must fullmatch rocm\\d+(?:\\.\\d+)?"
# The unanchored prefix must be gone from the family/flavor gates.
assert (
're.match(r"^rocm\\d"' not in text
), "install_python_stack.py must not gate a family on an unanchored re.match(^rocm\\d)"
def test_install_sh_rocm_side_effects_digit_gated(self):
"""The AMD bitsandbytes + 'repair ROCm torch' side effects must fire only on
an EXACT ROCm family (rocm7.2/gfx*), not a bare */rocm* whole-URL glob nor a
^rocm[0-9] prefix that catches a custom CPU/CUDA index like /rocm-current or a
suffixed /rocm7.2-private and force-repairs it from the wrong --default-index."""
text = INSTALL_SH.read_text(encoding = "utf-8")
assert (
'if _is_pip_rocm_family_leaf "$_torch_index_leaf"; then\n _torch_index_is_rocm_family=true'
in text
), "install.sh must set _torch_index_is_rocm_family from the exact-match helper"
assert (
'[ "$_torch_index_is_rocm_family" = true ]' in text
), "install.sh ROCm bnb/repair hooks must gate on _torch_index_is_rocm_family"
assert (
"*/rocm*|*/gfx*)\n _install_bnb_rocm" not in text
), "install.sh must not gate _install_bnb_rocm on a bare */rocm* whole-URL glob"
class TestPinnedIndexClearsUvEnvParity:
"""Every installer must neutralise the uv index env vars for a pinned torch
install (#6898). uv treats the default index (--index-url / --default-index) as
lowest priority, so an inherited UV_INDEX / UV_EXTRA_INDEX_URL mirror would win
under uv's first-index strategy and pull torch from the wrong index -- after
which the pinned wheel index is silently never used."""
UV_VARS = ("UV_DEFAULT_INDEX", "UV_INDEX_URL", "UV_INDEX", "UV_EXTRA_INDEX_URL")
def test_install_sh_clears_uv_index_vars(self):
text = INSTALL_SH.read_text(encoding = "utf-8")
assert (
"env -u UV_DEFAULT_INDEX -u UV_INDEX_URL -u UV_INDEX -u UV_EXTRA_INDEX_URL" in text
), "install.sh run_install_cmd must clear the uv index vars for --default-index installs"
def test_install_ps1_clears_uv_index_vars(self):
text = INSTALL_PS1.read_text(encoding = "utf-8")
for var in self.UV_VARS:
assert var in text, f"install.ps1 must clear {var} for pinned installs"
def test_setup_ps1_clears_uv_index_vars(self):
text = SETUP_PS1.read_text(encoding = "utf-8")
for var in self.UV_VARS:
assert var in text, f"setup.ps1 must clear {var} for pinned installs"
def test_stack_py_clears_uv_index_vars(self):
text = STACK_PY.read_text(encoding = "utf-8")
assert "_install_env_for_cmd" in text, (
"install_python_stack.py must scrub inherited uv index vars for pinned "
"installs via _install_env_for_cmd (parity with install.sh #6898)"
)
for var in self.UV_VARS:
assert var in text, f"install_python_stack.py must clear {var} for pinned installs"
def test_all_installers_clear_uv_torch_backend(self):
"""uv's torch backend redirects torch resolution to its own per-backend
index even against an explicit pin, so every installer's pinned-install
scrub must clear UV_TORCH_BACKEND too."""
sh = INSTALL_SH.read_text(encoding = "utf-8")
assert "-u UV_TORCH_BACKEND" in sh, "install.sh pinned scrub must clear UV_TORCH_BACKEND"
for path in (INSTALL_PS1, SETUP_PS1):
text = path.read_text(encoding = "utf-8")
assert (
"'UV_TORCH_BACKEND'" in text
), f"{path.name} pinned scrub must clear UV_TORCH_BACKEND"
stack = STACK_PY.read_text(encoding = "utf-8")
assert (
'"UV_TORCH_BACKEND",' in stack
), "install_python_stack.py strip tuple must include UV_TORCH_BACKEND"
def test_stack_py_strips_pip_extra_index_for_pip_fallback(self):
"""The pip fallback honours PIP_EXTRA_INDEX_URL (pip adds it IN ADDITION
to --index-url), so the pinned-command scrub must strip it."""
stack = STACK_PY.read_text(encoding = "utf-8")
assert (
'"PIP_EXTRA_INDEX_URL",' in stack
), "install_python_stack.py strip tuple must include PIP_EXTRA_INDEX_URL"
def test_all_installers_scrub_find_links(self):
"""uv's --find-links (env UV_FIND_LINKS) adds candidate locations that can
satisfy torch off a pinned index; every pinned-install scrub must clear it."""
sh = INSTALL_SH.read_text(encoding = "utf-8")
assert "-u UV_FIND_LINKS" in sh
for path in (INSTALL_PS1, SETUP_PS1):
assert "'UV_FIND_LINKS'" in path.read_text(encoding = "utf-8"), path.name
stack = STACK_PY.read_text(encoding = "utf-8")
assert '"UV_FIND_LINKS",' in stack and '"PIP_FIND_LINKS",' in stack
def test_setup_ps1_scrub_covers_pip_fallback(self):
"""setup.ps1's Fast-Install must keep the scrub active through the pip
fallback (pip honours PIP_EXTRA_INDEX_URL / PIP_FIND_LINKS in addition to
--index-url); restoring the vars before the fallback reopens the hole."""
text = SETUP_PS1.read_text(encoding = "utf-8")
fi = text[text.find("function Fast-Install") :][:2500]
assert "'PIP_EXTRA_INDEX_URL'" in fi and "'PIP_FIND_LINKS'" in fi
# the pip fallback must sit INSIDE the try whose finally restores the vars
assert fi.find("python -m pip install") < fi.find(
"finally"
), "pip fallback must run before the scrub is restored"
def test_all_installers_disable_uv_config_for_pinned_installs(self):
"""A DISCOVERED uv.toml / pyproject [tool.uv] outranks the CLI pin
(verified with uv 0.10: [pip] torch-backend = "cpu" and a non-default
[[index]] both resolve torch+cpu against an explicit --index-url /
--default-index cu126 pin; UV_NO_CONFIG=1 restores the pin). Every
installer's pinned scrub must set UV_NO_CONFIG=1 and drop UV_CONFIG_FILE."""
sh = INSTALL_SH.read_text(encoding = "utf-8")
assert "-u UV_CONFIG_FILE UV_NO_CONFIG=1" in sh, (
"install.sh run_install_cmd must set UV_NO_CONFIG=1 and drop "
"UV_CONFIG_FILE for --default-index installs"
)
for path in (INSTALL_PS1, SETUP_PS1):
text = path.read_text(encoding = "utf-8")
assert "'UV_CONFIG_FILE'" in text, f"{path.name} must drop UV_CONFIG_FILE"
assert (
"$env:UV_NO_CONFIG = '1'" in text
), f"{path.name} must set UV_NO_CONFIG=1 for pinned installs"
stack = STACK_PY.read_text(encoding = "utf-8")
assert (
'"UV_CONFIG_FILE",' in stack
), "install_python_stack.py strip tuple must include UV_CONFIG_FILE"
assert (
'env["UV_NO_CONFIG"] = "1"' in stack
), "_install_env_for_cmd must set UV_NO_CONFIG=1 for pinned installs"
def test_pip_fallbacks_disable_pip_config_files(self):
"""The pip FALLBACK (uv missing/failed) honours user/site pip config files
even with the PIP_* env vars stripped: `pip config set
global.extra-index-url` still adds indexes to a pinned install. pip loads
NO configuration files when PIP_CONFIG_FILE is the platform devnull, so
the two installers that HAVE a pip fallback (install_python_stack.py and
setup.ps1's Fast-Install) must set it in their pinned scrub. install.sh
and install.ps1 are uv-only (no python -m pip fallback) and need no
equivalent."""
stack = STACK_PY.read_text(encoding = "utf-8")
assert 'env["PIP_CONFIG_FILE"] = os.devnull' in stack, (
"_install_env_for_cmd must point PIP_CONFIG_FILE at os.devnull for "
"pinned installs (pip fallback isolation)"
)
setup = SETUP_PS1.read_text(encoding = "utf-8")
assert "$env:PIP_CONFIG_FILE = 'nul'" in setup, (
"setup.ps1 Fast-Install pinned scrub must point PIP_CONFIG_FILE at nul "
"(Windows devnull) so the pip fallback ignores user/site pip config"
)
assert (
"'PIP_CONFIG_FILE'" in setup
), "setup.ps1 must save/restore PIP_CONFIG_FILE around the pinned scrub"
def test_setup_ps1_bounds_unknown_leaf_pinned_torch(self):
"""A first-time/changed unknown-leaf custom pin routes through setup.ps1's
CUDA branch; install.ps1's fresh pinned install, install.sh, and the Python
verbatim path bound the WHOLE trio, so the Windows update path must too -- a
private mirror serving newer torch OR newer companions must not lift the venv
above the supported range under the pin."""
text = SETUP_PS1.read_text(encoding = "utf-8")
# The custom-leaf branch bounds torch AND both companions (parity with the
# other installers' custom-pin trio bounds), gated on a non-cu-family leaf.
for spec in (
'$cudaTorchSpec = "torch>=2.4,<2.11.0"',
'$cudaVisionSpec = "torchvision>=0.19,<0.26.0"',
'$cudaAudioSpec = "torchaudio>=2.4,<2.11.0"',
):
assert spec in text, f"setup.ps1 must bound the custom-leaf trio: {spec}"
assert (
"if ($TorchIndexPinned -and -not (Test-CudaFamilyLeaf $CuTag)) {" in text
), "the custom-leaf trio bounds must be gated on a pinned non-cu-family leaf"
assert (
"Fast-Install $cudaTorchSpec $cudaVisionSpec $cudaAudioSpec" in text
), "setup.ps1's CUDA branch must install via the bounded spec variables"
def test_setup_ps1_bounds_pinned_cpu_torch(self):
"""setup.ps1's CPU branch must bound the trio under an explicit pin (parity with
_CPU_TORCH_PKG_SPEC): the /cpu index serves newer torch, and _ensure_cpu_torch
keeps any CPU build, so a bare pinned trio could land an unsupported version.
An unpinned CPU host keeps the bare trio (pre-pin behavior unchanged)."""
text = SETUP_PS1.read_text(encoding = "utf-8")
for spec in (
'$cpuTorchSpec = "torch>=2.4,<2.12.0"',
'$cpuVisionSpec = "torchvision>=0.19,<0.27.0"',
'$cpuAudioSpec = "torchaudio>=2.4,<2.12.0"',
):
assert spec in text, f"setup.ps1 must bound the pinned CPU trio: {spec}"
assert (
"if ($TorchIndexPinned) {" in text
), "the CPU trio bounds must be gated on an explicit pin"
assert (
"Fast-Install $cpuTorchSpec $cpuVisionSpec $cpuAudioSpec @cpuForce" in text
), "setup.ps1's CPU branch must install via the spec variables"
# The ceilings mirror the Python repair spec exactly.
stack = STACK_PY.read_text(encoding = "utf-8")
spec_block = re.search(r"_CUDA_TORCH_PKG_SPEC[^(]*\(\s*(.*?)\)", stack, re.DOTALL)
assert spec_block and '"torch>=2.4,<2.12.0"' in spec_block.group(1), (
"_CPU_TORCH_PKG_SPEC (via _CUDA_TORCH_PKG_SPEC) must keep the torch<2.12 "
"ceiling the setup.ps1 pinned CPU branch mirrors"
)
def test_setup_ps1_stale_check_requires_rocm_digit(self):
"""The stale-venv check must use the same EXACT rocm/gfx gate as the install
selection (Test-PipRocmFamilyLeaf), or a custom rocm-* / suffixed rocm7.2-private
leaf is stale-compared as a family and force-reinstalls on every studio update."""
text = SETUP_PS1.read_text(encoding = "utf-8")
anchor = text.find("$_pinLeaf = Get-TorchIndexLeaf $_pinnedIdx")
assert anchor >= 0, "setup.ps1 stale check must classify the pinned leaf"
stale = text[anchor:][:2500]
assert (
"Test-PipRocmFamilyLeaf" in stale
), "setup.ps1 stale check must gate rocm leaves via the exact Test-PipRocmFamilyLeaf"
assert (
stale.count("-like 'rocm*'") == 0
), "setup.ps1 stale check must not use a bare -like 'rocm*' glob"
assert (
"-match '^rocm\\d'" not in stale
), "setup.ps1 stale check must not use an unanchored -match '^rocm\\d'"
class TestIndexPathSlashTrimParity:
"""Every installer must trim trailing PATH slashes only on the verbatim
UNSLOTH_TORCH_INDEX_URL override, preserving a ?query/#fragment token: a whole-URL
strip corrupts a base64 token ending in "/", a single strip leaves a double-slash leaf
empty. The helper must be DEFINED and WIRED into the override return in all four."""
def test_helper_defined_in_all_installers(self):
assert "def _trim_index_path_slashes(" in STACK_PY.read_text(encoding = "utf-8")
assert "_trim_index_path_slashes()" in INSTALL_SH.read_text(encoding = "utf-8")
assert "function Trim-IndexPathSlashes" in INSTALL_PS1.read_text(encoding = "utf-8")
assert "function Trim-IndexPathSlashes" in SETUP_PS1.read_text(encoding = "utf-8")
def test_helper_wired_into_override_in_all_installers(self):
assert "_trim_index_path_slashes(url)" in STACK_PY.read_text(encoding = "utf-8")
assert '_url=$(_trim_index_path_slashes "$_url")' in INSTALL_SH.read_text(encoding = "utf-8")
assert "Trim-IndexPathSlashes $env:UNSLOTH_TORCH_INDEX_URL" in INSTALL_PS1.read_text(
encoding = "utf-8"
)
assert "Trim-IndexPathSlashes $env:UNSLOTH_TORCH_INDEX_URL" in SETUP_PS1.read_text(
encoding = "utf-8"
)
class TestInstallOutputRedactionParity:
"""uv/pip failure text embeds the failing --index-url verbatim, so a captured install
log dumped on error can leak a user:token@ or ?token= secret. Every installer must
DEFINE a redaction helper and WIRE it into the captured-output print path."""
def test_helper_defined_in_all_installers(self):
assert "def _redact_install_output(" in STACK_PY.read_text(encoding = "utf-8")
assert "_redact_install_output()" in INSTALL_SH.read_text(encoding = "utf-8")
assert "function Redact-InstallOutput" in INSTALL_PS1.read_text(encoding = "utf-8")
assert "function Redact-InstallOutput" in SETUP_PS1.read_text(encoding = "utf-8")
def test_helper_wired_into_failure_print(self):
# install.sh dumps the captured log through the redactor on failure.
assert '_redact_install_output "$_log"' in INSTALL_SH.read_text(encoding = "utf-8")
# Both ps1 installers redact the captured $output before Write-Host on non-zero exit.
assert (
"Write-Host (Redact-InstallOutput $output) -ForegroundColor Red"
in INSTALL_PS1.read_text(encoding = "utf-8")
)
assert (
"Write-Host (Redact-InstallOutput $output) -ForegroundColor Red"
in SETUP_PS1.read_text(encoding = "utf-8")
)
# Python redacts the captured stdout before printing.
assert "_redact_install_output(" in STACK_PY.read_text(encoding = "utf-8")
class TestPipNoIndexScrubParity:
"""The plain-pip fallback honours PIP_*: PIP_NO_INDEX=1 makes it ignore ALL indexes
(defeating the pinned --index-url) and PIP_INDEX_URL replaces the pin. The two installers
that HAVE a plain-pip fallback (Python + setup.ps1) must scrub both for a pinned install.
install.sh / install.ps1 are uv-only (--default-index), which ignores pip config/env."""
def test_python_scrubs_pip_no_index_and_pip_index_url(self):
text = STACK_PY.read_text(encoding = "utf-8")
assert '"PIP_NO_INDEX"' in text
assert '"PIP_INDEX_URL"' in text
def test_setup_ps1_scrubs_pip_no_index_and_pip_index_url(self):
text = SETUP_PS1.read_text(encoding = "utf-8")
assert "'PIP_NO_INDEX'" in text
assert "'PIP_INDEX_URL'" in text

134
tests/python/test_install_python_stack.py
View file

@ -54,6 +54,24 @@ class TestBuildUvCmdTorchBackend:
a.startswith("--torch-backend") for a in cmd
), f"Empty UV_TORCH_BACKEND should not add flag, got: {cmd}"
def test_uv_torch_backend_skipped_for_pinned_index(self):
"""A pinned-index command must NOT get --torch-backend: uv's torch backend
redirects torch resolution to its own per-backend index even when
--index-url is given (verified: cu128 pin + backend cpu installs
torch+cpu), defeating the pin."""
for pin_flag in ("--index-url", "--default-index"):
with mock.patch.dict(os.environ, {"UV_TORCH_BACKEND": "cpu"}):
cmd = self._call(("torch", pin_flag, "https://download.pytorch.org/whl/cu128"))
assert not any(
a.startswith("--torch-backend") for a in cmd
), f"{pin_flag} command must not carry --torch-backend, got: {cmd}"
def test_uv_torch_backend_kept_for_unpinned(self):
"""Non-pinned commands still honour UV_TORCH_BACKEND."""
with mock.patch.dict(os.environ, {"UV_TORCH_BACKEND": "cpu"}):
cmd = self._call(("somepackage",))
assert "--torch-backend=cpu" in cmd
class TestUvSafePath:
"""_uv_safe_path hands uv a space-free `-c`/`-r` path (issue #6503)."""
@ -148,3 +166,119 @@ class TestUvSafePathHardening:
assert " " not in value
assert Path(value).read_text() == "transformers>=4.57.6\n"
class TestPinnedIndexClearsUvEnv:
"""A pinned torch install (--index-url / --default-index) must neutralise an
inherited UV_INDEX / UV_EXTRA_INDEX_URL so the pinned wheel index wins.
uv treats the default index (--index-url / --default-index) as LOWEST priority,
so an inherited UV_INDEX / UV_EXTRA_INDEX_URL (a corporate/CPU mirror) would be
searched first and, under uv's default first-index strategy, resolve torch from
the wrong mirror -- after which the marker records a wheel index that was never
used. install.sh (#6898), install.ps1 and setup.ps1 already clear these for
pinned installs; install_python_stack must match (parity across all installers).
"""
UV_VARS = ("UV_DEFAULT_INDEX", "UV_INDEX_URL", "UV_INDEX", "UV_EXTRA_INDEX_URL")
def test_pinned_index_url_strips_uv_index_vars(self):
cmd = [
"uv",
"pip",
"install",
"--force-reinstall",
"torch",
"torchvision",
"torchaudio",
"--index-url",
"https://download.pytorch.org/whl/cu128",
]
with mock.patch.dict(
os.environ,
{
"UV_INDEX": "https://mirror.corp/simple",
"UV_EXTRA_INDEX_URL": "https://mirror.corp/extra",
"UV_INDEX_URL": "https://mirror.corp/root",
"UV_DEFAULT_INDEX": "https://mirror.corp/default",
},
):
env = ips._install_env_for_cmd(cmd)
assert env is not None, "a --index-url install must run with a scrubbed env"
for var in self.UV_VARS:
assert var not in env, f"{var} must be cleared for a pinned-index install"
def test_pinned_default_index_strips_uv_index_vars(self):
# --default-index must be gated too (matches install.sh / install.ps1).
cmd = ["uv", "pip", "install", "torch", "--default-index", "https://x/cu126"]
with mock.patch.dict(os.environ, {"UV_INDEX": "https://mirror.corp/simple"}):
env = ips._install_env_for_cmd(cmd)
assert env is not None
assert "UV_INDEX" not in env
def test_non_pinned_install_keeps_user_mirror(self):
# A plain install (no --index-url) must NOT scrub the env, so a user's mirror
# still applies to base packages.
cmd = ["uv", "pip", "install", "unsloth", "unsloth-zoo"]
with mock.patch.dict(os.environ, {"UV_INDEX": "https://mirror.corp/simple"}):
env = ips._install_env_for_cmd(cmd)
assert env is None, "non-pinned installs must inherit the caller env unchanged"
def test_scrubbed_env_preserves_other_vars(self):
cmd = ["uv", "pip", "install", "torch", "--index-url", "https://x/cu128"]
with mock.patch.dict(
os.environ,
{"UV_INDEX": "https://mirror.corp/simple", "PATH_SENTINEL_XYZ": "keepme"},
):
env = ips._install_env_for_cmd(cmd)
assert env is not None
assert env.get("PATH_SENTINEL_XYZ") == "keepme", "only uv index vars are removed"
def test_pinned_cmd_strips_pip_extra_index_url(self):
"""PIP_EXTRA_INDEX_URL is stripped for pinned commands so the pip
fallback cannot satisfy torch from an inherited extra index."""
with mock.patch.dict(os.environ, {"PIP_EXTRA_INDEX_URL": "https://mirror/simple"}):
env = ips._install_env_for_cmd(
["pip", "install", "torch", "--index-url", "https://x/cu128"]
)
assert env is not None and "PIP_EXTRA_INDEX_URL" not in env
def test_pinned_cmd_strips_uv_torch_backend(self):
"""UV_TORCH_BACKEND is stripped for pinned commands so uv cannot read it
from the environment and reroute torch off the pinned index."""
with mock.patch.dict(os.environ, {"UV_TORCH_BACKEND": "cpu"}):
env = ips._install_env_for_cmd(
["uv", "pip", "install", "torch", "--index-url", "https://x/cu128"]
)
assert env is not None and "UV_TORCH_BACKEND" not in env
def test_pinned_cmd_disables_uv_config_discovery(self):
"""A DISCOVERED uv.toml / pyproject [tool.uv] outranks the CLI pin too
(verified with uv 0.10: [pip] torch-backend = "cpu" and a non-default
[[index]] both resolve torch+cpu against an explicit --index-url /
--default-index cu126 pin). Pinned commands must run with UV_NO_CONFIG=1
and without an inherited UV_CONFIG_FILE."""
with mock.patch.dict(os.environ, {"UV_CONFIG_FILE": "/etc/uv/uv.toml"}):
env = ips._install_env_for_cmd(
["uv", "pip", "install", "torch", "--index-url", "https://x/cu128"]
)
assert env is not None
assert env.get("UV_NO_CONFIG") == "1"
assert "UV_CONFIG_FILE" not in env
def test_pinned_cmd_disables_pip_config_files(self):
"""The pip FALLBACK honours user/site pip config files (pip config set
global.extra-index-url) even with the PIP_* env vars stripped; pip loads
NO configuration files when PIP_CONFIG_FILE is os.devnull. Harmless for
uv, decisive for the fallback."""
env = ips._install_env_for_cmd(
["uv", "pip", "install", "torch", "--index-url", "https://x/cu128"]
)
assert env is not None
assert env.get("PIP_CONFIG_FILE") == os.devnull
def test_non_pinned_cmd_keeps_uv_config_discovery(self):
"""Non-pinned installs inherit the caller env unchanged, so a user's uv
configuration still applies to base packages."""
env = ips._install_env_for_cmd(["uv", "pip", "install", "unsloth"])
assert env is None

1
tests/run_all.sh
View file

@ -15,6 +15,7 @@ sh "$TESTS_DIR/sh/test_resolve_cuda_archs.sh"
sh "$TESTS_DIR/sh/test_strixhalo_wsl_reroute.sh"
sh "$TESTS_DIR/sh/test_uninstall_shared_icon.sh"
sh "$TESTS_DIR/sh/test_torch_flavor.sh"
sh "$TESTS_DIR/sh/test_redact_install_output.sh"
sh "$TESTS_DIR/sh/test_install_uv_override_space.sh"
echo ""

57
tests/sh/test_get_torch_index_url.sh
View file

@ -23,6 +23,8 @@ _FAKE_SMI_DIR=$(mktemp -d)
echo ""
sed -n '/^_has_usable_nvidia_gpu()/,/^}/p' "$INSTALL_SH"
echo ""
sed -n '/^_trim_index_path_slashes()/,/^}/p' "$INSTALL_SH"
echo ""
sed -n '/^get_torch_index_url()/,/^}/p' "$INSTALL_SH"
} | sed "s|/usr/bin/nvidia-smi|$_FAKE_SMI_DIR/nvidia-smi-absent|g" \
> "$_FUNC_FILE"
@ -379,6 +381,61 @@ _result=$(run_func "$_dir" " -1 ")
assert_eq "CVD=' -1 ' hides NVIDIA -> cpu" "https://download.pytorch.org/whl/cpu" "$_result"
rm -rf "$_dir"
# --- explicit overrides (headless / container / CI; no GPU probing) ----------
# 39) UNSLOTH_TORCH_INDEX_FAMILY pins the family with no GPU present (not the cpu fallback).
_result=$(UNSLOTH_TORCH_INDEX_FAMILY="cu128" run_func "none")
assert_eq "family override (no GPU) -> cu128" "https://download.pytorch.org/whl/cu128" "$_result"
# 40) Family override beats real detection: an nvidia-smi 12.6 host still gets cu128
# (the Docker-build case -- builder sees the host driver but publishes a cu128 image).
_dir=$(make_mock_smi "12.6")
_result=$(UNSLOTH_TORCH_INDEX_FAMILY="cu128" run_func "$_dir")
assert_eq "family override beats detected 12.6 -> cu128" "https://download.pytorch.org/whl/cu128" "$_result"
rm -rf "$_dir"
# 41) UNSLOTH_TORCH_INDEX_URL is used verbatim and wins over detection.
_dir=$(make_mock_smi "12.6")
_result=$(UNSLOTH_TORCH_INDEX_URL="https://mirror.example.com/whl/cu999" run_func "$_dir")
assert_eq "url override beats detection -> verbatim" "https://mirror.example.com/whl/cu999" "$_result"
rm -rf "$_dir"
# 42) Family override is appended to UNSLOTH_PYTORCH_MIRROR (mirror still honoured).
_result=$(UNSLOTH_PYTORCH_MIRROR="https://mirror.example.com/whl" UNSLOTH_TORCH_INDEX_FAMILY="cu128" run_func "none")
assert_eq "mirror + family override -> mirror/cu128" "https://mirror.example.com/whl/cu128" "$_result"
# 43) Trailing slash in UNSLOTH_TORCH_INDEX_URL is stripped.
_result=$(UNSLOTH_TORCH_INDEX_URL="https://mirror.example.com/whl/cu128/" run_func "none")
assert_eq "url override trailing slash stripped" "https://mirror.example.com/whl/cu128" "$_result"
# 44) URL override takes precedence over family override.
_result=$(UNSLOTH_TORCH_INDEX_URL="https://mirror.example.com/whl/cu130" UNSLOTH_TORCH_INDEX_FAMILY="cu128" run_func "none")
assert_eq "url override beats family override -> url" "https://mirror.example.com/whl/cu130" "$_result"
# 45) An empty override is ignored (falls through to normal detection).
_result=$(UNSLOTH_TORCH_INDEX_FAMILY="" UNSLOTH_TORCH_INDEX_URL="" run_func "none")
assert_eq "empty overrides ignored -> detected cpu" "https://download.pytorch.org/whl/cpu" "$_result"
# 46) ALL trailing slashes are stripped from a URL override (not just one).
_result=$(UNSLOTH_TORCH_INDEX_URL="https://mirror.example.com/whl/cu128///" run_func "none")
assert_eq "url override double slash stripped" "https://mirror.example.com/whl/cu128" "$_result"
# 47) Leading and trailing slashes stripped from a family override.
_result=$(UNSLOTH_TORCH_INDEX_FAMILY="//cu128//" run_func "none")
assert_eq "family override slashes stripped" "https://download.pytorch.org/whl/cu128" "$_result"
# 48) A ?query token that ends in "/" is PRESERVED: only PATH slashes are trimmed, so a
# base64 token ending in "/" is not corrupted (path-only trim, not whole-URL rstrip).
_result=$(UNSLOTH_TORCH_INDEX_URL="https://mirror.example.com/whl/cu128?token=ab12cd/" run_func "none")
assert_eq "url override preserves query token slash" "https://mirror.example.com/whl/cu128?token=ab12cd/" "$_result"
# 49) Double PATH slash before a query is collapsed while the query survives intact.
_result=$(UNSLOTH_TORCH_INDEX_URL="https://mirror.example.com/whl/cu128//?token=ab12cd/" run_func "none")
assert_eq "url override path slash trimmed, query kept" "https://mirror.example.com/whl/cu128?token=ab12cd/" "$_result"
# 50) A #fragment ending in "/" is likewise preserved.
_result=$(UNSLOTH_TORCH_INDEX_URL="https://mirror.example.com/whl/cu128#anchor/" run_func "none")
assert_eq "url override preserves fragment slash" "https://mirror.example.com/whl/cu128#anchor/" "$_result"
rm -f "$_FUNC_FILE"
rm -rf "$_FAKE_SMI_DIR"
rm -rf "$_TOOLS_DIR"

89
tests/sh/test_redact_install_output.sh Executable file
View file

@ -0,0 +1,89 @@
#!/bin/bash
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
# Unit tests for install.sh's _redact_install_output helper. uv/pip failure text embeds the
# failing --index-url verbatim, so a captured install log dumped on error can leak a
# user:token@ or ?token= secret. The helper redacts both before printing. Mirrors
# _redact_install_output (install_python_stack.py) / Redact-InstallOutput (install.ps1 /
# setup.ps1).
set -e
SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
INSTALL_SH="$SCRIPT_DIR/../../install.sh"
PASS=0
FAIL=0
_FUNC_FILE=$(mktemp)
sed -n '/^_redact_install_output()/,/^}/p' "$INSTALL_SH" > "$_FUNC_FILE"
# shellcheck disable=SC1090
. "$_FUNC_FILE"
rm -f "$_FUNC_FILE"
assert_eq() {
_label="$1"; _expected="$2"; _actual="$3"
if [ "$_actual" = "$_expected" ]; then
echo " PASS: $_label"; PASS=$((PASS + 1))
else
echo " FAIL: $_label (expected '$_expected', got '$_actual')"; FAIL=$((FAIL + 1))
fi
}
# Redact from a file (the actual call site passes a captured-log tempfile).
redact_str() {
_rs_tmp=$(mktemp)
printf '%s\n' "$1" > "$_rs_tmp"
_rs_out=$(_redact_install_output "$_rs_tmp")
rm -f "$_rs_tmp"
printf '%s' "$_rs_out"
}
echo "=== _redact_install_output ==="
assert_eq "userinfo user:token@ redacted" \
"ERROR: failed https://<redacted>@download.pytorch.org/whl/cu128" \
"$(redact_str 'ERROR: failed https://alice:s3cr3t@download.pytorch.org/whl/cu128')"
assert_eq "bare-token@ userinfo redacted" \
"fetch https://<redacted>@host/whl/cu128 failed" \
"$(redact_str 'fetch https://ghp_deadbeef@host/whl/cu128 failed')"
assert_eq "single ?token= query redacted" \
"url https://host/whl/cu128?token=<redacted> unreachable" \
"$(redact_str 'url https://host/whl/cu128?token=abcd1234 unreachable')"
assert_eq "multiple query values redacted" \
"https://host/whl/cu128?token=<redacted>&channel=<redacted>" \
"$(redact_str 'https://host/whl/cu128?token=abcd1234&channel=beta')"
assert_eq "http (not https) userinfo redacted" \
"http://<redacted>@host/simple" \
"$(redact_str 'http://u:p@host/simple')"
assert_eq "fragment token redacted" \
"ERROR: could not fetch https://mirror.local/whl/cu128#<redacted> (403)" \
"$(redact_str 'ERROR: could not fetch https://mirror.local/whl/cu128#token=SECRET123 (403)')"
assert_eq "query and fragment both redacted" \
"https://host/whl/cu128?token=<redacted>#<redacted> done" \
"$(redact_str 'https://host/whl/cu128?token=abc#sig=xyz done')"
# Non-secret text is untouched (no false positives on ordinary log lines).
assert_eq "plain line untouched" \
"Resolved 42 packages in 1.2s" \
"$(redact_str 'Resolved 42 packages in 1.2s')"
assert_eq "plain url without creds untouched" \
"downloading https://download.pytorch.org/whl/cu128/torch-2.8.0.whl" \
"$(redact_str 'downloading https://download.pytorch.org/whl/cu128/torch-2.8.0.whl')"
assert_eq "bare hash comment untouched" \
"# retrying with --no-cache-dir" \
"$(redact_str '# retrying with --no-cache-dir')"
# Regression guard: no secret substring survives.
_leak=$(redact_str 'https://alice:s3cr3t@host/whl/cu128?token=SUPERSECRET#frag=ALSOSECRET')
case "$_leak" in
*s3cr3t*|*SUPERSECRET*|*ALSOSECRET*) assert_eq "no secret leak" "clean" "leaked:$_leak" ;;
*) assert_eq "no secret leak" "clean" "clean" ;;
esac
echo ""
echo "Results: $PASS passed, $FAIL failed"
[ "$FAIL" -eq 0 ]

74
tests/sh/test_torch_constraint.sh
View file

@ -108,6 +108,25 @@ assert_eq "\$TORCH_CONSTRAINT used in pip install" "yes" "$_has_var"
_hardcoded=$(grep -c '"torch>=2.4,<2.11.0"' "$INSTALL_SH" || true)
assert_eq "hardcoded torch>=2.4 appears exactly once" "1" "$_hardcoded"
# Companions must be bounded to torch's window everywhere: the <2.11 bound appears
# twice (default assignments + the pinned custom-leaf block), never bare. torchaudio
# 2.11 dropped its exact torch pin, so a bare companion next to a <2.11-capped torch
# resolves a mismatched 2.11 build.
_count=$(grep -c 'TORCHVISION_CONSTRAINT="torchvision>=0.19,<0.26.0"' "$INSTALL_SH" || true)
assert_eq "torchvision bounded (<0.26) at default + custom-leaf" "2" "$_count"
_count=$(grep -c 'TORCHAUDIO_CONSTRAINT="torchaudio>=2.4,<2.11.0"' "$INSTALL_SH" || true)
assert_eq "torchaudio bounded (<2.11) at default + custom-leaf" "2" "$_count"
_count=$(grep -c 'TORCHVISION_CONSTRAINT="torchvision"$' "$INSTALL_SH" || true)
assert_eq "no bare torchvision companion remains" "0" "$_count"
_count=$(grep -c 'TORCHAUDIO_CONSTRAINT="torchaudio"$' "$INSTALL_SH" || true)
assert_eq "no bare torchaudio companion remains" "0" "$_count"
# The cu* widen must carry the companions with it (torch <2.12 with torchaudio <2.11
# would cap a mismatched pair the other way).
assert_eq "cu widen pairs torchaudio (<2.12)" "1" "$(grep -c 'TORCHAUDIO_CONSTRAINT="torchaudio>=2.4,<2.12.0"' "$INSTALL_SH" || true)"
_gated=$(grep -c '_expected_torch_flavor_tag "$TORCH_INDEX_URL"' "$INSTALL_SH" || true)
_has_gate=$([ "$_gated" -ge 1 ] && echo "yes" || echo "no")
assert_eq "custom-companion bound gated on empty flavor tag" "yes" "$_has_gate"
# A fresh CUDA install widens the ceiling to <2.12.0 so cu12x/cu13x land torch
# 2.11.x (matches the base image and _CUDA_TORCH_PKG_SPEC).
_cuda_widen=$(grep -c 'TORCH_CONSTRAINT="torch>=2.4,<2.12.0"' "$INSTALL_SH" || true)
@ -285,6 +304,61 @@ bash -c "
_uv_got2=$(cat "$_UV_LOG2" 2>/dev/null || echo "")
assert_contains "mock uv arm64+py312 receives torch>=2.4" "$_uv_got2" "torch>=2.4,<2.11.0"
# ======================================================================
# ROCm 2.11 floor: leaf is lowercased before the gfx*/rocm* allowlist match
# ======================================================================
echo ""
echo "=== ROCm 2.11 floor case (leaf normalization) ==="
# Structural: install.sh lowercases _torch_index_leaf before the floor case, so the
# canonical gfx120X-all (capital X) matches gfx120x-all.
_has_lc=$(grep -c '_torch_index_leaf=$(printf .* | tr .\[:upper:\]. .\[:lower:\].)' "$INSTALL_SH" || true)
_has_lc_ok=$([ "$_has_lc" -ge 1 ] && echo "yes" || echo "no")
assert_eq "install.sh lowercases _torch_index_leaf" "yes" "$_has_lc_ok"
# Runtime: replicate install.sh's normalization + floor case and assert both gfx120X-all
# and gfx120x-all get the floor, while non-2.11 leaves keep the default.
run_floor_case() {
_url="$1"
bash -c '
TORCH_CONSTRAINT="torch>=2.4,<2.11.0"
TORCHVISION_CONSTRAINT="torchvision"
TORCHAUDIO_CONSTRAINT="torchaudio"
_torch_index_leaf="${1%/}"
_torch_index_leaf="${_torch_index_leaf##*/}"
_torch_index_leaf=$(printf "%s" "$_torch_index_leaf" | tr "[:upper:]" "[:lower:]")
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"
;;
esac
echo "$TORCH_CONSTRAINT"
' _ "$_url"
}
assert_eq "gfx120X-all (capital) -> 2.11 floor" "torch>=2.11.0,<2.12.0" \
"$(run_floor_case 'https://repo.amd.com/rocm/whl/gfx120X-all')"
assert_eq "gfx120X-all trailing slash -> 2.11 floor" "torch>=2.11.0,<2.12.0" \
"$(run_floor_case 'https://repo.amd.com/rocm/whl/gfx120X-all/')"
assert_eq "gfx120x-all (lowercase) -> 2.11 floor" "torch>=2.11.0,<2.12.0" \
"$(run_floor_case 'https://repo.amd.com/rocm/whl/gfx120x-all')"
assert_eq "gfx1151 -> 2.11 floor" "torch>=2.11.0,<2.12.0" \
"$(run_floor_case 'https://repo.amd.com/rocm/whl/gfx1151')"
assert_eq "gfx1150 -> 2.11 floor" "torch>=2.11.0,<2.12.0" \
"$(run_floor_case 'https://repo.amd.com/rocm/whl/gfx1150')"
assert_eq "rocm7.2 -> 2.11 floor" "torch>=2.11.0,<2.12.0" \
"$(run_floor_case 'https://download.pytorch.org/whl/rocm7.2')"
assert_eq "gfx110X-all -> default (no floor)" "torch>=2.4,<2.11.0" \
"$(run_floor_case 'https://repo.amd.com/rocm/whl/gfx110X-all')"
assert_eq "rocm6.4 -> default (no floor)" "torch>=2.4,<2.11.0" \
"$(run_floor_case 'https://download.pytorch.org/whl/rocm6.4')"
assert_eq "cu128 -> default (no floor)" "torch>=2.4,<2.11.0" \
"$(run_floor_case 'https://download.pytorch.org/whl/cu128')"
assert_eq "cpu -> default (no floor)" "torch>=2.4,<2.11.0" \
"$(run_floor_case 'https://download.pytorch.org/whl/cpu')"
# ======================================================================
# Summary
# ======================================================================

89
tests/sh/test_torch_flavor.sh
View file

@ -11,14 +11,21 @@ INSTALL_SH="$SCRIPT_DIR/../../install.sh"
PASS=0
FAIL=0
# Extract the three helper functions from install.sh and source them.
# Extract the helper functions from install.sh and source them
# (_torch_index_url_leaf is the shared leaf extractor the classifiers call).
_FUNC_FILE=$(mktemp)
{
sed -n '/^_torch_flavor_tag()/,/^}/p' "$INSTALL_SH"
echo ""
sed -n '/^_torch_index_url_leaf()/,/^}/p' "$INSTALL_SH"
echo ""
sed -n '/^_is_pip_rocm_family_leaf()/,/^}/p' "$INSTALL_SH"
echo ""
sed -n '/^_expected_torch_flavor_tag()/,/^}/p' "$INSTALL_SH"
echo ""
sed -n '/^_torch_index_repairable()/,/^}/p' "$INSTALL_SH"
echo ""
sed -n '/^_tauri_torch_index_family()/,/^}/p' "$INSTALL_SH"
} > "$_FUNC_FILE"
# shellcheck disable=SC1090
. "$_FUNC_FILE"
@ -56,6 +63,26 @@ assert_eq "amd gfx index" "rocm" "$(_expected_torch_flavor_tag 'https://re
assert_eq "mirror cu130 leaf" "cu130" "$(_expected_torch_flavor_tag 'https://my.mirror/pytorch/whl/cu130')"
assert_eq "unrecognized leaf" "" "$(_expected_torch_flavor_tag 'https://my.mirror/whl/simple')"
assert_eq "empty url" "" "$(_expected_torch_flavor_tag '')"
# Query/fragment dropped before classification: .../cu128?token=x classifies as cu128,
# not an opaque leaf that reinstalls every run.
assert_eq "query-bearing cu128" "cu128" "$(_expected_torch_flavor_tag 'https://m/whl/cu128?token=x')"
assert_eq "fragment-bearing cpu" "cpu" "$(_expected_torch_flavor_tag 'https://m/whl/cpu#frag')"
# A cu-suffixed CUSTOM leaf (cu128-private, cu128x) is NOT the cu128 family (exact
# cu+digits only). Mirrors Python re.fullmatch(cu[0-9]+) / PowerShell.
assert_eq "cu-suffix custom leaf" "" "$(_expected_torch_flavor_tag 'https://m/whl/cu128-private')"
assert_eq "cu-alnum custom leaf" "" "$(_expected_torch_flavor_tag 'https://m/whl/cu128x')"
assert_eq "bare cu digits stays" "cu126" "$(_expected_torch_flavor_tag 'https://m/whl/cu126')"
# A custom leaf merely STARTING with rocm (rocm-current, rocm-rel-7.2.1) is NOT a pip
# rocm family -> "" (custom); real families (rocm7.2) and gfx indexes stay "rocm".
assert_eq "custom rocm-current" "" "$(_expected_torch_flavor_tag 'https://mirror/whl/rocm-current')"
assert_eq "radeon rocm-rel leaf" "" "$(_expected_torch_flavor_tag 'https://repo.radeon.com/rocm/manylinux/rocm-rel-7.2.1')"
assert_eq "real rocm7.2 stays" "rocm" "$(_expected_torch_flavor_tag 'https://download.pytorch.org/whl/rocm7.2')"
# A rocm<digit>-SUFFIX private mirror (rocm7.2-private, rocm7-current) is a custom pin ->
# "" (custom); match the family exactly, not the prefix.
assert_eq "suffixed rocm7.2-private" "" "$(_expected_torch_flavor_tag 'https://co.internal/whl/rocm7.2-private')"
assert_eq "suffixed rocm7-current" "" "$(_expected_torch_flavor_tag 'https://co.internal/whl/rocm7-current')"
assert_eq "two-dot rocm7.2.1" "" "$(_expected_torch_flavor_tag 'https://co.internal/whl/rocm7.2.1')"
assert_eq "bare rocm7 stays" "rocm" "$(_expected_torch_flavor_tag 'https://download.pytorch.org/whl/rocm7')"
echo "=== _torch_index_repairable ==="
assert_eq "cu130 repairable" "yes" "$(_torch_index_repairable 'https://download.pytorch.org/whl/cu130')"
@ -64,6 +91,66 @@ assert_eq "gfx repairable" "yes" "$(_torch_index_repairable 'https://repo.
assert_eq "gfx1151 repairable" "yes" "$(_torch_index_repairable 'https://repo.amd.com/rocm/whl/gfx1151/')"
assert_eq "cpu NOT repairable" "no" "$(_torch_index_repairable 'https://download.pytorch.org/whl/cpu')"
assert_eq "unknown NOT repair" "no" "$(_torch_index_repairable 'https://my.mirror/whl/simple')"
# A suffixed rocm leaf is a verbatim pin, not a --default-index repairable family.
assert_eq "rocm-private NOT repair" "no" "$(_torch_index_repairable 'https://co.internal/whl/rocm7.2-private')"
echo "=== _is_pip_rocm_family_leaf ==="
assert_family() {
_label="$1"; _expected="$2"; _leaf="$3"
if _is_pip_rocm_family_leaf "$_leaf"; then _actual="yes"; else _actual="no"; fi
assert_eq "$_label" "$_expected" "$_actual"
}
assert_family "rocm7.2 family" "yes" "rocm7.2"
assert_family "rocm6.4 family" "yes" "rocm6.4"
assert_family "bare rocm7 family" "yes" "rocm7"
assert_family "gfx120x-all family" "yes" "gfx120x-all"
assert_family "gfx1151 family" "yes" "gfx1151"
assert_family "rocm7.2-private custom" "no" "rocm7.2-private"
assert_family "rocm7-current custom" "no" "rocm7-current"
assert_family "rocm-current custom" "no" "rocm-current"
assert_family "rocm-rel-7.2.1 custom" "no" "rocm-rel-7.2.1"
assert_family "rocm7.2.1 custom" "no" "rocm7.2.1"
# A trailing dot (rocm7.) or leading/double dot is NOT a family: both major and minor must
# be non-empty all-digits, matching Python re.fullmatch(rocm\d+(?:\.\d+)?). Bash previously
# accepted rocm7. via a bare %/-style trim while Python rejected it (validator asymmetry).
assert_family "rocm7. trailing-dot custom" "no" "rocm7."
assert_family "rocm.7 leading-dot custom" "no" "rocm.7"
assert_family "rocm7..2 double-dot custom" "no" "rocm7..2"
assert_family "cpu not rocm" "no" "cpu"
assert_family "cu128 not rocm" "no" "cu128"
assert_family "simple not rocm" "no" "simple"
echo "=== _torch_index_url_leaf (ALL trailing slashes stripped -> non-empty leaf) ==="
# A double (or triple) trailing slash must yield the real leaf, not an empty string that
# fails every classifier arm. Python .rstrip("/") drops them all; bash must match (a bare
# %/ left .../cu128// classifying as "").
assert_eq "double slash cu128 leaf" "cu128" "$(_torch_index_url_leaf 'https://m/whl/cu128//')"
assert_eq "triple slash rocm7.2 leaf" "rocm7.2" "$(_torch_index_url_leaf 'https://m/whl/rocm7.2///')"
assert_eq "double slash + token leaf" "cu128" "$(_torch_index_url_leaf 'https://m/whl/cu128//?token=x')"
assert_eq "single slash cu128 leaf" "cu128" "$(_torch_index_url_leaf 'https://m/whl/cu128/')"
# The classifier that consumes the leaf must therefore still tag a double-slash index.
assert_eq "double-slash cu128 tag" "cu128" "$(_expected_torch_flavor_tag 'https://m/whl/cu128//')"
assert_eq "double-slash rocm7.2 tag" "rocm" "$(_expected_torch_flavor_tag 'https://m/whl/rocm7.2//')"
echo "=== _tauri_torch_index_family (credential redaction) ==="
# A token/fragment must be stripped BEFORE classification so it never reaches the
# [TAURI:DIAG] line (the family is the last path segment, which else carries the query).
SKIP_TORCH=false
assert_eq "token stripped from rocm" "rocm7.2" "$(_tauri_torch_index_family 'https://mirror/whl/rocm7.2?token=SECRET')"
assert_eq "token-bearing cu classifies" "cu128" "$(_tauri_torch_index_family 'https://m/whl/cu128?token=x')"
assert_eq "fragment stripped cpu" "cpu" "$(_tauri_torch_index_family 'https://m/whl/cpu#frag')"
assert_eq "plain rocm7.2 unchanged" "rocm7.2" "$(_tauri_torch_index_family 'https://download.pytorch.org/whl/rocm7.2')"
# A trailing slash must be stripped too, or the */cu128 and */cpu arms miss .../cu128/
# and it falls through to "auto".
assert_eq "trailing slash cu128" "cu128" "$(_tauri_torch_index_family 'https://download.pytorch.org/whl/cu128/')"
assert_eq "slash + token cu128" "cu128" "$(_tauri_torch_index_family 'https://m/whl/cu128/?token=x')"
assert_eq "trailing slash cpu" "cpu" "$(_tauri_torch_index_family 'https://m/whl/cpu/')"
# Regression guard: no secret token substring may survive in any classification.
_leak=$(_tauri_torch_index_family 'https://mirror/whl/rocm7.2?token=SECRET')
case "$_leak" in
*SECRET*|*token*) assert_eq "no token leak in family" "clean" "leaked:$_leak" ;;
*) assert_eq "no token leak in family" "clean" "clean" ;;
esac
echo ""
echo "Results: $PASS passed, $FAIL failed"

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