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

Merge remote-tracking branch 'origin/main' into merge_to_fp8

Browse source 
# Conflicts:
#	unsloth/save.py
This commit is contained in:
Daniel Han 2026-06-30 10:35:37 +00:00
commit f637abc3d0
102 changed files with 5209 additions and 900 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -209,7 +209,7 @@ jobs:
'peft>=0.18,<0.20' 'accelerate>=0.34,<2' \
ipython
# torchvision: unsloth_zoo.vision_utils imports it at module scope.
pip install --index-url https://download.pytorch.org/whl/cpu \
pip install --index-url https://download.pytorch.org/whl/cpu --extra-index-url https://pypi.org/simple \
'torch>=2.4,<2.11' 'torchvision<0.26'
# transformers + trl from the matrix combo.
pip install "$RESOLVED_TRANSFORMERS_SPEC"
@ -2174,7 +2174,7 @@ jobs:
python -m pip install --upgrade pip
# Match the matrix job's torch path so unsloth_zoo's
# `import torch` resolves to the same CPU build.
pip install --index-url https://download.pytorch.org/whl/cpu \
pip install --index-url https://download.pytorch.org/whl/cpu --extra-index-url https://pypi.org/simple \
'torch>=2.4,<2.11' 'torchvision<0.26'
pip install \
'numpy<3' protobuf sentencepiece \

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

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

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

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

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

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

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

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

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

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

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

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

5
README.md
View file

@ -246,6 +246,11 @@ curl -fsSL https://unsloth.ai/install.sh | UNSLOTH_STUDIO_HOME=/abs/path sh
$env:UNSLOTH_STUDIO_HOME='C:\path'; irm https://unsloth.ai/install.ps1 | iex
```
On macOS, the installer defaults to the system certificate store (`UV_SYSTEM_CERTS=1`) so uv trusts the CAs in your Keychain, needed behind TLS-inspecting proxies (Cisco Umbrella, Zscaler, etc.). Opt out with:
```bash
curl -fsSL https://unsloth.ai/install.sh | UV_SYSTEM_CERTS=0 sh
```
Point the frontend build at a corporate npm mirror/proxy with `UNSLOTH_NPM_REGISTRY` (for the developer install behind a firewall that blocks `registry.npmjs.org`):
```bash
UNSLOTH_NPM_REGISTRY=https://artifactory.example.com/api/npm/npm/ ./install.sh --local

15
install.sh
View file

@ -1636,6 +1636,21 @@ export UV_HTTP_RETRIES
: "${UV_HTTP_TIMEOUT:=180}"
export UV_HTTP_TIMEOUT
# macOS: trust the system Keychain so uv uses SecureTransport instead of rustls.
# Required behind TLS-inspecting proxies (Cisco Umbrella, Zscaler, etc.) which
# present their own CA certificate. rustls (uv's default) ignores the Keychain
# and rejects intercepted connections with "invalid peer certificate: UnknownIssuer".
# Set both vars: UV_SYSTEM_CERTS is the modern one (uv >= 0.11), UV_NATIVE_TLS the
# legacy one understood by uv 0.8.16-0.10.x, which the installer keeps if already
# present (UV_MIN_VERSION) and which ignores UV_SYSTEM_CERTS. Mirror the choice onto
# both so it works on either uv. Opt out with UV_SYSTEM_CERTS=0.
if [ "$OS" = "macos" ]; then
: "${UV_SYSTEM_CERTS:=1}"
: "${UV_NATIVE_TLS:=$UV_SYSTEM_CERTS}"
fi
[ -n "${UV_SYSTEM_CERTS:-}" ] && export UV_SYSTEM_CERTS
[ -n "${UV_NATIVE_TLS:-}" ] && export UV_NATIVE_TLS
version_ge() {
# returns 0 if $1 >= $2
_a=$1

23
pyproject.toml
View file

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

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

@ -1271,6 +1271,9 @@ class LlamaCppBackend:
self._cache_type_kv: Optional[str] = None
# Whether --split-mode tensor was applied on the active load.
self._tensor_parallel: bool = False
# Layer load kept multi-GPU only to honor a downgraded tensor request, so a
# later explicit tensor-off reloads instead of deduping to it (#6659).
self._layer_preserves_tensor_intent: bool = False
self._reasoning_default: bool = True
self._speculative_type: Optional[str] = None
# Canonical UI-facing mode the user requested
@ -1643,6 +1646,11 @@ class LlamaCppBackend:
"""Whether --split-mode tensor is active on the loaded server."""
return self._tensor_parallel
@property
def layer_preserves_tensor_intent(self) -> bool:
"""True when a downgraded tensor request kept this layer load multi-GPU."""
return self._layer_preserves_tensor_intent
@property
def speculative_type(self) -> Optional[str]:
return self._speculative_type
@ -2430,6 +2438,37 @@ class LlamaCppBackend:
# aborts a --split-mode tensor load, so it's dropped for the tensor attempt.
_TENSOR_PARALLEL_KV_TYPES = frozenset({"f16", "bf16", "f32"})
# (binary, mtime, model) that aborted on --split-mode tensor this process (#6415
# geometry limit, e.g. MQA n_head_kv=1). Model-keyed so one model's abort doesn't
# skip tensor for others; tensor is tried by default, recorded only on a real abort.
_tensor_split_abort_keys: set[tuple[str, int, str]] = set()
@classmethod
def _tensor_split_cache_key(
cls, binary: Optional[str], model: Optional[str]
) -> Optional[tuple[str, int, str]]:
"""(path, mtime_ns, model) key; ns mtime re-probes a same-second binary swap."""
if not binary or not model:
return None
try:
mtime = Path(binary).stat().st_mtime_ns
except OSError:
mtime = 0
return (binary, mtime, model)
@classmethod
def _tensor_split_aborts(cls, binary: Optional[str], model: Optional[str]) -> bool:
"""True if (binary, model) aborted on --split-mode tensor this session."""
key = cls._tensor_split_cache_key(binary, model)
return key is not None and key in cls._tensor_split_abort_keys
@classmethod
def _record_tensor_split_abort(cls, binary: Optional[str], model: Optional[str]) -> None:
"""Remember a (binary, model) that aborts on --split-mode tensor."""
key = cls._tensor_split_cache_key(binary, model)
if key is not None:
cls._tensor_split_abort_keys.add(key)
@staticmethod
def _windows_pip_nvidia_dll_dirs(prefix: str) -> list[str]:
"""Return DLL dirs from pip-installed CUDA wheels under
@ -2569,9 +2608,13 @@ class LlamaCppBackend:
usable_fraction: Optional[float] = None,
total_by_idx: Optional[dict[int, int]] = None,
per_device_overhead_bytes: int = 0,
min_gpus: int = 1,
) -> tuple[Optional[list[int]], bool]:
"""Pick GPU(s) for a model from estimated VRAM and free memory.
``min_gpus`` (default 1, capped at ``len(gpus)``) keeps a downgraded
tensor/multi-GPU request spread instead of collapsing to one card.
``model_size_bytes`` should include weights and estimated KV cache.
``usable_fraction`` (default ``_GPU_PIN_VRAM_FRACTION``) provides
headroom for compute buffers, CUDA context, and other runtime
@ -2590,9 +2633,11 @@ class LlamaCppBackend:
if not gpus:
return None, True
min_gpus = max(1, min(min_gpus, len(gpus)))
model_size_mib = model_size_bytes / (1024 * 1024)
if usable_fraction is None:
usable_fraction = LlamaCppBackend._GPU_PIN_VRAM_FRACTION
overhead_mib = per_device_overhead_bytes / (1024 * 1024)
# Per-GPU usable budget: free - (1-frac)*total when total is known, else
# the legacy free*frac (also covers a total-0 two-column probe).
@ -2606,19 +2651,26 @@ class LlamaCppBackend:
# card can have less usable room than a less-used small one.
ranked = sorted(gpus, key = lambda g: _usable(g[0], g[1]), reverse = True)
# Try 1 GPU at the usable-VRAM threshold.
if _usable(ranked[0][0], ranked[0][1]) >= model_size_mib:
# Cap a downgraded multi-GPU request to the usable count so it doesn't pull
# in a near-full card to hit min_gpus. No-op for the default min_gpus == 1.
usable_count = sum(1 for idx, free_mib in ranked if _usable(idx, free_mib) > overhead_mib)
min_gpus = max(1, min(min_gpus, usable_count or 1))
# Try 1 GPU at the usable-VRAM threshold (only when one device is allowed).
if min_gpus <= 1 and _usable(ranked[0][0], ranked[0][1]) >= model_size_mib:
return [ranked[0][0]], False
# Try N GPUs (accumulate usable memory from most-free). Each GPU past the
# first adds a fixed per-device overhead the pool must hold.
overhead_mib = per_device_overhead_bytes / (1024 * 1024)
# Try N GPUs (most-free first); each past the first adds per-device overhead.
# Require at least min_gpus devices before accepting a fit.
cumulative = 0.0
selected = []
for idx, free_mib in ranked:
selected.append(idx)
cumulative += _usable(idx, free_mib)
if cumulative >= model_size_mib + (len(selected) - 1) * overhead_mib:
if (
len(selected) >= min_gpus
and cumulative >= model_size_mib + (len(selected) - 1) * overhead_mib
):
return sorted(selected), False
# Too large even for all GPUs; let --fit handle it
@ -3868,7 +3920,7 @@ class LlamaCppBackend:
logger.debug(f"Could not list repo files for {label}: {e}")
break
logger.debug(
f"Could not list repo files for {label} " f"(attempt {attempt + 1}/3): {e}"
f"Could not list repo files for {label} (attempt {attempt + 1}/3): {e}"
)
if attempt < 2:
self._cancel_event.wait(2**attempt)
@ -4332,6 +4384,17 @@ class LlamaCppBackend:
)
)
@staticmethod
def _is_tensor_split_assert(output: str) -> bool:
"""True only for the #6415 split-axis warmup assert (GGML_BACKEND_SPLIT_AXIS_*),
not any ggml assert/abort, so an unrelated invariant isn't cached. stderr is
merged into output."""
text = (output or "").lower()
if "ggml_assert" not in text and "ggml_abort" not in text:
return False
# the split-axis enum token, unique to this assert (not the source file).
return "split_axis" in text
@staticmethod
def _is_signal_crash(returncode: Optional[int]) -> bool:
"""True only on a hard fault (SIGSEGV/SIGABRT/SIGILL/SIGFPE/SIGBUS or a
@ -4344,6 +4407,20 @@ class LlamaCppBackend:
return True
return -returncode in (4, 6, 7, 8, 11) # SIGILL SIGABRT SIGBUS SIGFPE SIGSEGV
@staticmethod
def _is_abort_exit(returncode: Optional[int]) -> bool:
"""Windows CRT abort() exit code (3) from GGML_ASSERT on MSVC -- not a POSIX
signal or 0xC0000000+ NTSTATUS."""
return returncode == 3
@classmethod
def _should_record_tensor_split_abort(cls, returncode: Optional[int], output: str) -> bool:
"""The #6415 split-axis abort: the marker plus a hard crash (POSIX signal or
Windows abort exit). Marker required so a generic crash isn't cached."""
return cls._is_tensor_split_assert(output) and (
cls._is_signal_crash(returncode) or cls._is_abort_exit(returncode)
)
@staticmethod
def _with_flash_attn_off(cmd: list[str]) -> Optional[list[str]]:
"""Return cmd with flash attention forced off, or None when its effective
@ -4488,6 +4565,8 @@ class LlamaCppBackend:
n_gpu_layers: Optional[int] = None, # caller compat, unused
n_parallel: int = 1,
extra_args: Optional[List[str]] = None,
# Route-level tensor->layer fallback retry: keep the layer split multi-GPU.
preserve_multi_gpu_on_layer: bool = False,
) -> bool:
"""Start llama-server with a GGUF model.
@ -4518,6 +4597,8 @@ class LlamaCppBackend:
"n_gpu_layers": n_gpu_layers,
"n_parallel": n_parallel,
"extra_args": list(extra_args) if extra_args is not None else None,
# Replayed by _respawn_if_dead so a downgraded model stays multi-GPU.
"preserve_multi_gpu_on_layer": preserve_multi_gpu_on_layer,
}
# Serialise the whole load so concurrent /load calls never leave two
# llama-server processes alive (#5401 / #5161). Doesn't block /unload.
@ -4541,6 +4622,7 @@ class LlamaCppBackend:
chat_template_override = chat_template_override,
extra_args = extra_args,
is_vision = is_vision,
preserve_multi_gpu_on_layer = preserve_multi_gpu_on_layer,
):
logger.info(
f"load_model: backend already in target state for "
@ -4626,6 +4708,9 @@ class LlamaCppBackend:
# Block-diffusion GGUFs (DiffusionGemma) cannot run on llama-server;
# serve them with the diffusion runner (same OpenAI-compat interface).
if self._is_diffusion:
# Not a tensor/layer GGUF: clear any preserved-fallback flag from a
# prior load (this path skips the command builder that clears it).
self._layer_preserves_tensor_intent = False
with self._lock:
if self._cancel_event.is_set():
logger.info("Load cancelled before diffusion server start")
@ -4780,6 +4865,9 @@ class LlamaCppBackend:
"image input will be disabled for this session"
)
model_size = None # set in the fit try; used by the APU RAM guard
# Layer-fallback min GPUs; raised below on a tensor downgrade. Bound
# before the try so the --fit-on except path still has it (no UnboundLocal).
_layer_min_gpus = 1
try:
gguf_size = self._get_gguf_size_bytes(model_path)
# Include GPU-loaded mmproj in the fit budget (#5825).
@ -5064,10 +5152,8 @@ class LlamaCppBackend:
_apple_budget_mib = self._apple_metal_memory_budget_bytes() // (1024 * 1024)
def _restore_after_tensor_downgrade():
# Tensor mode dropped a quantized KV and stripped the cache
# extras (it rejects quantized); layer split supports them, so
# restore the original type + extras (minus --split-mode) and
# clear the env flag so the layer launch re-emits them.
# Restore the quantized KV + extras tensor dropped (layer
# split supports them), minus --split-mode.
nonlocal cache_type_kv, _cache_type_from_env, extra_args
if _tensor_dropped_cache_type_kv is not None:
cache_type_kv = _tensor_dropped_cache_type_kv
@ -5078,13 +5164,22 @@ class LlamaCppBackend:
else extra_args
)
if tensor_parallel and effective_is_vision:
# The route fallback retry is tensor-off; keep it multi-GPU.
if preserve_multi_gpu_on_layer:
_layer_min_gpus = max(_layer_min_gpus, len(gpus))
if tensor_parallel and self._tensor_split_aborts(binary, model_identifier):
# Aborted on tensor for this model this session (#6415); skip
# tensor upfront, layer split serves it.
logger.info(
"Tensor parallelism skipped for vision model: "
"--split-mode tensor is incompatible with --mmproj "
"in the current llama.cpp build; using layer split."
"Tensor parallelism skipped: this llama.cpp build aborted "
"on --split-mode tensor for this model earlier this "
"session; using layer split across %d GPU(s).",
len(gpus),
)
tensor_parallel = False
# Keep the multi-GPU request (gated on it, not the cache).
_layer_min_gpus = max(_layer_min_gpus, len(gpus))
_restore_after_tensor_downgrade()
# Tensor mode replicates a compute buffer on every GPU, so drop
@ -5124,6 +5219,11 @@ class LlamaCppBackend:
len(gpus),
)
tensor_parallel = False
# GPUs below tensor's compute-buffer reserve can still do layer
# split, so keep multi-GPU (mirrors the budget/geometry drops);
# _select_gpus caps unusable cards.
if len(gpus) >= 2:
_layer_min_gpus = max(_layer_min_gpus, len(gpus))
# Layer split supports a quantized KV the tensor attempt
# dropped; restore the original cache type + extras (minus
# --split-mode) so the layer launch re-emits them.
@ -5160,8 +5260,12 @@ class LlamaCppBackend:
"per-device compute buffers; falling back to layer split."
)
tensor_parallel = False
# Restore the dropped quantized KV + original cache extras
# (minus --split-mode); layer split supports them.
# Weights needed >1 card, so keep multi-GPU across the
# usable tensor GPUs.
if len(tp_gpus) >= 2:
_layer_min_gpus = max(_layer_min_gpus, len(tp_gpus))
# Restore the dropped quantized KV + cache extras (minus
# --split-mode); layer split supports them.
_restore_after_tensor_downgrade()
if tensor_parallel and tp_gpus:
@ -5263,6 +5367,7 @@ class LlamaCppBackend:
usable_fraction = _pin_fraction,
total_by_idx = total_by_idx,
per_device_overhead_bytes = _pipeline_overhead_bytes,
min_gpus = _layer_min_gpus,
)
# No silent shrink: effective_ctx stays == requested_ctx.
else:
@ -5273,7 +5378,22 @@ class LlamaCppBackend:
ranked = sorted(
gpus, key = lambda g: _gpu_usable(g, pin_fraction), reverse = True
)
for n_gpus in range(1, len(ranked) + 1):
# Skips _select_gpus, so apply its cap: count only cards
# whose usable VRAM clears the per-device layer overhead.
_pipeline_overhead_mib = _pipeline_overhead_bytes / (1024 * 1024)
_auto_min_gpus = max(
1,
min(
_layer_min_gpus,
sum(
1
for g in ranked
if _gpu_usable(g, pin_fraction) > _pipeline_overhead_mib
)
or 1,
),
)
for n_gpus in range(_auto_min_gpus, len(ranked) + 1):
subset = ranked[:n_gpus]
pool_budget = _pool_budget_mib(subset, pin_fraction)
_ms = _subset_model_size(n_gpus)
@ -5303,7 +5423,7 @@ class LlamaCppBackend:
# at 131k may pin fine with a 4096 KV (#5106).
effective_ctx = min(4096, effective_ctx)
if effective_ctx > 0:
for n_gpus in range(1, len(ranked) + 1):
for n_gpus in range(_auto_min_gpus, len(ranked) + 1):
subset = ranked[:n_gpus]
kv = self._estimate_kv_cache_bytes(
effective_ctx,
@ -5339,6 +5459,7 @@ class LlamaCppBackend:
usable_fraction = _pin_fraction,
total_by_idx = total_by_idx,
per_device_overhead_bytes = _pipeline_overhead_bytes,
min_gpus = _layer_min_gpus,
)
if use_fit and not explicit_ctx:
# Weights don't fit on any subset; default UI to 4096
@ -5476,6 +5597,15 @@ class LlamaCppBackend:
"--no-context-shift",
]
# Report a clean public model id (matching GET /v1/models) rather
# than the raw -m path in llama-server's own /v1/models and the
# "model" field of its chat/completions responses.
from core.inference.model_ids import public_model_id
_alias = public_model_id(self._model_identifier or model_path)
if _alias:
cmd.extend(["--alias", _alias])
fully_gpu_offloaded = False
if use_fit:
cmd.extend(["--fit", "on"])
@ -5569,12 +5699,15 @@ class LlamaCppBackend:
]
)
self._tensor_parallel = True
self._layer_preserves_tensor_intent = False
logger.info(
"Tensor parallelism: --split-mode tensor, --tensor-split %s",
tp_tensor_split,
)
else:
self._tensor_parallel = False
# > 1 only when a tensor request was downgraded but kept multi-GPU.
self._layer_preserves_tensor_intent = _layer_min_gpus > 1
# Speculative decoding. See _build_speculative_flags for the
# mode resolution, benchmarks, and llama.cpp references.
@ -5858,7 +5991,17 @@ class LlamaCppBackend:
_startup_crashed = (
self._process.poll() is not None and self._process.returncode != 0
)
if _spawn_attempt == 0 and _fit_retry_allowed and _startup_crashed:
# A split-axis abort (#6415) is fit-independent: skip the
# --fit off retry and let the caller latch it.
_split_axis_crash = self._is_tensor_split_assert(
"\n".join(self._stdout_lines[-50:])
)
if (
_spawn_attempt == 0
and _fit_retry_allowed
and _startup_crashed
and not _split_axis_crash
):
logger.warning(
"llama-server crashed during startup (exit code %s) "
"with the default memory-fit step enabled; Studio "
@ -5904,6 +6047,21 @@ class LlamaCppBackend:
)
healthy = _spawn_and_wait(cmd)
# #6415 split-mode tensor warmup abort. Latch it on THIS first spawn:
# the flash-attn-off retry below can't run tensor (needs flash_attn),
# so its output drops the marker and recording later would miss it,
# looping every load. Record and raise to the route's layer fallback,
# skipping the futile flash-attn/MTP retries.
if not healthy and self._tensor_parallel and not self._cancel_event.is_set():
_ts_out = "\n".join(self._stdout_lines[-50:])
_ts_rc = self._process.poll() if self._process is not None else None
if self._should_record_tensor_split_abort(_ts_rc, _ts_out):
LlamaCppBackend._record_tensor_split_abort(binary, model_identifier)
self._kill_process()
raise RuntimeError(
"llama-server aborted on --split-mode tensor "
"(split-axis geometry); retrying with layer split."
)
# Flash-attention kernels hard-crash at startup on some ROCm/GPU
# builds (frequently inside the vision tower). Disabling FA keeps
# both vision and MTP, so retry that way before dropping either.
@ -6048,6 +6206,7 @@ class LlamaCppBackend:
# Read the crash code before _kill_process() clears _process.
_crash_rc = self._process.poll() if self._process is not None else None
self._kill_process()
# The #6415 split-axis abort is latched earlier (first spawn).
# Skip if a cancel/unload is pending (mirrors the MTP guard).
if (
launched_with_mmproj
@ -6479,6 +6638,7 @@ class LlamaCppBackend:
spec_draft_n_max: Optional[int] = None,
tensor_parallel: bool = False,
mtp_draft_path: Optional[str] = None,
preserve_multi_gpu_on_layer: bool = False,
) -> bool:
"""True iff the live server already satisfies these load kwargs.
@ -6521,6 +6681,17 @@ class LlamaCppBackend:
# server. An identical request would downgrade the same way.
if not _tensor_parallel_matches_loaded(extra_args, tensor_parallel, self._tensor_parallel):
return False
# Preserved tensor->layer fallback + an EXPLICIT tensor drop: reload so
# placement re-selects instead of keeping the all-GPU mask (mirrors the route,
# #6659). preserve_multi_gpu_on_layer carries the route's carry-forward decision
# (True for an implicit same-settings reload), so those still dedupe -- the HF
# auto-pick / local-dir flows skip the route guard and only reach here.
if (
self._layer_preserves_tensor_intent
and not _effective_tensor_parallel(extra_args, tensor_parallel)
and not preserve_multi_gpu_on_layer
):
return False
# Compare on the canonical requested mode. With --spec-type in
# extra_args the backend stores None; mirror that here.
@ -6632,6 +6803,7 @@ class LlamaCppBackend:
self._supports_tools = False
self._cache_type_kv = None
self._tensor_parallel = False
self._layer_preserves_tensor_intent = False
self._speculative_type = None
self._requested_spec_mode = None
self._spec_draft_n_max = None

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

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

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

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

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

@ -299,6 +299,7 @@ class TrainingBackend:
# Build config dict for the subprocess
config = {
"model_name": kwargs["model_name"],
"project_name": kwargs.get("project_name"),
"training_type": kwargs.get("training_type", "LoRA/QLoRA"),
"hf_token": kwargs.get("hf_token", ""),
"load_in_4bit": kwargs.get("load_in_4bit", True),

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

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

90
studio/backend/main.py
View file

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

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

@ -106,8 +106,7 @@ class LoadRequest(BaseModel):
"Extra arguments forwarded verbatim to llama-server for GGUF models. "
"One token per list entry, e.g. ['--top-k', '20', '--seed', '42']. "
"Studio-managed flags (model identity, port, context length, GPU placement, "
"auth, --flash-attn, --no-context-shift, --jinja) are rejected. Ignored for "
"non-GGUF models."
"auth, UI/server mode) are rejected. Ignored for non-GGUF models."
),
)

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

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

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

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

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

@ -683,7 +683,9 @@ try:
detect_reasoning_flags,
)
from core.inference.llama_server_args import (
_effective_tensor_parallel,
_tensor_parallel_matches_loaded,
parse_split_mode_override,
resolve_tensor_parallel,
strip_shadowing_flags,
validate_extra_args,
@ -718,7 +720,9 @@ except ImportError:
detect_reasoning_flags,
)
from core.inference.llama_server_args import (
_effective_tensor_parallel,
_tensor_parallel_matches_loaded,
parse_split_mode_override,
resolve_tensor_parallel,
strip_shadowing_flags,
validate_extra_args,
@ -1107,6 +1111,7 @@ from auth.authentication import get_current_subject
from state.tool_approvals import resolve_tool_decision
from core.inference.key_exchange import decrypt_api_key
from core.inference.model_ids import public_model_id
from core.inference.api_monitor import api_monitor
from core.inference.llama_http import nonstreaming_client
from core.inference.providers import get_base_url
@ -2077,6 +2082,32 @@ def _should_strip_split_mode(request: LoadRequest, backend_extra: Optional[list[
)
def _carry_preserved_tensor_intent(
*, preserved: bool, same_model: bool, explicit_drop: bool
) -> bool:
"""Carry a preserved multi-GPU layer fallback forward only for a reload of the
SAME loaded model that doesn't explicitly drop tensor intent, so a fitting model
isn't collapsed to one GPU on a ctx-only change -- but an unrelated model switch
(without /unload) or an explicit tensor-off doesn't inherit it (#6659)."""
return preserved and same_model and not explicit_drop
def _is_explicit_tensor_drop(request: LoadRequest) -> bool:
"""True only when the request explicitly selects a non-tensor --split-mode (e.g.
layer/row/none), a deliberate departure from a preserved tensor->layer fallback.
A bare tensor_parallel field is NOT a drop: the Studio UI always sends it and echoes
the /load response's resolved value back, so after a fallback every reload carries
tensor_parallel=false even though the user never changed it -- treating that as a drop
would collapse the preserved multi-GPU placement on the next ctx/settings reload. An
empty clear is not a drop either (a fallback always stores --split-mode layer, never a
tensor split mode, so a clear never wipes tensor intent), nor is an unrelated extra
(--top-k) or inherit (None). tensor_parallel=true / --split-mode tensor re-engage
tensor. Shared by the already-loaded dedup and the load carry-forward (#6659)."""
override = parse_split_mode_override(request.llama_extra_args)
return override is not None and override.strip().lower() != "tensor"
def _request_matches_loaded_settings(
request: LoadRequest,
llama_backend: LlamaCppBackend,
@ -2115,6 +2146,13 @@ def _request_matches_loaded_settings(
effective_extra, request.tensor_parallel, llama_backend.tensor_parallel
):
return False
# Preserved tensor->layer fallback (both report tensor=off, so the check above
# matches): if the user now explicitly drops tensor intent, reload so placement
# re-selects instead of keeping the all-GPU mask (#6659). The effective check
# includes the env, so an env-only tensor (LLAMA_ARG_SPLIT_MODE=tensor) that
# can't actually be dropped falls through to the env-downgrade match, not a loop.
if llama_backend.layer_preserves_tensor_intent and _is_explicit_tensor_drop(request):
return False
# Spec decoding works on vision models too (MTP is mmproj-compatible,
# llama.cpp #22673; the old ``not is_vision`` gate is gone), so compare
# the real requested mode -- coercing vision to ``off`` here used to
@ -2809,6 +2847,48 @@ async def load_model(
hf_variant = config.gguf_variant,
)
# Tensor intent for this load: the request itself, or a preserved
# multi-GPU layer fallback carried across a reload of the SAME model that
# doesn't drop it (e.g. a ctx-only change), so a fitting model doesn't
# silently collapse to one GPU. Only an explicit non-tensor --split-mode
# override counts as the drop -- the tensor field echo / unrelated extras keep
# the preserved placement; the same-model guard stops a switch-without-unload
# inheriting the prior model's intent.
_explicit_tensor_drop = _is_explicit_tensor_drop(request)
# Compare the resolved config.identifier (what load_model stores), not the
# raw request id: from_identifier normalizes shorthands (adds unsloth/, fixes
# case), so a reload with the shorthand would otherwise miss the match and
# drop the carry-forward. #6659
_same_model_loaded = (
llama_backend.is_loaded
and (llama_backend.model_identifier or "").lower()
== (config.identifier or "").lower()
)
# model_identifier is variant-agnostic for HF repos and dir-level for a
# local multi-variant directory, so also require the loaded quant to match
# (path else variant, mirroring _already_in_target_state) -- otherwise a
# different variant inherits the prior one's preserved intent. #6659
if _same_model_loaded:
if config.gguf_file and llama_backend.gguf_path:
try:
_same_model_loaded = (
Path(llama_backend.gguf_path).resolve()
== Path(config.gguf_file).resolve()
)
except OSError:
_same_model_loaded = False
else:
_same_model_loaded = (llama_backend.hf_variant or "").lower() == (
config.gguf_variant or ""
).lower()
_tensor_intent_overall = _effective_tensor_parallel(
extra_llama_args, request.tensor_parallel
) or _carry_preserved_tensor_intent(
preserved = llama_backend.layer_preserves_tensor_intent,
same_model = _same_model_loaded,
explicit_drop = _explicit_tensor_drop,
)
# Run a single load attempt with the given tensor flag + extras.
async def _attempt_gguf_load(
tensor_parallel: bool, attempt_extra_args: Optional[list[str]]
@ -2822,6 +2902,12 @@ async def load_model(
**_source_load_kwargs,
**attempt_kwargs,
tensor_parallel = tensor_parallel,
# True on the layer fallback retry (tensor wanted overall but not on
# this attempt): keep multi-GPU. Mirrors the fallback's key.
preserve_multi_gpu_on_layer = bool(
_tensor_intent_overall
and not _effective_tensor_parallel(attempt_extra_args, tensor_parallel)
),
)
# Tensor parallelism is arch-gated in llama.cpp and crashes some loads
@ -3703,7 +3789,7 @@ async def generate_audio(
# Pick backend — both return (wav_bytes, sample_rate)
llama_backend = get_llama_cpp_backend()
if llama_backend.is_loaded and getattr(llama_backend, "_is_audio", False):
model_name = llama_backend.model_identifier
model_name = public_model_id(llama_backend.model_identifier)
gen = lambda: llama_backend.generate_audio_response(
text = text,
audio_type = llama_backend._audio_type,
@ -3721,7 +3807,7 @@ async def generate_audio(
model_info = backend.models.get(backend.active_model_name, {})
if not model_info.get("is_audio"):
raise HTTPException(status_code = 400, detail = "Active model is not an audio model.")
model_name = backend.active_model_name
model_name = public_model_id(backend.active_model_name)
gen = lambda: backend.generate_audio_response(
text = text,
temperature = payload.temperature,
@ -4837,7 +4923,8 @@ async def openai_chat_completions(
return response
if using_gguf:
model_name = llama_backend.model_identifier or payload.model
# Echo a clean public id in the response, never the absolute .gguf path.
model_name = public_model_id(llama_backend.model_identifier) or payload.model
if getattr(llama_backend, "_is_audio", False):
if _wants_multiple_choices(payload):
_raise_unsupported_n("GGUF audio chat completions")
@ -4852,7 +4939,9 @@ async def openai_chat_completions(
status_code = 400,
detail = "No model loaded. Call POST /inference/load first.",
)
model_name = backend.active_model_name or payload.model
# Clean public id so the response never echoes a local path; the audio
# branch below receives this sanitized label too.
model_name = public_model_id(backend.active_model_name) or payload.model
if _wants_multiple_choices(payload):
_raise_unsupported_n("non-GGUF chat completions")
@ -6387,6 +6476,9 @@ async def serve_sandbox_file(
# OpenAI-Compatible Models Listing (/models → /v1/models)
# =====================================================================
# `owned_by` marker on every /v1/models entry (loaded and available alike).
_OWNED_BY = "unsloth-studio"
def _openai_model_objects() -> list[dict]:
"""The model objects GET /v1/models exposes (one per loaded local backend).
@ -6401,10 +6493,12 @@ def _openai_model_objects() -> list[dict]:
llama_backend = get_llama_cpp_backend()
if llama_backend.is_loaded:
entry = {
"id": llama_backend.model_identifier,
# Public id, never the absolute .gguf path (which leaks the host
# filesystem layout); see core.inference.model_ids.public_model_id.
"id": public_model_id(llama_backend.model_identifier),
"object": "model",
"created": _created,
"owned_by": "local",
"owned_by": _OWNED_BY,
}
_ctx = _positive_int_or_none(getattr(llama_backend, "context_length", None))
if _ctx is not None:
@ -6422,10 +6516,10 @@ def _openai_model_objects() -> list[dict]:
if backend.active_model_name:
model_info = backend.models.get(backend.active_model_name, {})
entry = {
"id": backend.active_model_name,
"id": public_model_id(backend.active_model_name),
"object": "model",
"created": _created,
"owned_by": "local",
"owned_by": _OWNED_BY,
}
_ctx = _positive_int_or_none(model_info.get("context_length"))
if _ctx is None:
@ -6443,15 +6537,86 @@ def _openai_model_objects() -> list[dict]:
return models
# Brief cache for the local-model filesystem scan so repeated /v1/models calls
# don't rescan the HF cache and models dirs on every request.
_CATALOG_CACHE: dict = {"at": 0.0, "models": []}
_CATALOG_TTL_S = 30.0
_CATALOG_LOCK = asyncio.Lock()
async def _cached_local_catalog() -> list:
"""Locally available models (models dir + HF caches + LM Studio + scan
folders), cached for a few seconds. Returns a list of LocalModelInfo.
The scan walks several directories and stats many files, so it runs in a
worker thread (asyncio.to_thread) -- calling it inline would block the event
loop and stall every concurrent request and in-flight inference stream. A
lock with a double-check collapses a burst of simultaneous /v1/models calls
into a single scan instead of one per request."""
# Validity is keyed on "at" (set only after a scan), not on list contents, so
# an empty/errored scan is still cached instead of rescanning on every poll.
now = time.monotonic()
if _CATALOG_CACHE["at"] and (now - _CATALOG_CACHE["at"]) <= _CATALOG_TTL_S:
return _CATALOG_CACHE["models"]
async with _CATALOG_LOCK:
now = time.monotonic()
if _CATALOG_CACHE["at"] and (now - _CATALOG_CACHE["at"]) <= _CATALOG_TTL_S:
return _CATALOG_CACHE["models"]
try:
from routes.models import collect_local_models
_CATALOG_CACHE["models"] = await asyncio.to_thread(
collect_local_models, Path("./models").resolve()
)
except Exception as exc:
logger.debug("model catalog scan failed: %s", exc)
_CATALOG_CACHE["models"] = []
# Stamp after the scan, not the pre-scan "now": a scan slower than the TTL
# would otherwise leave the cache already expired, so every waiter rescans.
_CATALOG_CACHE["at"] = time.monotonic()
return _CATALOG_CACHE["models"]
async def _openai_catalog_objects() -> list[dict]:
"""Every model the server knows about for ``GET /v1/models``: the loaded
model(s) plus locally available (downloaded/cached) models discovered by
scanning. Loaded entries keep their context fields and are marked
``loaded: true``. All ids are clean public ids (never absolute paths)."""
_created = int(time.time())
# Loaded models first (clean ids + context fields), marked loaded.
by_id: dict[str, dict] = {}
for entry in _openai_model_objects():
by_id[entry["id"]] = {**entry, "loaded": True}
# Locally available (downloaded/cached) models that are not already loaded.
for info in await _cached_local_catalog():
cid = getattr(info, "model_id", None) or public_model_id(getattr(info, "id", None))
if not cid or cid in by_id:
continue
obj = {
"id": cid,
"object": "model",
"created": _created,
"owned_by": _OWNED_BY,
"loaded": False,
}
display = getattr(info, "display_name", None)
if display:
obj["display_name"] = display
by_id[cid] = obj
return list(by_id.values())
@router.get("/models")
async def openai_list_models(current_subject: str = Depends(get_current_subject)):
"""
OpenAI-compatible model listing endpoint.
OpenAI-compatible model listing endpoint (``GET /v1/models``).
Returns the currently loaded model in the format expected by
OpenAI-compatible clients (``GET /v1/models``).
Lists every model available on this server -- the loaded model(s) plus
locally available (downloaded/cached) models -- not only what is resident in
memory. Each entry carries a clean public id and a ``loaded`` flag.
"""
return {"object": "list", "data": _openai_model_objects()}
return {"object": "list", "data": await _openai_catalog_objects()}
@router.get("/models/{model_id:path}")
@ -6459,13 +6624,37 @@ async def openai_retrieve_model(model_id: str, current_subject: str = Depends(ge
"""
OpenAI-compatible single-model retrieval endpoint (``GET /v1/models/{id}``).
Returns the bare model object when ``model_id`` matches a loaded local
model, or 404 model_not_found otherwise. Defined after the LIST route so
it does not shadow it; ``{model_id:path}`` keeps ids with slashes intact.
Returns the bare model object when ``model_id`` matches a known model
(loaded or locally available), or 404 model_not_found otherwise. Defined
after the LIST route so it does not shadow it; ``{model_id:path}`` keeps ids
with slashes intact.
"""
for model in _openai_model_objects():
from core.inference.model_ids import model_id_matches
# Loaded models resolve without a catalog scan (the common case); only build
# the full catalog -- which may hit the filesystem -- for unloaded ids.
for entry in _openai_model_objects():
if entry["id"] == model_id:
return {**entry, "loaded": True}
objects = await _openai_catalog_objects()
for model in objects:
if model["id"] == model_id:
return model
# Backward compatibility: a client may still send the legacy raw identifier
# (e.g. an absolute .gguf path cached from an older /v1/models). Resolve it to
# the clean object so it keeps working, without ever echoing the path back.
llama_backend = get_llama_cpp_backend()
backend = get_inference_backend()
for raw in (
llama_backend.model_identifier if llama_backend.is_loaded else None,
backend.active_model_name or None,
):
if raw and model_id_matches(model_id, raw):
clean = public_model_id(raw)
for model in objects:
if model["id"] == clean:
return model
raise HTTPException(
status_code = 404,
detail = openai_error_body(
@ -7402,6 +7591,15 @@ async def _responses_stream(
target_url = f"{llama_backend.base_url}/v1/chat/completions"
async def event_generator():
# Clean public id for every response envelope. Prefer the loaded model's
# id so the stream agrees with /v1/models, chat/completions and the
# non-streaming twin; fall back to a sanitized payload.model (a legacy
# raw .gguf path is stripped, never echoed back).
_clean_model = (
public_model_id(getattr(llama_backend, "model_identifier", None))
or public_model_id(payload.model)
or payload.model
)
full_text = ""
full_reasoning = ""
input_tokens = 0
@ -7563,7 +7761,7 @@ async def _responses_stream(
"object": "response",
"created_at": created_at,
"status": "failed",
"model": payload.model,
"model": _clean_model,
"output": _snapshot_output(),
"usage": {
"input_tokens": input_tokens,
@ -7587,7 +7785,7 @@ async def _responses_stream(
"object": "response",
"created_at": created_at,
"status": "in_progress",
"model": payload.model,
"model": _clean_model,
"output": [],
"usage": {"input_tokens": 0, "output_tokens": 0, "total_tokens": 0},
},
@ -7627,7 +7825,7 @@ async def _responses_stream(
"object": "response",
"created_at": created_at,
"status": "failed",
"model": payload.model,
"model": _clean_model,
"output": [],
"error": {"code": 502, "message": _friendly_error(e)},
},
@ -7653,7 +7851,7 @@ async def _responses_stream(
"object": "response",
"created_at": created_at,
"status": "failed",
"model": payload.model,
"model": _clean_model,
"output": [],
"error": {
"code": resp.status_code,
@ -8002,7 +8200,7 @@ async def _responses_stream(
"object": "response",
"created_at": created_at,
"status": "completed",
"model": payload.model,
"model": _clean_model,
"output": _snapshot_output(),
"usage": {
"input_tokens": input_tokens,
@ -8274,7 +8472,13 @@ async def anthropic_messages(
),
)
model_name = getattr(llama_backend, "model_identifier", None) or payload.model
# Clean public id so /v1/messages never echoes the local .gguf path (and a
# legacy raw path sent as payload.model is sanitized rather than returned).
model_name = (
public_model_id(getattr(llama_backend, "model_identifier", None))
or public_model_id(payload.model)
or payload.model
)
message_id = f"msg_{uuid.uuid4().hex[:24]}"
# ── Translate Anthropic → OpenAI ──────────────────────────

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

@ -722,6 +722,94 @@ def _scan_ollama_dir(ollama_dir: Path, limit: Optional[int] = None) -> List[Loca
return found
def collect_local_models(models_root: Path) -> List[LocalModelInfo]:
"""Scan ``models_root``, the HF caches, LM Studio dirs, and user scan folders,
returning a deduplicated, hidden-filtered list of discovered local models.
Shared by ``GET /models/local`` (the model picker) and the OpenAI-compatible
catalog (``GET /v1/models``) so the UI and the API never drift. ``models_root``
must already be validated/trusted by the caller.
"""
from storage.studio_db import list_scan_folders
from utils.paths import (
hf_default_cache_dir,
legacy_hf_cache_dir,
lmstudio_model_dirs,
)
hf_cache_dir = _resolve_hf_cache_dir()
legacy_hf = legacy_hf_cache_dir()
hf_default = hf_default_cache_dir()
lm_dirs = lmstudio_model_dirs()
local_models = _scan_models_dir(models_root) + _scan_hf_cache(hf_cache_dir)
# Resolve once; an inaccessible aux cache must skip that scan, not 500.
hf_cache_real = _safe_resolve(hf_cache_dir)
legacy_real = _safe_resolve(legacy_hf)
default_real = _safe_resolve(hf_default)
# Scan legacy Unsloth HF cache for backward compatibility.
if _safe_is_dir(legacy_hf) and legacy_real != hf_cache_real:
local_models += _scan_hf_cache(legacy_hf)
# Scan HF system default cache (may differ under env overrides).
if _safe_is_dir(hf_default) and default_real != hf_cache_real and default_real != legacy_real:
local_models += _scan_hf_cache(hf_default)
# Scan LM Studio directories.
for lm_dir in lm_dirs:
local_models += _scan_lmstudio_dir(lm_dir)
# Scan user-added custom folders (per-folder cap).
_MAX_MODELS_PER_FOLDER = 200
try:
custom_folders = list_scan_folders()
except Exception as e:
logger.warning("Could not load custom scan folders: %s", e)
custom_folders = []
for folder in custom_folders:
folder_path = Path(folder["path"])
try:
# Filter Ollama .studio_links/ from generic scanners to
# avoid duplicates and leaking internal paths into the UI.
_generic = [
m
for m in (
_scan_models_dir(folder_path, limit = _MAX_MODELS_PER_FOLDER)
+ _scan_hf_cache(folder_path)
+ _scan_lmstudio_dir(folder_path)
)
if not any(p in (".studio_links", "ollama_links") for p in Path(m.path).parts)
]
custom_models = _generic
if len(custom_models) < _MAX_MODELS_PER_FOLDER:
custom_models += _scan_ollama_dir(
folder_path,
limit = _MAX_MODELS_PER_FOLDER - len(custom_models),
)
except OSError as e:
logger.warning("Skipping unreadable scan folder %s: %s", folder_path, e)
continue
local_models += [m.model_copy(update = {"source": "custom"}) for m in custom_models]
# Deduplicate, but always keep custom folder entries (keyed by
# (id, source)) so they show in the "Custom Folders" UI section
# even when the model is also in the HF cache.
deduped: dict[str, LocalModelInfo] = {}
for model in local_models:
key = f"{model.id}\x00custom" if model.source == "custom" else model.id
if key not in deduped:
deduped[key] = model
models = sorted(
deduped.values(),
key = lambda item: (item.updated_at or 0),
reverse = True,
)
return [m for m in models if not _is_hidden_model(m.id, m.path)]
@router.get("/local", response_model = LocalModelListResponse)
async def list_local_models(
models_dir: str = Query(
@ -770,78 +858,7 @@ async def list_local_models(
)
try:
local_models = _scan_models_dir(models_root) + _scan_hf_cache(hf_cache_dir)
# Resolve once; an inaccessible aux cache must skip that scan, not 500.
hf_cache_real = _safe_resolve(hf_cache_dir)
legacy_real = _safe_resolve(legacy_hf)
default_real = _safe_resolve(hf_default)
# Scan legacy Unsloth HF cache for backward compatibility.
if _safe_is_dir(legacy_hf) and legacy_real != hf_cache_real:
local_models += _scan_hf_cache(legacy_hf)
# Scan HF system default cache (may differ under env overrides).
if (
_safe_is_dir(hf_default)
and default_real != hf_cache_real
and default_real != legacy_real
):
local_models += _scan_hf_cache(hf_default)
# Scan LM Studio directories.
for lm_dir in lm_dirs:
local_models += _scan_lmstudio_dir(lm_dir)
# Scan user-added custom folders (per-folder cap).
from storage.studio_db import list_scan_folders
_MAX_MODELS_PER_FOLDER = 200
try:
custom_folders = list_scan_folders()
except Exception as e:
logger.warning("Could not load custom scan folders: %s", e)
custom_folders = []
for folder in custom_folders:
folder_path = Path(folder["path"])
try:
# Filter Ollama .studio_links/ from generic scanners to
# avoid duplicates and leaking internal paths into the UI.
_generic = [
m
for m in (
_scan_models_dir(folder_path, limit = _MAX_MODELS_PER_FOLDER)
+ _scan_hf_cache(folder_path)
+ _scan_lmstudio_dir(folder_path)
)
if not any(p in (".studio_links", "ollama_links") for p in Path(m.path).parts)
]
custom_models = _generic
if len(custom_models) < _MAX_MODELS_PER_FOLDER:
custom_models += _scan_ollama_dir(
folder_path,
limit = _MAX_MODELS_PER_FOLDER - len(custom_models),
)
except OSError as e:
logger.warning("Skipping unreadable scan folder %s: %s", folder_path, e)
continue
local_models += [m.model_copy(update = {"source": "custom"}) for m in custom_models]
# Deduplicate, but always keep custom folder entries (keyed by
# (id, source)) so they show in the "Custom Folders" UI section
# even when the model is also in the HF cache.
deduped: dict[str, LocalModelInfo] = {}
for model in local_models:
key = f"{model.id}\x00custom" if model.source == "custom" else model.id
if key not in deduped:
deduped[key] = model
models = sorted(
deduped.values(),
key = lambda item: (item.updated_at or 0),
reverse = True,
)
models = [m for m in models if not _is_hidden_model(m.id, m.path)]
models = collect_local_models(models_root)
return LocalModelListResponse(
models_dir = str(models_root),

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

@ -255,6 +255,7 @@ async def start_training(
# Convert request to backend kwargs.
training_kwargs = {
"model_name": request.model_name,
"project_name": request.project_name,
"training_type": request.training_type,
"hf_token": request.hf_token or "",
"load_in_4bit": request.load_in_4bit,

29
studio/backend/run.py
View file

@ -933,6 +933,9 @@ def run_server(
"""
global _server, _server_thread, _shutdown_event
boot_started = time.perf_counter()
logger.info("run_server startup begin api_only=%s host=%s port=%s", api_only, host, port)
# Reap every child if the parent dies abnormally (terminal close, Task
# Manager kill, SIGKILL); must run before any child can spawn.
from utils.process_lifetime import initialize_parent_lifetime
@ -984,7 +987,14 @@ def run_server(
from threading import Thread, Event
import uvicorn
import_started = time.perf_counter()
from main import app, setup_frontend, _IS_COLAB
logger.info(
"Imported FastAPI app in %.1fms",
(time.perf_counter() - import_started) * 1000,
)
from utils.paths import ensure_studio_directories
# Allow local stdio MCP servers on a loopback bind (the user's own machine),
@ -997,6 +1007,11 @@ def run_server(
# Create all standard directories on startup.
ensure_studio_directories()
logger.info(
"Ensured Studio directories in %.1fms",
(time.perf_counter() - boot_started) * 1000,
)
# Auto-find a free port if the requested one is in use.
if not _is_port_free(host, port):
original_port = port
@ -1060,6 +1075,11 @@ def run_server(
display_host = _resolve_external_ip() if host == "0.0.0.0" else host
_install_uvicorn_startup_log_rewrite(host, display_host)
logger.info(
"run_server pre-uvicorn setup completed in %.1fms",
(time.perf_counter() - boot_started) * 1000,
)
ready_event = Event()
startup_failed = Event()
startup_errors = []
@ -1068,6 +1088,10 @@ def run_server(
async def startup(self, *args, **kwargs):
await super().startup(*args, **kwargs)
if getattr(self, "started", False) and not self.should_exit:
logger.info(
"Uvicorn startup hook completed in %.1fms",
(time.perf_counter() - boot_started) * 1000,
)
ready_event.set()
# server_header=False suppresses uvicorn's "Server: uvicorn"; SecurityHeadersMiddleware sets its own.
@ -1150,6 +1174,11 @@ def run_server(
_shutdown_event.set()
raise
logger.info(
"run_server uvicorn ready after %.1fms",
(time.perf_counter() - boot_started) * 1000,
)
_write_pid_file()
import atexit

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

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

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

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

120
studio/backend/tests/test_data_recipe_seed.py
View file

@ -1,12 +1,126 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import asyncio
import importlib.util
from pathlib import Path
import pytest
def test_seed_inspect_load_kwargs_disables_remote_code_execution():
seed_route = (
def _seed_route_source() -> str:
return (
Path(__file__).resolve().parent.parent / "routes" / "data_recipe" / "seed.py"
).read_text()
assert '"trust_remote_code": False' in seed_route
def test_seed_inspect_load_kwargs_disables_remote_code_execution():
assert '"trust_remote_code": False' in _seed_route_source()
class _FakeUpload:
def __init__(self, filename: str, content: bytes):
self.filename = filename
self._content = content
async def read(self) -> bytes:
return self._content
def _load_seed_route(monkeypatch: pytest.MonkeyPatch, tmp_path: Path):
pytest.importorskip("fastapi")
pytest.importorskip("multipart")
pytest.importorskip("structlog")
backend_root = Path(__file__).resolve().parent.parent
monkeypatch.syspath_prepend(str(backend_root))
route_path = backend_root / "routes" / "data_recipe" / "seed.py"
spec = importlib.util.spec_from_file_location("seed_under_test", route_path)
assert spec is not None and spec.loader is not None
seed_route = importlib.util.module_from_spec(spec)
spec.loader.exec_module(seed_route)
seed_route.UNSTRUCTURED_UPLOAD_ROOT = tmp_path / "unstructured-uploads"
return seed_route
def _run_upload(
seed_route,
filename: str,
content: bytes,
block_id: str = "block",
):
return asyncio.run(
seed_route.upload_unstructured_file(_FakeUpload(filename, content), block_id)
)
def _block_files(seed_route, block_id: str = "block") -> list[str]:
block_dir = seed_route.UNSTRUCTURED_UPLOAD_ROOT / block_id
if not block_dir.exists():
return []
return sorted(path.name for path in block_dir.iterdir())
def _raise(exc: BaseException):
def raise_exc(*args, **kwargs):
raise exc
return raise_exc
@pytest.mark.parametrize(
("filename", "package"),
[
("paper.pdf", "pymupdf4llm"),
("notes.docx", "mammoth"),
],
)
def test_unstructured_upload_names_missing_extractor_dependency(
monkeypatch, tmp_path, filename, package
):
seed_route = _load_seed_route(monkeypatch, tmp_path)
monkeypatch.setattr(
seed_route,
"_extract_text_from_file",
_raise(ModuleNotFoundError(f"No module named {package!r}", name = package)),
)
result = _run_upload(seed_route, filename, b"%PDF-1.7")
assert result.status == "error"
assert (
result.error
== f"Cannot read {Path(filename).suffix} files: the '{package}' package is not installed."
)
assert _block_files(seed_route) == []
def test_unstructured_upload_keeps_txt_path_working(monkeypatch, tmp_path):
seed_route = _load_seed_route(monkeypatch, tmp_path)
result = _run_upload(seed_route, "notes.txt", b"hello")
assert result.status == "ok"
assert result.error is None
assert any(name.endswith(".txt") for name in _block_files(seed_route))
assert any(name.endswith(".extracted.txt") for name in _block_files(seed_route))
@pytest.mark.parametrize(
"exc",
[
ImportError("cannot import internal symbol"),
ModuleNotFoundError(
"No module named 'missing_transitive_pkg'",
name = "missing_transitive_pkg",
),
],
)
def test_unstructured_upload_import_errors_stay_generic(monkeypatch, tmp_path, exc):
seed_route = _load_seed_route(monkeypatch, tmp_path)
monkeypatch.setattr(seed_route, "_extract_text_from_file", _raise(exc))
result = _run_upload(seed_route, "paper.pdf", b"%PDF-1.7")
assert result.status == "error"
assert result.error == "Text extraction failed."
assert _block_files(seed_route) == []

62
studio/backend/tests/test_model_ids.py 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
import sys
from pathlib import Path
_BACKEND = Path(__file__).resolve().parents[1]
if str(_BACKEND) not in sys.path:
sys.path.insert(0, str(_BACKEND))
from core.inference.model_ids import model_id_matches, public_model_id # noqa: E402
def test_local_gguf_path_becomes_clean_stem():
assert public_model_id("/srv/models/Qwen3-30B-A3B-Q4_K_M.gguf") == "Qwen3-30B-A3B-Q4_K_M"
assert public_model_id("/home/u/.cache/models/llama.gguf") == "llama"
def test_hf_repo_id_unchanged():
assert public_model_id("unsloth/Qwen3-30B-A3B-GGUF") == "unsloth/Qwen3-30B-A3B-GGUF"
assert public_model_id("Qwen3-30B-A3B") == "Qwen3-30B-A3B"
def test_none_and_empty_passthrough():
assert public_model_id(None) is None
assert public_model_id("") == ""
def test_windows_path():
assert public_model_id("C:\\models\\foo.gguf") == "foo"
assert public_model_id("models\\sub\\bar.gguf") == "bar"
def test_directory_path_uses_basename():
assert public_model_id("/opt/models/MyModelDir") == "MyModelDir"
# A 3+ segment relative path is a local path, not an org/model repo id.
assert public_model_id("a/b/c") == "c"
def test_relative_and_home_paths_are_sanitized():
# ./ ../ ~ prefixed paths are local and must not be echoed raw.
assert public_model_id("./model.gguf") == "model"
assert public_model_id("../models/foo.gguf") == "foo"
assert public_model_id("~/models/baz.gguf") == "baz"
assert public_model_id("./mistral") == "mistral"
assert public_model_id("~/mistral") == "mistral"
assert public_model_id(".\\models\\foo.gguf") == "foo"
def test_dotted_repo_id_not_mistaken_for_relative_path():
# A leading dot that is not ./ or ../ is an ordinary clean name.
assert public_model_id(".hidden-model") == ".hidden-model"
assert public_model_id("org/.config") == "org/.config"
def test_matches_clean_and_legacy():
path = "/srv/models/Qwen3-Q4.gguf"
assert model_id_matches("Qwen3-Q4", path) # clean public id
assert model_id_matches(path, path) # legacy raw path
assert not model_id_matches("other", path)
assert not model_id_matches(None, path)
assert not model_id_matches("x", None)

181
studio/backend/tests/test_openai_catalog.py Normal file
View file

@ -0,0 +1,181 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""GET /v1/models lists the full server catalog (loaded + locally available)."""
import asyncio
import json
import sys
from pathlib import Path
_BACKEND = Path(__file__).resolve().parents[1]
if str(_BACKEND) not in sys.path:
sys.path.insert(0, str(_BACKEND))
import routes.inference as inf # noqa: E402
class _Info:
def __init__(
self,
id,
display_name,
model_id = None,
):
self.id = id
self.display_name = display_name
self.model_id = model_id
class _FakeLlama:
is_loaded = True
model_identifier = "/srv/models/Qwen3-Q4.gguf"
context_length = 4096
max_context_length = None
native_context_length = None
def __init__(self, loaded = True):
self.is_loaded = loaded
class _FakeUnsloth:
active_model_name = None
models: dict = {}
context_length = None
max_seq_length = None
def test_catalog_lists_loaded_and_available(monkeypatch):
monkeypatch.setattr(inf, "get_llama_cpp_backend", lambda: _FakeLlama())
monkeypatch.setattr(inf, "get_inference_backend", lambda: _FakeUnsloth())
async def _fake_catalog():
return [
_Info("/data/models/Qwen3-Q4.gguf", "Qwen3-Q4"), # same as loaded -> dedup
_Info("/data/models/Llama-8B-Q8.gguf", "Llama-8B-Q8"), # available, not loaded
_Info("models--org--Foo", "Foo", model_id = "org/Foo"), # hf cache repo id
]
monkeypatch.setattr(inf, "_cached_local_catalog", _fake_catalog)
data = asyncio.run(inf._openai_catalog_objects())
ids = {m["id"]: m for m in data}
# Loaded model is present, marked loaded, and keeps context fields.
assert ids["Qwen3-Q4"]["loaded"] is True
assert ids["Qwen3-Q4"]["context_length"] == 4096
# Available-but-not-loaded models are listed too.
assert ids["Llama-8B-Q8"]["loaded"] is False
assert ids["org/Foo"]["loaded"] is False
# The loaded gguf and the on-disk copy collapse to one clean id.
assert [m["id"] for m in data].count("Qwen3-Q4") == 1
# No absolute paths or .gguf suffixes leak anywhere.
blob = json.dumps(data)
assert ".gguf" not in blob
assert "/srv/" not in blob
assert "/data/" not in blob
def test_empty_and_errored_scans_are_cached(monkeypatch):
# Cache validity is keyed on the timestamp, not list contents, so an empty
# (fresh install / no local models) or errored scan is still cached for the
# TTL instead of rescanning the filesystem on every /v1/models poll.
import routes.models as models_mod
for outcome in ("empty", "error"):
calls = {"n": 0}
def _scan(_root, _outcome = outcome):
calls["n"] += 1
if _outcome == "error":
raise RuntimeError("scan blew up")
return []
monkeypatch.setattr(models_mod, "collect_local_models", _scan)
monkeypatch.setattr(inf, "_CATALOG_CACHE", {"at": 0.0, "models": []})
async def _run():
return [await inf._cached_local_catalog() for _ in range(3)]
results = asyncio.run(_run())
assert results == [[], [], []], outcome
assert calls["n"] == 1, f"{outcome} scan ran {calls['n']}x (TTL not honored)"
def test_catalog_ttl_starts_after_scan_completes(monkeypatch):
# The cache timestamp must be taken AFTER the scan, not before it. A scan that
# outlives the TTL would otherwise leave the cache born-expired, so the next
# caller rescans instead of reusing the just-computed catalog.
import routes.models as models_mod
clock = {"t": 1000.0}
monkeypatch.setattr(inf.time, "monotonic", lambda: clock["t"])
monkeypatch.setattr(inf, "_CATALOG_CACHE", {"at": 0.0, "models": []})
calls = {"n": 0}
def _slow_scan(_root):
calls["n"] += 1
clock["t"] += inf._CATALOG_TTL_S + 10 # the scan itself outlives the TTL
return [_Info("/m/A.gguf", "A")]
monkeypatch.setattr(models_mod, "collect_local_models", _slow_scan)
async def _run():
first = await inf._cached_local_catalog()
second = await inf._cached_local_catalog() # clock unchanged since scan end
return first, second
first, second = asyncio.run(_run())
assert [i.id for i in first] == ["/m/A.gguf"]
assert calls["n"] == 1, "TTL started before the scan -> cache born expired, rescanned"
def test_retrieve_loaded_model_skips_catalog_scan(monkeypatch):
# Retrieving a loaded id must resolve from the loaded set alone, never paying
# for the filesystem scan that _cached_local_catalog drives.
monkeypatch.setattr(inf, "get_llama_cpp_backend", lambda: _FakeLlama())
monkeypatch.setattr(inf, "get_inference_backend", lambda: _FakeUnsloth())
async def _boom():
raise AssertionError("catalog scan must not run for a loaded id")
monkeypatch.setattr(inf, "_cached_local_catalog", _boom)
model = asyncio.run(inf.openai_retrieve_model("Qwen3-Q4", current_subject = "t"))
assert model["id"] == "Qwen3-Q4"
assert model["loaded"] is True
def test_cached_local_catalog_offloads_and_caches(monkeypatch):
# The filesystem scan must run off the event loop (asyncio.to_thread) and be
# cached, so a burst of /v1/models calls does not re-scan or block.
calls = {"scan": 0, "threaded": 0}
def _fake_collect(_root):
calls["scan"] += 1
return [_Info("/data/models/A.gguf", "A")]
import routes.models as models_mod
monkeypatch.setattr(models_mod, "collect_local_models", _fake_collect)
real_to_thread = inf.asyncio.to_thread
async def _counting_to_thread(fn, *a, **k):
calls["threaded"] += 1
return await real_to_thread(fn, *a, **k)
monkeypatch.setattr(inf.asyncio, "to_thread", _counting_to_thread)
# Fresh cache for a deterministic count.
monkeypatch.setattr(inf, "_CATALOG_CACHE", {"at": 0.0, "models": []})
async def _run():
first = await inf._cached_local_catalog()
second = await inf._cached_local_catalog() # within TTL -> cached
return first, second
first, second = asyncio.run(_run())
assert [i.id for i in first] == ["/data/models/A.gguf"]
assert second is first or [i.id for i in second] == [i.id for i in first]
assert calls["scan"] == 1 # cached: scanned once for two calls
assert calls["threaded"] == 1 # offloaded to a worker thread

45
studio/backend/tests/test_openai_models_path_leak.py Normal file
View file

@ -0,0 +1,45 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""GET /v1/models must report a clean public id, never the on-disk .gguf path."""
import json
import sys
from pathlib import Path
_BACKEND = Path(__file__).resolve().parents[1]
if str(_BACKEND) not in sys.path:
sys.path.insert(0, str(_BACKEND))
import routes.inference as inf # noqa: E402
class _FakeLlama:
is_loaded = True
model_identifier = "/srv/models/Qwen3-30B-A3B-Q4_K_M.gguf"
context_length = 4096
max_context_length = None
native_context_length = None
class _FakeUnsloth:
active_model_name = None
models: dict = {}
context_length = None
max_seq_length = None
def test_openai_models_returns_clean_id_without_path(monkeypatch):
monkeypatch.setattr(inf, "get_llama_cpp_backend", lambda: _FakeLlama())
monkeypatch.setattr(inf, "get_inference_backend", lambda: _FakeUnsloth())
objs = inf._openai_model_objects()
assert len(objs) == 1
assert objs[0]["id"] == "Qwen3-30B-A3B-Q4_K_M"
# The serialized payload must not leak the absolute path or the .gguf suffix.
blob = json.dumps(objs)
assert "/srv/models" not in blob
assert ".gguf" not in blob
# Context fields still flow through.
assert objs[0]["context_length"] == 4096

805
studio/backend/tests/test_tp_vision_regression.py Normal file
View file

@ -0,0 +1,805 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Regression guards for silent tensor-parallel downgrades in load_model.
PR #6416 blanket-disabled tensor parallelism for vision models to dodge a
--split-mode tensor + --mmproj GGML_ASSERT (#6415), which silently single-GPU'd
any mmproj/MTP GGUF that fit on one card. The fix makes the skip self-healing:
tensor is tried by default and recorded per (binary, model) only on a real abort.
load_model is too entangled to drive end-to-end, so these tests inspect the
source / drive the pure helpers. The headline test pins the set of TP-drop
conditions, so a new silent drop fails CI. No GPU; fully deterministic.
"""
from __future__ import annotations
import ast
import importlib.util
import inspect
import os
import sys
import textwrap
import types as _types
from pathlib import Path
_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
if _BACKEND_DIR not in sys.path:
sys.path.insert(0, _BACKEND_DIR)
# External-dep stubs so importing the backend doesn't require structlog / httpx /
# loggers -- but only when the real module is missing, so a lightweight stub never
# shadows the real package (or `loggers.handlers` submodule) for tests collected
# later in the same pytest process.
try:
import structlog # noqa: F401
except ImportError:
_structlog_stub = _types.ModuleType("structlog")
_structlog_stub.get_logger = lambda *a, **k: __import__("logging").getLogger("stub")
sys.modules["structlog"] = _structlog_stub
try:
import loggers # noqa: F401
except ImportError:
_loggers_stub = _types.ModuleType("loggers")
_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name)
sys.modules["loggers"] = _loggers_stub
try:
import httpx as _httpx_real # noqa: F401
except ImportError:
_httpx_stub = _types.ModuleType("httpx")
for _exc in (
"ConnectError",
"TimeoutException",
"ReadTimeout",
"ReadError",
"RemoteProtocolError",
"CloseError",
"HTTPError",
"RequestError",
):
setattr(_httpx_stub, _exc, type(_exc, (Exception,), {}))
_httpx_stub.Timeout = type("T", (), {"__init__": lambda s, *a, **k: None})
_httpx_stub.Response = type("Response", (), {})
_httpx_stub.Client = type(
"C",
(),
{
"__init__": lambda s, **kw: None,
"__enter__": lambda s: s,
"__exit__": lambda s, *a: None,
},
)
sys.modules["httpx"] = _httpx_stub
from core.inference.llama_cpp import LlamaCppBackend # noqa: E402
_GB = 1024**3
def _load_inference_routes_module():
"""Load routes/inference.py directly, bypassing routes/__init__.py (which imports
every router, dragging in unrelated deps like python-multipart) (Codex #6659)."""
route_path = Path(_BACKEND_DIR) / "routes" / "inference.py"
spec = importlib.util.spec_from_file_location(
"tp_vision_regression_inference_routes", route_path
)
assert spec is not None and spec.loader is not None
module = importlib.util.module_from_spec(spec)
sys.modules[spec.name] = module
spec.loader.exec_module(module)
return module
def _load_model_ast() -> ast.FunctionDef:
"""Parse load_model into an AST FunctionDef (no import side effects)."""
src = textwrap.dedent(inspect.getsource(LlamaCppBackend.load_model))
return ast.parse(src).body[0]
def _tensor_parallel_false_drop_guards() -> list[str]:
"""Source of the guard expression for every `if ...: tensor_parallel = False`
(the LOCAL variable, not self._tensor_parallel) inside load_model."""
fn = _load_model_ast()
def _body_drops_tp(body) -> bool:
for n in body:
if (
isinstance(n, ast.Assign)
and any(isinstance(t, ast.Name) and t.id == "tensor_parallel" for t in n.targets)
and isinstance(n.value, ast.Constant)
and n.value.value is False
):
return True
return False
return [
ast.unparse(node.test)
for node in ast.walk(fn)
if isinstance(node, ast.If) and _body_drops_tp(node.body)
]
# Every condition that may flip a requested tensor_parallel back to False. Adding
# one must be conscious: update this allowlist and keep multi-GPU where possible.
_ALLOWED_TP_DROP_GUARDS = {
# Capability: --split-mode tensor aborted for this (binary, model) (#6415).
# Self-healing -- tried by default, skipped only after a real abort (vs #6416).
"tensor_parallel and self._tensor_split_aborts(binary, model_identifier)",
# Capacity: tensor needs >= 2 GPUs clearing the compute-buffer reserve.
"tensor_parallel and len(tp_gpus) < 2",
# Capacity: pooled usable VRAM can't hold weights + MTP reserve -> layer split.
"_tp_weight_budget_mib <= _tp_required_mib",
}
def test_tensor_parallel_drop_sites_match_allowlist():
"""The set of reasons a requested TP can be dropped is fixed and reviewed: a new
drop site fails this set-equality until consciously allowlisted (would catch #6416)."""
found = set(_tensor_parallel_false_drop_guards())
assert found == _ALLOWED_TP_DROP_GUARDS, (
"tensor_parallel drop sites changed.\n"
f" unexpected (new) : {sorted(found - _ALLOWED_TP_DROP_GUARDS)}\n"
f" missing (removed): {sorted(_ALLOWED_TP_DROP_GUARDS - found)}\n"
"A new drop means a user's TP request is ignored for a new reason -- "
"review it, keep multi-GPU where possible, surface it, then update "
"_ALLOWED_TP_DROP_GUARDS."
)
def test_every_tp_drop_is_logged_not_silent():
"""Each tensor_parallel downgrade must log why, so it never disappears silently."""
fn = _load_model_ast()
def _body_drops_tp(body):
return any(
isinstance(n, ast.Assign)
and any(isinstance(t, ast.Name) and t.id == "tensor_parallel" for t in n.targets)
and isinstance(n.value, ast.Constant)
and n.value.value is False
for n in body
)
def _body_logs(body) -> bool:
for n in ast.walk(ast.Module(body = list(body), type_ignores = [])):
if (
isinstance(n, ast.Call)
and isinstance(n.func, ast.Attribute)
and isinstance(n.func.value, ast.Name)
and n.func.value.id == "logger"
):
return True
return False
for node in ast.walk(fn):
if isinstance(node, ast.If) and _body_drops_tp(node.body):
assert _body_logs(node.body), (
f"TP drop under `{ast.unparse(node.test)}` has no logger call -- "
"downgrades must explain themselves."
)
def test_tensor_split_gate_is_self_healing_not_blanket():
"""Skip is conditional on a recorded (binary, model) abort, not a blanket
is_vision disable (the #6416 regression)."""
src = inspect.getsource(LlamaCppBackend.load_model)
assert "self._tensor_split_aborts(binary, model_identifier)" in src
assert "if tensor_parallel and is_vision:" not in src
assert "if tensor_parallel and effective_is_vision:" not in src
def test_tensor_split_skip_documents_layer_split_fallback():
"""When the skip fires (known-bad binary+model), it states the fallback."""
src = inspect.getsource(LlamaCppBackend.load_model)
gate = src.find("self._tensor_split_aborts(binary, model_identifier)")
assert gate != -1
block = src[gate : gate + 600]
assert "layer split" in block, "the skip should state it falls back to layer split"
def test_tensor_split_abort_recorded_early_on_first_spawn():
"""Recorded on the first spawn showing the marker, before the flash-attn-off
retry (which can't run tensor so drops the marker) -- else it loops (oobabooga, #6659)."""
src = inspect.getsource(LlamaCppBackend.load_model)
idx = src.find("_record_tensor_split_abort(binary, model_identifier)")
assert idx != -1, "load_model must record a (binary, model) tensor-split abort"
guard = src[max(0, idx - 600) : idx]
assert "self._tensor_parallel" in guard
assert (
"_should_record_tensor_split_abort" in guard
), "record must be gated on the marker-plus-hard-crash decision helper"
# Recorded before the flash-attn-off retry, not after the full ladder.
fa_off = src.find("_with_flash_attn_off")
assert 0 <= idx < fa_off, "recording must latch on the first spawn, before flash-off"
def test_vision_downgrade_preserves_multi_gpu_intent():
"""The vision downgrade raises _layer_min_gpus and threads it into both the
_select_gpus and auto-context layer paths, so a fitting model still spreads."""
src = inspect.getsource(LlamaCppBackend.load_model)
assert "_layer_min_gpus = max(_layer_min_gpus, len(gpus))" in src
assert src.count("min_gpus = _layer_min_gpus") >= 2
assert "range(_auto_min_gpus, len(ranked) + 1)" in src
auto = src.find("_auto_min_gpus = max(")
assert auto != -1 and "_layer_min_gpus" in src[auto : auto + 200]
# ── per-binary capability cache (pure) ───────────────────────────────
def test_tensor_attempted_by_default_for_unknown_binary():
"""A (binary, model) not seen to abort -> tensor is attempted (not skipped)."""
assert LlamaCppBackend._tensor_split_aborts("/never/seen/llama-server", "m") is False
assert LlamaCppBackend._tensor_split_aborts(None, "m") is False
assert LlamaCppBackend._tensor_split_aborts("/x", None) is False
def test_recorded_tensor_abort_is_per_model():
"""A recorded (binary, model) abort trips the gate for that model only -- a
different model on the same binary still attempts tensor (oobabooga, #6659)."""
b = f"/tmp/llama-server-{id(object())}"
try:
assert LlamaCppBackend._tensor_split_aborts(b, "model-a") is False
LlamaCppBackend._record_tensor_split_abort(b, "model-a")
assert LlamaCppBackend._tensor_split_aborts(b, "model-a") is True
# a different model on the same binary is unaffected
assert LlamaCppBackend._tensor_split_aborts(b, "model-b") is False
finally:
LlamaCppBackend._tensor_split_abort_keys.discard(
LlamaCppBackend._tensor_split_cache_key(b, "model-a")
)
# ── _select_gpus: single-GPU collapse vs honored multi-GPU intent (pure) ──
def test_select_gpus_collapses_to_single_gpu_when_model_fits():
"""Default (min_gpus=1): a 39 GB model on four 183 GB GPUs pins ONE GPU -- the
'single GPU' symptom once TP drops, and why the downgrade needs min_gpus."""
gpus = [(0, 180000), (1, 180000), (2, 180000), (3, 180000)] # (idx, free MiB)
gpu_indices, _use_fit = LlamaCppBackend._select_gpus(int(39 * _GB), gpus)
assert gpu_indices is not None and len(gpu_indices) == 1
def test_select_gpus_min_gpus_keeps_multi_gpu_for_fitting_model():
"""min_gpus>=2 must NOT collapse to one GPU for a model that fits on one."""
gpus = [(0, 180000), (1, 180000), (2, 180000), (3, 180000)]
gpu_indices, _ = LlamaCppBackend._select_gpus(int(39 * _GB), gpus, min_gpus = 2)
assert gpu_indices is not None and len(gpu_indices) >= 2
def test_select_gpus_min_gpus_capped_to_available():
"""min_gpus larger than the GPU count is capped, not an error."""
gpus = [(0, 180000), (1, 180000)]
gi, _ = LlamaCppBackend._select_gpus(int(10 * _GB), gpus, min_gpus = 8)
assert gi is not None and len(gi) == 2
def test_select_gpus_uses_multiple_gpus_when_model_does_not_fit():
"""Sanity: selection spreads across GPUs when one card can't hold the model."""
gpus = [(0, 40000), (1, 40000), (2, 40000), (3, 40000)] # 40 GB free each
gpu_indices, _use_fit = LlamaCppBackend._select_gpus(int(120 * _GB), gpus)
assert gpu_indices is not None and len(gpu_indices) >= 2
def test_select_gpus_min_gpus_excludes_unusable_gpu():
"""min_gpus caps to usable cards: 2 free + 1 nearly-full -> 2-GPU split, not
forcing the full card (OOM) or tripping --fit (#6659)."""
gpus = [(0, 180000), (1, 180000), (2, 500)] # GPU 2 is nearly full
total = {0: 180000, 1: 180000, 2: 180000}
gi, _ = LlamaCppBackend._select_gpus(
int(39 * _GB),
gpus,
min_gpus = 3,
total_by_idx = total,
per_device_overhead_bytes = int(1 * _GB),
)
assert gi is not None
assert 2 not in gi, "a nearly-full GPU must not be forced in to satisfy min_gpus"
assert len(gi) == 2
def test_tensor_abort_cache_invalidated_on_binary_mtime_change(tmp_path):
"""Cache keys on (path, mtime, model), so a binary swapped in place (in-app
update, no restart) is re-probed instead of inheriting the old abort (#6659)."""
binp = tmp_path / "llama-server"
binp.write_text("v1")
p = str(binp)
try:
LlamaCppBackend._record_tensor_split_abort(p, "m")
assert LlamaCppBackend._tensor_split_aborts(p, "m") is True
# Simulate an in-place update bumping the binary's mtime.
st = binp.stat()
os.utime(p, (st.st_atime, st.st_mtime + 10))
assert (
LlamaCppBackend._tensor_split_aborts(p, "m") is False
), "a binary swapped in place (new mtime) must be re-probed"
# A same-second replacement (sub-second mtime bump) must also re-probe:
# second-resolution mtime would inherit the stale abort (reviewer.py P2).
sec_ns = (binp.stat().st_mtime_ns // 1_000_000_000) * 1_000_000_000
os.utime(p, ns = (sec_ns, sec_ns))
LlamaCppBackend._record_tensor_split_abort(p, "m")
binp.write_text("v2")
os.utime(p, ns = (sec_ns, sec_ns + 1))
assert (
LlamaCppBackend._tensor_split_aborts(p, "m") is False
), "a same-second in-place swap (ns mtime bump) must be re-probed"
finally:
for key in list(LlamaCppBackend._tensor_split_abort_keys):
if key and key[0] == p:
LlamaCppBackend._tensor_split_abort_keys.discard(key)
def test_tensor_split_abort_raises_early_to_layer_fallback():
"""The first-spawn abort raises to the route's layer fallback (not the text-only
mmproj strip), before the flash-attn-off retry, preserving the projector (#6659)."""
src = inspect.getsource(LlamaCppBackend.load_model)
raise_idx = src.find("(split-axis geometry); retrying with layer split")
assert raise_idx != -1, "the split-axis abort must raise to trigger a layer retry"
# raises before both the flash-attn-off retry and the text-only mmproj strip
assert raise_idx < src.find("_with_flash_attn_off")
assert raise_idx < src.find("_strip_mmproj_args(_last_spawn_cmd)")
# gated on the marker-plus-crash helper, which also drives the record just above
guard = src[max(0, raise_idx - 600) : raise_idx]
assert "_should_record_tensor_split_abort" in guard
rec_idx = src.find("_record_tensor_split_abort(binary, model_identifier)")
assert rec_idx != -1 and rec_idx < raise_idx
def test_budget_downgrade_preserves_multi_gpu_intent():
"""The pooled-VRAM downgrade raises _layer_min_gpus from the usable tensor GPUs
too, symmetric with the vision downgrade (reviewer.py asymmetric fix, #6659)."""
src = inspect.getsource(LlamaCppBackend.load_model)
budget = src.find("_tp_weight_budget_mib <= _tp_required_mib")
assert budget != -1
block = src[budget : budget + 1000]
assert "tensor_parallel = False" in block
assert (
"_layer_min_gpus = max(_layer_min_gpus, len(tp_gpus))" in block
), "the budget downgrade must preserve multi-GPU intent like the vision gate"
def test_compute_buffer_downgrade_preserves_multi_gpu_intent():
"""The len(tp_gpus) < 2 compute-buffer downgrade raises _layer_min_gpus from the
full GPU set too, so it is symmetric with the budget/geometry downgrades and
doesn't collapse a multi-GPU layer load to one card (reviewer.py P1 on #6659)."""
src = inspect.getsource(LlamaCppBackend.load_model)
gate = src.find("tensor_parallel and len(tp_gpus) < 2")
assert gate != -1
# Bound to exactly this block: from its gate to the next (budget) downgrade.
nxt = src.find("_tp_weight_budget_mib <= _tp_required_mib", gate)
assert nxt != -1
block = src[gate:nxt]
assert "tensor_parallel = False" in block
assert (
"_layer_min_gpus = max(_layer_min_gpus, len(gpus))" in block
), "the compute-buffer downgrade must preserve multi-GPU intent like the others"
def test_tensor_split_layer_min_gpus_bump_requires_tensor_request():
"""Every guard that bumps _layer_min_gpus off the abort cache also tests
tensor_parallel, so a non-tensor load on a known-bad binary doesn't grab every
GPU for a fitting model (#6659)."""
fn = _load_model_ast()
checked = 0
for node in ast.walk(fn):
if isinstance(node, ast.If):
test_src = ast.unparse(node.test)
if "self._tensor_split_aborts(binary, model_identifier)" not in test_src:
continue
body = "\n".join(ast.unparse(n) for n in node.body)
if "_layer_min_gpus" in body:
checked += 1
assert "tensor_parallel" in test_src, (
"the cached _layer_min_gpus bump must require a current tensor "
f"request, but fires under `{test_src}`"
)
assert checked >= 1, "expected an abort-cache guard that bumps _layer_min_gpus"
# ── round-2 follow-up: route-fallback retry + auto-context cap + assert marker ──
def test_layer_fallback_retry_preserves_multi_gpu_intent():
"""load_model takes a preserve_multi_gpu_on_layer hint and raises _layer_min_gpus
for it, so the tensor-off fallback retry still spreads a fitting model (#6659)."""
sig = inspect.signature(LlamaCppBackend.load_model)
assert "preserve_multi_gpu_on_layer" in sig.parameters
assert sig.parameters["preserve_multi_gpu_on_layer"].default is False
fn = _load_model_ast()
found = any(
isinstance(n, ast.If)
and "preserve_multi_gpu_on_layer" in ast.unparse(n.test)
and "_layer_min_gpus" in "\n".join(ast.unparse(b) for b in n.body)
for n in ast.walk(fn)
)
assert found, "preserve_multi_gpu_on_layer must raise _layer_min_gpus"
def test_auto_context_layer_loops_capped_to_usable_gpus():
"""The auto-context loops bypass _select_gpus, so they apply its cap: a card
counts only if usable VRAM clears the per-device layer overhead (#6659)."""
src = inspect.getsource(LlamaCppBackend.load_model)
assert (
"range(max(1, _layer_min_gpus), len(ranked) + 1)" not in src
), "auto-context loops must cap _layer_min_gpus to usable GPUs, not use it raw"
assert "_auto_min_gpus" in src
assert "range(_auto_min_gpus, len(ranked) + 1)" in src
# the eligibility threshold is the per-device layer overhead, not bare > 0
auto = src.find("_auto_min_gpus = max(")
assert auto != -1
block = src[auto : auto + 400]
assert "_pipeline_overhead_mib" in block, (
"a card must clear the per-device layer overhead to count, mirroring "
"_select_gpus, so a nearly-full GPU is not exposed and OOMs"
)
def test_fallback_hint_uses_effective_tensor_request_not_just_toggle():
"""Tensor intent keys off _effective_tensor_parallel (toggle + extras + env), not
just the toggle, so extra/env-driven tensor users keep multi-GPU (#6659)."""
route = Path(_BACKEND_DIR) / "routes" / "inference.py"
src = route.read_text()
idx = src.find("_tensor_intent_overall = _effective_tensor_parallel(")
assert idx != -1, "the GGUF load closure must compute tensor intent"
block = src[idx : idx + 300]
assert "extra_llama_args, request.tensor_parallel" in block
pres = src.find("preserve_multi_gpu_on_layer = bool(")
assert (
"_effective_tensor_parallel(attempt_extra_args, tensor_parallel)" in src[pres : pres + 200]
)
# not the toggle-only form this replaced
assert (
"bool(\n request.tensor_parallel and not tensor_parallel" not in src
)
def test_carry_preserved_tensor_intent_truth_table():
"""Behavioral check of the carry-forward decision: carried only for the SAME
model, preserved, and not an explicit drop. Catches a `not` inversion (ctx-only
collapse) and a missing same-model guard (cross-model leak) (#6659)."""
inference_routes = _load_inference_routes_module()
f = inference_routes._carry_preserved_tensor_intent
assert f(preserved = True, same_model = True, explicit_drop = False) is True
assert f(preserved = True, same_model = True, explicit_drop = True) is False # explicit drop
assert f(preserved = True, same_model = False, explicit_drop = False) is False # model switch
assert f(preserved = False, same_model = True, explicit_drop = False) is False # not a fallback
def test_preserved_fallback_carried_across_non_drop_reload():
"""The hint carries the preserved fallback via _carry_preserved_tensor_intent,
gated on the same model loaded, so a ctx-only reload keeps multi-GPU but a model
switch / explicit drop doesn't inherit it (#6659)."""
route = Path(_BACKEND_DIR) / "routes" / "inference.py"
src = route.read_text()
idx = src.find("_tensor_intent_overall = _effective_tensor_parallel(")
assert idx != -1
block = src[idx : idx + 400]
assert "_carry_preserved_tensor_intent(" in block
assert "preserved = llama_backend.layer_preserves_tensor_intent" in block
assert "same_model = _same_model_loaded" in block
assert "explicit_drop = _explicit_tensor_drop" in block
def test_same_model_guard_checks_path_and_variant():
"""The same-model guard matches the resolved config.identifier (what load_model
stores, after from_identifier normalizes shorthands) -- not the raw request id --
and also matches the loaded quant by path (local multi-variant dir) else variant (HF
repo), so a reload keeps the carry-forward and a different variant doesn't inherit
the prior one's preserved tensor intent (#6659)."""
route = Path(_BACKEND_DIR) / "routes" / "inference.py"
src = route.read_text()
idx = src.find("_same_model_loaded = (")
assert idx != -1
block = src[idx : idx + 1300]
# Identity compares the normalized config.identifier, not the raw model_identifier.
head = src[idx : idx + 200]
assert "config.identifier" in head and "== (model_identifier" not in head
assert "llama_backend.gguf_path" in block and "config.gguf_file" in block
assert "llama_backend.hf_variant" in block and "config.gguf_variant" in block
def test_diffusion_load_clears_preserved_tensor_flag():
"""The diffusion early-return path (skips the command builder) clears the
preserved-fallback flag, so a prior tensor fallback doesn't churn it (#6659)."""
src = inspect.getsource(LlamaCppBackend.load_model)
diff = src.find("if self._is_diffusion:")
assert diff != -1
start = src.find("return self._start_diffusion_server", diff)
assert start != -1
assert "self._layer_preserves_tensor_intent = False" in src[diff:start]
def test_is_tensor_split_assert_marker():
"""Matches the specific #6415 split-axis assert, not any ggml assert/abort, so
an unrelated invariant a corrupt GGUF/projector trips isn't cached (#6659)."""
f = LlamaCppBackend._is_tensor_split_assert
# the real #6415 warmup assert (split-axis enum, in ggml-backend-meta)
assert (
f(
"ggml-backend-meta.cpp:541: GGML_ASSERT(src_ss[0].axis != "
"GGML_BACKEND_SPLIT_AXIS_0) failed"
)
is True
)
# the split-axis token alone (file path elided / reworded) still matches
assert f("GGML_ASSERT(x.axis != GGML_BACKEND_SPLIT_AXIS_1) failed") is True
# UNRELATED asserts must NOT match -- including a different invariant from the
# same multi-assert source file (matched on the token, not the file name).
assert f("ggml-backend-meta.cpp:99: GGML_ASSERT(buf != NULL) failed") is False
assert f("/x/ggml.c:1234: GGML_ASSERT(ne == 1) failed") is False
assert f("ggml_abort: something else entirely") is False
assert f("Segmentation fault (core dumped)") is False
assert f("") is False
assert f(None) is False
def test_layer_preserve_hint_replayed_on_respawn():
"""The preserve hint is in the replay snapshot (_pending_load_kwargs), so a
respawn keeps the downgraded model multi-GPU (Codex review on #6659)."""
src = inspect.getsource(LlamaCppBackend.load_model)
pend = src.find("_pending_load_kwargs = {")
assert pend != -1
block = src[pend : src.find("}", pend) + 1]
assert '"preserve_multi_gpu_on_layer": preserve_multi_gpu_on_layer' in block, (
"the layer-preserve hint must be in the replay snapshot so _respawn_if_dead "
"keeps the multi-GPU placement"
)
def test_should_record_tensor_split_abort_decision():
"""Behavioral check of marker AND (signal crash OR Windows abort), so an
or->and typo or caching a generic crash fails here, not just the source pins."""
f = LlamaCppBackend._should_record_tensor_split_abort
marker = "ggml-backend-meta.cpp:541: GGML_ASSERT(x.axis != GGML_BACKEND_SPLIT_AXIS_0) failed"
# marker + a hard crash records, across every platform's abort encoding
assert f(-6, marker) is True # POSIX SIGABRT
assert f(-11, marker) is True # POSIX SIGSEGV
assert f(3, marker) is True # Windows CRT abort() exit (not a signal)
assert f(0xC0000005, marker) is True # Windows NTSTATUS access violation
# marker present but no hard crash -> not recorded
assert f(0, marker) is False # clean exit
assert f(-9, marker) is False # SIGKILL (OOM / unload), not a fault
assert f(None, marker) is False # still running
# hard crash but not the split-axis marker -> not recorded (no over-caching)
assert f(3, "some other failure") is False
assert f(-6, "GGML_ASSERT(buf != NULL) failed") is False
assert f(0xC0000005, "") is False
def test_fit_off_retry_skipped_on_split_axis_abort():
"""The fit-independent --fit off retry is skipped on the split-axis marker, else
the model crashes a second time before the latch records it (reviewer.py, #6659)."""
src = inspect.getsource(LlamaCppBackend.load_model)
retry = src.find('run_cmd = [*run_cmd, "--fit", "off"]')
assert retry != -1
guard = src[max(0, retry - 1000) : retry]
assert "_fit_retry_allowed" in guard and "_startup_crashed" in guard
assert (
"not _split_axis_crash" in guard
), "the fit-off retry must be skipped when the crash is a split-axis abort"
def test_is_abort_exit_recognizes_windows_crt_abort():
"""exit code 3 (MSVC abort()) counts as a crash; signals / clean exits do not."""
f = LlamaCppBackend._is_abort_exit
assert f(3) is True
assert f(0) is False
assert f(-6) is False # POSIX SIGABRT is handled by _is_signal_crash, not here
assert f(None) is False
# ── tensor-off after a multi-GPU fallback forces a reload (route dedup) ─
class _NoopProcess:
"""Stand-in for Popen so is_loaded is True and atexit cleanup doesn't crash."""
def terminate(self):
pass
def wait(self, timeout = None):
return 0
def kill(self):
pass
def poll(self):
return 0
def _fallback_loaded_backend(layer_preserves_tensor_intent: bool) -> LlamaCppBackend:
"""A loaded backend in the tensor->layer fallback state (tensor off, --split-mode
layer stored), differing only in the preserved-multi-GPU flag."""
b = LlamaCppBackend()
b._model_identifier = "owner/repo"
b._requested_n_ctx = 0
b._cache_type_kv = None
b._tensor_parallel = False
b._layer_preserves_tensor_intent = layer_preserves_tensor_intent
b._extra_args = ["--split-mode", "layer"]
b._requested_spec_mode = "auto"
b._chat_template_override = None
b._gguf_path = None
return b
def test_tensor_off_echo_preserves_multi_gpu_fallback():
"""The Studio UI always sends tensor_parallel and echoes the /load response's
resolved value, so after a fallback a ctx/settings reload carries tensor_parallel=
false even though the user never changed it. That echo must NOT collapse the
preserved multi-GPU placement -- it dedupes (Codex #6659)."""
from models.inference import LoadRequest
inference_routes = _load_inference_routes_module()
req = LoadRequest(model_path = "owner/repo", tensor_parallel = False)
assert "tensor_parallel" in req.model_fields_set, "the UI always sends the field"
# Preserved fallback + bare tensor=false echo: dedupe, keep multi-GPU (no collapse).
assert (
inference_routes._request_matches_loaded_settings(
req, _fallback_loaded_backend(layer_preserves_tensor_intent = True)
)
is True
)
# A genuine layer load (no preserved intent): tensor-off also dedupes, no churn.
assert (
inference_routes._request_matches_loaded_settings(
req, _fallback_loaded_backend(layer_preserves_tensor_intent = False)
)
is True
)
def test_explicit_split_mode_layer_extras_reloads_after_multi_gpu_fallback():
"""Tensor intent can be dropped via extras too: an explicit --split-mode layer
matches the stored fallback extras but must still reload (reviewer.py P1, #6659)."""
from models.inference import LoadRequest
inference_routes = _load_inference_routes_module()
req = LoadRequest(model_path = "owner/repo", llama_extra_args = ["--split-mode", "layer"])
assert "llama_extra_args" in req.model_fields_set
assert (
inference_routes._request_matches_loaded_settings(
req, _fallback_loaded_backend(layer_preserves_tensor_intent = True)
)
is False
)
def test_tensor_off_reload_requires_explicit_toggle():
"""An Apply that doesn't touch the toggle (e.g. a context change) isn't churned
by the preserved-fallback reload -- the working server is kept (Codex #6659)."""
from models.inference import LoadRequest
inference_routes = _load_inference_routes_module()
req = LoadRequest(model_path = "owner/repo") # tensor_parallel left unset
assert "tensor_parallel" not in req.model_fields_set
assert (
inference_routes._request_matches_loaded_settings(
req, _fallback_loaded_backend(layer_preserves_tensor_intent = True)
)
is True
)
def test_tensor_off_under_env_tensor_does_not_reload_loop(monkeypatch):
"""With LLAMA_ARG_SPLIT_MODE=tensor set, a tensor-off request can't drop tensor
intent, so the env-aware guard dedupes instead of reload-looping (Codex #6659)."""
from models.inference import LoadRequest
inference_routes = _load_inference_routes_module()
monkeypatch.setenv("LLAMA_ARG_SPLIT_MODE", "tensor")
req = LoadRequest(model_path = "owner/repo", tensor_parallel = False)
assert "tensor_parallel" in req.model_fields_set
# env still forces tensor -> not a real drop -> dedupe (no reload loop).
assert (
inference_routes._request_matches_loaded_settings(
req, _fallback_loaded_backend(layer_preserves_tensor_intent = True)
)
is True
)
def test_is_explicit_tensor_drop_truth_table():
"""Only an explicit non-tensor --split-mode override is a drop. A bare
tensor_parallel field (the UI always sends it and echoes the fallback's false), an
empty clear, an unrelated extra (--top-k), or inherit (None) must NOT collapse a
preserved fallback; --split-mode tensor / tensor_parallel=true re-engage (Codex
#6659)."""
from models.inference import LoadRequest
f = _load_inference_routes_module()._is_explicit_tensor_drop
# A non-tensor split-mode override is the one deliberate departure -> drop.
assert (
f(LoadRequest(model_path = "owner/repo", llama_extra_args = ["--split-mode", "layer"])) is True
)
# tensor / retry re-engages, never a drop.
assert (
f(LoadRequest(model_path = "owner/repo", llama_extra_args = ["--split-mode", "tensor"]))
is False
)
# A bare tensor_parallel field is the UI echo, not a drop (would collapse on reload).
assert f(LoadRequest(model_path = "owner/repo", tensor_parallel = False)) is False
assert f(LoadRequest(model_path = "owner/repo", tensor_parallel = True)) is False
# Unrelated extra / empty clear / inherit all keep the preserved placement.
assert f(LoadRequest(model_path = "owner/repo", llama_extra_args = ["--top-k", "20"])) is False
assert f(LoadRequest(model_path = "owner/repo", llama_extra_args = [])) is False
assert f(LoadRequest(model_path = "owner/repo")) is False
def test_explicit_tensor_drop_uses_shared_helper_in_both_readers():
"""Both the already-loaded dedup and the load carry-forward derive the drop from
_is_explicit_tensor_drop, so they agree on what counts as a drop -- a reload for
an unrelated extra still carries the preserved intent rather than collapsing to one
GPU (Codex #6659)."""
src = (Path(_BACKEND_DIR) / "routes" / "inference.py").read_text()
# Dedup reader (the preserved-fallback reload guard).
assert "layer_preserves_tensor_intent and _is_explicit_tensor_drop(request)" in src
# Load carry-forward reader feeds the same decision into the carry-forward.
assert "_explicit_tensor_drop = _is_explicit_tensor_drop(request)" in src
def test_layer_preserves_tensor_intent_set_only_on_preserved_downgrade():
"""load_model latches the flag from _layer_min_gpus (raised only when a tensor
request is downgraded but kept multi-GPU), and clears it when tensor stays on."""
src = inspect.getsource(LlamaCppBackend.load_model)
on = src.find("self._tensor_parallel = True")
off = src.find("self._tensor_parallel = False")
assert 0 <= on and 0 <= off
assert "self._layer_preserves_tensor_intent = False" in src[on : on + 120]
assert "self._layer_preserves_tensor_intent = _layer_min_gpus > 1" in src[off : off + 400]
def test_layer_min_gpus_bound_before_gpu_selection_try():
"""_layer_min_gpus is bound before the GPU-selection try, so the --fit-on except
path can't UnboundLocalError when the command builder reads it (Codex #6659)."""
src = inspect.getsource(LlamaCppBackend.load_model)
assert src.count("_layer_min_gpus = 1") == 1, "exactly one init, before the try"
init = src.find("_layer_min_gpus = 1")
try_body = src.find("gguf_size = self._get_gguf_size_bytes")
fit_except = src.find("GPU selection failed")
use_after = src.find("self._layer_preserves_tensor_intent = _layer_min_gpus > 1")
assert (
-1 < init < try_body < fit_except < use_after
), "the init must precede the try body, the except, and the command-builder use"
def test_already_in_target_state_reloads_on_tensor_off_after_fallback():
"""The backend fast path mirrors the route dedup: a preserved fallback reloads on
an EXPLICIT tensor-off request, but an implicit same-settings reload (carry-forward
preserve_multi_gpu_on_layer=True) still dedupes (Codex #6659)."""
def _backend(layer_preserves: bool) -> LlamaCppBackend:
b = _fallback_loaded_backend(layer_preserves_tensor_intent = layer_preserves)
b._process = _NoopProcess()
b._healthy = True
return b
kwargs = dict(
gguf_path = None,
mtp_draft_path = None,
model_identifier = "owner/repo",
hf_variant = None,
n_ctx = 0,
cache_type_kv = None,
speculative_type = None,
spec_draft_n_max = None,
tensor_parallel = False,
chat_template_override = None,
extra_args = ["--split-mode", "layer"],
is_vision = False,
)
# Preserved fallback + EXPLICIT tensor drop -> reload (not already in target state).
assert _backend(True)._already_in_target_state(**kwargs) is False
# Same preserved fallback but an implicit reload that carries the intent forward
# (HF auto-pick / local-dir flows skip the route guard and reach here) -> dedupe.
assert (
_backend(True)._already_in_target_state(**kwargs, preserve_multi_gpu_on_layer = True) is True
)
# A genuine layer load (no preserved intent) -> dedupe, no churn.
assert _backend(False)._already_in_target_state(**kwargs) is True

103
studio/backend/tests/test_training_runs.py Normal file
View file

@ -0,0 +1,103 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import json
from storage.studio_db import _extract_project_name_from_config_json
from utils.training_runs import (
build_default_output_dir_name,
model_segment_from_default_output_dir_name,
normalize_project_name,
slugify_project_name,
)
def test_normalize_project_name_trims_and_collapses_whitespace():
assert normalize_project_name(" Customer Support LoRA ") == "Customer Support LoRA"
def test_normalize_project_name_returns_none_for_empty_or_invalid_values():
assert normalize_project_name(" ") is None
assert normalize_project_name(None) is None
def test_slugify_project_name_makes_safe_suffix():
assert slugify_project_name("Customer Support / LoRA v2") == "customer-support-lora-v2"
def test_slugify_project_name_rejects_path_only_or_separator_only_values():
assert slugify_project_name("..") is None
assert slugify_project_name("///") is None
def test_build_default_output_dir_name_appends_project_slug():
output_dir = build_default_output_dir_name(
"unsloth/Llama-3.2-3B-Instruct",
"Customer Support",
timestamp = 1771227800,
)
assert output_dir == "unsloth_Llama-3.2-3B-Instruct__project-customer-support_1771227800"
def test_build_default_output_dir_name_caps_final_component(tmp_path):
output_dir = build_default_output_dir_name(
"a" * 240,
"b" * 80,
timestamp = 1771227800,
)
assert len(output_dir.encode()) <= 255
(tmp_path / output_dir).mkdir()
def test_build_default_output_dir_name_skips_invalid_project_slug():
output_dir = build_default_output_dir_name(
"unsloth/Llama-3.2-3B-Instruct",
"..",
timestamp = 1771227800,
)
assert output_dir == "unsloth_Llama-3.2-3B-Instruct_1771227800"
def test_model_segment_from_default_output_dir_name_strips_project_slug():
assert (
model_segment_from_default_output_dir_name(
"unsloth_Llama-3.2-3B-Instruct__project-customer-support_1771227800"
)
== "unsloth_Llama-3.2-3B-Instruct"
)
def test_model_segment_preserves_project_marker_text_in_model_name():
output_dir = build_default_output_dir_name(
"org/foo__project-bar",
timestamp = 1771227800,
)
assert output_dir == "org_foo__project--bar_1771227800"
assert model_segment_from_default_output_dir_name(output_dir) == "org_foo__project-bar"
def test_model_segment_strips_project_slug_after_escaped_model_marker():
output_dir = build_default_output_dir_name(
"org/foo__project-bar",
"Customer Support",
timestamp = 1771227800,
)
assert output_dir == "org_foo__project--bar__project-customer-support_1771227800"
assert model_segment_from_default_output_dir_name(output_dir) == "org_foo__project-bar"
def test_extract_project_name_from_config_json_returns_normalized_name():
config_json = json.dumps({"project_name": " Sales Assistant "})
assert _extract_project_name_from_config_json(config_json) == "Sales Assistant"
def test_extract_project_name_from_config_json_handles_missing_or_invalid_payload():
assert _extract_project_name_from_config_json(None) is None
assert _extract_project_name_from_config_json("not-json") is None
assert _extract_project_name_from_config_json(json.dumps({"project_name": " "})) is None

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

@ -195,6 +195,16 @@ def test_hf_dataset_rejects_unsafe_values(bad_hf_dataset):
)
def test_project_name_rejects_values_over_ui_limit():
with pytest.raises(ValidationError):
TrainingStartRequest(
model_name = "unsloth/test",
project_name = "x" * 81,
training_type = "LoRA/QLoRA",
format_type = "alpaca",
)
# --- Start-route streaming compatibility guards ---

101
studio/backend/utils/models/checkpoints.py
View file

@ -9,6 +9,12 @@ import structlog
from loggers import get_logger
from pathlib import Path
from typing import List, Optional, Tuple
from storage.studio_db import get_connection
from utils.training_runs import (
build_default_output_dir_name,
extract_project_name,
model_segment_from_default_output_dir_name,
)
from utils.paths import outputs_root, resolve_output_dir
logger = get_logger(__name__)
@ -30,6 +36,93 @@ def _checkpoint_sort_key(checkpoint_path: Path) -> tuple[int, int, str]:
return (1, 0, str(checkpoint_path))
def _infer_base_model_from_history(checkpoint_dir: Path) -> Optional[str]:
"""Best-effort base-model lookup using persisted Studio run metadata."""
checkpoint_name = checkpoint_dir.name
resolved_checkpoint_dir = str(checkpoint_dir.resolve())
try:
conn = get_connection()
except Exception:
return None
try:
exact_rows = conn.execute(
"""
SELECT model_name
FROM training_runs
WHERE output_dir IN (?, ?)
ORDER BY started_at DESC
""",
(
resolved_checkpoint_dir,
str(checkpoint_dir),
),
).fetchall()
for row in exact_rows:
model_name = row["model_name"]
if model_name:
return model_name
suffix_rows = conn.execute(
"""
SELECT model_name, output_dir
FROM training_runs
WHERE output_dir IS NOT NULL
ORDER BY started_at DESC
"""
).fetchall()
for row in suffix_rows:
output_dir = str(row["output_dir"] or "").rstrip("/\\")
if not (
output_dir.endswith(f"/{checkpoint_name}")
or output_dir.endswith(f"\\{checkpoint_name}")
):
continue
model_name = row["model_name"]
if model_name:
return model_name
parts = checkpoint_name.rsplit("_", 1)
if len(parts) != 2 or not parts[1].isdigit():
return None
timestamp = int(parts[1])
generated_rows = conn.execute(
"""
SELECT model_name, config_json
FROM training_runs
ORDER BY started_at DESC
"""
).fetchall()
for row in generated_rows:
model_name = row["model_name"]
if not model_name:
continue
project_name = None
config_json = row["config_json"]
if config_json:
try:
project_name = extract_project_name(json.loads(config_json))
except (TypeError, json.JSONDecodeError):
project_name = None
expected_dir_name = build_default_output_dir_name(
model_name,
project_name,
timestamp = timestamp,
)
if expected_dir_name == checkpoint_name:
return model_name
except Exception:
return None
finally:
conn.close()
return None
def _read_checkpoint_loss(checkpoint_path: Path) -> Optional[float]:
"""Read loss from the last log_history entry of trainer_state.json, or None."""
trainer_state = checkpoint_path / "trainer_state.json"
@ -106,9 +199,11 @@ def scan_checkpoints(
# Fallback: extract base model name from the folder name, e.g.
# "unsloth_Llama-3.2-3B-Instruct_1771227800" → "unsloth/Llama-3.2-3B-Instruct"
if not metadata.get("base_model"):
parts = item.name.rsplit("_", 1)
if len(parts) == 2 and parts[1].isdigit():
name_part = parts[0]
metadata["base_model"] = _infer_base_model_from_history(item)
if not metadata.get("base_model"):
name_part = model_segment_from_default_output_dir_name(item.name)
if name_part:
idx = name_part.find("_")
if idx > 0:
metadata["base_model"] = name_part[:idx] + "/" + name_part[idx + 1 :]

104
studio/backend/utils/training_runs.py Normal file
View file

@ -0,0 +1,104 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Helpers for naming and describing Studio training runs."""
from __future__ import annotations
import re
import time
from typing import Any, Optional
_INVALID_SEGMENT_CHARS = re.compile(r"[^A-Za-z0-9._-]+")
_MAX_RUN_DIR_NAME_CHARS = 255
_PROJECT_MARKER = "__project-"
_PROJECT_MARKER_ESCAPE = f"{_PROJECT_MARKER}-"
def _trim_segment(segment: str, max_chars: int) -> str:
if max_chars <= 0:
return ""
return segment[:max_chars].strip("._-")
def _escape_project_marker(segment: str) -> str:
return segment.replace(_PROJECT_MARKER, _PROJECT_MARKER_ESCAPE)
def _unescape_project_marker(segment: str) -> str:
return segment.replace(_PROJECT_MARKER_ESCAPE, _PROJECT_MARKER)
def _appended_project_marker_index(segment: str) -> int:
marker_index = segment.rfind(_PROJECT_MARKER)
while marker_index >= 0 and segment.startswith(_PROJECT_MARKER_ESCAPE, marker_index):
marker_index = segment.rfind(_PROJECT_MARKER, 0, marker_index)
return marker_index
def normalize_project_name(project_name: Any) -> Optional[str]:
"""Return a trimmed project name, or None when empty/invalid."""
if not isinstance(project_name, str):
return None
normalized = " ".join(project_name.strip().split())
return normalized or None
def slugify_project_name(project_name: Any) -> Optional[str]:
"""Convert a project name into a filesystem-safe suffix."""
normalized = normalize_project_name(project_name)
if normalized is None:
return None
slug = _INVALID_SEGMENT_CHARS.sub("-", normalized).strip("-._")
if not slug:
return None
return slug.lower()
def build_default_output_dir_name(
model_name: str,
project_name: Any = None,
*,
timestamp: Optional[int] = None,
) -> str:
"""Build the default training output folder name."""
from utils.paths import default_run_dir_name
timestamp_part = str(int(time.time() if timestamp is None else timestamp))
timestamp_suffix = f"_{timestamp_part}"
model_segment = _escape_project_marker(default_run_dir_name(model_name))
project_slug = slugify_project_name(project_name)
if not project_slug:
max_model_chars = _MAX_RUN_DIR_NAME_CHARS - len(timestamp_suffix)
model_segment = _trim_segment(model_segment, max_model_chars) or "model"
return f"{model_segment}{timestamp_suffix}"
max_project_chars = (
_MAX_RUN_DIR_NAME_CHARS - len("model") - len(_PROJECT_MARKER) - len(timestamp_suffix)
)
project_slug = _trim_segment(project_slug, max_project_chars) or "project"
project_suffix = f"{_PROJECT_MARKER}{project_slug}{timestamp_suffix}"
max_model_chars = _MAX_RUN_DIR_NAME_CHARS - len(project_suffix)
model_segment = _trim_segment(model_segment, max_model_chars) or "model"
return f"{model_segment}{project_suffix}"
def model_segment_from_default_output_dir_name(output_dir_name: str) -> Optional[str]:
"""Return the encoded model segment from a default run folder name."""
parts = str(output_dir_name or "").rsplit("_", 1)
if len(parts) != 2 or not parts[1].isdigit():
return None
model_segment = parts[0]
marker_index = _appended_project_marker_index(model_segment)
if marker_index >= 0:
model_segment = model_segment[:marker_index]
model_segment = _unescape_project_marker(model_segment)
return model_segment or None
def extract_project_name(config: Any) -> Optional[str]:
"""Read and normalize a project name from a stored config dict."""
if not isinstance(config, dict):
return None
return normalize_project_name(config.get("project_name"))

35
studio/frontend/package-lock.json generated
View file

@ -1704,6 +1704,7 @@
"os": [
"android"
],
"peer": true,
"engines": {
"node": ">= 10"
},
@ -1724,6 +1725,7 @@
"os": [
"darwin"
],
"peer": true,
"engines": {
"node": ">= 10"
},
@ -1744,6 +1746,7 @@
"os": [
"darwin"
],
"peer": true,
"engines": {
"node": ">= 10"
},
@ -1764,6 +1767,7 @@
"os": [
"linux"
],
"peer": true,
"engines": {
"node": ">= 10"
},
@ -1784,6 +1788,7 @@
"os": [
"linux"
],
"peer": true,
"engines": {
"node": ">= 10"
},
@ -1804,6 +1809,7 @@
"os": [
"linux"
],
"peer": true,
"engines": {
"node": ">= 10"
},
@ -1824,6 +1830,7 @@
"os": [
"linux"
],
"peer": true,
"engines": {
"node": ">= 10"
},
@ -1844,6 +1851,7 @@
"os": [
"linux"
],
"peer": true,
"engines": {
"node": ">= 10"
},
@ -1864,6 +1872,7 @@
"os": [
"linux"
],
"peer": true,
"engines": {
"node": ">= 10"
},
@ -1884,6 +1893,7 @@
"os": [
"win32"
],
"peer": true,
"engines": {
"node": ">= 10"
},
@ -1904,6 +1914,7 @@
"os": [
"win32"
],
"peer": true,
"engines": {
"node": ">= 10"
},
@ -5669,9 +5680,6 @@
"cpu": [
"arm64"
],
"libc": [
"glibc"
],
"license": "MIT",
"optional": true,
"os": [
@ -5688,9 +5696,6 @@
"cpu": [
"arm64"
],
"libc": [
"musl"
],
"license": "MIT",
"optional": true,
"os": [
@ -5707,9 +5712,6 @@
"cpu": [
"ppc64"
],
"libc": [
"glibc"
],
"license": "MIT",
"optional": true,
"os": [
@ -5726,9 +5728,6 @@
"cpu": [
"s390x"
],
"libc": [
"glibc"
],
"license": "MIT",
"optional": true,
"os": [
@ -5745,9 +5744,6 @@
"cpu": [
"x64"
],
"libc": [
"glibc"
],
"license": "MIT",
"optional": true,
"os": [
@ -5764,9 +5760,6 @@
"cpu": [
"x64"
],
"libc": [
"musl"
],
"license": "MIT",
"optional": true,
"os": [
@ -10282,9 +10275,9 @@
}
},
"node_modules/hono": {
"version": "4.12.21",
"resolved": "https://registry.npmjs.org/hono/-/hono-4.12.21.tgz",
"integrity": "sha512-uV63apnb0kyPtAUwoWgaGh9HyIFcv8lgmzPZSiTBQAFOFGIzka5EZ1dZocmGnn0XdX0+XTqJ6Tqv7selMuGLRQ==",
"version": "4.12.25",
"resolved": "https://registry.npmjs.org/hono/-/hono-4.12.25.tgz",
"integrity": "sha512-2NFaIyNVgJmBs/ecmtGzlmluTFs5cHEWGTdu0t1HBwYzoGXOL5nUQBRMXsXWla5i4KkG//QMzVP88m1+I3fdAQ==",
"license": "MIT",
"engines": {
"node": ">=16.9.0"

2
studio/frontend/package.json
View file

@ -86,7 +86,7 @@
"@tanstack/router-core": "1.169.2",
"@tanstack/history": "1.161.6",
"mermaid": "11.15.0",
"hono": "4.12.21",
"hono": "4.12.25",
"qs": "6.15.2",
"ip-address": "10.1.1",
"brace-expansion@5.0.5": "5.0.6"

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

@ -354,12 +354,11 @@ function TauriWrapper({ children }: { children: ReactNode }) {
);
}
const showApp = status === "running" && desktopAuthReady;
const showApp = status === "running";
const desktopBooting = status === "running" && !desktopAuthReady;
const showInteractiveApp = showApp && desktopAuthReady;
const startupStatus = status === "running" ? "starting" : status;
const startupProgressDetail =
status === "running" && !desktopAuthReady
? "Signing in to desktop session..."
: progressDetail;
const startupProgressDetail = progressDetail;
const usesCustomTitlebar = shouldUseCustomWindowTitlebar();
const usesNativeMacTitlebar = shouldUseNativeMacWindowTitlebar();
const hidesTitlebarSidebar = HIDDEN_TITLEBAR_SIDEBAR_ROUTES.has(pathname);
@ -369,12 +368,23 @@ function TauriWrapper({ children }: { children: ReactNode }) {
<TauriUpdateLayer isExternalServer={isExternalServer}>
<LlamaUpdateBanner
positioned={false}
enabled={!hidesTitlebarSidebar}
enabled={showInteractiveApp && !hidesTitlebarSidebar}
/>
<DownloadManagerPanel positioned={false} />
{showInteractiveApp ? <DownloadManagerPanel positioned={false} /> : null}
</TauriUpdateLayer>
<NativeIntentDrain />
{children}
{showInteractiveApp ? <NativeIntentDrain /> : null}
{showInteractiveApp ? children : null}
{desktopBooting ? (
<div className="pointer-events-none fixed inset-x-0 bottom-5 z-[9999] flex justify-center px-4">
<div className="absolute inset-x-4 bottom-16 mx-auto flex max-w-[520px] flex-col items-center gap-2 rounded-2xl border border-border/70 bg-background/95 px-6 py-5 text-center shadow-xl">
<div className="font-medium text-sm">Preparing Studio</div>
<div className="text-muted-foreground text-xs">The local backend is ready. Signing in to your desktop session before loading chats.</div>
</div>
<div className="rounded-full border border-border/70 bg-background/95 px-4 py-2 text-xs text-muted-foreground shadow-lg">
Signing in to desktop session...
</div>
</div>
) : null}
</>
) : (
<StartupScreen

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

@ -123,6 +123,7 @@ import {
deleteTrainingRun,
emitTrainingRunDeleted,
emitTrainingRunUpdated,
getTrainingRunDisplayTitle,
removeTrainingUnloadGuard,
renameTrainingRun,
useTrainingCompletionWatch,
@ -261,19 +262,6 @@ function NavItem({
);
}
// TEMP DEV override: preview the update card on installs with no real update
// (e.g. an editable checkout). In the browser console run
// `localStorage.setItem("unsloth_force_update_card", "1")` and reload. Remove
// before merge.
function devForceUpdateCard(): boolean {
if (typeof window === "undefined") return false;
try {
return window.localStorage.getItem("unsloth_force_update_card") === "1";
} catch {
return false;
}
}
export function AppSidebar() {
const t = useT();
const { isDark, toggleTheme, anchorRef } = useAnimatedThemeToggle();
@ -290,13 +278,10 @@ export function AppSidebar() {
// Web update detection: `webUpdate` is non-null only when the installed
// (PyPI) version is behind the latest release, so the card is hidden by
// default. `forceUpdateCard` is a TEMP dev override to preview it on installs
// with no real update (e.g. an editable checkout); remove before merge.
// default.
const { status: webUpdate } = useWebUpdateCheck();
const [forceUpdateCard] = useState(devForceUpdateCard);
const showUpdateCard = Boolean(webUpdate) || forceUpdateCard;
const updateVersion =
webUpdate?.latestVersion ?? (forceUpdateCard ? "0.0.0" : null);
const showUpdateCard = Boolean(webUpdate);
const updateVersion = webUpdate?.latestVersion ?? null;
// Auto-close mobile Sheet after navigation
const closeMobileIfOpen = () => {
@ -592,7 +577,7 @@ export function AppSidebar() {
setRenamingTarget({ kind: "chat", item, current: item.title });
}
function openRenameRun(run: TrainingRunSummary) {
const current = run.display_name ?? run.model_name;
const current = getTrainingRunDisplayTitle(run);
setRenameDraft(current);
setRenamingTarget({ kind: "run", run, current });
}
@ -1377,7 +1362,7 @@ export function AppSidebar() {
aria-hidden
/>
<span className="truncate">
{run.display_name ?? run.model_name}
{getTrainingRunDisplayTitle(run)}
</span>
<span className="ml-auto mr-0.5 shrink-0 text-[10px] text-muted-foreground">
{formatRelativeShort(run.started_at)}
@ -1653,8 +1638,7 @@ export function AppSidebar() {
renderEmphasizedTranslation(
t,
"shell.dialog.deleteRun.description",
confirmingDelete.run.display_name ??
confirmingDelete.run.model_name,
getTrainingRunDisplayTitle(confirmingDelete.run),
)
) : confirmingDelete?.kind === "chat" ? (
renderEmphasizedTranslation(

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

@ -1114,6 +1114,17 @@ function localModelIsGguf(m: LocalModelInfo): boolean {
);
}
function localPathTooltip(name: string, path: string): ReactNode {
return (
<>
<span className="block break-words">{name}</span>
<span className="block mt-1 text-[10px] text-muted-foreground break-all">
{path}
</span>
</>
);
}
/** Whether a local model is an MLX build (name hint). MLX runs on Mac only, so
* callers gate visibility on the host being a Mac. */
function localModelIsMlx(m: LocalModelInfo): boolean {
@ -2928,6 +2939,10 @@ export function HubModelPicker({
<ModelRow
label={m.model_id ?? m.display_name}
meta={isGguf ? "GGUF" : "Local"}
tooltipText={localPathTooltip(
m.model_id ?? m.display_name,
m.path,
)}
selected={value === m.id}
optionProps={hubModelList.getOptionProps(
optionKey,
@ -3017,6 +3032,10 @@ export function HubModelPicker({
<ModelRow
label={m.model_id ?? m.display_name}
meta={isGguf ? "GGUF" : "Local"}
tooltipText={localPathTooltip(
m.model_id ?? m.display_name,
m.path,
)}
selected={value === m.id}
optionProps={hubModelList.getOptionProps(
optionKey,
@ -3098,6 +3117,10 @@ export function HubModelPicker({
<ModelRow
label={m.model_id ?? m.display_name}
meta={isGguf ? "GGUF" : "Local"}
tooltipText={localPathTooltip(
m.model_id ?? m.display_name,
m.path,
)}
selected={value === m.id}
optionProps={hubModelList.getOptionProps(
optionKey,

39
studio/frontend/src/components/assistant-ui/model-selector/recommended-fit.ts
View file

@ -64,38 +64,19 @@ export function matchesFormatFilter(
}
}
// Model-size extraction from repo id, matching the backend's 3-regex priority:
// active params (MoE "A3B") > effective params (Gemma "E4B") > total ("8B").
// Bounded by separators so we never read "16" from "bf16" or "2" from "Kimi-K2".
// Examples: "Qwen3.5-35B-A3B" -> 3, "gemma-4-E4B" -> 4, "Llama-3-8B" -> 8.
const ACTIVE_PARAM_RE = /(?:^|[-_/. ])a(\d+(?:\.\d+)?)\s*[bB](?=$|[-_/. ])/i;
const EFFECTIVE_PARAM_RE = /(?:^|[-_/. ])e(\d+(?:\.\d+)?)\s*[bB](?=$|[-_/. ])/i;
const TOTAL_PARAM_RE = /(?:^|[-_/. ])(\d+(?:\.\d+)?)\s*[bB](?=$|[-_/. ])/;
function paramsFromMatch(match: RegExpExecArray | null): number | undefined {
if (!match) return undefined;
const billions = parseFloat(match[1]);
return Number.isFinite(billions) && billions > 0
? billions * 1e9
: undefined;
}
/** Active/effective parameter count parsed from a repo id, if it uses explicit
* MoE/Gemma-style notation such as A3B or E4B. */
export function activeOrEffectiveParamsFromId(id: string): number | undefined {
return (
paramsFromMatch(ACTIVE_PARAM_RE.exec(id)) ??
paramsFromMatch(EFFECTIVE_PARAM_RE.exec(id))
);
}
// First "<n>B" token in a repo id, e.g. "Qwen3-4B-GGUF" -> 4, "gpt-oss-20b" ->
// 20, "Qwen3-30B-A3B" -> 30 (MoE total), "gemma-4-E4B" -> 4 (effective-param
// "E" series). The digits must be bounded by a separator so we never read "16"
// from "bf16" or the "2" in "Kimi-K2".
const PARAM_RE = /(?:^|[-_/. ])[eE]?(\d+(?:\.\d+)?)\s*[bB](?=$|[-_./ ])/;
/** Parameter count (absolute, e.g. 4e9) parsed from a repo id, or undefined
* when the id has no size token (so callers can treat the size as unknown).
* Prefers MoE active-param notation (A3B) over effective (E4B) over total. */
* when the id has no size token (so callers can treat the size as unknown). */
export function paramsFromId(id: string): number | undefined {
return (
activeOrEffectiveParamsFromId(id) ?? paramsFromMatch(TOTAL_PARAM_RE.exec(id))
);
const match = PARAM_RE.exec(id);
if (!match) return undefined;
const billions = parseFloat(match[1]);
return Number.isFinite(billions) && billions > 0 ? billions * 1e9 : undefined;
}
// Smallest practical GGUF/MLX quant (~Q2_K, low-bit). The fit check asks whether

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

@ -47,6 +47,7 @@ import {
useChatRuntimeStore,
} from "../stores/chat-runtime-store";
import { useExternalProvidersStore } from "../stores/external-providers-store";
import type { ModelType } from "../types";
import { isMultimodalResponse } from "../types/api";
import type {
GgufVariantDetail,
@ -142,6 +143,11 @@ interface ServerTimings {
type RunMessages = Parameters<ChatModelAdapter["run"]>[0]["messages"];
type RunMessage = RunMessages[number];
type OpenAIStreamAdapterOptions = {
modelType?: ModelType;
pairId?: string;
};
/** Tracks which user messages were sent with an audio file (messageId → filename). */
export const sentAudioNames = new Map<string, string>();
@ -1182,7 +1188,17 @@ export function findLatestUserAudioBase64(
async function resolveUseAdapter(
threadId: string | undefined,
options: OpenAIStreamAdapterOptions = {},
): Promise<boolean | undefined> {
if (options.modelType === "model1" || options.modelType === "model2") {
return undefined;
}
if (
options.pairId &&
(options.modelType === "base" || options.modelType === "lora")
) {
return options.modelType === "lora";
}
if (!threadId) {
return undefined;
}
@ -1629,7 +1645,9 @@ async function autoLoadSmallestModel(): Promise<{
}
}
export function createOpenAIStreamAdapter(): ChatModelAdapter {
export function createOpenAIStreamAdapter(
options: OpenAIStreamAdapterOptions = {},
): ChatModelAdapter {
return {
async *run({ messages, abortSignal, unstable_threadId }) {
await useChatRuntimeStore.getState().hydratePersistedSettings();
@ -2076,7 +2094,7 @@ export function createOpenAIStreamAdapter(): ChatModelAdapter {
}
runtime.clearPendingAudio();
}
const useAdapter = await resolveUseAdapter(resolvedThreadId);
const useAdapter = await resolveUseAdapter(resolvedThreadId, options);
// ── Audio model path (non-streaming) ─────────────────────
const activeModel = runtime.models.find(

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

@ -606,7 +606,7 @@ const LoraCompareContent = memo(function LoraCompareContent({
handleName="lora"
borderClassName="border-t border-border/60 md:border-t-0 md:border-l"
header={
<div className="shrink-0 px-3 py-1.5 text-start md:text-end">
<div className="shrink-0 px-3 py-1.5 text-start md:text-end md:pr-[calc(4rem+var(--studio-chat-header-right-inset,var(--studio-window-control-inset,0px)))]">
<span className="text-[10px] font-semibold uppercase tracking-wider text-primary">
Fine-tuned
</span>
@ -654,7 +654,7 @@ function GeneralCompareHeader({
return (
<div
className={cn(
"pointer-events-none relative z-[65] flex h-[48px] shrink-0 items-start gap-2 bg-background pt-[var(--studio-chat-header-padding-top,11px)]",
"pointer-events-none relative z-40 flex h-[48px] shrink-0 items-start gap-2 bg-background pt-[var(--studio-chat-header-padding-top,11px)]",
side === "left"
? pinned
? "pl-12 pr-3 md:pl-2"
@ -2388,7 +2388,7 @@ export function ChatPage({
)}
<div
className={cn(
"pointer-events-none absolute top-[var(--studio-content-top-inset,0px)] left-0 right-[10px] z-[66] flex h-[var(--studio-chat-header-height,48px)] shrink-0 items-start bg-background pt-[var(--studio-chat-header-padding-top,11px)] pr-[calc(0.5rem+var(--studio-chat-header-right-inset,var(--studio-window-control-inset,0px)))]",
"pointer-events-none absolute top-[var(--studio-content-top-inset,0px)] left-0 right-[10px] z-40 flex h-[var(--studio-chat-header-height,48px)] shrink-0 items-start bg-background pt-[var(--studio-chat-header-padding-top,11px)] pr-[calc(0.5rem+var(--studio-chat-header-right-inset,var(--studio-window-control-inset,0px)))]",
isMobile
? "pl-12"
: pinned

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

@ -1045,16 +1045,17 @@ function useStudioRuntimeAdapters(
return adapters;
}
const chatAdapter = createOpenAIStreamAdapter();
function useRuntimeHook(
modelType: ModelType,
pairId?: string,
): ReturnType<typeof useLocalRuntime> {
const adapters = useStudioRuntimeAdapters(modelType, pairId);
const persistedChatAdapter = useMemo(
() => createPersistedRunAdapter(chatAdapter),
[],
() =>
createPersistedRunAdapter(
createOpenAIStreamAdapter({ modelType, pairId }),
),
[modelType, pairId],
);
return useLocalRuntime(persistedChatAdapter, { adapters });
}

19
studio/frontend/src/features/chat/utils/delete-thread-message.ts
View file

@ -112,7 +112,26 @@ export async function deleteThreadMessage(args: {
const exported = thread.export();
const repo = new MessageRepository();
repo.import(exported);
const target = exported.messages.find(
({ message }) => message.id === messageId,
);
const assistantReplyIds =
target?.message.role === "user"
? exported.messages
.filter(
({ parentId, message }) =>
parentId === messageId && message.role === "assistant",
)
.map(({ message }) => message.id)
: [];
// Delete the prompt first; that relinks its replies up to the prompt's parent
repo.deleteMessage(messageId);
for (const replyId of assistantReplyIds) {
repo.deleteMessage(replyId);
}
const next = repo.export();
if (remoteId) {
await syncExportedRepositoryToBackend(remoteId, next, {

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

@ -75,16 +75,35 @@ import { exportTourSteps } from "./tour";
const SEARCH_INPUT_REASONS = new Set(["input-change", "input-paste", "input-clear"]);
type SourceTab = "local" | "checkpoint" | "hf";
type SourceMode = "checkpoint" | "model";
function safePathSegment(
value: string | null | undefined,
fallback = "model",
maxLength = 250,
): string {
const safe = (value ?? "")
.replace(/[^a-zA-Z0-9._-]/g, "-")
.replace(/^[._-]+|[._-]+$/g, "")
.slice(0, maxLength)
.replace(/[._-]+$/g, "");
return safe || fallback;
}
function buildRelativeSaveDirectory(
exportMethod: ExportMethod | null,
sourceMode: SourceMode,
sourceBaseModelName: string,
selectedModelIdx: string | null,
checkpoint: string | null,
): string {
if (exportMethod === "gguf") {
return `${(sourceBaseModelName.split("/").pop() ?? selectedModelIdx ?? "model")
.replace(/[^a-zA-Z0-9._-]/g, "-")}-GGUF`;
const rawName =
sourceMode === "checkpoint"
? checkpoint ?? selectedModelIdx ?? sourceBaseModelName
: sourceBaseModelName;
return `${safePathSegment(rawName)}-GGUF`;
}
return `${selectedModelIdx ?? "model"}/${checkpoint}`;
}
@ -125,9 +144,7 @@ export function ExportPage() {
const [selectedModelIdx, setSelectedModelIdx] = useState<string | null>(null);
const [checkpoint, setCheckpoint] = useState<string | null>(null);
const [sourceMode, setSourceMode] = useState<"checkpoint" | "model">(
"checkpoint",
);
const [sourceMode, setSourceMode] = useState<SourceMode>("checkpoint");
const [modelSource, setModelSource] = useState<"hf" | "local">("hf");
const [modelInput, setModelInput] = useState("");
const [selectedSourceModel, setSelectedSourceModel] = useState<string | null>(
@ -449,6 +466,7 @@ export function ExportPage() {
const defaultSaveDirectory = useMemo(() => {
const relative = buildRelativeSaveDirectory(
exportMethod,
sourceMode,
sourceBaseModelName,
selectedModelIdx,
checkpoint,

26
studio/frontend/src/features/hub/catalog/model-card.tsx
View file

@ -11,7 +11,7 @@ import { ownerPaletteColor } from "@/features/hub/lib/avatar-theme";
import { buildAdaptiveCardAccentStyle } from "@/features/hub/lib/card-accent";
import { useDominantColor } from "@/features/hub/lib/use-dominant-color";
import { formatModelParamLabel } from "@/features/hub/lib/view-models";
import { cn, formatCompact } from "@/lib/utils";
import { formatCompact } from "@/lib/utils";
import { Download01Icon, FavouriteIcon } from "@hugeicons/core-free-icons";
import { type CSSProperties, memo, useMemo } from "react";
import type { DiscoverRow } from "../types";
@ -339,25 +339,11 @@ export const ModelCard = memo(function ModelCard({
value={formatCompact(row.result.likes)}
/>
</div>
<div className="flex shrink-0 items-center gap-1.5">
{row.fitLevel && (
<span
className={cn(
"text-[9px] font-medium px-1.5 py-0.5 rounded leading-none uppercase shrink-0",
row.fitLevel === "comfortable" && "text-emerald-700 bg-emerald-50 dark:text-emerald-300 dark:bg-emerald-500/15",
row.fitLevel === "fits" && "text-amber-700 bg-amber-50 dark:text-amber-300 dark:bg-amber-500/15",
row.fitLevel === "oom" && "text-red-700 bg-red-50 dark:text-red-300 dark:bg-red-500/15"
)}
>
{row.fitLevel === "comfortable" ? "Comfortable" : row.fitLevel === "fits" ? "Fits GPU" : "OOM"}
</span>
)}
{hasSize ? (
<span className="hub-chip shrink-0">{sizeLabel}</span>
) : topCapability ? (
<CapabilityPill capability={topCapability} iconOnly={true} />
) : null}
</div>
{hasSize ? (
<span className="hub-chip shrink-0">{sizeLabel}</span>
) : topCapability ? (
<CapabilityPill capability={topCapability} iconOnly={true} />
) : null}
</div>
</button>
);

46
studio/frontend/src/features/hub/catalog/models-table.tsx
View file

@ -570,22 +570,6 @@ export const ResultCard = memo(function ResultCard({
node: <span className="shrink-0">{sizeLabel}</span>,
});
}
if (row.fitLevel) {
const label = row.fitLevel === "comfortable" ? "Comfortable" : row.fitLevel === "fits" ? "Fits GPU" : "OOM";
const toneClass = row.fitLevel === "comfortable"
? "text-emerald-700 bg-emerald-50 dark:text-emerald-300 dark:bg-emerald-500/15"
: row.fitLevel === "fits"
? "text-amber-700 bg-amber-50 dark:text-amber-300 dark:bg-amber-500/15"
: "text-red-700 bg-red-50 dark:text-red-300 dark:bg-red-500/15";
textParts.push({
key: "gpuFit",
node: (
<span className={cn("shrink-0 uppercase text-[9px] font-medium px-1.5 py-0.5 rounded leading-none", toneClass)}>
{label}
</span>
),
});
}
if (row.result.updatedAt) {
textParts.push({
key: "updated",
@ -745,21 +729,6 @@ export const ResultGridRow = memo(function ResultGridRow({
</div>
<span className="mt-0.5 flex min-w-0 items-center gap-1 text-[11.5px] leading-[15px] text-muted-foreground/80">
<VerifiedOwner owner={row.owner} />
{row.fitLevel && (
<>
<span aria-hidden="true" className="text-muted-foreground/35">•</span>
<span
className={cn(
"text-[9px] font-semibold uppercase shrink-0",
row.fitLevel === "comfortable" && "text-emerald-600 dark:text-emerald-400",
row.fitLevel === "fits" && "text-amber-600 dark:text-amber-400",
row.fitLevel === "oom" && "text-red-500"
)}
>
{row.fitLevel === "comfortable" ? "Comfortable" : row.fitLevel === "fits" ? "Fits GPU" : "OOM"}
</span>
</>
)}
</span>
</div>
</div>
@ -882,21 +851,6 @@ export const ResultSplitRow = memo(function ResultSplitRow({
</div>
<span className="mt-0.5 flex min-w-0 items-center gap-1 text-[10.5px] leading-[14px] text-muted-foreground/80">
<VerifiedOwner owner={row.owner} />
{row.fitLevel && (
<>
<span aria-hidden="true" className="text-muted-foreground/35">•</span>
<span
className={cn(
"text-[9px] font-semibold uppercase shrink-0",
row.fitLevel === "comfortable" && "text-emerald-600 dark:text-emerald-400",
row.fitLevel === "fits" && "text-amber-600 dark:text-amber-400",
row.fitLevel === "oom" && "text-red-500"
)}
>
{row.fitLevel === "comfortable" ? "Comfortable" : row.fitLevel === "fits" ? "Fits GPU" : "OOM"}
</span>
</>
)}
</span>
</div>
<div className="flex shrink-0 flex-col items-end gap-0.5 text-[10.5px] tabular-nums text-muted-foreground/70">

24
studio/frontend/src/features/hub/catalog/models-toolbar.tsx
View file

@ -21,7 +21,6 @@ import { HugeiconsIcon } from "@hugeicons/react";
import type { HfSortKey } from "@/features/hub/hooks/use-hub-model-search";
import type {
CapabilityFilter,
GpuFitFilter,
ModelFormatFilter,
ModelsTab,
ResourceTypeFilter,
@ -29,7 +28,6 @@ import type {
import {
CAPABILITY_FILTER_OPTIONS,
FORMAT_FILTER_OPTIONS,
GPU_FIT_FILTER_OPTIONS,
} from "../lib/view-models";
import { HubOptionMenu, type HubOption } from "./hub-option-menu";
import {
@ -70,8 +68,6 @@ export const ModelsToolbar = memo(function ModelsToolbar({
onFormatFilterChange,
capabilityFilter,
onCapabilityFilterChange,
gpuFitFilter,
onGpuFitFilterChange,
onManageLocalFolders,
onOpenFineTune,
}: {
@ -88,8 +84,6 @@ export const ModelsToolbar = memo(function ModelsToolbar({
onFormatFilterChange: (value: ModelFormatFilter) => void;
capabilityFilter: CapabilityFilter;
onCapabilityFilterChange: (value: CapabilityFilter) => void;
gpuFitFilter: GpuFitFilter;
onGpuFitFilterChange: (value: GpuFitFilter) => void;
onManageLocalFolders: () => void;
/** Opens the curated "Fine-tune ready" channel (discover only). Exposed as a
* format-dropdown option rather than a standalone feed section. */
@ -159,14 +153,6 @@ export const ModelsToolbar = memo(function ModelsToolbar({
})),
[],
);
const gpuFitOptions = useMemo<HubOption<GpuFitFilter>[]>(
() =>
GPU_FIT_FILTER_OPTIONS.map((option) => ({
value: option.value,
label: option.label,
})),
[],
);
const sortOptions = useMemo<HubOption<HfSortKey>[]>(
() =>
SORT_OPTIONS.map((option) => ({
@ -357,16 +343,6 @@ export const ModelsToolbar = memo(function ModelsToolbar({
/>
)}
{tab === "discover" && !isDataset && (
<HubOptionMenu
value={gpuFitFilter}
options={gpuFitOptions}
onValueChange={onGpuFitFilterChange}
ariaLabel="GPU fit filter"
className={cn(triggerBase, "w-[128px]")}
/>
)}
{tab === "discover" && (
<HubOptionMenu
value={sortBy}

20
studio/frontend/src/features/hub/catalog/on-device-folders-dialog.tsx
View file

@ -180,7 +180,7 @@ export function OnDeviceFoldersDialog({
<>
<Dialog open={open} onOpenChange={onOpenChange}>
<DialogContent
className="max-w-[600px] gap-0 overflow-hidden p-0"
className="gap-0 overflow-hidden p-0 sm:max-w-[620px] lg:max-w-[660px] xl:max-w-[680px] [&_[data-slot=dialog-close]]:right-3 [&_[data-slot=dialog-close]]:top-3"
overlayClassName="bg-black/20 backdrop-blur-none"
>
<DialogHeader className="border-b border-border/60 px-5 py-4">
@ -319,7 +319,12 @@ export function OnDeviceFoldersDialog({
return (
<div
key={folder.id}
className="flex min-h-12 items-center gap-3 border-b border-border/50 px-3 py-2 last:border-b-0"
className={cn(
"grid min-h-12 w-full items-center gap-3 border-b border-border/50 px-3 py-2 last:border-b-0",
isTauri
? "grid-cols-[2rem_minmax(0,1fr)_2rem_2rem]"
: "grid-cols-[2rem_minmax(0,1fr)_2rem]",
)}
>
<div className="flex size-8 shrink-0 items-center justify-center rounded-[9px] bg-muted text-muted-foreground">
<HugeiconsIcon
@ -328,13 +333,18 @@ export function OnDeviceFoldersDialog({
className="size-4"
/>
</div>
<div className="min-w-0 flex-1">
<p className="truncate text-[12.5px] font-medium text-foreground">
<div className="min-w-0 overflow-hidden">
<p
className="block w-full truncate text-[12.5px] font-medium text-foreground"
title={pathTail(folder.path)}
>
{pathTail(folder.path)}
</p>
<Tooltip>
<TooltipTrigger asChild={true}>
<p className="truncate font-mono text-[10.5px] text-muted-foreground">
<p
className="block w-full truncate font-mono text-[10.5px] text-muted-foreground"
>
{folder.path}
</p>
</TooltipTrigger>

85
studio/frontend/src/features/hub/hub-page.tsx
View file

@ -85,17 +85,12 @@ import type {
CachedInventoryRow,
CapabilityFilter,
DiscoverRow,
GpuFitFilter,
LocalInventoryRow,
ModelFormatFilter,
ModelsTab,
ResourceTypeFilter,
SelectedModelView,
} from "./types";
import {
classifyGpuFit,
matchesGpuFitFilter,
} from "./lib/gpu-fit-filter";
const MODELS_TAB_STORAGE_KEY = "unsloth.hub.modelsTab";
const ALL_MODELS_VIEW_STORAGE_KEY = "unsloth.hub.allModelsView";
@ -418,7 +413,6 @@ export function ModelsPage() {
);
const [capabilityFilter, setCapabilityFilter] =
useState<CapabilityFilter>("all");
const [gpuFitFilter, setGpuFitFilter] = useState<GpuFitFilter>("all");
const [allModelsView, setAllModelsViewState] = useState<AllModelsView>(
readAllModelsViewPreference,
);
@ -567,7 +561,6 @@ export function ModelsPage() {
const apiHfToken = hfApiToken(debouncedHfToken);
const deferredFormatFilter = useDeferredValue(formatFilter);
const deferredCapabilityFilter = useDeferredValue(capabilityFilter);
const deferredGpuFitFilter = useDeferredValue(gpuFitFilter);
const hasQuery = deferredDebouncedQuery.trim() !== "";
const mode: DiscoverMode = !isModelDiscover
@ -697,59 +690,20 @@ export function ModelsPage() {
const discoverRows = isDatasetMode ? datasetDiscoverRows : modelDiscoverRows;
// Pre-compute GPU fit level for every discover row so filteredDiscoverRows
// and model cards can both consume the same classification.
const gpuFitLevelById = useMemo(() => {
const map = new Map<string, ReturnType<typeof classifyGpuFit>>();
for (const row of discoverRows) {
map.set(
row.id,
classifyGpuFit({
totalParams: row.result.totalParams,
estimatedSizeBytes: row.result.estimatedSizeBytes,
repoId: row.id,
gpu,
}),
);
}
return map;
}, [discoverRows, gpu]);
const addGpuFitLevel = useCallback(
(row: DiscoverRow): DiscoverRow => ({
...row,
fitLevel: classifyGpuFit({
totalParams: row.result.totalParams,
estimatedSizeBytes: row.result.estimatedSizeBytes,
repoId: row.id,
gpu,
}),
}),
[gpu],
);
const filteredDiscoverRows = useMemo(() => {
if (isDatasetMode) return discoverRows;
return discoverRows
.filter(
(row) =>
!isHiddenModelId(row.id) &&
matchesFormat(detectResultFormat(row.result), effectiveDiscoverFormat) &&
matchesCapability(row.capabilities, deferredCapabilityFilter) &&
matchesGpuFitFilter(gpuFitLevelById.get(row.id) ?? null, deferredGpuFitFilter) &&
(!activeChannel?.finetunableOnly || isUnslothFinetunable(row.result)),
)
.map((row) => ({
...row,
fitLevel: gpuFitLevelById.get(row.id) ?? null,
}));
return discoverRows.filter(
(row) =>
!isHiddenModelId(row.id) &&
matchesFormat(detectResultFormat(row.result), effectiveDiscoverFormat) &&
matchesCapability(row.capabilities, deferredCapabilityFilter) &&
(!activeChannel?.finetunableOnly || isUnslothFinetunable(row.result)),
);
}, [
discoverRows,
isDatasetMode,
effectiveDiscoverFormat,
deferredCapabilityFilter,
deferredGpuFitFilter,
gpuFitLevelById,
activeChannel,
]);
@ -770,17 +724,8 @@ export function ModelsPage() {
effectiveLocalRows,
)
.filter((row) => !isHiddenModelId(row.id))
.filter((row) => matchesFormat(row.result.isGguf, "gguf"))
.map(addGpuFitLevel)
.filter((row) =>
matchesGpuFitFilter(row.fitLevel ?? null, deferredGpuFitFilter),
),
[
hubFeed.trending.results,
modelDiscoveryInventorySignature,
addGpuFitLevel,
deferredGpuFitFilter,
],
.filter((row) => matchesFormat(row.result.isGguf, "gguf")),
[hubFeed.trending.results, modelDiscoveryInventorySignature],
);
const feedRows = useMemo(() => {
if (!isFeedMode) return [];
@ -897,7 +842,6 @@ export function ModelsPage() {
resourceType,
deferredFormatFilter,
deferredCapabilityFilter,
deferredGpuFitFilter,
effectiveSort,
effectiveDirection,
activeChannelId,
@ -908,7 +852,6 @@ export function ModelsPage() {
resourceType,
deferredFormatFilter,
deferredCapabilityFilter,
deferredGpuFitFilter,
effectiveSort,
effectiveDirection,
activeChannelId,
@ -929,7 +872,6 @@ export function ModelsPage() {
setDownloadedFormat("all");
}
setCapabilityFilter("all");
setGpuFitFilter("all");
}, [isDiscoverTab, urlSection, navigate]);
const handleDiscoverFetchIntent = useCallback(() => {
setDiscoverFetchIntent((value) => value + 1);
@ -1295,10 +1237,8 @@ export function ModelsPage() {
hasMore,
manualFetchAvailable: discoverManualFetchAvailable,
hasActiveFilters:
deferredGpuFitFilter !== "all" ||
(!isFeedMode &&
(deferredFormatFilter !== "all" ||
deferredCapabilityFilter !== "all")),
!isFeedMode &&
(deferredFormatFilter !== "all" || deferredCapabilityFilter !== "all"),
}),
[
tab,
@ -1324,7 +1264,6 @@ export function ModelsPage() {
discoverManualFetchAvailable,
deferredFormatFilter,
deferredCapabilityFilter,
deferredGpuFitFilter,
],
);
@ -1509,8 +1448,6 @@ export function ModelsPage() {
onFormatFilterChange={setFormatFilter}
capabilityFilter={capabilityFilter}
onCapabilityFilterChange={setCapabilityFilter}
gpuFitFilter={gpuFitFilter}
onGpuFitFilterChange={setGpuFitFilter}
onManageLocalFolders={handleManageLocalFolders}
onOpenFineTune={() => handleOpenList("finetune")}
/>

80
studio/frontend/src/features/hub/lib/gpu-fit-filter.ts
View file

@ -1,80 +0,0 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
// GPU-aware model-fit filtering: classifies whether a model fits the device
// and provides a filter predicate for the Hub page and model selector.
import type { GpuInfo } from "@/hooks/use-gpu-info";
import {
activeOrEffectiveParamsFromId,
estimateQuantBytes,
paramsFromId,
} from "@/components/assistant-ui/model-selector/recommended-fit";
/** The three filter states exposed in the toolbar dropdown. */
export type GpuFitFilter = "all" | "fits" | "comfortable";
/** Per-model fit classification. */
export type GpuFitLevel = "comfortable" | "fits" | "oom";
/**
* Classify whether a model fits the device.
*
* - "comfortable": estimated size ≤ 70% of GPU VRAM (runs fully in VRAM)
* - "fits": estimated size ≤ 70% GPU + 70% system RAM (runs with CPU offload)
* - "oom": exceeds both budgets
*
* Returns null when we can't determine the size (unknown → no badge).
*/
export function classifyGpuFit(opts: {
totalParams?: number;
estimatedSizeBytes?: number;
repoId: string;
gpu: GpuInfo;
}): GpuFitLevel | null {
const { totalParams, estimatedSizeBytes, repoId, gpu } = opts;
const gpuGb = gpu.memoryTotalGb;
const ramGb = gpu.systemRamAvailableGb;
if (gpuGb <= 0 && ramGb <= 0) return null; // no budget info
// Active/effective model tokens (for example MoE A3B) describe runnable size
// better than HF total-parameter metadata; otherwise prefer exact metadata.
const activeOrEffectiveParams = activeOrEffectiveParamsFromId(repoId);
const params = activeOrEffectiveParams ?? totalParams ?? paramsFromId(repoId);
const sizeBytes =
activeOrEffectiveParams
? estimateQuantBytes(activeOrEffectiveParams)
: estimatedSizeBytes ?? (params ? estimateQuantBytes(params) : undefined);
if (!sizeBytes || sizeBytes <= 0) return null; // can't determine
const sizeGb = sizeBytes / 1024 ** 3;
let comfortBudget: number;
let fitBudget: number;
if (!gpu.available || gpuGb <= 0) {
// Unified memory system (no discrete GPU)
comfortBudget = ramGb * 0.7;
fitBudget = ramGb * 0.7;
} else {
// Discrete GPU
comfortBudget = gpuGb * 0.7;
fitBudget = gpuGb * 0.7 + ramGb * 0.7;
}
if (sizeGb <= comfortBudget) return "comfortable";
if (sizeGb <= fitBudget) return "fits";
return "oom";
}
/** Whether a row passes the given GPU fit filter. */
export function matchesGpuFitFilter(
level: GpuFitLevel | null,
filter: GpuFitFilter,
): boolean {
if (filter === "all") return true;
if (level === null) return false;
if (filter === "comfortable") return level === "comfortable";
// "fits" shows both comfortable and fits
return level === "comfortable" || level === "fits";
}

10
studio/frontend/src/features/hub/lib/view-models.ts
View file

@ -10,7 +10,6 @@ import type {
import type {
CapabilityFilter,
DiscoverRow,
GpuFitFilter,
ModelFormatFilter,
} from "../types";
import {
@ -52,15 +51,6 @@ export const FORMAT_FILTER_OPTIONS: ReadonlyArray<{
{ value: "mlx", label: "MLX" },
];
export const GPU_FIT_FILTER_OPTIONS: ReadonlyArray<{
value: GpuFitFilter;
label: string;
}> = [
{ value: "all", label: "All sizes" },
{ value: "fits", label: "Fits GPU" },
{ value: "comfortable", label: "Comfortable" },
];
const BILLION = 1_000_000_000;
export function formatParamCount(totalParams: number | undefined): string {

4
studio/frontend/src/features/hub/types.ts
View file

@ -28,9 +28,6 @@ export type ModelFormatFilter = "all" | "gguf" | "checkpoint" | "mlx";
export type CapabilityFilter = "all" | CapabilityKey;
import type { GpuFitFilter, GpuFitLevel } from "./lib/gpu-fit-filter";
export type { GpuFitFilter, GpuFitLevel };
export interface DiscoverRow {
id: string;
owner: string;
@ -40,7 +37,6 @@ export interface DiscoverRow {
isPartialOnDevice: boolean;
summary: string;
capabilities: Capability[];
fitLevel?: GpuFitLevel | null;
}
export type SelectedResourceSource = "huggingface" | "hub_cache" | LocalSource;

25
studio/frontend/src/features/onboarding/components/steps/hyperparameters-step.tsx
View file

@ -7,6 +7,7 @@ import {
FieldLegend,
FieldSet,
} from "@/components/ui/field";
import { Input } from "@/components/ui/input";
import {
Select,
SelectContent,
@ -59,6 +60,8 @@ function stepLR(value: number, direction: 1 | -1): number {
export function HyperparametersStep() {
const {
trainingMethod,
projectName,
setProjectName,
maxSteps,
setMaxSteps,
epochs,
@ -79,6 +82,8 @@ export function HyperparametersStep() {
} = useTrainingConfigStore(
useShallow((s) => ({
trainingMethod: s.trainingMethod,
projectName: s.projectName,
setProjectName: s.setProjectName,
maxSteps: s.maxSteps,
setMaxSteps: s.setMaxSteps,
epochs: s.epochs,
@ -125,6 +130,26 @@ export function HyperparametersStep() {
<FieldSet>
<FieldLegend variant="label">Choose your training parameters</FieldLegend>
<div className="flex flex-col gap-4">
<div className="flex flex-col gap-2">
<div className="flex items-center justify-between">
<FieldLabel className="flex items-center gap-1.5 !text-sm text-muted-foreground">
Project Name
<span className="text-xs font-normal text-muted-foreground/70">
Optional
</span>
</FieldLabel>
</div>
<Input
value={projectName || ""}
onChange={(e) => setProjectName(e.target.value)}
placeholder="customer-support-lora"
maxLength={80}
/>
<p className="text-xs text-muted-foreground">
Used in training output folder names, export defaults, and history.
</p>
</div>
<div
key={useEpochs ? "epochs" : "steps"}
className="flex flex-col gap-2 animate-in fade-in-1 slide-in-from-bottom-1 duration-200"

4
studio/frontend/src/features/onboarding/components/steps/summary-step.tsx
View file

@ -50,6 +50,7 @@ export function SummaryStep() {
const {
modelType,
selectedModel,
projectName,
trainingMethod,
datasetSource,
datasetFormat,
@ -68,6 +69,7 @@ export function SummaryStep() {
({
modelType,
selectedModel,
projectName,
trainingMethod,
datasetSource,
datasetFormat,
@ -84,6 +86,7 @@ export function SummaryStep() {
}) => ({
modelType,
selectedModel,
projectName,
trainingMethod,
datasetSource,
datasetFormat,
@ -152,6 +155,7 @@ export function SummaryStep() {
<Separator className="my-2" />
<div className="space-y-1 text-sm">
<Row label="Type" value={modelType} capitalize />
<Row label="Project" value={projectName || "--"} />
<Row label="Method" value={trainingMethodLabel} />
</div>
</CardContent>

1
studio/frontend/src/features/studio/historical-training-view.tsx
View file

@ -78,6 +78,7 @@ function mapToViewData(
error: run.status === "error" ? run.error_message : null,
isTrainingRunning: false,
modelName: run.display_name ?? run.model_name,
projectName: run.project_name,
trainingMethod: parseBackendTrainingMethod(
detail.config?.training_type,
detail.config?.load_in_4bit,

26
studio/frontend/src/features/studio/history-card-grid.tsx
View file

@ -15,6 +15,8 @@ import { Button } from "@/components/ui/button";
import type { TrainingRunSummary } from "@/features/training";
import {
deleteTrainingRun,
getTrainingRunDisplayTitle,
getTrainingRunModelSubtitle,
emitTrainingRunDeleted,
listTrainingRuns,
onTrainingRunDeleted,
@ -387,6 +389,12 @@ export function HistoryCardGrid({
const isRunning = run.status === "running";
const canResume = run.can_resume && !wasContinued;
const isResuming = resumeTarget === run.id;
const title = getTrainingRunDisplayTitle(run);
const modelSubtitle = getTrainingRunModelSubtitle(run);
const projectSubtitle =
run.project_name && title !== run.project_name ? run.project_name : null;
// Backend /p ref + its capability token. Both are required: the link
// is useless (404s) without the signature, so don't offer to copy it.
const canCopyPreview = !!run.preview_ref && !!run.preview_sig;
@ -476,16 +484,16 @@ export function HistoryCardGrid({
<div className="min-w-0">
<p
className="truncate text-sm font-medium"
title={run.display_name ?? run.model_name}
title={title}
>
{run.display_name ?? run.model_name}
{title}
</p>
{run.display_name && (
{modelSubtitle && (
<p
className="truncate text-xs text-muted-foreground"
title={run.model_name}
title={modelSubtitle}
>
{run.model_name}
{modelSubtitle}
</p>
)}
<p
@ -494,6 +502,14 @@ export function HistoryCardGrid({
>
{run.dataset_name}
</p>
{projectSubtitle && (
<p
className="truncate text-xs text-muted-foreground/80"
title={projectSubtitle}
>
{projectSubtitle}
</p>
)}
</div>
{run.loss_sparkline && run.loss_sparkline.length >= 2 && (
<div className={cn((canResume || canCopyPreview) && "h-7 overflow-hidden")}>

11
studio/frontend/src/features/studio/live-training-view.tsx
View file

@ -33,6 +33,8 @@ export function LiveTrainingView(): ReactElement {
evalEnabled: state.evalEnabled,
outputDir: state.outputDir,
isTrainingRunning: state.isTrainingRunning,
startModelName: state.startModelName,
startProjectName: state.startProjectName,
lossHistory: state.lossHistory,
lrHistory: state.lrHistory,
gradNormHistory: state.gradNormHistory,
@ -45,10 +47,16 @@ export function LiveTrainingView(): ReactElement {
const config = useTrainingConfigStore(
useShallow((state) => ({
selectedModel: state.selectedModel,
projectName: state.projectName,
trainingMethod: state.trainingMethod,
})),
);
const activeProjectName =
runtime.startProjectName !== null
? runtime.startProjectName.trim() || null
: (config.projectName || "").trim() || null;
const viewData: TrainingViewData = {
phase: runtime.phase,
currentStep: runtime.currentStep,
@ -66,7 +74,8 @@ export function LiveTrainingView(): ReactElement {
message: runtime.message,
error: runtime.error,
isTrainingRunning: runtime.isTrainingRunning,
modelName: config.selectedModel ?? "",
modelName: runtime.startModelName ?? config.selectedModel ?? "",
projectName: activeProjectName,
trainingMethod: config.trainingMethod ?? "",
lossHistory: runtime.lossHistory,
lrHistory: runtime.lrHistory,

18
studio/frontend/src/features/studio/sections/params-section.tsx
View file

@ -229,6 +229,24 @@ export function ParamsSection(): ReactElement {
: "h-studio-config-column"} duration-150`}
>
<div className="flex flex-col gap-4">
<div className="flex flex-col gap-2">
<span className="flex items-center gap-1.5 text-xs font-medium text-muted-foreground">
{t("studio.params.projectName")}
<span className="text-[10px] font-normal text-muted-foreground/70">
{t("studio.params.optional")}
</span>
</span>
<Input
value={store.projectName || ""}
onChange={(event) => store.setProjectName(event.target.value)}
placeholder="customer-support-lora"
maxLength={80}
/>
<p className="text-[10px] text-muted-foreground">
{t("studio.params.projectNameDescription")}
</p>
</div>
{/* Max Steps / Epochs */}
<div className="flex flex-col gap-2">
<div

19
studio/frontend/src/features/studio/sections/progress-section.tsx
View file

@ -270,6 +270,11 @@ export function ProgressSection({
>
{t(phaseLabelKeys[data.phase])}
</span>
{data.projectName && (
<span className="rounded-full border border-border/60 px-2.5 py-1 text-[10px] font-medium text-foreground/80">
{data.projectName}
</span>
)}
<span className="text-[10px] tabular-nums text-muted-foreground">
{t("studio.progress.epoch", {
value: formatNumber(data.currentEpoch, 2),
@ -290,7 +295,7 @@ export function ProgressSection({
</span>
<span>{pct}%</span>
</div>
<Progress value={pct} className="h-2 bg-foreground/[0.05]" />
<Progress value={pct} className="h-2 bg-foreground/5" />
</div>
{!isHistorical && (
@ -307,7 +312,12 @@ export function ProgressSection({
</p>
)}
<div className="grid gap-x-4 gap-y-3 pt-1 sm:grid-cols-2 xl:grid-cols-5">
<div
className={cn(
"grid gap-x-4 gap-y-3 pt-1 sm:grid-cols-2",
data.projectName ? "xl:grid-cols-6" : "xl:grid-cols-5",
)}
>
<MetricStat
label={t("studio.progress.loss")}
valueClassName="text-2xl font-bold tracking-tight"
@ -318,6 +328,11 @@ export function ProgressSection({
<MetricStat label={t("studio.progress.gradNorm")}>
{formatNumber(stoppedGradNorm, 3)}
</MetricStat>
{data.projectName && (
<MetricStat label={t("studio.progress.project")} valueClassName="truncate">
{data.projectName}
</MetricStat>
)}
<MetricStat label={t("studio.progress.model")} valueClassName="truncate">
{data.modelName || "--"}
</MetricStat>

1
studio/frontend/src/features/training/api/mappers.ts
View file

@ -73,6 +73,7 @@ export function buildTrainingStartPayload(
return {
model_name: config.selectedModel ?? "",
project_name: (config.projectName || "").trim() || null,
training_type: toBackendTrainingType(config.trainingMethod),
hf_token: config.hfToken.trim() || null,
load_in_4bit: (adapterMethod && isQloraMethod) || (isCpt && isFourBitModel),

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

@ -60,6 +60,7 @@ export function useTrainingActions() {
config.selectedModel ?? null,
getHfDatasetName(config),
false,
config.projectName || "",
);
runtimeStore.setStarting(true);
@ -152,7 +153,12 @@ export function useTrainingActions() {
// Re-read config after potential store updates from dataset check
const payload = buildTrainingStartPayload(useTrainingConfigStore.getState());
runtimeStore.setStartResources(payload.model_name, payload.hf_dataset, false);
runtimeStore.setStartResources(
payload.model_name,
payload.hf_dataset,
false,
payload.project_name ?? "",
);
const response = await startTraining(payload);
if (response.status === "error") {
@ -196,7 +202,7 @@ export function useTrainingActions() {
const resumeTrainingRunFromHistory = useCallback(async (runId: string): Promise<boolean> => {
const runtimeStore = useTrainingRuntimeStore.getState();
runtimeStore.setStartError(null);
runtimeStore.setStartResources(null, null, true);
runtimeStore.setStartResources(null, null, true, null);
runtimeStore.setStarting(true);
try {
@ -220,7 +226,12 @@ export function useTrainingActions() {
resume_from_checkpoint: outputDir,
} as TrainingStartRequest;
runtimeStore.setStartResources(payload.model_name, payload.hf_dataset, true);
runtimeStore.setStartResources(
payload.model_name,
payload.hf_dataset,
true,
payload.project_name ?? "",
);
// Resume goes straight to startTraining, so it runs the same consent gate as a
// fresh start; otherwise a resumed custom-code run hits the worker block with no dialog.

5
studio/frontend/src/features/training/index.ts
View file

@ -7,6 +7,11 @@ export {
useTrainingRuntimeStore,
} from "./stores/training-runtime-store";
export { useTrainingActions } from "./hooks/use-training-actions";
export {
getTrainingRunDisplayTitle,
getTrainingRunModelSubtitle,
} from "./lib/run-display";
export { useTrainingHistorySidebarItems } from "./hooks/use-training-history-sidebar";
export { useTrainingRuntimeLifecycle } from "./hooks/use-training-runtime-lifecycle";
export { useTrainingCompletionWatch } from "./hooks/use-training-completion-watch";

24
studio/frontend/src/features/training/lib/run-display.ts Normal file
View file

@ -0,0 +1,24 @@
// 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 type { TrainingRunSummary } from "../types/history";
type TrainingRunTitleFields = Pick<
TrainingRunSummary,
"display_name" | "project_name" | "model_name"
>;
function nonEmpty(value: string | null | undefined): string | null {
const trimmed = value?.trim();
return trimmed ? trimmed : null;
}
export function getTrainingRunDisplayTitle(run: TrainingRunTitleFields): string {
return nonEmpty(run.display_name) ?? nonEmpty(run.project_name) ?? run.model_name;
}
export function getTrainingRunModelSubtitle(
run: TrainingRunTitleFields,
): string | null {
return getTrainingRunDisplayTitle(run) === run.model_name ? null : run.model_name;
}

2
studio/frontend/src/features/training/stores/training-config-store.ts
View file

@ -58,6 +58,7 @@ const initialState: TrainingConfigState = {
currentStep: MIN_STEP,
modelType: null,
selectedModel: null,
projectName: "",
trainingMethod: "qlora",
hfToken: "",
datasetSource: "huggingface",
@ -613,6 +614,7 @@ export const useTrainingConfigStore = create<TrainingConfigStore>()(
if (state.modelDefaultsAppliedFor === state.selectedModel) return;
void loadAndApplyModelDefaults(state.selectedModel);
},
setProjectName: (projectName) => set({ projectName }),
setTrainingMethod: (trainingMethod) => {
const state = get();
set(

9
studio/frontend/src/features/training/stores/training-runtime-store.ts
View file

@ -24,6 +24,7 @@ const initialState: TrainingRuntimeState = {
startError: null,
startModelName: null,
startDatasetName: null,
startProjectName: null,
startFromResume: false,
sseConnected: false,
firstStepReceived: false,
@ -125,8 +126,12 @@ export const useTrainingRuntimeStore = create<TrainingRuntimeStore>()((set) => (
setHasHydrated: (value) => set({ hasHydrated: value }),
setStarting: (value) => set({ isStarting: value }),
setStartError: (value) => set({ startError: value }),
setStartResources: (startModelName, startDatasetName, startFromResume = false) =>
set({ startModelName, startDatasetName, startFromResume }),
setStartResources: (
startModelName,
startDatasetName,
startFromResume = false,
startProjectName = null,
) => set({ startModelName, startDatasetName, startProjectName, startFromResume }),
setSseConnected: (value) => set({ sseConnected: value }),
setLastEventId: (value) => set({ lastEventId: value }),

1
studio/frontend/src/features/training/types/api.ts
View file

@ -5,6 +5,7 @@ import type { S3Config } from "@/types/training";
export interface TrainingStartRequest {
model_name: string;
project_name: string | null;
training_type: string;
hf_token: string | null;
load_in_4bit: boolean;

2
studio/frontend/src/features/training/types/config.ts
View file

@ -21,6 +21,7 @@ export interface TrainingConfigState {
currentStep: StepNumber;
modelType: ModelType | null;
selectedModel: string | null;
projectName: string;
trainingMethod: TrainingMethod;
hfToken: string;
datasetSource: DatasetSource;
@ -95,6 +96,7 @@ export interface TrainingConfigActions {
prevStep: () => void;
setModelType: (type: ModelType) => void;
setSelectedModel: (model: string | null) => void;
setProjectName: (value: string) => void;
ensureModelDefaultsLoaded: () => void;
ensureDatasetChecked: () => void;
setTrainingMethod: (method: TrainingMethod) => void;

1
studio/frontend/src/features/training/types/history.ts
View file

@ -5,6 +5,7 @@ export interface TrainingRunSummary {
id: string;
status: "running" | "completed" | "stopped" | "error";
model_name: string;
project_name: string | null;
dataset_name: string;
display_name: string | null;
started_at: string;

3
studio/frontend/src/features/training/types/runtime.ts
View file

@ -83,6 +83,7 @@ export interface TrainingRuntimeState {
startError: string | null;
startModelName: string | null;
startDatasetName: string | null;
startProjectName: string | null;
startFromResume: boolean;
sseConnected: boolean;
firstStepReceived: boolean;
@ -121,6 +122,7 @@ export interface TrainingRuntimeActions {
modelName: string | null,
datasetName: string | null,
fromResume?: boolean,
projectName?: string | null,
) => void;
setSseConnected: (value: boolean) => void;
setLastEventId: (value: number | null) => void;
@ -160,6 +162,7 @@ export interface TrainingViewData {
// Config summary
modelName: string;
projectName: string | null;
trainingMethod: string;
// Time-series (for ChartsSection)

20
studio/frontend/src/hooks/use-tauri-backend.ts
View file

@ -75,8 +75,7 @@ function externalConflictMessage(preflight: DesktopPreflightResult) {
: "A Unsloth server for this install is already running from a terminal. Stop that server, or run `unsloth studio update` from that terminal before using the desktop app.";
}
async function waitForManagedServerReady(
invoke: TauriInvoke,
async function waitForManagedServerPort(
getPort: () => number | null,
shouldContinue: () => boolean,
): Promise<ManagedStartupResult> {
@ -91,15 +90,7 @@ async function waitForManagedServerReady(
continue;
}
const healthy = await invoke<boolean>("check_health", { port });
if (!shouldContinue()) {
return { status: "aborted" };
}
if (healthy && getPort() === port) {
return { status: "ready", port };
}
await wait(MANAGED_STARTUP_POLL_MS);
return { status: "ready", port };
}
}
@ -280,10 +271,9 @@ export function useTauriBackend() {
// backend/run.py keeps the 8888-8908 fallback via server-port/TAURI_PORT.
await invoke("start_managed_server", { port: 8888 });
// Wait for the owned backend's server-port event. Don't attach to an
// external backend if the managed start doesn't report a port.
const startupResult = await waitForManagedServerReady(
invoke,
// Rust emits server-port only after validating the desktop-owned process.
// Treat that as the UI handoff point instead of doing a second health poll.
const startupResult = await waitForManagedServerPort(
() => portRef.current,
() => startingRef.current,
);

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

@ -609,6 +609,10 @@ export const en = {
params: {
title: "Parameters",
description: "Configure training hyperparameters",
projectName: "Project Name",
optional: "Optional",
projectNameDescription:
"Used in training output folder names, export defaults, and history.",
loraSettings: "LoRA Settings",
trainingHyperparameters: "Training Hyperparameters",
maxSteps: "Max Steps",
@ -850,6 +854,7 @@ export const en = {
loss: "Loss",
lr: "LR",
gradNorm: "Grad Norm",
project: "Project",
model: "Model",
method: "Method",
elapsed: "Elapsed: {value}",

4
studio/frontend/src/i18n/locales/zh-CN.ts
View file

@ -541,6 +541,9 @@ export const zhCN = {
params: {
title: "参数",
description: "配置训练超参数",
projectName: "项目名称",
optional: "可选",
projectNameDescription: "用于训练输出文件夹名称、导出默认值和历史记录。",
loraSettings: "LoRA 设置",
trainingHyperparameters: "训练超参数",
maxSteps: "最大步数",
@ -766,6 +769,7 @@ export const zhCN = {
loss: "Loss",
lr: "LR",
gradNorm: "梯度范数",
project: "项目",
model: "模型",
method: "方法",
elapsed: "已用时间：{value}",

140
studio/install_llama_prebuilt.py
View file

@ -7,6 +7,7 @@
from __future__ import annotations
import argparse
import atexit
import errno
import fnmatch
import hashlib
@ -5203,7 +5204,8 @@ def ldconfig_runtime_dirs(required_libraries: Iterable[str]) -> list[str]:
def linux_runtime_dirs(binary_path: Path) -> list[str]:
missing = linux_missing_libraries(binary_path)
# ldd may execute the binary, so probe it with a secret-free env.
missing = linux_missing_libraries(binary_path, env = scrubbed_environ())
if not missing:
return []
return linux_runtime_dirs_for_required_libraries(missing)
@ -5499,6 +5501,140 @@ def _wsl_system_rocm_lib_dirs() -> list[str]:
return out
# Secrets a downloaded llama.cpp binary never needs; keep them out of binary_env().
# The installer's own API calls read os.environ directly, so auth is unaffected.
_SECRET_ENV_EXACT_NAMES = frozenset(
{
"HF_TOKEN",
"HUGGING_FACE_HUB_TOKEN",
"GH_TOKEN",
"GITHUB_TOKEN",
"WANDB_API_KEY",
"OPENAI_API_KEY",
"ANTHROPIC_API_KEY",
"AWS_ACCESS_KEY_ID",
"AWS_SECRET_ACCESS_KEY",
"AWS_SESSION_TOKEN",
"GOOGLE_APPLICATION_CREDENTIALS",
"AZURE_CLIENT_SECRET",
"ACTIONS_ID_TOKEN_REQUEST_TOKEN",
"ACTIONS_ID_TOKEN_REQUEST_URL",
"ACTIONS_RUNTIME_TOKEN",
# Credential pointers (cluster / remote-host access).
"KUBECONFIG",
"SSH_AUTH_SOCK",
}
)
# Case-insensitive substring markers for names we do not enumerate (no bare "KEY",
# which would hit benign runtime vars).
_SECRET_ENV_MARKERS = (
"TOKEN",
"SECRET",
"PASSWORD",
"PASSWD",
"PASSPHRASE",
"CREDENTIAL",
"PRIVATE_KEY",
"API_KEY",
)
# Proxy / index URLs embed creds in their value; the offline binaries never need them.
_SECRET_ENV_URL_NAMES = frozenset(
{
"HTTP_PROXY",
"HTTPS_PROXY",
"ALL_PROXY",
"FTP_PROXY",
"RSYNC_PROXY",
"PIP_INDEX_URL",
"PIP_EXTRA_INDEX_URL",
"UV_INDEX_URL",
"UV_DEFAULT_INDEX",
"UV_EXTRA_INDEX_URL",
}
)
# Also drop values with URL userinfo creds (scheme://user:secret@host or token@host).
_URL_USERINFO_CREDENTIAL_RE = re.compile(r"://[^/@\s]+@")
def is_secret_env_name(name: str) -> bool:
upper = name.upper()
return (
upper in _SECRET_ENV_EXACT_NAMES
or upper in _SECRET_ENV_URL_NAMES
or any(marker in upper for marker in _SECRET_ENV_MARKERS)
)
def scrub_env(env: dict[str, str]) -> dict[str, str]:
"""Drop secret-bearing variables before handing an env to a downloaded binary."""
return {
key: value
for key, value in env.items()
if not is_secret_env_name(key) and not _URL_USERINFO_CREDENTIAL_RE.search(value or "")
}
# Home / cache pointers to on-disk token stores (~/.cache/huggingface/token,
# ~/.aws/credentials, ...). Stripping env tokens is not enough; point these at an
# empty home so the binary cannot read those files via $HOME.
_RUNTIME_HOME_POINTER_VARS = (
"HOME",
"USERPROFILE",
"APPDATA",
"LOCALAPPDATA",
"XDG_CACHE_HOME",
"XDG_CONFIG_HOME",
"XDG_DATA_HOME",
"HF_HOME",
"HUGGINGFACE_HUB_CACHE",
"HF_HUB_CACHE",
)
# Credential / config file pointers outside HOME; drop so lookups fall back to the
# empty home.
_CREDENTIAL_FILE_POINTER_VARS = (
"NETRC",
"PIP_CONFIG_FILE",
"DOCKER_CONFIG",
"GIT_CONFIG_GLOBAL",
)
# GitHub Actions command files: appending to these injects PATH/env into later steps.
_CI_COMMAND_FILE_VARS = (
"GITHUB_ENV",
"GITHUB_PATH",
"GITHUB_OUTPUT",
"GITHUB_STEP_SUMMARY",
"BASH_ENV",
)
_isolated_runtime_home_dir: str | None = None
def isolated_runtime_home() -> str:
# Empty dir, created lazily and removed at exit. (A binary resolving the real
# home via getpwuid is out of scope; that needs OS sandboxing.)
global _isolated_runtime_home_dir
if _isolated_runtime_home_dir is None:
path = tempfile.mkdtemp(prefix = "unsloth-prebuilt-home-")
atexit.register(shutil.rmtree, path, ignore_errors = True)
_isolated_runtime_home_dir = path
return _isolated_runtime_home_dir
def scrubbed_environ() -> dict[str, str]:
# os.environ minus secrets, with home / credential pointers neutralised. Used for
# the binary env and any probe (e.g. ldd) that runs the untrusted binary.
env = scrub_env(os.environ.copy())
runtime_home = isolated_runtime_home()
for pointer in _RUNTIME_HOME_POINTER_VARS:
env[pointer] = runtime_home
# Windows rebuilds the profile from %HOMEDRIVE%%HOMEPATH% (no-op pair on POSIX).
drive, tail = os.path.splitdrive(runtime_home)
env["HOMEDRIVE"], env["HOMEPATH"] = drive, tail or runtime_home
for pointer in (*_CREDENTIAL_FILE_POINTER_VARS, *_CI_COMMAND_FILE_VARS):
env.pop(pointer, None)
return env
def binary_env(
binary_path: Path,
install_dir: Path,
@ -5506,7 +5642,7 @@ def binary_env(
*,
runtime_line: str | None = None,
) -> dict[str, str]:
env = os.environ.copy()
env = scrubbed_environ()
if host.is_windows:
path_dirs = [
str(binary_path.parent),

16
studio/src-tauri/src/commands.rs
View file

@ -65,10 +65,18 @@ pub async fn desktop_preflight(
shutdown: tauri::State<'_, ShutdownFlag>,
diagnostics: tauri::State<'_, DiagnosticsState>,
) -> Result<crate::preflight::DesktopPreflightResult, String> {
let started = Instant::now();
let (result, adopted_watchdog_generation) =
crate::preflight::desktop_preflight_result_with_state(state.inner()).await?;
diagnostics::record_preflight(&diagnostics, &result);
info!(
"desktop_preflight completed disposition={:?} port={:?} in {}ms",
result.disposition,
result.port,
started.elapsed().as_millis()
);
if let Some((generation, newly_adopted)) = adopted_watchdog_generation {
if newly_adopted {
if let Some(port) = result.port {
@ -205,9 +213,17 @@ pub async fn start_managed_server(
port: u16,
) -> Result<(), String> {
info!("start_managed_server command called with port {}", port);
let started = Instant::now();
let diagnostics_state = diagnostics.inner().clone();
let generation = process::start_backend(&app, &state, port, &shutdown, &diagnostics_state)?;
info!(
"start_managed_server spawned generation={} in {}ms",
generation,
started.elapsed().as_millis()
);
let watchdog_state = state.inner().clone();
let watchdog_shutdown = shutdown.inner().clone();
let watchdog_app = app.clone();

147
studio/src-tauri/src/desktop_backend_owner.rs
View file

@ -93,7 +93,7 @@ enum PreviousAppPidStatus {
Uncertain,
}
#[derive(Debug, Deserialize)]
#[derive(Clone, Debug, Deserialize)]
struct HealthDesktopOwner {
kind: Option<String>,
token_sha256: Option<String>,
@ -101,15 +101,7 @@ struct HealthDesktopOwner {
#[derive(Debug, Deserialize)]
struct HealthResponse {
status: Option<String>,
service: Option<String>,
version: Option<String>,
desktop_protocol_version: Option<u16>,
desktop_manageability_version: Option<u16>,
supports_desktop_auth: Option<bool>,
supports_desktop_backend_ownership: Option<bool>,
studio_root_id: Option<String>,
desktop_owner: Option<HealthDesktopOwner>,
}
#[derive(Debug)]
@ -123,6 +115,18 @@ struct DesktopLoginPayload<'a> {
secret: &'a str,
}
#[derive(Clone, Debug, Deserialize)]
pub(crate) struct DesktopLiveness {
status: Option<String>,
service: Option<String>,
desktop_protocol_version: Option<u16>,
desktop_manageability_version: Option<u16>,
supports_desktop_auth: Option<bool>,
supports_desktop_backend_ownership: Option<bool>,
studio_root_id: Option<String>,
desktop_owner: Option<HealthDesktopOwner>,
}
#[derive(Deserialize)]
struct TokenResponse {
access_token: String,
@ -290,10 +294,10 @@ impl BackendOwnerState {
}
pub(crate) fn verifies_exact_port_blocking(&self, port: u16) -> bool {
match fetch_health_blocking(port) {
Ok(Some(health)) => {
health_verifies_metadata(&health, &self.metadata)
&& lifecycle_control_block_reason(&health).is_none()
match fetch_liveness_blocking(port) {
Ok(Some(liveness)) => {
liveness_verifies_metadata(&liveness, &self.metadata)
&& lifecycle_control_block_reason(&liveness).is_none()
}
_ => false,
}
@ -498,66 +502,110 @@ pub(crate) fn test_owner_state(root_id: &str, token: &str, port: u16) -> Backend
}
}
fn health_verifies_metadata(health: &HealthResponse, metadata: &DesktopBackendMetadata) -> bool {
let healthy = health.status.as_deref() == Some("healthy")
&& health.service.as_deref() == Some("Unsloth UI Backend");
let Some(owner) = health.desktop_owner.as_ref() else {
fn liveness_verifies_metadata(
liveness: &DesktopLiveness,
metadata: &DesktopBackendMetadata,
) -> bool {
let alive = matches!(liveness.status.as_deref(), Some("alive") | Some("healthy"))
&& liveness.service.as_deref() == Some("Unsloth UI Backend");
let Some(owner) = liveness.desktop_owner.as_ref() else {
return false;
};
healthy
alive
&& owner_matches_metadata(
metadata,
health.studio_root_id.as_deref(),
liveness.studio_root_id.as_deref(),
owner.kind.as_deref(),
owner.token_sha256.as_deref(),
)
}
fn lifecycle_control_block_reason(health: &HealthResponse) -> Option<String> {
if health.desktop_protocol_version != Some(crate::preflight::DESKTOP_PROTOCOL_VERSION) {
fn lifecycle_control_block_reason(liveness: &DesktopLiveness) -> Option<String> {
if liveness.desktop_protocol_version != Some(crate::preflight::DESKTOP_PROTOCOL_VERSION) {
return Some("desktop_protocol_incompatible".to_string());
}
if health.supports_desktop_auth != Some(true) {
if liveness.supports_desktop_auth != Some(true) {
return Some("desktop_auth_unsupported".to_string());
}
if health.desktop_manageability_version.unwrap_or(0)
if liveness.desktop_manageability_version.unwrap_or(0)
< crate::preflight::DESKTOP_MANAGEABILITY_VERSION
{
return Some("desktop_manageability_unsupported".to_string());
}
if health.supports_desktop_backend_ownership != Some(true) {
if liveness.supports_desktop_backend_ownership != Some(true) {
return Some("desktop_backend_ownership_unsupported".to_string());
}
None
}
fn ready_for_use_status(health: &HealthResponse) -> OwnedBackendReadiness {
match crate::preflight::backend_version_stale_reason(health.version.as_deref()) {
fn ready_for_use_status(health: Option<&HealthResponse>) -> OwnedBackendReadiness {
let version = health
.and_then(|h| h.version.as_deref())
.filter(|v| !v.is_empty());
match crate::preflight::backend_version_stale_reason(version) {
Some(reason) => OwnedBackendReadiness::Stale { reason },
None => OwnedBackendReadiness::Ready,
}
}
async fn fetch_health(port: u16) -> Result<Option<HealthResponse>, reqwest::Error> {
async fn health_ready_status(port: u16) -> OwnedBackendReadiness {
match fetch_health(port).await {
Ok(health) => ready_for_use_status(health.as_ref()),
Err(reason) => OwnedBackendReadiness::Stale { reason },
}
}
async fn fetch_liveness(port: u16) -> Result<Option<DesktopLiveness>, reqwest::Error> {
let client = reqwest::Client::builder()
.timeout(LOCAL_HTTP_TIMEOUT)
.build()?;
for path in ["/api/liveness", "/api/health"] {
let response = client
.get(format!("http://127.0.0.1:{port}{path}"))
.send()
.await?;
if response.status() == reqwest::StatusCode::NOT_FOUND && path == "/api/liveness" {
continue;
}
if !response.status().is_success() {
return Ok(None);
}
return response.json::<DesktopLiveness>().await.map(Some);
}
Ok(None)
}
fn fetch_liveness_blocking(port: u16) -> Result<Option<DesktopLiveness>, String> {
for path in ["/api/liveness", "/api/health"] {
let response = http_request_blocking(port, "GET", path, &[], &[])?;
if response.status == 404 && path == "/api/liveness" {
continue;
}
if !(200..300).contains(&response.status) {
return Ok(None);
}
return serde_json::from_slice::<DesktopLiveness>(&response.body)
.map(Some)
.map_err(|e| e.to_string());
}
Ok(None)
}
async fn fetch_health(port: u16) -> Result<Option<HealthResponse>, String> {
let client = reqwest::Client::builder()
.timeout(LOCAL_HTTP_TIMEOUT)
.build()
.map_err(|e| e.to_string())?;
let response = client
.get(format!("http://127.0.0.1:{port}/api/health"))
.send()
.await?;
.await
.map_err(|e| e.to_string())?;
if !response.status().is_success() {
return Ok(None);
}
response.json::<HealthResponse>().await.map(Some)
}
fn fetch_health_blocking(port: u16) -> Result<Option<HealthResponse>, String> {
let response = http_request_blocking(port, "GET", "/api/health", &[], &[])?;
if !(200..300).contains(&response.status) {
return Ok(None);
}
serde_json::from_slice::<HealthResponse>(&response.body)
response
.json::<HealthResponse>()
.await
.map(Some)
.map_err(|e| e.to_string())
}
@ -618,21 +666,21 @@ pub(crate) async fn probe_owned_backend_state(
};
let mut verified = Vec::new();
for port in ports {
let health = match fetch_health(port).await {
Ok(Some(health)) => health,
let liveness = match fetch_liveness(port).await {
Ok(Some(liveness)) => liveness,
Ok(None) => continue,
Err(error) => {
warn!(
"Desktop-owned backend probe skipped port {} after health error: {}",
"Desktop-owned backend probe skipped port {} after liveness error: {}",
port, error
);
continue;
}
};
if !health_verifies_metadata(&health, &owner.metadata) {
if !liveness_verifies_metadata(&liveness, &owner.metadata) {
continue;
}
if let Some(reason) = lifecycle_control_block_reason(&health) {
if let Some(reason) = lifecycle_control_block_reason(&liveness) {
return OwnedBackendProbe::Unmanageable { port, reason };
}
if !desktop_login_route_compatible(port).await {
@ -646,7 +694,7 @@ pub(crate) async fn probe_owned_backend_state(
return OwnedBackendProbe::Unmanageable { port, reason };
}
}
verified.push((port, ready_for_use_status(&health)));
verified.push((port, health_ready_status(port).await));
}
if verified.len() != 1 {
@ -967,12 +1015,11 @@ mod tests {
}
#[test]
fn health_verification_requires_root_kind_and_token_sha() {
fn liveness_verification_requires_root_kind_and_token_sha() {
let metadata = metadata(1, Some(8888));
let health = HealthResponse {
status: Some("healthy".to_string()),
let liveness = DesktopLiveness {
status: Some("alive".to_string()),
service: Some("Unsloth UI Backend".to_string()),
version: Some("2026.5.2".to_string()),
desktop_protocol_version: Some(1),
desktop_manageability_version: Some(1),
supports_desktop_auth: Some(true),
@ -983,12 +1030,12 @@ mod tests {
token_sha256: Some(token_sha256(TOKEN)),
}),
};
assert!(health_verifies_metadata(&health, &metadata));
assert!(liveness_verifies_metadata(&liveness, &metadata));
let mut wrong_root = health;
let mut wrong_root = liveness;
wrong_root.studio_root_id =
Some("bbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbbb".to_string());
assert!(!health_verifies_metadata(&wrong_root, &metadata));
assert!(!liveness_verifies_metadata(&wrong_root, &metadata));
}
#[tokio::test]

16
studio/src-tauri/src/preflight.rs
View file

@ -193,6 +193,9 @@ pub async fn desktop_preflight_result_with_state(
if let Some(snapshot) = crate::process::owned_backend_snapshot(state)? {
let Some(owner) = snapshot.owner.clone() else {
// TAURI_PORT is emitted only after uvicorn lifespan completes; keep
// this ownerless path on full health so auth/bootstrap are ready.
let probe = match snapshot.port {
Some(port) => backend::probe_ownerless_spawned_backend(port).await,
None => backend,
@ -494,9 +497,22 @@ mod tests {
FakeCli { bin, dir }
}
#[cfg(unix)]
fn remove_managed_capability_cache() {
let _ = std::fs::remove_file(
dirs::home_dir()
.unwrap()
.join(".unsloth")
.join("studio")
.join("desktop_capability_cache.json"),
);
}
#[cfg(unix)]
#[tokio::test]
async fn managed_cli_capability_probe_classifies_core_cases() {
remove_managed_capability_cache();
for (name, script, stale_reason) in [
(
"cap-missing",

12
studio/src-tauri/src/preflight/backend.rs
View file

@ -3,7 +3,10 @@ use super::version::{
backend_version_stale_reason, DESKTOP_MANAGEABILITY_VERSION, DESKTOP_PROTOCOL_VERSION,
};
use serde::{Deserialize, Serialize};
use log::info;
use std::time::Duration;
use std::time::Instant;
#[derive(Debug, Deserialize)]
struct DesktopOwnerHealth {
@ -24,6 +27,7 @@ pub(super) struct BackendHealth {
}
pub(super) async fn backend_health(client: &reqwest::Client, port: u16) -> Option<BackendHealth> {
let started = Instant::now();
let url = format!("http://127.0.0.1:{port}/api/health");
let response = client.get(url).send().await.ok()?;
if !response.status().is_success() {
@ -40,6 +44,14 @@ pub(super) async fn backend_health(client: &reqwest::Client, port: u16) -> Optio
.and_then(|v| v.as_str())
.map(|s| s == "Unsloth UI Backend")
.unwrap_or(false);
info!(
"Desktop preflight: health probe on port {} healthy={} service={} in {}ms",
port,
healthy,
service,
started.elapsed().as_millis()
);
if !healthy || !service {
return None;
}

333
studio/src-tauri/src/preflight/managed.rs
View file

@ -2,14 +2,30 @@ use super::types::ManagedProbe;
use super::version::{
backend_version_stale_reason, DESKTOP_MANAGEABILITY_VERSION, DESKTOP_PROTOCOL_VERSION,
};
use serde::Deserialize;
use log::{info, warn};
use serde::{Deserialize, Serialize};
use std::fs;
use std::path::{Path, PathBuf};
use std::process::Stdio;
use std::time::Duration;
use std::time::{Duration, Instant, UNIX_EPOCH};
use tokio::io::AsyncReadExt;
use tokio::process::Command;
#[derive(Debug, Deserialize)]
const MANAGED_CAPABILITY_CACHE_SCHEMA: u16 = 2;
const FNV64_OFFSET_BASIS: u64 = 0xcbf29ce484222325;
const FNV64_PRIME: u64 = 0x100000001b3;
const HASHED_MARKER_MAX_BYTES: u64 = 64 * 1024;
const FALLBACK_MARKER_NAMES: &[&str] = &[
"pyvenv.cfg",
"uv.lock",
"requirements.txt",
"python.exe",
"python",
];
#[derive(Debug, Clone, Deserialize, Serialize)]
struct DesktopCapability {
desktop_protocol_version: Option<u16>,
desktop_manageability_version: Option<u16>,
@ -20,7 +36,225 @@ struct DesktopCapability {
version: Option<String>,
}
#[derive(Debug, Clone, Deserialize, Serialize)]
struct ManagedCapabilityCache {
schema: u16,
bin_path: String,
bin_size: u64,
bin_mtime_ms: u64,
studio_root_id: Option<String>,
marker_path: Option<String>,
marker_size: Option<u64>,
marker_mtime_ms: Option<u64>,
desktop_protocol_version: u16,
desktop_manageability_version: u16,
capability: DesktopCapability,
}
#[derive(Debug, Clone)]
struct MarkerFingerprint {
path: String,
size: u64,
mtime_ms: u64,
content_hash: Option<u64>,
}
#[derive(Debug, Clone, PartialEq, Eq)]
struct ManagedBinFingerprint {
bin_path: String,
bin_size: u64,
bin_mtime_ms: u64,
studio_root_id: Option<String>,
marker_path: Option<String>,
marker_size: Option<u64>,
marker_mtime_ms: Option<u64>,
}
fn modified_ms(metadata: &fs::Metadata) -> Option<u64> {
metadata
.modified()
.ok()?
.duration_since(UNIX_EPOCH)
.ok()
.and_then(|duration| u64::try_from(duration.as_millis()).ok())
}
fn hash_bytes(hash: u64, bytes: &[u8]) -> u64 {
bytes.iter().fold(hash, |mut next, byte| {
next ^= u64::from(*byte);
next.wrapping_mul(FNV64_PRIME)
})
}
fn marker_content_hash(path: &Path, metadata: &fs::Metadata) -> Option<u64> {
if metadata.len() > HASHED_MARKER_MAX_BYTES {
return None;
}
fs::read(path)
.ok()
.map(|bytes| hash_bytes(FNV64_OFFSET_BASIS, &bytes))
}
fn marker_candidates_for_bin(bin: &Path) -> Vec<PathBuf> {
let Some(scripts_dir) = bin.parent() else {
return Vec::new();
};
let Some(venv_dir) = scripts_dir.parent() else {
return Vec::new();
};
let mut out = Vec::new();
#[cfg(unix)]
{
if let Ok(lib_dir) = fs::read_dir(venv_dir.join("lib")) {
for entry in lib_dir.flatten() {
out.push(
entry
.path()
.join("site-packages")
.join("unsloth_cli")
.join("commands")
.join("studio.py"),
);
}
}
}
for marker_name in FALLBACK_MARKER_NAMES {
out.push(venv_dir.join(marker_name));
out.push(scripts_dir.join(marker_name));
}
out.push(
venv_dir
.join("Lib")
.join("site-packages")
.join("unsloth_cli")
.join("commands")
.join("studio.py"),
);
out
}
fn managed_bin_fingerprint(bin: &Path) -> Option<ManagedBinFingerprint> {
let bin_metadata = fs::metadata(bin).ok()?;
let bin_path = bin
.canonicalize()
.unwrap_or_else(|_| bin.to_path_buf())
.to_string_lossy()
.into_owned();
let studio_root_id = crate::desktop_backend_owner::read_expected_studio_root_id();
let mut marker_entries: Vec<MarkerFingerprint> = marker_candidates_for_bin(bin)
.into_iter()
.filter_map(|path| {
let metadata = fs::metadata(&path).ok()?;
Some(MarkerFingerprint {
path: path
.canonicalize()
.unwrap_or(path.clone())
.to_string_lossy()
.into_owned(),
size: metadata.len(),
mtime_ms: modified_ms(&metadata)?,
content_hash: marker_content_hash(&path, &metadata),
})
})
.collect();
marker_entries.sort_by(|left, right| left.path.cmp(&right.path));
let marker_hash = marker_entries
.iter()
.fold(FNV64_OFFSET_BASIS, |hash, marker| {
let next = hash_bytes(hash, marker.path.as_bytes());
let next = hash_bytes(next, &marker.size.to_le_bytes());
let next = hash_bytes(next, &marker.mtime_ms.to_le_bytes());
if let Some(content_hash) = marker.content_hash {
hash_bytes(next, &content_hash.to_le_bytes())
} else {
next
}
});
let marker_path = (!marker_entries.is_empty()).then(|| "markers".to_string());
let marker_size = (!marker_entries.is_empty()).then(|| marker_entries.len() as u64);
let marker_mtime_ms = (!marker_entries.is_empty()).then_some(marker_hash);
Some(ManagedBinFingerprint {
bin_path,
bin_size: bin_metadata.len(),
bin_mtime_ms: modified_ms(&bin_metadata)?,
studio_root_id,
marker_path,
marker_size,
marker_mtime_ms,
})
}
fn capability_cache_path() -> Option<PathBuf> {
dirs::home_dir().map(|home| {
home.join(".unsloth")
.join("studio")
.join("desktop_capability_cache.json")
})
}
fn cache_matches(cache: &ManagedCapabilityCache, fingerprint: &ManagedBinFingerprint) -> bool {
cache.schema == MANAGED_CAPABILITY_CACHE_SCHEMA
&& cache.desktop_protocol_version == DESKTOP_PROTOCOL_VERSION
&& cache.desktop_manageability_version == DESKTOP_MANAGEABILITY_VERSION
&& cache.bin_path == fingerprint.bin_path
&& cache.bin_size == fingerprint.bin_size
&& cache.bin_mtime_ms == fingerprint.bin_mtime_ms
&& cache.studio_root_id == fingerprint.studio_root_id
&& cache.marker_path == fingerprint.marker_path
&& cache.marker_size == fingerprint.marker_size
&& cache.marker_mtime_ms == fingerprint.marker_mtime_ms
&& desktop_capability_ready(&cache.capability)
}
fn read_cached_capability(fingerprint: &ManagedBinFingerprint) -> Option<DesktopCapability> {
let path = capability_cache_path()?;
let bytes = fs::read(path).ok()?;
let cache = serde_json::from_slice::<ManagedCapabilityCache>(&bytes).ok()?;
if cache_matches(&cache, fingerprint) {
Some(cache.capability)
} else {
None
}
}
fn write_cached_capability(fingerprint: &ManagedBinFingerprint, capability: &DesktopCapability) {
let Some(path) = capability_cache_path() else {
return;
};
let cache = ManagedCapabilityCache {
schema: MANAGED_CAPABILITY_CACHE_SCHEMA,
bin_path: fingerprint.bin_path.clone(),
bin_size: fingerprint.bin_size,
bin_mtime_ms: fingerprint.bin_mtime_ms,
studio_root_id: fingerprint.studio_root_id.clone(),
marker_path: fingerprint.marker_path.clone(),
marker_size: fingerprint.marker_size,
marker_mtime_ms: fingerprint.marker_mtime_ms,
desktop_protocol_version: DESKTOP_PROTOCOL_VERSION,
desktop_manageability_version: DESKTOP_MANAGEABILITY_VERSION,
capability: capability.clone(),
};
if let Some(parent) = path.parent() {
if fs::create_dir_all(parent).is_err() {
return;
}
}
let Ok(bytes) = serde_json::to_vec_pretty(&cache) else {
return;
};
if let Err(error) = fs::write(&path, bytes) {
warn!(
"Managed preflight: could not write capability cache: {}",
error
);
}
}
async fn run_cli_probe(bin: &Path, args: &[&str]) -> bool {
let started = Instant::now();
let mut cmd = Command::new(bin);
cmd.args(args).stdout(Stdio::null()).stderr(Stdio::null());
@ -43,20 +277,33 @@ async fn run_cli_probe(bin: &Path, args: &[&str]) -> bool {
}
let Ok(mut child) = cmd.spawn() else {
info!(
"Managed preflight probe {:?} failed to spawn in {}ms",
args,
started.elapsed().as_millis()
);
return false;
};
match tokio::time::timeout(Duration::from_secs(10), child.wait()).await {
let ok = match tokio::time::timeout(Duration::from_secs(10), child.wait()).await {
Ok(Ok(status)) => status.success(),
_ => {
let _ = child.kill().await;
let _ = child.wait().await;
false
}
}
};
info!(
"Managed preflight probe {:?} finished ok={} in {}ms",
args,
ok,
started.elapsed().as_millis()
);
ok
}
async fn probe_cli_capability(bin: &Path) -> Option<DesktopCapability> {
let started = Instant::now();
let mut cmd = Command::new(bin);
cmd.args(["studio", "desktop-capabilities", "--json"])
.stdout(Stdio::piped())
@ -81,6 +328,10 @@ async fn probe_cli_capability(bin: &Path) -> Option<DesktopCapability> {
}
let Ok(mut child) = cmd.spawn() else {
info!(
"Managed desktop-capabilities probe failed to spawn in {}ms",
started.elapsed().as_millis()
);
return None;
};
let Some(mut stdout) = child.stdout.take() else {
@ -92,9 +343,19 @@ async fn probe_cli_capability(bin: &Path) -> Option<DesktopCapability> {
Err(_) => {
let _ = child.kill().await;
let _ = child.wait().await;
info!(
"Managed desktop-capabilities probe timed out in {}ms",
started.elapsed().as_millis()
);
return None;
}
_ => {
info!(
"Managed desktop-capabilities probe exited unsuccessfully in {}ms",
started.elapsed().as_millis()
);
return None;
}
_ => return None,
}
let mut output = Vec::new();
@ -102,7 +363,13 @@ async fn probe_cli_capability(bin: &Path) -> Option<DesktopCapability> {
return None;
}
serde_json::from_slice::<DesktopCapability>(&output).ok()
let capability = serde_json::from_slice::<DesktopCapability>(&output).ok();
info!(
"Managed desktop-capabilities probe finished ok={} in {}ms",
capability.is_some(),
started.elapsed().as_millis()
);
capability
}
fn desktop_capability_stale_reason(capability: &DesktopCapability) -> Option<String> {
@ -132,18 +399,48 @@ fn desktop_capability_ready(capability: &DesktopCapability) -> bool {
}
pub(super) async fn probe_managed_bin(bin: PathBuf) -> ManagedProbe {
let started = Instant::now();
if !run_cli_probe(&bin, &["-h"]).await {
info!(
"Managed preflight: cli unusable for {:?} in {}ms",
bin,
started.elapsed().as_millis()
);
return ManagedProbe::Stale {
bin,
reason: "cli_unusable".to_string(),
};
}
let capability = probe_cli_capability(&bin).await;
if let Some(capability) = capability {
if desktop_capability_ready(&capability) {
if let Some(fingerprint) = managed_bin_fingerprint(&bin) {
if read_cached_capability(&fingerprint).is_some() {
info!(
"Managed preflight: using cached desktop capability for {:?} in {}ms",
bin,
started.elapsed().as_millis()
);
return ManagedProbe::Ready { bin };
}
}
let capability = probe_cli_capability(&bin).await;
if let Some(capability) = capability {
if let Some(fingerprint) = managed_bin_fingerprint(&bin) {
write_cached_capability(&fingerprint, &capability);
}
if desktop_capability_ready(&capability) {
info!(
"Managed preflight: cli ready for {:?} in {}ms",
bin,
started.elapsed().as_millis()
);
return ManagedProbe::Ready { bin };
}
info!(
"Managed preflight: cli stale for {:?} in {}ms",
bin,
started.elapsed().as_millis()
);
return ManagedProbe::Stale {
bin,
reason: desktop_capability_stale_reason(&capability)
@ -151,6 +448,11 @@ pub(super) async fn probe_managed_bin(bin: PathBuf) -> ManagedProbe {
};
}
info!(
"Managed preflight: desktop capability probe failed for {:?} in {}ms",
bin,
started.elapsed().as_millis()
);
ManagedProbe::Stale {
bin,
reason: "desktop_capability_probe_failed".to_string(),
@ -158,10 +460,17 @@ pub(super) async fn probe_managed_bin(bin: PathBuf) -> ManagedProbe {
}
pub(super) async fn probe_managed_install() -> ManagedProbe {
match crate::process::find_unsloth_binary() {
let started = Instant::now();
let result = match crate::process::find_unsloth_binary() {
Some(bin) => probe_managed_bin(bin).await,
None => ManagedProbe::Missing,
}
};
info!(
"Managed preflight: install probe result {:?} in {}ms",
result,
started.elapsed().as_millis()
);
result
}
pub async fn managed_install_ready() -> bool {

86
studio/src-tauri/src/process.rs
View file

@ -735,6 +735,7 @@ pub fn start_backend(
}
async fn generic_backend_health_ok(port: u16) -> bool {
let started = std::time::Instant::now();
let client = match reqwest::Client::builder()
.timeout(Duration::from_secs(2))
.build()
@ -745,49 +746,75 @@ async fn generic_backend_health_ok(port: u16) -> bool {
return false;
}
};
let response = match client
.get(format!("http://127.0.0.1:{port}/api/health"))
.send()
.await
{
Ok(response) => response,
Err(error) => {
let mut last_status = None;
let mut json = None;
for path in ["/api/liveness", "/api/health"] {
let response = match client
.get(format!("http://127.0.0.1:{port}{path}"))
.send()
.await
{
Ok(response) => response,
Err(error) => {
warn!(
"Backend port candidate {} failed health request: {}",
port, error
);
return false;
}
};
if response.status() == reqwest::StatusCode::NOT_FOUND && path == "/api/liveness" {
last_status = Some(response.status());
continue;
}
if !response.status().is_success() {
warn!(
"Backend port candidate {} failed health request: {}",
port, error
"Backend port candidate {} returned HTTP {} from health",
port,
response.status()
);
return false;
}
};
if !response.status().is_success() {
json = match response.json::<serde_json::Value>().await {
Ok(json) => Some(json),
Err(error) => {
warn!(
"Backend port candidate {} returned invalid health JSON: {}",
port, error
);
return false;
}
};
break;
}
let Some(json) = json else {
warn!(
"Backend port candidate {} returned HTTP {} from health",
port,
response.status()
last_status
.map(|status| status.to_string())
.unwrap_or_else(|| "unknown".to_string())
);
return false;
}
let json = match response.json::<serde_json::Value>().await {
Ok(json) => json,
Err(error) => {
warn!(
"Backend port candidate {} returned invalid health JSON: {}",
port, error
);
return false;
}
};
let healthy = json
let live = json
.get("status")
.and_then(|v| v.as_str())
.map(|s| s == "healthy")
.map(|s| s == "alive" || s == "healthy")
.unwrap_or(false);
let service = json
.get("service")
.and_then(|v| v.as_str())
.map(|s| s == "Unsloth UI Backend")
.unwrap_or(false);
healthy && service
info!(
"Backend port candidate {} liveness live={} service={} in {}ms",
port,
live,
service,
started.elapsed().as_millis()
);
live && service
}
async fn validate_candidate_port(
@ -798,6 +825,7 @@ async fn validate_candidate_port(
generation: u64,
port: u16,
) {
let started = std::time::Instant::now();
let owner = {
let proc = match state.lock() {
Ok(proc) => proc,
@ -852,6 +880,14 @@ async fn validate_candidate_port(
}
};
info!(
"Validated backend port candidate {} valid={} emit={} in {}ms",
port,
valid,
should_emit,
started.elapsed().as_millis()
);
if should_emit {
diagnostics::record_backend_port(&diagnostics_state, &session_id, port);
info!("Validated backend port: {}", port);

76
tests/python/test_change_system_message.py Normal file
View file

@ -0,0 +1,76 @@
import ast
import re
import types
from pathlib import Path
import pytest
def _load_change_system_message():
# Extract just _change_system_message from chat_templates.py so the test runs
# without importing unsloth (which needs unsloth_zoo / a GPU). Same pattern as
# tests/saving/test_is_gpt_oss_detection.py.
source = Path(__file__).parents[2] / "unsloth" / "chat_templates.py"
tree = ast.parse(source.read_text(encoding = "utf-8"))
funcs = [
node
for node in tree.body
if isinstance(node, ast.FunctionDef) and node.name == "_change_system_message"
]
namespace = {
"re": re,
"logger": types.SimpleNamespace(warning_once = lambda *a, **k: None),
"DEFAULT_SYSTEM_MESSAGE": {"unsloth": "You are a helpful assistant to the user"},
}
module = ast.Module(body = funcs, type_ignores = [])
ast.fix_missing_locations(module)
exec(compile(module, str(source), "exec"), namespace)
return namespace["_change_system_message"]
CUSTOM = "mycustom" # not in DEFAULT_SYSTEM_MESSAGE -> no predefined default
def test_custom_template_fills_placeholder():
# A custom template with a {system_message} placeholder must be filled, not
# left with the literal placeholder.
fn = _load_change_system_message()
template, used = fn("System: {system_message}\nUser:", CUSTOM, "You are a pirate")
assert template == "System: You are a pirate\nUser:"
assert "{system_message}" not in template
assert used == "You are a pirate"
def test_custom_template_preserves_backslashes():
# Why str.replace and not re.sub: a system message with backslashes (Windows
# paths, LaTeX, group-like text) must be inserted verbatim. re.sub treats the
# replacement specially -- r"C:\Users" raises bad-escape, r"\1" is a group ref.
fn = _load_change_system_message()
for msg in (r"C:\Users\me", r"\frac{a}{b}", r"see \1 here"):
template, used = fn("System: {system_message}", CUSTOM, msg)
assert template == f"System: {msg}"
assert used == msg
def test_custom_template_requires_system_message():
# A custom template with a placeholder but no system message must raise,
# rather than silently leaving the placeholder in.
fn = _load_change_system_message()
with pytest.raises(ValueError):
fn("System: {system_message}", CUSTOM, None)
def test_custom_template_without_placeholder_unchanged():
fn = _load_change_system_message()
template, used = fn("System: fixed", CUSTOM, "ignored")
assert template == "System: fixed"
def test_predefined_template_uses_default_then_override():
# Predefined templates with a default are unaffected by the change.
fn = _load_change_system_message()
t1, u1 = fn("System: {system_message}", "unsloth", None)
assert t1 == "System: You are a helpful assistant to the user"
t2, u2 = fn("System: {system_message}", "unsloth", "Custom override")
assert t2 == "System: Custom override"
assert u2 == "Custom override"

193
tests/python/test_v100_fullft_precision.py Normal file
View file

@ -0,0 +1,193 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
"""Regression tests for full finetuning precision on no-bf16 GPUs (V100/T4).
Full finetuning upcasts trainable weights to float32, so the model dtype is
float32 (not bfloat16). The SFTTrainer mixed-precision template in
unsloth/models/rl.py must then:
- run the forward pass under float16 autocast for normal models,
- keep FORCE_FLOAT32 models (Gemma3, gpt_oss, ...) in pure float32,
- never select bf16 on hardware without bf16.
We execute the REAL template block extracted from rl.py source (no heavy unsloth
import) against mocked inputs. See issue #4082.
"""
from __future__ import annotations
import os
import sys
import types
from pathlib import Path
import pytest
torch = pytest.importorskip("torch")
RL_PY = Path(__file__).resolve().parents[2] / "unsloth" / "models" / "rl.py"
def _extract_mixed_precision_code() -> str:
lines = RL_PY.read_text().split("\n")
try:
start = next(i for i, l in enumerate(lines) if "mixed_precision = (" in l)
except StopIteration:
pytest.skip("mixed_precision template not found in rl.py")
body, k = [], start + 1
while lines[k].strip() != ")":
body.append(lines[k])
k += 1
return eval("(\n" + "\n".join(body) + "\n)") # only string literals + comments
CODE = _extract_mixed_precision_code()
def _restore(mapping, saved):
"""Restore a dict-like to its saved snapshot: pop keys that were absent."""
for k, v in saved.items():
if v is None:
mapping.pop(k, None)
else:
mapping[k] = v
def _decide(dtype, *, bf16_supported, force_float32, full_finetuning, mixed_precision, fp16, bf16):
"""Run the template block; return (args.fp16, args.bf16, ACCELERATE_MP, raised).
Stubs (sys.modules, env vars, torch.cuda.is_bf16_supported) are restored on
exit so a decision can't leak into later tests in the same process.
"""
uzu = types.ModuleType("unsloth_zoo.utils")
uzu._get_dtype = lambda x: x
uzd = types.ModuleType("unsloth_zoo.device_type")
uzd.device_is_bf16_supported = lambda: bf16_supported # device-aware signal stub
env_keys = (
"UNSLOTH_FORCE_FLOAT32",
"UNSLOTH_ENABLE_FULL_FINETUNING",
"UNSLOTH_MIXED_PRECISION",
"ACCELERATE_MIXED_PRECISION",
)
mod_keys = ("unsloth_zoo", "unsloth_zoo.utils", "unsloth_zoo.device_type")
saved_env = {k: os.environ.get(k) for k in env_keys}
saved_mods = {k: sys.modules.get(k) for k in mod_keys}
orig_bf16 = torch.cuda.is_bf16_supported
try:
sys.modules.setdefault("unsloth_zoo", types.ModuleType("unsloth_zoo"))
sys.modules["unsloth_zoo.utils"] = uzu
sys.modules["unsloth_zoo.device_type"] = uzd
for k in env_keys:
os.environ.pop(k, None)
os.environ["UNSLOTH_FORCE_FLOAT32"] = "1" if force_float32 else "0"
os.environ["UNSLOTH_ENABLE_FULL_FINETUNING"] = "1" if full_finetuning else "0"
os.environ["UNSLOTH_MIXED_PRECISION"] = mixed_precision
torch.cuda.is_bf16_supported = lambda *a, **k: bf16_supported
args = types.SimpleNamespace(fp16 = fp16, bf16 = bf16, mixed_precision = None)
emb = types.SimpleNamespace(weight = types.SimpleNamespace(dtype = dtype))
model = types.SimpleNamespace(
config = types.SimpleNamespace(dtype = dtype, torch_dtype = dtype),
get_input_embeddings = lambda: emb,
)
raised = None
try:
exec(CODE, {"torch": torch, "os": os}, {"args": args, "model": model})
except TypeError:
raised = "TypeError"
return args.fp16, args.bf16, os.environ.get("ACCELERATE_MIXED_PRECISION"), raised
finally:
torch.cuda.is_bf16_supported = orig_bf16
_restore(os.environ, saved_env)
_restore(sys.modules, saved_mods)
def test_v100_normal_fullft_fp16_explicit():
# Normal model, full FT (weights upcast to float32), V100, fp16=True.
fp16, bf16, amp, raised = _decide(
torch.float32,
bf16_supported = False,
force_float32 = False,
full_finetuning = True,
mixed_precision = "float32",
fp16 = True,
bf16 = False,
)
assert raised is None
assert (fp16, bf16) == (True, False) # float32 weights + fp16 forward
def test_v100_normal_fullft_precision_unset():
# Same, but user left precision unset -> must pick fp16, never bf16.
fp16, bf16, amp, raised = _decide(
torch.float32,
bf16_supported = False,
force_float32 = False,
full_finetuning = True,
mixed_precision = "float32",
fp16 = False,
bf16 = False,
)
assert raised is None
assert (fp16, bf16) == (True, False)
assert amp == "fp16"
def test_force_float32_model_fullft_is_pure_float32():
# FORCE_FLOAT32 model (Gemma3, gpt_oss, ...) in full FT -> pure float32, no autocast.
fp16, bf16, amp, raised = _decide(
torch.float32,
bf16_supported = False,
force_float32 = True,
full_finetuning = True,
mixed_precision = "float32",
fp16 = True,
bf16 = False,
)
assert raised is None
assert (fp16, bf16) == (False, False)
assert amp in (None, "no")
def test_no_bf16_on_volta_in_auto_branch():
# bf16 model dtype but no bf16 HW, precision unset -> fp16, never bf16.
fp16, bf16, amp, raised = _decide(
torch.bfloat16,
bf16_supported = False,
force_float32 = False,
full_finetuning = False,
mixed_precision = "float32",
fp16 = False,
bf16 = False,
)
assert bf16 is False
def test_bf16_gpu_unchanged_auto_branch():
# Regression guard: on a bf16 GPU, a float32 model with unset precision
# still selects bf16 autocast (behavior must not change for bf16 hardware).
fp16, bf16, amp, raised = _decide(
torch.float32,
bf16_supported = True,
force_float32 = False,
full_finetuning = True,
mixed_precision = "float32",
fp16 = False,
bf16 = False,
)
assert raised is None
assert (fp16, bf16) == (False, True)
def test_genuine_bf16_model_with_fp16_still_raises():
# A real bfloat16 model on bf16 HW with fp16 requested is a genuine mismatch.
_, _, _, raised = _decide(
torch.bfloat16,
bf16_supported = True,
force_float32 = False,
full_finetuning = False,
mixed_precision = "float32",
fp16 = True,
bf16 = False,
)
assert raised == "TypeError"

52
tests/saving/test_is_gpt_oss_detection.py Normal file
View file

@ -0,0 +1,52 @@
import ast
import types
from pathlib import Path
def _load_is_gpt_oss():
# Extract just the helper from save.py so the test runs without importing
# unsloth (which requires unsloth_zoo / a GPU), matching the pattern used by
# test_qwen3_5_vlm_full_finetune_key_remap.py.
source = Path(__file__).parents[2] / "unsloth" / "save.py"
tree = ast.parse(source.read_text(encoding = "utf-8"))
helpers = [
node
for node in tree.body
if isinstance(node, ast.FunctionDef) and node.name == "_is_gpt_oss"
]
module = ast.Module(body = helpers, type_ignores = [])
ast.fix_missing_locations(module)
namespace = {}
exec(compile(module, str(source), "exec"), namespace)
return namespace["_is_gpt_oss"]
def _model(architectures = None, model_type = None):
config = types.SimpleNamespace()
if architectures is not None:
config.architectures = architectures
if model_type is not None:
config.model_type = model_type
return types.SimpleNamespace(config = config)
def test_detects_gpt_oss_by_architecture():
# config.architectures is a list, so detection must use membership, not ==.
# A model that declares GptOssForCausalLM but has no matching model_type must
# still be routed to the mxfp4 save path.
is_gpt_oss = _load_is_gpt_oss()
assert is_gpt_oss(_model(architectures = ["GptOssForCausalLM"])) is True
assert is_gpt_oss(_model(architectures = ["GptOssForCausalLM"], model_type = "gpt_oss")) is True
def test_detects_gpt_oss_by_model_type():
is_gpt_oss = _load_is_gpt_oss()
assert is_gpt_oss(_model(architectures = ["SomethingElse"], model_type = "gpt-oss")) is True
assert is_gpt_oss(_model(architectures = ["SomethingElse"], model_type = "gpt_oss")) is True
def test_non_gpt_oss_is_false():
is_gpt_oss = _load_is_gpt_oss()
assert is_gpt_oss(_model(architectures = ["LlamaForCausalLM"], model_type = "llama")) is False
assert is_gpt_oss(_model()) is False
assert is_gpt_oss(types.SimpleNamespace()) is False

256
tests/studio/install/test_install_llama_prebuilt_logic.py
View file

@ -22,6 +22,9 @@ SPEC.loader.exec_module(INSTALL_LLAMA_PREBUILT)
PrebuiltFallback = INSTALL_LLAMA_PREBUILT.PrebuiltFallback
extract_archive = INSTALL_LLAMA_PREBUILT.extract_archive
binary_env = INSTALL_LLAMA_PREBUILT.binary_env
is_secret_env_name = INSTALL_LLAMA_PREBUILT.is_secret_env_name
scrub_env = INSTALL_LLAMA_PREBUILT.scrub_env
isolated_runtime_home = INSTALL_LLAMA_PREBUILT.isolated_runtime_home
HostInfo = INSTALL_LLAMA_PREBUILT.HostInfo
AssetChoice = INSTALL_LLAMA_PREBUILT.AssetChoice
ApprovedArtifactHash = INSTALL_LLAMA_PREBUILT.ApprovedArtifactHash
@ -779,6 +782,259 @@ def test_binary_env_linux_includes_binary_parent_in_ld_library_path(
assert str(install_dir) in ld_dirs
def test_scrub_env_drops_secrets_and_keeps_runtime_vars():
raw = {
# secrets
"HF_TOKEN": "hf_x",
"HUGGING_FACE_HUB_TOKEN": "hf_y",
"GH_TOKEN": "gh_x",
"GITHUB_TOKEN": "gh_y",
"WANDB_API_KEY": "wandb_x",
"AWS_SECRET_ACCESS_KEY": "aws_x",
"ACTIONS_ID_TOKEN_REQUEST_TOKEN": "oidc_x",
"ACTIONS_ID_TOKEN_REQUEST_URL": "https://oidc",
"SOME_VENDOR_API_KEY": "vendor_x",
"DB_PASSWORD": "pw",
"MY_PRIVATE_KEY": "pk",
"KUBECONFIG": "/home/runner/.kube/config",
"SSH_AUTH_SOCK": "/tmp/ssh-agent.sock",
"SSH_PASSPHRASE": "ssh_pass",
# runtime vars to keep
"PATH": "/usr/bin",
"LD_LIBRARY_PATH": "/opt/lib",
"DYLD_LIBRARY_PATH": "/opt/dyld",
"HOME": "/home/runner",
"TMPDIR": "/tmp",
"CUDA_VISIBLE_DEVICES": "0",
"HSA_OVERRIDE_GFX_VERSION": "11.0.0",
}
cleaned = scrub_env(raw)
for secret in (
"HF_TOKEN",
"HUGGING_FACE_HUB_TOKEN",
"GH_TOKEN",
"GITHUB_TOKEN",
"WANDB_API_KEY",
"AWS_SECRET_ACCESS_KEY",
"ACTIONS_ID_TOKEN_REQUEST_TOKEN",
"ACTIONS_ID_TOKEN_REQUEST_URL",
"SOME_VENDOR_API_KEY",
"DB_PASSWORD",
"MY_PRIVATE_KEY",
"KUBECONFIG",
"SSH_AUTH_SOCK",
"SSH_PASSPHRASE",
):
assert secret not in cleaned, f"{secret} must be stripped from binary env"
for keep in (
"PATH",
"LD_LIBRARY_PATH",
"DYLD_LIBRARY_PATH",
"HOME",
"TMPDIR",
"CUDA_VISIBLE_DEVICES",
"HSA_OVERRIDE_GFX_VERSION",
):
assert cleaned[keep] == raw[keep], f"{keep} must be preserved for the binary"
# no bare "KEY" marker: benign KEY-containing names survive
assert is_secret_env_name("API_KEY") is True
assert is_secret_env_name("SSH_KEYFILE_PATH") is False
assert is_secret_env_name("PATH") is False
def test_scrub_env_drops_proxy_index_and_embedded_url_credentials():
raw = {
# proxy / package-index URLs whose values commonly embed credentials
"HTTPS_PROXY": "https://user:secret@proxy:8080",
"https_proxy": "https://user:secret@proxy:8080", # lower-case variant
"ALL_PROXY": "socks5://user:secret@proxy:1080",
"PIP_INDEX_URL": "https://u:p@pypi.internal/simple",
"UV_INDEX_URL": "https://u:p@index.internal/simple",
# credentials embedded in an otherwise benign-named variable's value
"MY_DB_DSN": "postgres://admin:secret@db:5432/app",
# benign vars the binary needs, including a URL with no userinfo
"PATH": "/usr/bin",
"CUDA_VISIBLE_DEVICES": "0",
"NO_PROXY": "localhost,127.0.0.1",
"SOME_ENDPOINT": "https://example.com:8080/v1",
}
cleaned = scrub_env(raw)
for secret in (
"HTTPS_PROXY",
"https_proxy",
"ALL_PROXY",
"PIP_INDEX_URL",
"UV_INDEX_URL",
"MY_DB_DSN",
):
assert secret not in cleaned, f"{secret} must be stripped from binary env"
for keep in ("PATH", "CUDA_VISIBLE_DEVICES", "NO_PROXY", "SOME_ENDPOINT"):
assert cleaned[keep] == raw[keep], f"{keep} must be preserved for the binary"
assert is_secret_env_name("HTTPS_PROXY") is True
assert is_secret_env_name("https_proxy") is True
assert is_secret_env_name("NO_PROXY") is False
def test_binary_env_strips_secrets_from_downloaded_binary_environment(
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
):
install_dir = tmp_path / "llama.cpp"
bin_dir = install_dir / "build" / "bin"
bin_dir.mkdir(parents = True)
binary_path = bin_dir / "llama-server"
binary_path.write_bytes(b"fake")
host = HostInfo(
system = "Linux",
machine = "x86_64",
is_windows = False,
is_linux = True,
is_macos = False,
is_x86_64 = True,
is_arm64 = False,
nvidia_smi = None,
driver_cuda_version = None,
compute_caps = [],
visible_cuda_devices = None,
has_physical_nvidia = False,
has_usable_nvidia = False,
)
monkeypatch.setattr(INSTALL_LLAMA_PREBUILT, "linux_runtime_dirs", lambda _bp: [])
monkeypatch.setenv("HF_TOKEN", "hf_secret_from_ci")
monkeypatch.setenv("GITHUB_TOKEN", "gh_secret_from_ci")
monkeypatch.setenv("GH_TOKEN", "gh_secret_from_ci")
monkeypatch.setenv("WANDB_API_KEY", "wandb_secret_from_ci")
monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "1")
env = binary_env(binary_path, install_dir, host)
assert "HF_TOKEN" not in env
assert "GITHUB_TOKEN" not in env
assert "GH_TOKEN" not in env
assert "WANDB_API_KEY" not in env
# library/runtime resolution unaffected
assert str(bin_dir) in env["LD_LIBRARY_PATH"].split(os.pathsep)
assert env["CUDA_VISIBLE_DEVICES"] == "1"
def test_binary_env_redirects_home_away_from_real_credential_stores(
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
):
install_dir = tmp_path / "llama.cpp"
bin_dir = install_dir / "build" / "bin"
bin_dir.mkdir(parents = True)
binary_path = bin_dir / "llama-server"
binary_path.write_bytes(b"fake")
host = HostInfo(
system = "Linux",
machine = "x86_64",
is_windows = False,
is_linux = True,
is_macos = False,
is_x86_64 = True,
is_arm64 = False,
nvidia_smi = None,
driver_cuda_version = None,
compute_caps = [],
visible_cuda_devices = None,
has_physical_nvidia = False,
has_usable_nvidia = False,
)
monkeypatch.setattr(INSTALL_LLAMA_PREBUILT, "linux_runtime_dirs", lambda _bp: [])
real_home = str(tmp_path / "real_home")
monkeypatch.setenv("HOME", real_home)
monkeypatch.setenv("HF_HOME", real_home + "/.cache/huggingface")
env = binary_env(binary_path, install_dir, host)
# HOME and the cache pointers are redirected to a single empty, existing dir.
assert env["HOME"] != real_home
assert env["HF_HOME"] == env["HOME"]
assert env["HOME"] == isolated_runtime_home()
assert os.path.isdir(env["HOME"])
assert os.listdir(env["HOME"]) == []
# Windows reconstructs the profile from HOMEDRIVE + HOMEPATH.
assert env["HOMEDRIVE"] + env["HOMEPATH"] == env["HOME"]
def test_scrub_env_drops_token_only_url_userinfo():
raw = {
"GENERIC_REPO": "https://ghp_tokenonly@github.com/org/repo",
"GENERIC_OK": "https://example.com:8080/v1",
}
cleaned = scrub_env(raw)
assert "GENERIC_REPO" not in cleaned
assert cleaned["GENERIC_OK"] == raw["GENERIC_OK"]
def test_binary_env_drops_explicit_credential_file_pointers(
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
):
host = HostInfo(
system = "Linux",
machine = "x86_64",
is_windows = False,
is_linux = True,
is_macos = False,
is_x86_64 = True,
is_arm64 = False,
nvidia_smi = None,
driver_cuda_version = None,
compute_caps = [],
visible_cuda_devices = None,
has_physical_nvidia = False,
has_usable_nvidia = False,
)
monkeypatch.setattr(INSTALL_LLAMA_PREBUILT, "linux_runtime_dirs", lambda _bp: [])
dropped = (
"NETRC",
"PIP_CONFIG_FILE",
"DOCKER_CONFIG",
"GIT_CONFIG_GLOBAL",
"GITHUB_ENV",
"GITHUB_PATH",
"GITHUB_OUTPUT",
"GITHUB_STEP_SUMMARY",
"BASH_ENV",
)
for var in dropped:
monkeypatch.setenv(var, "/home/realuser/secret")
env = binary_env(tmp_path / "llama-server", tmp_path, host)
for var in dropped:
assert var not in env
def test_linux_runtime_dirs_probes_with_secret_free_env(monkeypatch: pytest.MonkeyPatch):
captured: dict[str, object] = {}
def fake_missing(binary_path, *, env = None):
captured["env"] = env
return []
monkeypatch.setattr(INSTALL_LLAMA_PREBUILT, "linux_missing_libraries", fake_missing)
monkeypatch.setenv("HF_TOKEN", "hf_secret")
monkeypatch.setenv("GITHUB_TOKEN", "gh_secret")
INSTALL_LLAMA_PREBUILT.linux_runtime_dirs(Path("/fake/llama-server"))
probe_env = captured["env"]
assert probe_env is not None
assert "HF_TOKEN" not in probe_env
assert "GITHUB_TOKEN" not in probe_env
def test_install_prebuilt_falls_back_to_older_release_plan(
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
):

12
tests/studio/playwright_chat_ui.py
View file

@ -1225,15 +1225,17 @@ with sync_playwright() as p:
# ─────────────────────────────────────────────────────
step("Shutdown via account menu")
# Re-login with NEW2 for a valid /api/shutdown token (CLI rotation
# invalidated the old one). The stale token can make the SPA auth
# guard abort this goto with ERR_ABORTED; resolve on
# domcontentloaded and tolerate it -- the pw-field wait confirms /login.
# invalidated the old one). The stale token can make the SPA auth guard
# abort this goto with ERR_ABORTED, or redirect to the same /login URL
# ("interrupted by another navigation"); resolve on domcontentloaded and
# tolerate either -- the pw-field wait below confirms we are on /login.
_tolerated_nav = ("ERR_ABORTED", "interrupted by another navigation")
try:
page.goto(f"{BASE}/login", wait_until = "domcontentloaded", timeout = 60_000)
except Exception as exc:
if "ERR_ABORTED" not in str(exc):
if not any(t in str(exc) for t in _tolerated_nav):
raise
info(f"goto /login aborted ({exc!r}); password-field wait will confirm /login")
info(f"goto /login interrupted ({exc!r}); password-field wait will confirm /login")
pw_field = page.locator("#password")
pw_field.wait_for(state = "visible", timeout = 60_000)
pw_field.fill(NEW2)

166
tests/test_vllm_broken_detection.py Normal file
View file

@ -0,0 +1,166 @@
# Unsloth - 2x faster, 60% less VRAM LLM training and finetuning
# Copyright 2023-present Daniel Han-Chen, Michael Han-Chen & the Unsloth team. All rights reserved.
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Lesser General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU Lesser General Public License for more details.
"""Regression test for #6590: modern vLLM lazy-loads its compiled extensions, so
a bare ``import vllm`` succeeds even when ``vllm._C`` (or a sibling) is ABI-broken
and ``disable_broken_vllm`` missed it. GPU-free, via a synthetic vLLM."""
from __future__ import annotations
import contextlib
import importlib.abc
import importlib.machinery
import importlib.util
import sys
import types
import pytest
_LIBCUDART_ERROR = "libcudart.so.13: cannot open shared object file: No such file or directory"
class _ExtensionLoader(importlib.abc.Loader):
"""A compiled extension that loads cleanly or fails on dlopen."""
def __init__(self, broken, error):
self.broken = broken
self.error = error
def create_module(self, spec):
return None
def exec_module(self, module):
if self.broken:
raise ImportError(self.error)
class _FakeVllmFinder(importlib.abc.MetaPathFinder):
"""Lazy vLLM: ``import vllm`` succeeds; each ``vllm._*`` ext is healthy,
ABI-broken, or absent, as real vLLM only loads ``_C`` & friends on use."""
def __init__(self, present, broken, error):
self.present = present
self.broken = broken
self.error = error
def find_spec(
self,
fullname,
path = None,
target = None,
):
if fullname in self.present:
return importlib.machinery.ModuleSpec(
name = fullname,
loader = _ExtensionLoader(broken = fullname in self.broken, error = self.error),
is_package = False,
)
return None # absent -> ModuleNotFoundError, which the guard ignores
@contextlib.contextmanager
def _fake_vllm(
present,
broken,
error = _LIBCUDART_ERROR,
):
"""Install a synthetic lazy vLLM, restoring VLLM_BROKEN, find_spec,
meta_path, and the vllm* sys.modules entries on exit."""
from unsloth import import_fixes
submodules = import_fixes._VLLM_COMPILED_EXTENSIONS
saved_meta_path = list(sys.meta_path)
saved_find_spec = importlib.util.find_spec
saved_broken = import_fixes.VLLM_BROKEN
saved_modules = {n: sys.modules.get(n) for n in ("vllm", *submodules)}
try:
import_fixes.VLLM_BROKEN = False
fake_vllm = types.ModuleType("vllm")
fake_vllm.__path__ = []
fake_vllm.__spec__ = importlib.machinery.ModuleSpec("vllm", loader = None, is_package = True)
sys.modules["vllm"] = fake_vllm
for name in submodules:
sys.modules.pop(name, None)
sys.meta_path.insert(0, _FakeVllmFinder(present, broken, error))
yield import_fixes
finally:
import_fixes.VLLM_BROKEN = saved_broken
sys.meta_path[:] = saved_meta_path
importlib.util.find_spec = saved_find_spec
for name, module in saved_modules.items():
if module is None:
sys.modules.pop(name, None)
else:
sys.modules[name] = module
@pytest.mark.parametrize(
"broken_ext",
["vllm._C", "vllm._C_stable_libtorch"],
ids = ["core_C", "sibling_C_stable_libtorch"],
)
def test_disable_broken_vllm_detects_lazy_loaded_broken_extension(broken_ext):
# A CUDA-major mismatch breaks every ext; whichever one loads first must trip detection.
present = {"vllm._C", "vllm._C_stable_libtorch"}
with _fake_vllm(present = present, broken = {broken_ext}) as import_fixes:
detected = import_fixes.disable_broken_vllm()
assert detected is True, (
f"disable_broken_vllm missed an ABI-broken {broken_ext} behind a "
"lazily-importable vllm package — issue #6590 would resurface."
)
assert import_fixes.VLLM_BROKEN is True
# Once disabled, vLLM must look absent so callers fall back cleanly.
assert importlib.util.find_spec("vllm") is None
@pytest.mark.parametrize(
"error",
[
"libnccl.so.2: cannot open shared object file: No such file or directory",
"libcuda.so.1: cannot open shared object file: No such file or directory",
],
ids = ["libnccl", "libcuda"],
)
def test_disable_broken_vllm_detects_non_cudart_so_failure(error):
# A CUDA mismatch can surface through a non-libcudart .so (libnccl, libcuda),
# which the old libcudart/libcublas/libnvrtc allow-list let slip through.
with _fake_vllm(present = {"vllm._C"}, broken = {"vllm._C"}, error = error) as import_fixes:
detected = import_fixes.disable_broken_vllm()
assert detected is True, (
f"disable_broken_vllm missed a present-but-broken vllm._C raising "
f"{error!r} — vLLM would be left enabled and crash later."
)
assert import_fixes.VLLM_BROKEN is True
@pytest.mark.parametrize(
"present",
[{"vllm._C"}, {"vllm._C", "vllm._C_stable_libtorch", "vllm._moe_C"}],
ids = ["core_only", "all_present"],
)
def test_disable_broken_vllm_keeps_healthy_vllm_enabled(present):
# Healthy install: an absent sibling (ModuleNotFoundError) or an extra present
# ext that loads cleanly must NOT be mistaken for an ABI break.
with _fake_vllm(present = present, broken = set()) as import_fixes:
detected = import_fixes.disable_broken_vllm()
assert detected is False
assert import_fixes.VLLM_BROKEN is False
assert importlib.util.find_spec("vllm") is not None
if __name__ == "__main__":
raise SystemExit(pytest.main([__file__, "-v"]))

19
unsloth/_gpu_init.py
View file

@ -36,16 +36,15 @@ from .import_fixes import (
fix_huggingface_hub,
)
# Redirect a read-only Hugging Face cache before anything below can import
# huggingface_hub / transformers / vllm (disable_broken_vllm probes
# `import vllm`, check_fbgemm_gpu_version imports transformers, and
# fix_huggingface_hub imports huggingface_hub itself), all of which can
# freeze Hub's cache constants with the un-redirected paths. unsloth_zoo
# runs the same redirect at import, but that happens after these probes.
# hf_cache.py is stdlib-only, so load it straight from its file without
# triggering the full unsloth_zoo package init this early; the zoo's own
# call later is an idempotent no-op. Older unsloth_zoo without hf_cache.py
# is skipped silently.
# Redirect a read-only Hugging Face cache before anything below imports
# huggingface_hub / transformers / vllm (disable_broken_vllm probes `import vllm`
# and its compiled extensions, check_fbgemm_gpu_version imports transformers,
# fix_huggingface_hub imports huggingface_hub) -- any of which would freeze Hub's
# cache constants with the un-redirected paths. unsloth_zoo runs the same redirect
# at import, but only after these probes. hf_cache.py is stdlib-only, so load it
# straight from its file without triggering the full unsloth_zoo init this early;
# the zoo's later call is an idempotent no-op. Older unsloth_zoo without it is
# skipped silently.
try:
import importlib.util as _importlib_util
from pathlib import Path as _Path

45
unsloth/chat_templates.py
View file

@ -1804,10 +1804,19 @@ CHAT_TEMPLATES["yi-chat"] = (yi_chat_template, yi_chat_template_eos_token, False
DEFAULT_SYSTEM_MESSAGE["yi-chat"] = None
def _change_system_message(template: str, type_chat_template: str, system_message: str = None):
system_message_pattern = r"\{system_message\}"
# For predefined templates, check if default system message exists
default_system_message = DEFAULT_SYSTEM_MESSAGE.get(f"{type_chat_template}", None)
# Custom templates have no predefined default, but may still carry a
# {system_message} placeholder. Handle it before the no-default early return
# below, which would otherwise leave the literal "{system_message}" in the
# template. A placeholder with no system message is an error, not a no-op.
if default_system_message is None and "{system_message}" in template:
if system_message is None:
raise ValueError("Unsloth: You need to provide a system message for custom templates.")
new_template = template.replace("{system_message}", system_message)
return new_template, system_message
if default_system_message is None:
if system_message is not None:
logger.warning_once(
@ -1817,21 +1826,9 @@ def _change_system_message(template: str, type_chat_template: str, system_messag
)
return template, system_message
# For custom templates
if type_chat_template is None:
has_placeholder = re.search(system_message_pattern, template) is not None
if has_placeholder:
if system_message is None:
raise ValueError("Unsloth: You need to provide a system message for custom templates.")
new_template = re.sub(system_message_pattern, system_message, template)
return new_template, system_message
return template, system_message
# For predefined templates with default system message
message_to_use = system_message if system_message is not None else default_system_message
new_template = re.sub(system_message_pattern, message_to_use, template)
new_template = template.replace("{system_message}", message_to_use)
return new_template, message_to_use
@ -2279,28 +2276,28 @@ def get_ollama_eos_tokens(tokenizer, extra_eos_tokens = []):
if getattr(tokenizer, "bos_token", None) is not None:
added_tokens_decoder = [x for x in added_tokens_decoder if x != tokenizer.bos_token]
repeatted_tokens = []
repeated_tokens = []
# Join all vocab
joined_text = "\x01\x00".join(added_tokens_decoder)
for token in added_tokens_decoder:
n = len(token)
repeatted_counts = joined_text.count(token[:n//2])
repeated_counts = joined_text.count(token[:n//2])
# Try finding longer than 1/2 of the token in the rest
# For eg <|reserved_special_token_0|>, <|reserved_special_token_1|>
if repeatted_counts > 2:
if repeated_counts > 2:
for j in range(n//2+1, n):
if joined_text.count(token[:j]) < repeatted_counts:
if joined_text.count(token[:j]) < repeated_counts:
j -= 1
# Remove repeatted tokens to reduce search space
# Remove repeated tokens to reduce search space
joined_text = joined_text.replace(token[:j], "")
repeatted_tokens.append(token[:j])
repeated_tokens.append(token[:j])
break
# Remove duplicates
splitted = joined_text.split("\x01\x00")
final_eos_tokens = [old for old, new in zip(added_tokens_decoder, splitted) if old == new]
split = joined_text.split("\x01\x00")
final_eos_tokens = [old for old, new in zip(added_tokens_decoder, split) if old == new]
final_eos_tokens += extra_eos_tokens
final_eos_tokens += repeatted_tokens
final_eos_tokens += repeated_tokens
# Remove new lines, spaces and HTML tags
filtered_eos_tokens = []

41
unsloth/dataprep/raw_text.py
View file

@ -223,48 +223,63 @@ class RawTextDataLoader:
return "\n\n".join(texts)
return ""
# Cache text fields/columns for better performance
_TEXT_FIELDS = ("text", "content", "message", "body", "description", "prompt")
_TEXT_COLUMNS = _TEXT_FIELDS
def _extract_text_from_json(self, data):
"""Extract text from JSON object using common field names."""
text_fields = ["text", "content", "message", "body", "description", "prompt"]
for field in text_fields:
for field in self._TEXT_FIELDS:
if field in data and isinstance(data[field], str):
return data[field]
return ""
def _extract_text_from_csv_row(self, row):
"""Extract text from CSV row using common column names."""
text_columns = ["text", "content", "message", "body", "description", "prompt"]
for column in text_columns:
for column in self._TEXT_COLUMNS:
if column in row and row[column]:
return row[column]
return ""
class TextPreprocessor:
# Compile regex patterns once for better performance
_WHITESPACE_PATTERN = re.compile(r"[^\S\n]+")
_INVALID_CHARS_PATTERN = re.compile(r"[^\x20-\x7E\n]")
_MULTIPLE_SPACES_PATTERN = re.compile(r"[ ]{2,}")
_NEWLINE_SPACES_PATTERN = re.compile(r" *\n *")
_MULTIPLE_NEWLINES_PATTERN = re.compile(r"\n{3,}")
_CHAPTER_PATTERN = re.compile(r"^# (.+)$", re.MULTILINE)
_SECTION_PATTERN = re.compile(r"^## (.+)$", re.MULTILINE)
_SUBSECTION_PATTERN = re.compile(r"^### (.+)$", re.MULTILINE)
_CODE_BLOCK_PATTERN = re.compile(r"```(\w*)\n(.*?)\n```", re.DOTALL)
def clean_text(self, text):
"""Remove unwanted characters, normalize whitespace"""
text = text.replace("\r\n", "\n").replace("\r", "\n")
text = re.sub(r"[^\S\n]+", " ", text)
text = re.sub(r"[^\x20-\x7E\n]", "", text)
text = re.sub(r"[ ]{2,}", " ", text)
text = re.sub(r" *\n *", "\n", text)
text = re.sub(r"\n{3,}", "\n\n", text)
text = self._WHITESPACE_PATTERN.sub(" ", text)
text = self._INVALID_CHARS_PATTERN.sub("", text)
text = self._MULTIPLE_SPACES_PATTERN.sub(" ", text)
text = self._NEWLINE_SPACES_PATTERN.sub("\n", text)
text = self._MULTIPLE_NEWLINES_PATTERN.sub("\n\n", text)
return text.strip()
def extract_sections(self, text, patterns):
"""Extract specific sections (e.g., code blocks, quotes)"""
sections = []
for pattern in patterns:
# Compile pattern on first use and cache? Well, patterns are user-provided,
# so just use re.findall with compiled flags
matches = re.findall(pattern, text, re.MULTILINE | re.DOTALL)
sections.extend(matches)
return sections
def add_structure_tokens(self, text):
"""Add special tokens for structure (chapters, sections)"""
text = re.sub(r"^# (.+)$", r"<|chapter|>\1<|/chapter|>", text, flags = re.MULTILINE)
text = re.sub(r"^## (.+)$", r"<|section|>\1<|/section|>", text, flags = re.MULTILINE)
text = re.sub(r"^### (.+)$", r"<|subsection|>\1<|/subsection|>", text, flags = re.MULTILINE)
text = re.sub(r"```(\w*)\n(.*?)\n```", r"<|code|\1|>\2<|/code|>", text, flags = re.DOTALL)
text = self._CHAPTER_PATTERN.sub(r"<|chapter|>\1<|/chapter|>", text)
text = self._SECTION_PATTERN.sub(r"<|section|>\1<|/section|>", text)
text = self._SUBSECTION_PATTERN.sub(r"<|subsection|>\1<|/subsection|>", text)
text = self._CODE_BLOCK_PATTERN.sub(r"<|code|\1|>\2<|/code|>", text)
return text
def validate_dataset(self, dataset):

32
unsloth/import_fixes.py
View file

@ -158,6 +158,11 @@ if not UNSLOTH_ENABLE_LOGGING:
logging.getLogger("torchao").addFilter(
HideLoggingMessage("Skipping import of cpp extensions due to incompatible torch version")
)
# torch >= 2.11 path: torchao dlopens each prebuilt _C*.so and logs "Failed to load
# .../_C*.so" when one can't (ABI tag mismatch in the wheel, e.g. a cp310 .so under a
# cp312 runtime on Colab, or an arch-specific kernel the GPU lacks). It falls back to
# non-cpp paths and Unsloth doesn't use these kernels, so drop the cosmetic record.
logging.getLogger("torchao").addFilter(HideLoggingMessage("Failed to load "))
# SyntaxWarning: invalid escape sequence '\.'
warnings.filterwarnings("ignore", message = "invalid escape sequence", category = SyntaxWarning)
# PYTORCH_CUDA_ALLOC_CONF is deprecated warning from torch
@ -2349,11 +2354,9 @@ def _is_broken_vllm_error(error) -> bool:
)
) or ("vllm" in message and "undefined symbol" in message):
return True
# Also catch CUDA shared library mismatches during vllm import
# e.g. "libcudart.so.12: cannot open shared object file"
if (
"libcudart" in message or "libcublas" in message or "libnvrtc" in message
) and "cannot open shared object file" in message:
# Forced extension load raises the bare loader error (no "vllm._C"
# wrapper); match any .so failure as callers feed only vLLM imports.
if "cannot open shared object file" in message:
return True
current = getattr(current, "__cause__", None) or getattr(current, "__context__", None)
return False
@ -2545,6 +2548,16 @@ def _clear_vllm_modules():
sys.modules.pop(module_name, None)
# vLLM's compiled extensions. A CUDA-major ABI break hits all of them, so
# probing the eagerly-loaded _C and its siblings reliably trips it.
_VLLM_COMPILED_EXTENSIONS = (
"vllm._C",
"vllm._C_stable_libtorch",
"vllm._moe_C",
"vllm._rocm_C",
)
def disable_broken_vllm(error = None):
"""Disable vLLM dynamically when its shared library is ABI-broken."""
global VLLM_BROKEN
@ -2562,6 +2575,15 @@ def disable_broken_vllm(error = None):
try:
import vllm # noqa: F401
# Lazy vLLM lets a bare `import vllm` succeed even when an extension
# is ABI-broken; force-load each to surface the .so failure here.
# A missing one raises ModuleNotFoundError (skipped below).
for _ext in _VLLM_COMPILED_EXTENSIONS:
try:
importlib.import_module(_ext)
except ModuleNotFoundError:
pass
return False
except Exception as import_error:
failure = import_error

2
unsloth/kernels/geglu.py
View file

@ -97,7 +97,7 @@ def _exact_backward_kernel(
e_row = tl.load(e + offsets, mask = mask, other = 0).to(tl.float32)
g_row = tl.load(g + offsets, mask = mask, other = 0) # .to(tl.float32)
# Break e_row away for re-use
# Break e_row away for reuse
# f = 1/2 * e * (1 + erf(1/sqrt(2) * e))
f_partial_row = 0.5 * (tl.math.erf(tl.math.rsqrt(2.0) * e_row) + 1.0)
f_row = f_partial_row * e_row

2
unsloth/kernels/moe/grouped_gemm/kernels/forward.py
View file

@ -111,7 +111,7 @@ def _grouped_gemm_forward_kernel(
while tidx >= processed_tiles and tidx < processed_tiles + num_tiles_per_expert:
tile_idx = tidx - processed_tiles
# Check if L2 cache re-use for this order is optimal
# Check if L2 cache reuse for this order is optimal
tile_m_idx = tile_idx % num_m_tiles
tile_n_idx = tile_idx // num_m_tiles

83
unsloth/models/_utils.py
View file

@ -86,6 +86,8 @@ __all__ = [
"is_moe_model",
"get_moe_target_parameters",
"make_fast_generate_wrapper",
"_mark_unsloth_disable_data_parallel",
"_patch_transformers_trainer_data_parallel",
]
import torch
@ -160,6 +162,85 @@ from unsloth_zoo.training_utils import (
)
def _iter_wrapped_models(model):
seen = set()
current = model
while current is not None and id(current) not in seen:
yield current
seen.add(id(current))
next_model = getattr(current, "model", None)
if next_model is None:
next_model = getattr(current, "base_model", None)
if next_model is None:
next_model = getattr(current, "module", None)
current = next_model
def _patch_transformers_trainer_data_parallel():
try:
from transformers.trainer import Trainer
except (ImportError, ModuleNotFoundError):
return False
original_wrap_model = getattr(Trainer, "_wrap_model", None)
if original_wrap_model is None:
return False
if getattr(original_wrap_model, "_unsloth_data_parallel_patched", False):
return True
try:
supports_dataloader = "dataloader" in inspect.signature(original_wrap_model).parameters
except (TypeError, ValueError):
supports_dataloader = True
def _call_original_wrap_model(self, model, wrap_args, wrap_kwargs):
if supports_dataloader:
return original_wrap_model(self, model, *wrap_args, **wrap_kwargs)
if "dataloader" in wrap_kwargs:
wrap_kwargs = {k: v for k, v in wrap_kwargs.items() if k != "dataloader"}
return original_wrap_model(self, model, *wrap_args, **wrap_kwargs)
@functools.wraps(original_wrap_model)
def _unsloth_wrap_model(self, model, *wrap_args, **wrap_kwargs):
args = getattr(self, "args", None)
disable_data_parallel = getattr(model, "_unsloth_disable_data_parallel", False)
is_real_8bit = getattr(model, "is_loaded_in_8bit", False)
if (
args is None
or not disable_data_parallel
or is_real_8bit
or getattr(args, "n_gpu", 0) <= 1
):
return _call_original_wrap_model(self, model, wrap_args, wrap_kwargs)
had_n_gpu = hasattr(args, "_n_gpu")
old_n_gpu = getattr(args, "_n_gpu", None)
args._n_gpu = 1
try:
return _call_original_wrap_model(self, model, wrap_args, wrap_kwargs)
finally:
if had_n_gpu:
args._n_gpu = old_n_gpu
else:
try:
delattr(args, "_n_gpu")
except AttributeError:
pass
_unsloth_wrap_model._unsloth_data_parallel_patched = True
_unsloth_wrap_model._unsloth_original_wrap_model = original_wrap_model
Trainer._wrap_model = _unsloth_wrap_model
return True
def _mark_unsloth_disable_data_parallel(model, disable = True):
if disable:
_patch_transformers_trainer_data_parallel()
for module in _iter_wrapped_models(model):
setattr(module, "_unsloth_disable_data_parallel", bool(disable))
return model
def resolve_hip_gpu_stats_name(gpu_stats):
name = str(getattr(gpu_stats, "name", "") or "").strip()
name = re.sub(r"\s*\([^)]*\)\s*$", "", name).strip()
@ -2997,7 +3078,7 @@ class TorchAOConfig:
def _untie_input_output_embeddings(model: torch.nn.Module) -> None:
"""
Utility to untie input/output embeddings in a HuggingFace model.
This is useful if we want to quantize the input/ouput embeddings differently.
This is useful if we want to quantize the input/output embeddings differently.
Model is modified in-place.
"""

1
unsloth/models/falcon_h1.py
View file

@ -38,7 +38,6 @@ try:
FalconH1Model,
FalconH1ForCausalLM,
FalconH1RMSNorm,
FalconH1RMSNormGated,
FalconHybridMambaAttentionDynamicCache,
)
except:

12
unsloth/models/llama.py
View file

@ -2832,13 +2832,11 @@ class FastLlamaModel:
internal_model = model
while hasattr(internal_model, "model"):
internal_model._saved_temp_tokenizer = tokenizer
# Also set is_loaded_in_8bit to disable incorrect DDP
internal_model.is_loaded_in_8bit = True
internal_model = internal_model.model
internal_model._saved_temp_tokenizer = tokenizer
# Also set is_loaded_in_8bit to disable incorrect DDP
internal_model.is_loaded_in_8bit = True
# Prevent Transformers Trainer from auto-wrapping Unsloth LoRA models in DP.
_mark_unsloth_disable_data_parallel(model)
# For transformers > 4.47.1, we need to add rotary_emb to all attention layers
if IS_ATTENTION_REFACTOR or hasattr(model.model, "rotary_emb"):
@ -3379,13 +3377,11 @@ class FastLlamaModel:
while hasattr(internal_model, "model"):
if hasattr(internal_model, "_saved_temp_tokenizer"):
internal_model._saved_temp_tokenizer.padding_side = "right"
# Also set is_loaded_in_8bit to disable incorrect DDP
internal_model.is_loaded_in_8bit = True
internal_model = internal_model.model
if hasattr(internal_model, "_saved_temp_tokenizer"):
internal_model._saved_temp_tokenizer.padding_side = "right"
# Also set is_loaded_in_8bit to disable incorrect DDP
internal_model.is_loaded_in_8bit = True
# Prevent Transformers Trainer from auto-wrapping Unsloth LoRA models in DP.
_mark_unsloth_disable_data_parallel(model)
# Clear deleted GPU items
for _ in range(3):

26
unsloth/models/rl.py
View file

@ -994,8 +994,18 @@ def _patch_trl_rl_trainers_impl(trainer_file = "grpo_trainer"):
"use_fp16 = getattr(args, 'fp16', False)\n"
"if type(use_fp16) is not bool: use_fp16 = False\n"
"force_float32 = False\n"
# device-aware bf16 check (CUDA/XPU/HIP), so V100/T4 never pick bf16
# but AMD/Intel are unaffected; fall back on older unsloth_zoo.
"try:\n"
" from unsloth_zoo.device_type import device_is_bf16_supported as _bf16_supported\n"
"except Exception:\n"
" _bf16_supported = torch.cuda.is_bf16_supported\n"
# FORCE_FLOAT32 models (Gemma3, gpt_oss, ...) cannot use float16. On a GPU without
# bf16 (V100/T4) keep them in float32 so they never autocast to fp16. On a bf16 GPU,
# full finetuning can still use bf16 autocast (master weights stay float32), which is
# faster and uses less memory; LoRA/QLoRA keep float32 when forced.
"full_finetuning = os.environ.get('UNSLOTH_ENABLE_FULL_FINETUNING', '0') == '1'\n"
"if not full_finetuning and (os.environ.get('UNSLOTH_FORCE_FLOAT32', '0') == '1'):\n"
"if os.environ.get('UNSLOTH_FORCE_FLOAT32', '0') == '1' and not (full_finetuning and _bf16_supported()):\n"
" print('Unsloth: Switching to float32 training since model cannot work with float16')\n"
" force_float32 = True\n"
"mixed_precision_dtype = os.environ.get('UNSLOTH_MIXED_PRECISION', 'float32')\n"
@ -1004,8 +1014,9 @@ def _patch_trl_rl_trainers_impl(trainer_file = "grpo_trainer"):
"from unsloth_zoo.utils import _get_dtype\n"
"dtype = _get_dtype(dtype)\n"
"float16 = dtype == torch.float16\n"
"bfloat16 = dtype == torch.bfloat16\n"
"if not force_float32 and (float16 and use_bf16): raise TypeError('Unsloth: Model is in float16 precision but you want to use bfloat16 precision. Set fp16 to `True` and bf16 to `False`')\n"
"if not force_float32 and (not float16 and use_fp16): raise TypeError('Unsloth: Model is in bfloat16 precision but you want to use float16 precision. Set fp16 to `False` and bf16 to `True`')\n"
"if not force_float32 and (bfloat16 and use_fp16): raise TypeError('Unsloth: Model is in bfloat16 precision but you want to use float16 precision. Set fp16 to `False` and bf16 to `True`')\n"
"if force_float32:\n"
" # Forced float32 training\n"
" args.fp16 = False\n"
@ -1014,11 +1025,12 @@ def _patch_trl_rl_trainers_impl(trainer_file = "grpo_trainer"):
" if hasattr(args, 'mixed_precision'): args.mixed_precision = 'no'\n"
" # args.mixed_precision is a new argument which needs to be set now\n"
"elif (not use_bf16 and not use_fp16) and mixed_precision_dtype == 'float32':\n"
" # Mixed precision training\n"
" args.fp16 = float16\n"
" args.bf16 = not float16\n"
" os.environ['ACCELERATE_MIXED_PRECISION'] = 'fp16' if float16 else 'bf16'\n"
" if hasattr(args, 'mixed_precision'): args.mixed_precision = 'fp16' if float16 else 'bf16'\n"
" # Mixed precision training. bf16 only if the GPU supports it; V100/T4 use fp16.\n"
" use_bf16_amp = (not float16) and _bf16_supported()\n"
" args.fp16 = not use_bf16_amp\n"
" args.bf16 = use_bf16_amp\n"
" os.environ['ACCELERATE_MIXED_PRECISION'] = 'bf16' if use_bf16_amp else 'fp16'\n"
" if hasattr(args, 'mixed_precision'): args.mixed_precision = 'bf16' if use_bf16_amp else 'fp16'\n"
" # args.mixed_precision is a new argument which needs to be set now\n"
"elif mixed_precision_dtype == 'bfloat16':\n"
" # Both False since bfloat16 full finetuning doesn't do any autocasting.\n"

12
unsloth/models/vision.py
View file

@ -1440,16 +1440,14 @@ class FastBaseModel:
while hasattr(m, "model"):
m.max_seq_length = max_seq_length
m._saved_temp_tokenizer = tokenizer
# Also set is_loaded_in_8bit to disable incorrect DDP
m.is_loaded_in_8bit = True if not full_finetuning else False
m = m.model
m.max_seq_length = max_seq_length
# Save to modules as well
for module in model.modules():
module.max_seq_length = max_seq_length
m._saved_temp_tokenizer = tokenizer
# Also set is_loaded_in_8bit to disable incorrect DDP
m.is_loaded_in_8bit = True if not full_finetuning else False
# Prevent Transformers Trainer from auto-wrapping Unsloth LoRA models in DP.
_mark_unsloth_disable_data_parallel(model, disable = not full_finetuning)
# Patch generate
if os.environ.get("UNSLOTH_DISABLE_FAST_GENERATION", "0") == "0" and hasattr(
@ -1826,14 +1824,12 @@ class FastBaseModel:
if hasattr(m, "_saved_temp_tokenizer"):
if hasattr(m._saved_temp_tokenizer, "tokenizer"):
m._saved_temp_tokenizer.tokenizer.padding_side = "left"
# Also set is_loaded_in_8bit to disable incorrect DDP
m.is_loaded_in_8bit = True if not full_finetuning else False
m = m.model
if hasattr(m, "_saved_temp_tokenizer"):
if hasattr(m._saved_temp_tokenizer, "tokenizer"):
m._saved_temp_tokenizer.tokenizer.padding_side = "left"
# Also set is_loaded_in_8bit to disable incorrect DDP
m.is_loaded_in_8bit = True if not full_finetuning else False
# Prevent Transformers Trainer from auto-wrapping Unsloth LoRA models in DP.
_mark_unsloth_disable_data_parallel(model, disable = not full_finetuning)
# Clear deleted GPU items
for _ in range(3):

113
unsloth/save.py
View file

@ -226,6 +226,29 @@ def _normalize_compressed_method(save_method):
return None
def _is_cmake_only_llama_cpp(llama_cpp_dir: str = "llama.cpp") -> bool:
"""
True if llama.cpp's Makefile is the post-CMake-migration deprecation stub,
so `make` cannot build it. A genuinely missing/empty checkout returns False
so it isn't treated as CMake-only: the caller then probes make and fails
loudly on a real error rather than silently assuming a CMake build.
"""
makefile_path = os.path.join(llama_cpp_dir, "Makefile")
if not os.path.exists(makefile_path):
# No Makefile: only CMake-only if a real CMake project is present
return os.path.exists(os.path.join(llama_cpp_dir, "CMakeLists.txt"))
try:
with open(makefile_path, "r", encoding = "utf-8", errors = "ignore") as f:
content = f.read(4096).lower()
if "cmake" in content and "deprecated" in content:
return True
if "build system changed" in content:
return True
except (IOError, OSError):
pass
return False
def print_quantization_methods():
for key, value in ALLOWED_QUANTS.items():
print(f'"{key}" ==> {value}')
@ -565,6 +588,17 @@ def _is_qwen3_5_vlm(model):
) or getattr(config, "model_type", None) in ("qwen3_5", "qwen3_5_moe")
def _is_gpt_oss(model):
config = getattr(model, "config", None)
if config is None:
return False
architectures = getattr(config, "architectures", None) or ()
return "GptOssForCausalLM" in architectures or getattr(config, "model_type", None) in (
"gpt-oss",
"gpt_oss",
)
def _qwen3_5_vlm_state_dict_for_save(state_dict):
remapped_state_dict = {}
for key, value in state_dict.items():
@ -1269,14 +1303,27 @@ def install_llama_cpp_make_non_blocking():
# https://github.com/ggerganov/llama.cpp/issues/7062
# Weirdly GPU conversion for GGUF breaks??
# env = { **os.environ, "LLAMA_CUDA": "1", }
# Force make clean
check = os.system("make clean -C llama.cpp")
IS_CMAKE = False
if check == 0:
# Skip the make-clean probe on CMake-only checkouts (its error output is misleading)
IS_CMAKE = _is_cmake_only_llama_cpp("llama.cpp")
if not IS_CMAKE:
# Confirm make still works, silently
try:
result = subprocess.run(
["make", "clean", "-C", "llama.cpp"],
stdout = subprocess.DEVNULL,
stderr = subprocess.DEVNULL,
)
IS_CMAKE = result.returncode != 0
except FileNotFoundError:
# No make executable; use CMake
IS_CMAKE = True
if not IS_CMAKE:
# Uses old MAKE
n_jobs = max(int((psutil.cpu_count() or 1) * 1.5), 1)
full_command = ["make", "all", "-j" + str(n_jobs), "-C", "llama.cpp"]
IS_CMAKE = False
else:
# Uses new CMAKE
n_jobs = max(int(psutil.cpu_count() or 1), 1) # Use less CPUs since 1.5x faster
@ -1299,7 +1346,6 @@ def install_llama_cpp_make_non_blocking():
"--clean-first",
"--target",
] + LLAMA_CPP_TARGETS
IS_CMAKE = True
# https://github.com/ggerganov/llama.cpp/issues/7062
# Weirdly GPU conversion for GGUF breaks??
# run_installer = subprocess.Popen(full_command, env = env, stdout = subprocess.DEVNULL, stderr = subprocess.STDOUT)
@ -1468,20 +1514,25 @@ def install_llama_cpp_old(version = -10):
]
try_execute(commands)
# Try using MAKE
commands = [
"make clean -C llama.cpp",
f"make all -j{(psutil.cpu_count() or 1)*2} -C llama.cpp",
]
if try_execute(commands) == "CMAKE":
# Instead use CMAKE
# Detect CMake-only build system before trying make
use_cmake = _is_cmake_only_llama_cpp("llama.cpp")
if not use_cmake:
# Try using MAKE
commands = [
"make clean -C llama.cpp",
f"make all -j{(psutil.cpu_count() or 1)*2} -C llama.cpp",
]
use_cmake = try_execute(commands) == "CMAKE"
if use_cmake:
# Use CMAKE
commands = [
f"cmake llama.cpp -B llama.cpp/build -DBUILD_SHARED_LIBS=OFF -DGGML_CUDA=OFF {CURL_FLAG}",
f"cmake --build llama.cpp/build --config Release -j{(psutil.cpu_count() or 1)*2} --clean-first --target {' '.join(LLAMA_CPP_TARGETS)}",
"cp llama.cpp/build/bin/llama-* llama.cpp",
"rm -rf llama.cpp/build",
]
try_execute(commands)
# Check if successful
@ -1513,15 +1564,21 @@ def install_llama_cpp_blocking(use_cuda = False):
return
try_execute(commands)
commands = [
"make clean -C llama.cpp",
# https://github.com/ggerganov/llama.cpp/issues/7062
# Weirdly GPU conversion for GGUF breaks??
# f"{use_cuda} make all -j{(psutil.cpu_count() or 1)*2} -C llama.cpp",
f"make all -j{(psutil.cpu_count() or 1)*2} -C llama.cpp",
]
if try_execute(commands) == "CMAKE":
# Instead use CMAKE
# Detect CMake-only build system before trying make
use_cmake = _is_cmake_only_llama_cpp("llama.cpp")
if not use_cmake:
commands = [
"make clean -C llama.cpp",
# https://github.com/ggerganov/llama.cpp/issues/7062
# Weirdly GPU conversion for GGUF breaks??
# f"{use_cuda} make all -j{(psutil.cpu_count() or 1)*2} -C llama.cpp",
f"make all -j{(psutil.cpu_count() or 1)*2} -C llama.cpp",
]
use_cmake = try_execute(commands) == "CMAKE"
if use_cmake:
# Use CMAKE
commands = [
f"cmake llama.cpp -B llama.cpp/build -DBUILD_SHARED_LIBS=OFF -DGGML_CUDA=OFF {CURL_FLAG}",
f"cmake --build llama.cpp/build --config Release -j{(psutil.cpu_count() or 1)*2} --clean-first --target {' '.join(LLAMA_CPP_TARGETS)}",
@ -2541,15 +2598,7 @@ def unsloth_save_pretrained_gguf(
is_processor = is_vlm and isinstance(tokenizer, ProcessorMixin)
is_gpt_oss = (
True
if (
hasattr(self.config, "architectures")
and self.config.architectures == "GptOssForCausalLM"
)
or (hasattr(self.config, "model_type") and self.config.model_type in ["gpt-oss", "gpt_oss"])
else False
)
is_gpt_oss = _is_gpt_oss(self)
# Step 2: Prepare arguments for model saving
arguments = dict(locals())
arguments["model"] = self

108
unsloth_cli/_inference.py
View file

@ -3,11 +3,12 @@
"""Model loading and streaming shared by `inference` and `chat`."""
import asyncio
import os
import re
import sys
from pathlib import Path
from typing import Optional
from typing import List, Optional
import typer
@ -211,28 +212,65 @@ def resolve_model_config(model: str, *, hf_token: Optional[str]):
return model_config
def _load_gguf_backend(model_config, *, hf_token, max_seq_length):
def _validate_llama_extra_args_or_exit(llama_extra_args: Optional[List[str]]) -> list[str]:
from core.inference.llama_server_args import validate_extra_args
try:
return validate_extra_args(llama_extra_args)
except ValueError as exc:
typer.echo(f"Error: {exc}", err = True)
raise typer.Exit(code = 1)
def _load_gguf_backend(
model_config,
*,
hf_token,
max_seq_length,
tensor_parallel: bool = False,
llama_extra_args: Optional[List[str]] = None,
):
ensure_studio_backend_path()
from core.inference.llama_cpp import LlamaCppBackend
from core.inference.tensor_fallback import load_with_tensor_fallback
llama_backend = LlamaCppBackend()
extra_args = _validate_llama_extra_args_or_exit(llama_extra_args)
common = dict(
hf_variant = model_config.gguf_variant,
model_identifier = model_config.identifier,
is_vision = model_config.is_vision,
n_ctx = max_seq_length,
)
if model_config.gguf_hf_repo:
loaded = llama_backend.load_model(
hf_repo = model_config.gguf_hf_repo, hf_token = hf_token, **common
async def _attempt_gguf_load(
requested_tensor_parallel: bool, attempt_extra_args: Optional[List[str]]
) -> bool:
attempt_common = dict(
common,
tensor_parallel = requested_tensor_parallel,
extra_args = attempt_extra_args,
)
else:
loaded = llama_backend.load_model(
if model_config.gguf_hf_repo:
return llama_backend.load_model(
hf_repo = model_config.gguf_hf_repo,
hf_token = hf_token,
**attempt_common,
)
return llama_backend.load_model(
gguf_path = model_config.gguf_file,
mmproj_path = model_config.gguf_mmproj_file,
mtp_draft_path = model_config.gguf_mtp_file,
**common,
**attempt_common,
)
loaded = asyncio.run(
load_with_tensor_fallback(
_attempt_gguf_load,
requested_tensor = tensor_parallel,
extra_args = extra_args,
label = model_config.identifier,
)
)
if not loaded:
typer.echo("Model load failed", err = True)
raise typer.Exit(code = 1)
@ -245,6 +283,8 @@ def load_chat_backend(
hf_token: Optional[str],
max_seq_length: int,
load_in_4bit: bool,
tensor_parallel: bool = False,
llama_extra_args: Optional[List[str]] = None,
model_config = None,
fresh_backend: bool = False,
):
@ -259,7 +299,13 @@ def load_chat_backend(
typer.echo(f"Loading {model}", err = True)
if model_config.is_gguf:
return _load_gguf_backend(model_config, hf_token = hf_token, max_seq_length = max_seq_length)
return _load_gguf_backend(
model_config,
hf_token = hf_token,
max_seq_length = max_seq_length,
tensor_parallel = tensor_parallel,
llama_extra_args = llama_extra_args,
)
if fresh_backend:
ensure_studio_backend_path()
@ -447,18 +493,31 @@ class HttpChatBackend:
# No redirects: this carries a bearer token (see urlopen_no_redirect).
return urlopen_no_redirect(request, timeout = timeout)
def ensure_loaded(self, model: str, *, hf_token, max_seq_length, load_in_4bit) -> None:
def ensure_loaded(
self,
model: str,
*,
hf_token,
max_seq_length,
load_in_4bit,
tensor_parallel: bool = False,
llama_extra_args: Optional[List[str]] = None,
) -> None:
typer.echo(f"Loading {model} on the Studio server", err = True)
payload = {
"model_path": model,
"hf_token": hf_token,
"max_seq_length": max_seq_length,
"load_in_4bit": load_in_4bit,
"tensor_parallel": tensor_parallel,
}
if llama_extra_args:
payload["llama_extra_args"] = llama_extra_args
try:
self._request(
"POST",
"/api/inference/load",
{
"model_path": model,
"hf_token": hf_token,
"max_seq_length": max_seq_length,
"load_in_4bit": load_in_4bit,
},
payload,
).close()
except Exception as exc:
typer.echo(f"Model load failed: {exc}", err = True)
@ -538,7 +597,15 @@ class HttpChatBackend:
pass
def connect_studio_server(model: str, *, hf_token, max_seq_length, load_in_4bit):
def connect_studio_server(
model: str,
*,
hf_token,
max_seq_length,
load_in_4bit,
tensor_parallel: bool = False,
llama_extra_args: Optional[List[str]] = None,
):
"""Backend on a running Studio server, or None (caller loads locally)."""
base_url = find_studio_server()
if not base_url:
@ -576,6 +643,11 @@ def connect_studio_server(model: str, *, hf_token, max_seq_length, load_in_4bit)
return _refuse("couldn't self-issue a Studio token (is Studio set up here?).")
backend = HttpChatBackend(base_url, token)
backend.ensure_loaded(
model, hf_token = hf_token, max_seq_length = max_seq_length, load_in_4bit = load_in_4bit
model,
hf_token = hf_token,
max_seq_length = max_seq_length,
load_in_4bit = load_in_4bit,
tensor_parallel = tensor_parallel,
llama_extra_args = llama_extra_args,
)
return backend

26
unsloth_cli/commands/chat.py
View file

@ -1,7 +1,7 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
from typing import Optional
from typing import List, Optional
import typer
from rich.console import Console
@ -153,6 +153,22 @@ def chat(
),
max_seq_length: int = typer.Option(4096, "--max-seq-length"),
load_in_4bit: bool = typer.Option(True, "--load-in-4bit/--no-load-in-4bit"),
tensor_parallel: bool = typer.Option(
False,
"--tensor-parallel/--no-tensor-parallel",
help = (
"Split a GGUF across GPUs by tensor (--split-mode tensor) instead "
"of by layer. Ignored for non-GGUF models."
),
),
llama_extra_args: Optional[List[str]] = typer.Option(
None,
"--llama-extra-arg",
help = (
"Extra llama-server arg for GGUF models. Repeat for multiple "
"tokens, e.g. --llama-extra-arg=--top-k --llama-extra-arg 20."
),
),
think: bool = typer.Option(
False,
"--think/--no-think",
@ -190,7 +206,13 @@ def chat(
err.print(f"--compare unavailable: {compare_blocked}", style = "red", markup = False)
raise typer.Exit(code = 1)
load_opts = dict(hf_token = hf_token, max_seq_length = max_seq_length, load_in_4bit = load_in_4bit)
load_opts = dict(
hf_token = hf_token,
max_seq_length = max_seq_length,
load_in_4bit = load_in_4bit,
tensor_parallel = tensor_parallel,
llama_extra_args = llama_extra_args,
)
# Prefer a running Studio server: instant starts, model shared with the UI.
chat_backend = None if no_server else connect_studio_server(model, **load_opts)

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