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

Merge branch 'main' into pip

Browse source
This commit is contained in:
Daniel Han 2026-05-31 07:10:55 -07:00
commit 5688072af6
168 changed files with 27831 additions and 2342 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -0,0 +1,57 @@
#!/usr/bin/env bash
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
#
# Assert Studio installed a llama.cpp that loads and runs on THIS macOS. Tests
# the contract that matters (binaries load and their minimum-OS is <= this host)
# instead of the old "did install.sh fall back to a source build?" grep, since a
# source build with a correct deployment target is a valid outcome.
set -uo pipefail
UNSLOTH_HOME="${STUDIO_HOME:-$HOME/.unsloth}"
LLAMA_DIR="${LLAMA_CPP_DIR:-$UNSLOTH_HOME/llama.cpp}"
BIN_DIR="$LLAMA_DIR/build/bin"
fail() {
echo "::error::$*"
if [ -f logs/install.log ]; then
echo "---- install.log (llama.cpp lines) ----"
grep -E "llama-prebuilt|llama\.cpp|macos prebuilt|falling back" logs/install.log | tail -80 || true
fi
exit 1
}
SERVER="$(find "$LLAMA_DIR" -type f -name 'llama-server' 2>/dev/null | head -1)"
QUANT="$(find "$LLAMA_DIR" -type f -name 'llama-quantize' 2>/dev/null | head -1)"
[ -n "$SERVER" ] || fail "llama-server not found under $LLAMA_DIR after install"
[ -n "$QUANT" ] || fail "llama-quantize not found under $LLAMA_DIR after install"
HOST_VER="$(sw_vers -productVersion 2>/dev/null || echo '0')"
HOST_MAJOR="${HOST_VER%%.*}"
# Static minimum-OS check on every Mach-O we ship. vtool ships with the Xcode
# command line tools, which GitHub macOS runners always have; if it is somehow
# missing we skip the static check and rely on the runtime launch below.
if command -v vtool >/dev/null 2>&1; then
while IFS= read -r macho; do
[ -n "$macho" ] || continue
minos="$(vtool -show-build "$macho" 2>/dev/null | awk '/minos/{print $2; exit}')"
[ -n "$minos" ] || continue
min_major="${minos%%.*}"
if [ "$min_major" -gt "$HOST_MAJOR" ] 2>/dev/null; then
fail "$(basename "$macho") is built for macOS $minos but this runner is macOS $HOST_VER (prebuilt is newer than the host)"
fi
done < <(find "$BIN_DIR" -type f \( -name '*.dylib' -o -name 'llama-server' -o -name 'llama-quantize' \) 2>/dev/null)
fi
# Runtime launch: --version forces dyld to load every linked dylib (including
# libggml-metal.dylib). A missing Metal symbol or too-new binary fails here.
if ! "$SERVER" --version >/tmp/llama-server-version.txt 2>&1; then
echo "---- llama-server --version output ----"
cat /tmp/llama-server-version.txt || true
fail "llama-server failed to launch on macOS $HOST_VER (dyld load / symbol error)"
fi
echo "llama.cpp load validation passed on macOS $HOST_VER"
echo " server: $SERVER"
sed -n '1,4p' /tmp/llama-server-version.txt 2>/dev/null || true

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

@ -79,6 +79,56 @@ jobs:
run: |
ruff check unsloth unsloth_cli studio tests cli.py unsloth-cli.py
- name: Import-hoist verifier self-test
# scripts/verify_import_hoist.py is a scope-aware (LEGB) AST
# resolver that gates import-hoisting / alias-rename refactors
# against two bugs ruff and pyflakes both miss:
# 1. dangling alias -- `from a import b as _b` hoisted to
# `from a import b` but a leftover `_b` reference now
# resolves to nothing (or to some other module-level `_b`).
# 2. rename clash -- `_b -> b` silently re-points at a
# different object already named `b` in that scope.
# This step runs the tool's 8 negative-control cases so a
# regression in the verifier itself fails before we trust it on
# a diff. Hermetic, stdlib-only, sub-second. Hard gate.
run: |
python scripts/verify_import_hoist.py --self-test
- name: Import-hoist / alias-rename safety (changed Python files)
# Runs the verifier in compare mode on every in-place-modified
# .py in the PR: parses each file BEFORE (base branch) and AFTER
# (this diff), resolves every name load, and fails on a BLOCKER
# (dangling alias / rename clash / re-pointed import). INFO
# findings (a helper relocated to another file) do not fail.
#
# --diff-filter=M (in-place edits only) is deliberate: that is
# exactly where a hoist refactor lives, and it skips brand-new
# files whose re-export imports would otherwise look "unused".
#
# actions/checkout uses fetch-depth: 1, so the base branch is not
# present locally. Fetch the single base commit with an explicit
# refspec so origin/<base> is reliably created (a bare
# `git fetch origin <ref>` only updates FETCH_HEAD in some
# configs). Two-dot diff avoids needing a merge-base on a shallow
# clone.
if: github.event_name == 'pull_request'
run: |
git fetch --no-tags --depth=1 origin \
"${{ github.base_ref }}:refs/remotes/origin/${{ github.base_ref }}"
mapfile -t CHANGED < <(
git diff --name-only --diff-filter=M \
"origin/${{ github.base_ref }}" HEAD -- '*.py' \
| grep -vE '(^|/)(unsloth_compiled_cache|node_modules|build|dist)/' || true
)
if [ "${#CHANGED[@]}" -eq 0 ]; then
echo "no in-place-modified Python files to check"
exit 0
fi
printf 'checking %d file(s):\n' "${#CHANGED[@]}"
printf ' %s\n' "${CHANGED[@]}"
python scripts/verify_import_hoist.py \
--before "origin/${{ github.base_ref }}" --after HEAD "${CHANGED[@]}"
- name: No leftover debugger / pdb / breakpoint calls
# Catches the "I'll just stick a breakpoint() here" mistake
# before it ships. AST-based so commented-out debugger

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

@ -285,7 +285,15 @@ jobs:
# The PR-time CI must validate the code in this PR; PyPI unsloth
# may lag the in-repo CPU-torch fallback in unsloth/kernels/utils.py
# (lines 162-170) that handles missing torch._C._cuda_getCurrentRawStream.
pip install --no-deps unsloth_zoo
# unsloth_zoo from git main mirrors every other CI (Core / MLX /
# install.sh) so PR-time validation sees the same zoo HEAD.
for attempt in 1 2 3; do
if pip install --no-deps "unsloth_zoo @ git+https://github.com/unslothai/unsloth-zoo"; then
break
fi
[ "$attempt" -eq 3 ] && { echo "::error::unsloth_zoo install failed after 3 attempts"; exit 1; }
sleep $((5 * attempt))
done
pip install --no-deps -e ./unsloth
- name: Convert notebooks for AST scan

126
.github/workflows/security-audit.yml vendored
View file

@ -72,6 +72,31 @@ concurrency:
permissions:
contents: read
# ──────────────────────────────────────────────────────────────────────
# Network-resilience knobs, applied to every job/step. These add retries
# and backoff ONLY; they do not relax a single integrity check. cargo
# still resolves against Cargo.lock (--locked), pip still verifies the
# wheels it downloads, npm still enforces package-lock integrity, the
# harden-runner egress allowlists below are unchanged, and every action
# stays SHA-pinned. The advisory-audit run on 2026-05-29 red-failed when
# one crates.io tarball fetch hit "Recv failure: Connection reset by
# peer" (curl 56); cargo's default of 3 retries over an HTTP/2-multiplexed
# connection did not recover. The settings below make that class of
# transient fault self-heal instead of failing the whole run.
env:
# pip: raise the built-in retry count and per-connection timeout.
PIP_RETRIES: "10"
PIP_DEFAULT_TIMEOUT: "60"
# cargo: retry network ops and disable HTTP/2 multiplexing -- the
# documented mitigation for the curl-56 connection resets above.
CARGO_NET_RETRY: "10"
CARGO_HTTP_MULTIPLEXING: "false"
CARGO_NET_GIT_FETCH_WITH_CLI: "true"
# npm: retry registry fetches with capped exponential backoff.
NPM_CONFIG_FETCH_RETRIES: "5"
NPM_CONFIG_FETCH_RETRY_MINTIMEOUT: "2000"
NPM_CONFIG_FETCH_RETRY_MAXTIMEOUT: "60000"
jobs:
# ─────────────────────────────────────────────────────────────────────
# Combined advisory-DB audit: pip-audit + npm audit + cargo audit
@ -140,7 +165,7 @@ jobs:
- uses: dtolnay/rust-toolchain@29eef336d9b2848a0b548edc03f92a220660cdb8 # stable @ 2026-03-27
- uses: swatinem/rust-cache@e18b497796c12c097a38f9edb9d0641fb99eee32 # v2.9.1
- uses: swatinem/rust-cache@c19371144df3bb44fab255c43d04cbc2ab54d1c4 # v2.9.1
with:
workspaces: studio/src-tauri -> target
@ -153,8 +178,23 @@ jobs:
# crashes with a TOML parse error on that file.
# npm audit is bundled with the node toolchain, no install.
run: |
python -m pip install --upgrade pip 'pip-audit>=2.7'
cargo install --locked --version '^0.22' cargo-audit
retry() { # retry <max-attempts> <command...> with exponential backoff
local max="$1"; shift
local n=1 delay=5
until "$@"; do
if [ "$n" -ge "$max" ]; then
echo "::error::command failed after ${n} attempts: $*" >&2
return 1
fi
echo "attempt ${n}/${max} failed; retrying in ${delay}s: $*" >&2
sleep "$delay"; n=$((n + 1)); delay=$((delay * 2))
done
}
retry 5 python -m pip install --upgrade pip 'pip-audit>=2.7'
# --locked keeps the resolved tree identical to Cargo.lock; the
# CARGO_NET_* env above plus this outer loop survive transient
# crates.io connection resets without weakening that guarantee.
retry 5 cargo install --locked --version '^0.22' cargo-audit
# ─────────────────────────────────────────────────────────────
# Python: pip-audit
@ -330,32 +370,60 @@ jobs:
# ─────────────────────────────────────────────────────────────
# OSV-Scanner: cross-ecosystem advisory DB (PyPI + npm + cargo)
# ─────────────────────────────────────────────────────────────
- name: Download + verify OSV-Scanner
# Split out from the scan below so binary integrity is a HARD gate:
# a checksum mismatch (swapped release asset, the Trivy-style pivot
# this workflow refuses) fails the job instead of being swallowed by
# the scan step's continue-on-error. A download still failing after
# retries is transient, so we skip the scan rather than red-fail.
# SHA-256 verified BEFORE chmod +x / exec. Bump OSV_SHA256 in lockstep
# with OSV_VERSION (value from the release's osv-scanner_SHA256SUMS).
run: |
set -euo pipefail
OSV_VERSION="v2.0.2"
OSV_SHA256="3abcfd7126c453a00421487e721b296e0cb68085bd431d6cef60872774170fc8"
if ! curl --proto '=https' --tlsv1.2 -fsSL \
--retry 5 --retry-delay 3 --retry-connrefused --retry-all-errors \
-o /tmp/osv-scanner \
"https://github.com/google/osv-scanner/releases/download/${OSV_VERSION}/osv-scanner_linux_amd64"; then
echo "::warning::osv-scanner download failed after retries; skipping scan" >&2
rm -f /tmp/osv-scanner
exit 0 # transient availability: do not red-fail the job
fi
if ! echo "${OSV_SHA256} /tmp/osv-scanner" | sha256sum -c -; then
echo "::error::osv-scanner checksum mismatch; refusing to execute" >&2
rm -f /tmp/osv-scanner
exit 1 # integrity failure: hard-fail
fi
chmod +x /tmp/osv-scanner
/tmp/osv-scanner --version
- name: OSV-Scanner (PyPI + npm + cargo, cross-ecosystem advisories)
# OSV's advisory feed is a superset of GitHub-Advisory + RustSec
# + npm advisories; running it alongside the per-ecosystem audit
# tools catches CVEs that haven't propagated to the per-ecosystem
# DBs yet (e.g. langchain-core CVE-2025-68664 was on OSV before
# GitHub Advisory). Single binary, one transitive resolver, all
# three lockfile types in one pass. Non-blocking until baselines
# close.
# three lockfile types in one pass. Binary is checksum-verified in
# the step above; only the advisory scan stays non-blocking until
# baselines close.
continue-on-error: true
run: |
set +e
# OSV-Scanner ships a raw binary (no tarball) in v2.x.
curl -fsSL -o /tmp/osv-scanner \
https://github.com/google/osv-scanner/releases/download/v2.0.2/osv-scanner_linux_amd64
chmod +x /tmp/osv-scanner
/tmp/osv-scanner --version
/tmp/osv-scanner scan source \
--lockfile=studio/frontend/package-lock.json \
--lockfile=studio/src-tauri/Cargo.lock \
--lockfile=requirements.txt:audit-reqs/unsloth-deps.txt \
--lockfile=requirements.txt:audit-reqs/studio.txt \
--lockfile=requirements.txt:audit-reqs/no-torch-runtime.txt \
--lockfile=requirements.txt:audit-reqs/overrides.txt \
--lockfile=requirements.txt:audit-reqs/extras.txt \
--lockfile=requirements.txt:audit-reqs/extras-no-deps.txt \
--format=table 2>&1 | tee logs-osv-scanner.txt
if [ ! -x /tmp/osv-scanner ]; then
echo "osv-scanner unavailable this run; skipping scan" | tee logs-osv-scanner.txt
else
/tmp/osv-scanner scan source \
--lockfile=studio/frontend/package-lock.json \
--lockfile=studio/src-tauri/Cargo.lock \
--lockfile=requirements.txt:audit-reqs/unsloth-deps.txt \
--lockfile=requirements.txt:audit-reqs/studio.txt \
--lockfile=requirements.txt:audit-reqs/no-torch-runtime.txt \
--lockfile=requirements.txt:audit-reqs/overrides.txt \
--lockfile=requirements.txt:audit-reqs/extras.txt \
--lockfile=requirements.txt:audit-reqs/extras-no-deps.txt \
--format=table 2>&1 | tee logs-osv-scanner.txt
fi
{
echo "## OSV-Scanner (cross-ecosystem)"
echo
@ -1075,7 +1143,23 @@ jobs:
# new-install-script gate below protects against, and we must
# not run any third-party hook to set up the audit.
working-directory: studio/frontend
run: npm ci --ignore-scripts
run: |
retry() { # retry <max-attempts> <command...> with exponential backoff
local max="$1"; shift
local n=1 delay=5
until "$@"; do
if [ "$n" -ge "$max" ]; then
echo "::error::command failed after ${n} attempts: $*" >&2
return 1
fi
echo "attempt ${n}/${max} failed; retrying in ${delay}s: $*" >&2
sleep "$delay"; n=$((n + 1)); delay=$((delay * 2))
done
}
# --ignore-scripts is mandatory here (no third-party hook runs);
# the retry only re-attempts the registry fetch, it never relaxes
# that flag or the package-lock integrity check npm ci enforces.
retry 5 npm ci --ignore-scripts
- name: npm audit signatures (informational)
# Surfaces unsigned / mis-signed packages from the npm

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

@ -144,9 +144,19 @@ jobs:
# versions ship a CPU build that imports cleanly on Linux.
pip install 'bitsandbytes>=0.45'
# unsloth.device_type imports unsloth_zoo.utils.Version at module
# scope, so the conftest preload needs unsloth_zoo even though
# it is an optional dep of unsloth.
pip install 'unsloth_zoo>=2026.5.1'
# scope, so the conftest preload needs unsloth_zoo. Pull from
# git main so this job sees the same zoo HEAD as Core / MLX /
# install.sh do (otherwise a fix on zoo main hides until release).
# No --no-deps: matches prior `pip install 'unsloth_zoo>=2026.5.1'`
# behaviour so triton etc. still come in for the Repo tests CPU
# collection imports.
for attempt in 1 2 3; do
if pip install "unsloth_zoo @ git+https://github.com/unslothai/unsloth-zoo"; then
break
fi
[ "$attempt" -eq 3 ] && { echo "::error::unsloth_zoo install failed after 3 attempts"; exit 1; }
sleep $((5 * attempt))
done
pip install -e . --no-deps
- name: Repo tests (CPU, auto-discovered)
@ -212,6 +222,7 @@ jobs:
for s in \
tests/sh/test_get_torch_index_url.sh \
tests/sh/test_mac_intel_compat.sh \
tests/sh/test_nvcc_meets_llama_minimum.sh \
tests/sh/test_tauri_install_exit_order.sh \
tests/sh/test_torch_constraint.sh; do
echo "::group::$s"

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

@ -89,13 +89,8 @@ jobs:
set -o pipefail
bash install.sh --local --no-torch 2>&1 | tee logs/install.log
- name: Assert install.sh used the Mac llama.cpp prebuilt
run: |
if grep -q "falling back to source build" logs/install.log; then
echo "::error::install.sh fell back to source-build llama.cpp on Mac. Studio must install the prebuilt llama-bNNNN-bin-macos-arm64 on Apple Silicon."
grep -E "llama-prebuilt|llama.cpp" logs/install.log | tail -60
exit 1
fi
- name: Assert llama.cpp loads on this macOS
run: bash .github/scripts/assert-llama-loads.sh
- name: Install pyjwt for the JWT-expiry forge test
run: pip install 'pyjwt>=2.6'

27
.github/workflows/studio-mac-inference-smoke.yml vendored
View file

@ -114,13 +114,8 @@ jobs:
set -o pipefail
bash install.sh --local --no-torch 2>&1 | tee logs/install.log
- name: Assert install.sh used the Mac llama.cpp prebuilt
run: |
if grep -q "falling back to source build" logs/install.log; then
echo "::error::install.sh fell back to source-build llama.cpp on Mac. Studio must install the prebuilt llama-bNNNN-bin-macos-arm64 on Apple Silicon."
grep -E "llama-prebuilt|llama.cpp" logs/install.log | tail -60
exit 1
fi
- name: Assert llama.cpp loads on this macOS
run: bash .github/scripts/assert-llama-loads.sh
- name: Install OpenAI + Anthropic Python SDKs
run: pip install 'openai>=1.50' 'anthropic>=0.40'
@ -369,13 +364,8 @@ jobs:
set -o pipefail
bash install.sh --local --no-torch 2>&1 | tee logs/install.log
- name: Assert install.sh used the Mac llama.cpp prebuilt
run: |
if grep -q "falling back to source build" logs/install.log; then
echo "::error::install.sh fell back to source-build llama.cpp on Mac. Studio must install the prebuilt llama-bNNNN-bin-macos-arm64 on Apple Silicon."
grep -E "llama-prebuilt|llama.cpp" logs/install.log | tail -60
exit 1
fi
- name: Assert llama.cpp loads on this macOS
run: bash .github/scripts/assert-llama-loads.sh
- name: Reset auth + boot Studio (API-only, default tool policy)
# We deliberately use the API-only mode rather than
@ -760,13 +750,8 @@ jobs:
set -o pipefail
bash install.sh --local --no-torch 2>&1 | tee logs/install.log
- name: Assert install.sh used the Mac llama.cpp prebuilt
run: |
if grep -q "falling back to source build" logs/install.log; then
echo "::error::install.sh fell back to source-build llama.cpp on Mac. Studio must install the prebuilt llama-bNNNN-bin-macos-arm64 on Apple Silicon."
grep -E "llama-prebuilt|llama.cpp" logs/install.log | tail -60
exit 1
fi
- name: Assert llama.cpp loads on this macOS
run: bash .github/scripts/assert-llama-loads.sh
- name: Install OpenAI + Anthropic Python SDKs
run: pip install 'openai>=1.50' 'anthropic>=0.40'

80
.github/workflows/studio-mac-install-matrix.yml vendored Normal file
View file

@ -0,0 +1,80 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
# Proves Studio's llama.cpp install loads on every supported macOS. The heavy
# app smokes stay single-OS; this matrix covers the OS-version dimension cheaply
# (install.sh + binary-load assert). Regression guard for the macOS-version
# selection in studio/install_llama_prebuilt.py.
name: Mac Studio Install Matrix CI
on:
pull_request:
paths:
- 'studio/install_llama_prebuilt.py'
- 'studio/setup.sh'
- 'install.sh'
- '.github/scripts/assert-llama-loads.sh'
- '.github/workflows/studio-mac-install-matrix.yml'
push:
branches: [main, pip]
workflow_dispatch:
concurrency:
group: ${{ github.workflow }}-${{ github.ref }}
cancel-in-progress: true
permissions:
contents: read
jobs:
install-load:
name: Install + load (${{ matrix.os }})
runs-on: ${{ matrix.os }}
timeout-minutes: 25
continue-on-error: ${{ matrix.experimental }}
strategy:
fail-fast: false
matrix:
include:
- os: macos-14 # Apple Silicon, macOS 14 Sonoma
experimental: false
- os: macos-15 # Apple Silicon, macOS 15 Sequoia
experimental: false
- os: macos-26 # Apple Silicon, macOS 26 Tahoe
experimental: false
- os: macos-15-intel # Intel x86_64, macOS 15 (informational)
experimental: true
- os: macos-26-intel # Intel x86_64, macOS 26 (last Intel macOS)
experimental: true
steps:
- uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
with:
persist-credentials: false
- uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0
with:
node-version: '22'
- uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0
with:
python-version: '3.12'
- name: Install Studio (--local, --no-torch)
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
run: |
mkdir -p logs
set -o pipefail
bash install.sh --local --no-torch 2>&1 | tee logs/install.log
- name: Assert llama.cpp loads on this macOS
run: bash .github/scripts/assert-llama-loads.sh
- name: Upload install log
if: always()
uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
with:
name: mac-install-matrix-${{ matrix.os }}-log
path: logs/install.log
retention-days: 7

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

@ -89,13 +89,8 @@ jobs:
set -o pipefail
bash install.sh --local --no-torch 2>&1 | tee logs/install.log
- name: Assert install.sh used the Mac llama.cpp prebuilt
run: |
if grep -q "falling back to source build" logs/install.log; then
echo "::error::install.sh fell back to source-build llama.cpp on Mac. Studio must install the prebuilt llama-bNNNN-bin-macos-arm64 on Apple Silicon."
grep -E "llama-prebuilt|llama.cpp" logs/install.log | tail -60
exit 1
fi
- name: Assert llama.cpp loads on this macOS
run: bash .github/scripts/assert-llama-loads.sh
- name: Install Playwright + Chromium
# No --with-deps on Mac: that flag installs Linux apt packages.

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

@ -67,21 +67,8 @@ jobs:
set -o pipefail
bash install.sh --local --no-torch 2>&1 | tee logs/install.log
- name: Assert install.sh used the Mac llama.cpp prebuilt
run: |
# Mac install must take the prebuilt path. Source-build
# fallback here is an Unsloth bug.
if grep -q "falling back to source build" logs/install.log; then
echo "::error::install.sh fell back to source-build llama.cpp on Mac. Studio must install the prebuilt llama-bNNNN-bin-macos-arm64 on Apple Silicon."
grep -E "llama-prebuilt|llama.cpp" logs/install.log | tail -60
exit 1
fi
if ! grep -qE "prebuilt installed and validated|prebuilt up to date and validated|bin-macos-arm64" logs/install.log; then
echo "::error::no Mac prebuilt llama.cpp marker in install.log."
grep -E "llama-prebuilt|llama.cpp" logs/install.log | tail -60
exit 1
fi
echo "install.sh installed the Mac prebuilt llama.cpp"
- name: Assert llama.cpp loads on this macOS
run: bash .github/scripts/assert-llama-loads.sh
- name: First update should be a no-op (prebuilt already validated)
env:

5
README.md
View file

@ -202,7 +202,7 @@ unsloth studio -p 8888
#### Nightly: Windows:
Run in Windows Powershell:
```bash
```powershell
git clone https://github.com/unslothai/unsloth.git
cd unsloth
git checkout nightly
@ -215,6 +215,9 @@ Then to launch every time:
unsloth studio -p 8888
```
#### Advanced launch options
Cap Studio's native CPU thread pools on high-core hosts: `UNSLOTH_CPU_THREADS=8 unsloth studio -p 8888`. Explicit `OMP_NUM_THREADS` / `MKL_NUM_THREADS` / `OPENBLAS_NUM_THREADS` / `NUMEXPR_NUM_THREADS` still take precedence.
#### Uninstall
The recommended way to fully remove Unsloth Studio is the matching uninstall script for your OS. It stops any running servers, removes the install dir, the launcher data dir, the desktop shortcut, and any platform-specific entries (macOS `.app` bundle + Launch Services on Mac; Start Menu, `HKCU\Software\Unsloth` registry key and user `PATH` entries on Windows):

321
install.ps1
View file

@ -887,12 +887,19 @@ shell.Run cmd, 0, False
}
# ── Check winget ──
# winget is only needed to install Python or uv. If both are
# already on PATH (Windows ARM64 GitHub-hosted runners, manual
# python.org + Astral uv installs, corporate locked-down hosts
# without the Store, etc.) the script can proceed without it.
# We defer the hard failure to the Python / uv install branches
# below, where winget is actually invoked.
Write-TauriLog "STEP" "Checking system dependencies"
if (-not (Get-Command winget -ErrorAction SilentlyContinue)) {
step "winget" "not available" "Red"
substep "Install it from https://aka.ms/getwinget" "Yellow"
substep "or install Python $PythonVersion and uv manually, then re-run." "Yellow"
return (Exit-InstallFailure "winget is not available")
$script:WingetAvailable = [bool](Get-Command winget -ErrorAction SilentlyContinue)
if ($script:WingetAvailable) {
step "winget" "available"
} else {
step "winget" "not available -- will require Python + uv to be already installed" "Yellow"
substep "Get it from https://aka.ms/getwinget if Python / uv are not already on PATH." "Yellow"
}
# ── Helper: detect a working Python 3.11-3.13 on the system ──
@ -969,10 +976,17 @@ shell.Run cmd, 0, False
# Find-CompatiblePython returns @{ Version = "3.13"; Path = "C:\...\python.exe" } or $null.
Write-TauriLog "STEP" "Installing Python"
$DetectedPython = Find-CompatiblePython
if ($DetectedPython) {
step "python" "Python $($DetectedPython.Version) already installed"
}
if (-not $DetectedPython) {
if (-not $script:WingetAvailable) {
Write-Host "[ERROR] No compatible Python (3.11-3.13) found and winget is unavailable on this host." -ForegroundColor Red
Write-Host " Install Python $PythonVersion from https://www.python.org/downloads/" -ForegroundColor Yellow
Write-Host " and re-run this installer (make sure 'Add Python to PATH' is checked)." -ForegroundColor Yellow
return (Exit-InstallFailure "winget required to install Python on this host")
}
substep "installing Python ${PythonVersion}..."
$pythonPackageId = "Python.Python.$PythonVersion"
# Temporarily lower ErrorActionPreference so that winget stderr
@ -1024,14 +1038,19 @@ shell.Run cmd, 0, False
Write-TauriLog "STEP" "Installing uv package manager"
if (-not (Get-Command uv -ErrorAction SilentlyContinue)) {
substep "installing uv package manager..."
$prevEAP = $ErrorActionPreference
$ErrorActionPreference = "Continue"
try { winget install --id=astral-sh.uv -e --accept-package-agreements --accept-source-agreements } catch {}
$ErrorActionPreference = $prevEAP
Refresh-SessionPath
# Fallback: if winget didn't put uv on PATH, try the PowerShell installer
if ($script:WingetAvailable) {
$prevEAP = $ErrorActionPreference
$ErrorActionPreference = "Continue"
try { winget install --id=astral-sh.uv -e --accept-package-agreements --accept-source-agreements } catch {}
$ErrorActionPreference = $prevEAP
Refresh-SessionPath
}
# Fallback: if winget is unavailable or didn't put uv on PATH,
# use Astral's official PowerShell installer. This is the only
# supported path on hosts without winget (Windows ARM64 runners,
# corporate machines without the Store, etc.).
if (-not (Get-Command uv -ErrorAction SilentlyContinue)) {
substep "trying alternative uv installer..." "Yellow"
substep "installing uv via https://astral.sh/uv/install.ps1..." "Yellow"
Invoke-Expression (Invoke-RestMethod -Uri "https://astral.sh/uv/install.ps1")
Refresh-SessionPath
}
@ -1221,11 +1240,196 @@ shell.Run cmd, 0, False
}
}
}
# ── AMD ROCm detection (Windows) — mirrors setup.ps1 ──
$HasROCm = $false
$HipSdkInstalled = $false # HIP SDK binary found (independent of device accessibility)
$ROCmGpuLabel = $null
$ROCmVersion = $null
$ROCmGfxArch = $null
if (-not $HasNvidiaSmi) {
# hipinfo: PATH first, then HIP_PATH/ROCM_PATH bin fallback (mirrors NVIDIA smi path resolution).
# AMD HIP SDK sets HIP_PATH but may not add the bin dir to PATH depending on install type.
$hipinfoExe = Get-Command hipinfo -ErrorAction SilentlyContinue
if (-not $hipinfoExe) {
$hipRoot = if ($env:HIP_PATH) { $env:HIP_PATH } elseif ($env:ROCM_PATH) { $env:ROCM_PATH } else { $null }
$hipEnvLabel = if ($env:HIP_PATH) { "HIP_PATH" } else { "ROCM_PATH" }
if ($hipRoot) {
$hipinfoCandidate = Join-Path $hipRoot "bin\hipinfo.exe"
if (Test-Path $hipinfoCandidate) {
Write-Host " [WARN] hipinfo not on PATH -- located via ${hipEnvLabel}: $hipinfoCandidate" -ForegroundColor Yellow
Write-Host " Add '$(Join-Path $hipRoot 'bin')' to your PATH to suppress this warning" -ForegroundColor Yellow
Write-Host " Quick fix: [Environment]::SetEnvironmentVariable('PATH',`$env:PATH+';$(Join-Path $hipRoot 'bin')','User')" -ForegroundColor Yellow
$hipinfoExe = [PSCustomObject]@{ Source = $hipinfoCandidate }
} else {
Write-Host " [WARN] ${hipEnvLabel}=$hipRoot is set but hipinfo.exe not found at $hipinfoCandidate" -ForegroundColor Yellow
Write-Host " HIP SDK install may be incomplete -- re-install from:" -ForegroundColor Yellow
Write-Host " https://rocm.docs.amd.com/en/latest/deploy/windows/index.html" -ForegroundColor Yellow
}
}
}
if ($hipinfoExe) {
$HipSdkInstalled = $true # binary found → SDK is installed regardless of device state
try {
$hipOut = & $hipinfoExe.Source 2>&1 | Out-String
if ($LASTEXITCODE -eq 0 -and $hipOut -match "(?i)gcnArchName") {
$HasROCm = $true
$_hipAllArches = @([regex]::Matches($hipOut, "(?im)^\s*gcnArchName\s*:\s*(\S+)") | ForEach-Object { ($_.Groups[1].Value -split ':')[0].Trim().ToLower() })
$_hipVisIdx = if ($env:HIP_VISIBLE_DEVICES -match '^\d') { [int]($env:HIP_VISIBLE_DEVICES -split ',')[0] } elseif ($env:ROCR_VISIBLE_DEVICES -match '^\d') { [int]($env:ROCR_VISIBLE_DEVICES -split ',')[0] } else { 0 }
if ($_hipAllArches.Count -gt 0) {
$ROCmGfxArch = if ($_hipVisIdx -lt $_hipAllArches.Count) { $_hipAllArches[$_hipVisIdx] } else { $_hipAllArches[0] }
$ROCmGpuLabel = "AMD ROCm ($ROCmGfxArch)"
} else {
$ROCmGpuLabel = "AMD ROCm"
}
} elseif ($LASTEXITCODE -ne 0) {
# hipinfo ran but returned a HIP runtime error (e.g. "no ROCm-capable device detected")
$firstLine = ($hipOut -split '\r?\n' | Where-Object { $_.Trim() } | Select-Object -First 1)
Write-Host " [WARN] hipinfo returned a HIP runtime error (exit $LASTEXITCODE)" -ForegroundColor Yellow
Write-Host " $firstLine" -ForegroundColor Yellow
Write-Host " Ensure ROCm drivers are installed: https://rocm.docs.amd.com/en/latest/deploy/windows/index.html" -ForegroundColor Yellow
}
} catch {}
}
if (-not $HasROCm) {
$amdSmiExe = Get-Command "amd-smi" -ErrorAction SilentlyContinue
if ($amdSmiExe) {
try {
$smiOut = & $amdSmiExe.Source list 2>&1 | Out-String
if ($LASTEXITCODE -eq 0 -and $smiOut -match "(?im)^GPU\s*[:\[]\s*\d") {
$HasROCm = $true
# Mirror the hipinfo path: collect all gfx tokens in enumeration
# order and pick the runtime-visible one via HIP_VISIBLE_DEVICES.
$_smiVisIdx = if ($env:HIP_VISIBLE_DEVICES -match '^\d') { [int]($env:HIP_VISIBLE_DEVICES -split ',')[0] } elseif ($env:ROCR_VISIBLE_DEVICES -match '^\d') { [int]($env:ROCR_VISIBLE_DEVICES -split ',')[0] } else { 0 }
# Attempt 1: newer amd-smi versions embed the gfx arch in list output.
$_smiGfxTokens = @([regex]::Matches($smiOut, "(?i)\b(gfx\d+[a-z]?)\b") | ForEach-Object { $_.Groups[1].Value.ToLower() })
if ($_smiGfxTokens.Count -gt 0) {
$ROCmGfxArch = if ($_smiVisIdx -lt $_smiGfxTokens.Count) { $_smiGfxTokens[$_smiVisIdx] } else { $_smiGfxTokens[0] }
$ROCmGpuLabel = "AMD ROCm ($ROCmGfxArch)"
} else {
# Attempt 2: 'static --asic' exposes ASIC details on ROCm 6+,
# including the GFX target needed for wheel index selection.
$smiAsicOut = ""
try { $smiAsicOut = & $amdSmiExe.Source static --asic 2>&1 | Out-String } catch {}
$_asicGfxTokens = @([regex]::Matches($smiAsicOut, "(?i)\b(gfx\d+[a-z]?)\b") | ForEach-Object { $_.Groups[1].Value.ToLower() })
if ($_asicGfxTokens.Count -gt 0) {
$ROCmGfxArch = if ($_smiVisIdx -lt $_asicGfxTokens.Count) { $_asicGfxTokens[$_smiVisIdx] } else { $_asicGfxTokens[0] }
$ROCmGpuLabel = "AMD ROCm ($ROCmGfxArch)"
} elseif ($smiAsicOut -match "(?im)Market.?Name\s*[:\|]\s*([^\r\n]+)") {
$ROCmGpuLabel = "AMD ROCm ($($Matches[1].Trim()))"
} else {
$ROCmGpuLabel = "AMD ROCm"
}
}
}
} catch {}
}
}
if (-not $HasROCm) {
try {
$wmiGpu = Get-WmiObject Win32_VideoController -ErrorAction SilentlyContinue |
Where-Object { $_.Name -match "AMD|Radeon" } |
Select-Object -First 1
if ($wmiGpu) { $ROCmGpuLabel = $wmiGpu.Name }
} catch {}
}
# ── Arch resolution: env-var override → name inference ──────────────
# Covers users whose amd-smi is too old to report the GFX target and
# who don't have hipinfo (HIP-runtime-only, common on Strix Halo / iGPU).
if ($HasROCm -and -not $ROCmGfxArch) {
# 1. Manual override: set UNSLOTH_ROCM_GFX_ARCH=gfx1151 before running.
if ($env:UNSLOTH_ROCM_GFX_ARCH) {
$ROCmGfxArch = $env:UNSLOTH_ROCM_GFX_ARCH.Trim().ToLower()
$ROCmGpuLabel = "AMD ROCm ($ROCmGfxArch)"
substep "gfx arch from UNSLOTH_ROCM_GFX_ARCH env override: $ROCmGfxArch" "Cyan"
}
# 2. Best-effort name → arch lookup from marketing name (amd-smi / WMI).
elseif ($ROCmGpuLabel) {
$nameArchTable = @(
@{ P = "9070 XT|9080"; A = "gfx1201" } # RDNA 4
@{ P = "9070|9060"; A = "gfx1200" } # RDNA 4
@{ P = "8060S|890M|Strix Halo|HX 37[05]|HX 38[05]|AI 9 HX"; A = "gfx1151" } # RDNA 3.5 iGPU (Strix Halo / Radeon 8060S retail)
@{ P = "880M|Strix Point|AI 9 36[05]|AI 7 35[05]|AI 5 34[05]"; A = "gfx1150" } # RDNA 3.5 iGPU (Strix Point)
@{ P = "RX 7900|RX 7800|RX 7700(?! S)"; A = "gfx1100" } # RDNA 3 desktop
@{ P = "RX 7600"; A = "gfx1102" } # RDNA 3
@{ P = "780M|760M|740M|Phoenix"; A = "gfx1103" } # RDNA 3 iGPU (Phoenix)
)
foreach ($row in $nameArchTable) {
if ($ROCmGpuLabel -match $row.P) {
$ROCmGfxArch = $row.A
$ROCmGpuLabel = "AMD ROCm ($ROCmGfxArch)"
substep "gfx arch inferred from GPU name: $ROCmGfxArch" "Cyan"
substep "Tip: set UNSLOTH_ROCM_GFX_ARCH=$ROCmGfxArch to skip inference next time" "Cyan"
break
}
}
}
}
# Capture ROCm version for wheel selection (hipconfig, then amd-smi).
# Run whenever the HIP SDK binary is present, not just when the device is accessible --
# hipconfig --version works even when hipinfo reports no ROCm device (driver issue).
if ($HasROCm -or $HipSdkInstalled) {
$hipConfigExe = Get-Command hipconfig -ErrorAction SilentlyContinue
if (-not $hipConfigExe) {
$hipRoot = if ($env:HIP_PATH) { $env:HIP_PATH } elseif ($env:ROCM_PATH) { $env:ROCM_PATH } else { $null }
if ($hipRoot) {
$hipConfigCandidate = Join-Path $hipRoot "bin\hipconfig.exe"
if (Test-Path $hipConfigCandidate) {
$hipConfigEnvLabel = if ($env:HIP_PATH) { "HIP_PATH" } else { "ROCM_PATH" }
Write-Host " [WARN] hipconfig not on PATH -- located via ${hipConfigEnvLabel}: $hipConfigCandidate" -ForegroundColor Yellow
$hipConfigExe = [PSCustomObject]@{ Source = $hipConfigCandidate }
}
}
}
if ($hipConfigExe) {
try {
$hipVerOut = & $hipConfigExe.Source --version 2>&1 | Out-String
if ($LASTEXITCODE -eq 0) {
$hipVerLine = ($hipVerOut -split '\r?\n' | Where-Object { $_.Trim() } | Select-Object -First 1).Trim()
if ($hipVerLine -match '(\d+\.\d+)') {
$ROCmVersion = $Matches[1]
$ROCmVersionFull = $hipVerLine
}
}
} catch {}
}
if (-not $ROCmVersion) {
$amdSmiVer = Get-Command "amd-smi" -ErrorAction SilentlyContinue
if ($amdSmiVer) {
try {
$smiVerOut = & $amdSmiVer.Source version 2>&1 | Out-String
if ($LASTEXITCODE -eq 0 -and $smiVerOut -match 'ROCm version:\s*(\d+\.\d+)') {
$ROCmVersion = $Matches[1]
}
} catch {}
}
}
}
}
if ($HasNvidiaSmi) {
step "gpu" "NVIDIA GPU detected"
} elseif ($HasROCm) {
step "gpu" $ROCmGpuLabel
$hipSdkPath = if ($env:HIP_PATH) { $env:HIP_PATH } elseif ($env:ROCM_PATH) { $env:ROCM_PATH } else { "on system PATH" }
substep "HIP SDK: $hipSdkPath"
if ($ROCmVersionFull) { substep "hipconfig: $ROCmVersionFull" }
} elseif ($HipSdkInstalled -and $ROCmGpuLabel) {
# HIP SDK is installed but ROCm can't see the device (driver issue, not SDK issue)
$sdkVer = if ($ROCmVersionFull) { " (HIP $ROCmVersionFull)" } else { "" }
step "gpu" "AMD GPU detected -- not ROCm-accessible$sdkVer" "Yellow"
substep "Detected: $ROCmGpuLabel" "Yellow"
substep "[WARN] HIP SDK is installed but hipinfo reports no ROCm-capable device." "Yellow"
substep " This is a driver issue, not an SDK issue." "Yellow"
substep " Ensure the ROCm compute driver is installed alongside the display driver:" "Yellow"
substep " https://rocm.docs.amd.com/en/latest/deploy/windows/index.html" "Yellow"
} elseif ($ROCmGpuLabel) {
step "gpu" "AMD GPU detected -- HIP SDK not found" "Yellow"
substep "Detected: $ROCmGpuLabel" "Yellow"
substep "Install the HIP SDK for ROCm GPU inference:" "Yellow"
substep "https://rocm.docs.amd.com/en/latest/deploy/windows/index.html" "Yellow"
} else {
step "gpu" "none (chat-only / GGUF)" "Yellow"
substep "Training and GPU inference require an NVIDIA GPU with drivers installed." "Yellow"
substep "Training and GPU inference require an NVIDIA or AMD ROCm GPU." "Yellow"
}
# ── Choose the correct PyTorch index URL based on driver CUDA version ──
@ -1235,7 +1439,10 @@ shell.Run cmd, 0, False
if (-not $NvidiaSmiExe) { return "$baseUrl/cpu" }
try {
$output = & $NvidiaSmiExe 2>&1 | Out-String
if ($output -match 'CUDA Version:\s+(\d+)\.(\d+)') {
# Newer NVIDIA drivers (e.g. 610.x on Windows) print
# "CUDA UMD Version: X.Y" instead of the legacy "CUDA Version: X.Y".
# Accept both spellings so we don't fall through to the cu126 default.
if ($output -match 'CUDA(?: UMD)? Version:\s+(\d+)\.(\d+)') {
$major = [int]$Matches[1]; $minor = [int]$Matches[2]
if ($major -ge 13) { return "$baseUrl/cu130" }
if ($major -eq 12 -and $minor -ge 8) { return "$baseUrl/cu128" }
@ -1249,14 +1456,73 @@ shell.Run cmd, 0, False
return "$baseUrl/cu126"
}
$TorchIndexUrl = Get-TorchIndexUrl
$TorchIndexFamily = Get-TauriTorchIndexFamily $TorchIndexUrl
# ── GPU arch → newest compatible Windows ROCm wheel release ──
# Wheels bundle their own ROCm runtime; the installed HIP SDK version does
# not constrain which release to use. Always picks the newest release that
# supports the GPU architecture.
# ── AMD Windows ROCm: arch-aware pip index (repo.amd.com) ──
# Wheels bundle their own ROCm runtime and support all Python versions.
# Override with UNSLOTH_ROCM_WINDOWS_MIRROR for air-gapped / mirror installs.
$ROCmIndexUrl = $null
$ROCmTorchFloor = $null
if ($HasROCm -and $TorchIndexUrl -like "*/cpu" -and -not $SkipTorch) {
$amdIndexBase = if ($env:UNSLOTH_ROCM_WINDOWS_MIRROR) { $env:UNSLOTH_ROCM_WINDOWS_MIRROR.TrimEnd('/') } else { "https://repo.amd.com/rocm/whl" }
$archFamilyMap = @{
"gfx1201" = "gfx120X-all"; "gfx1200" = "gfx120X-all" # RDNA 4
"gfx1151" = "gfx1151"; "gfx1150" = "gfx1150" # RDNA 3.5 (Strix Halo/Point)
"gfx1103" = "gfx110X-all"; "gfx1102" = "gfx110X-all" # RDNA 3
"gfx1101" = "gfx110X-all"; "gfx1100" = "gfx110X-all"
"gfx90a" = "gfx90a"; "gfx908" = "gfx908" # MI200/MI100
}
# gfx120X (RDNA 4) and gfx1151/gfx1150 (Strix) have a null-pointer bug in
# torch._C._grouped_mm on torch <2.11.0 (rocm7.12 and rocm7.1 respectively).
# TheRock issues #5284 and #3284. Force torch>=2.11.0 so pip never resolves
# to the broken 2.10.0 wheels even though they exist on the AMD index.
# The <2.12.0 ceiling matches the Linux install_python_stack.py constraint
# for the same arches: AMD actively publishes new versions on their index,
# so without a ceiling a future 2.12.0+rocmX.Y wheel would be pulled in
# automatically before it has been validated on these architectures.
# Bump the ceiling here (and in install_python_stack.py) when 2.12.x is
# confirmed working on gfx120X / Strix.
$torchFloorMap = @{
"gfx1201" = "torch>=2.11.0,<2.12.0"; "gfx1200" = "torch>=2.11.0,<2.12.0"
"gfx1151" = "torch>=2.11.0,<2.12.0"; "gfx1150" = "torch>=2.11.0,<2.12.0"
}
$archFamily = if ($ROCmGfxArch -and $archFamilyMap.ContainsKey($ROCmGfxArch)) { $archFamilyMap[$ROCmGfxArch] } else { $null }
if ($archFamily) {
$ROCmIndexUrl = "$amdIndexBase/$archFamily/"
$ROCmTorchFloor = if ($ROCmGfxArch -and $torchFloorMap.ContainsKey($ROCmGfxArch)) { $torchFloorMap[$ROCmGfxArch] } else { $null }
$archLabel = if ($ROCmGfxArch) { $ROCmGfxArch } else { "AMD GPU" }
substep "$archLabel -- AMD repo.amd.com index selected" "Cyan"
if ($ROCmTorchFloor) {
substep " enforcing $ROCmTorchFloor (known _grouped_mm bug in older wheels)" "Cyan"
}
} elseif ($ROCmGfxArch) {
substep "AMD GPU ($ROCmGfxArch) not in supported arch list -- falling back to CPU-only PyTorch" "Yellow"
} else {
substep "AMD GPU detected but arch unknown -- falling back to CPU-only PyTorch" "Yellow"
}
}
if ($ROCmIndexUrl) {
$TorchIndexFamily = "rocm"
} else {
$TorchIndexFamily = Get-TauriTorchIndexFamily $TorchIndexUrl
}
$GpuBranch = Get-TauriGpuBranch $TorchIndexFamily
Write-TauriDiag -GpuBranch $GpuBranch -TorchIndexFamily $TorchIndexFamily -PythonVersionForDiag $DetectedPython.Version
# ── Print CPU-only hint when no GPU detected ──
if (-not $SkipTorch -and $TorchIndexUrl -like "*/cpu") {
if (-not $SkipTorch -and -not $ROCmIndexUrl -and $TorchIndexUrl -like "*/cpu") {
Write-Host ""
substep "No NVIDIA GPU detected." "Yellow"
if ($HipSdkInstalled -and -not $HasROCm) {
substep "Installing CPU-only PyTorch (HIP SDK found but GPU not ROCm-accessible)." "Yellow"
} elseif ($ROCmGpuLabel) {
substep "Installing CPU-only PyTorch (ROCm wheels require the HIP SDK)." "Yellow"
} else {
substep "No NVIDIA GPU detected." "Yellow"
}
substep "Installing CPU-only PyTorch. If you only need GGUF chat/inference," "Yellow"
substep "re-run with --no-torch for a faster, lighter install:" "Yellow"
substep ".\install.ps1 --no-torch" "Yellow"
@ -1300,7 +1566,7 @@ shell.Run cmd, 0, False
if ($SkipTorch) {
# No-torch: install unsloth + unsloth-zoo with --no-deps, then
# runtime deps (typer, safetensors, transformers, etc.) with --no-deps.
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.5.7" unsloth-zoo }
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.5.8" unsloth-zoo }
if ($baseInstallExit -eq 0) {
# Resolve pydantic WITH deps so pip pins pydantic-core
# to the matching version (no-torch-runtime.txt below
@ -1314,7 +1580,7 @@ shell.Run cmd, 0, False
}
}
} else {
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.5.7" unsloth-zoo }
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --reinstall-package unsloth --reinstall-package unsloth-zoo "unsloth>=2026.5.8" unsloth-zoo }
}
if ($baseInstallExit -ne 0) {
Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red
@ -1334,9 +1600,18 @@ shell.Run cmd, 0, False
return (Exit-InstallFailure "Failed to overlay unsloth-zoo (exit code $zooOverlayExit)" $zooOverlayExit)
}
}
} elseif ($TorchIndexUrl) {
} elseif ($TorchIndexUrl -or $ROCmIndexUrl) {
if ($SkipTorch) {
substep "skipping PyTorch (--no-torch flag set)." "Yellow"
} elseif ($ROCmIndexUrl) {
Write-TauriLog "STEP" "Installing PyTorch (AMD ROCm Windows)"
substep "installing PyTorch from $ROCmIndexUrl..."
$torchSpec = if ($ROCmTorchFloor) { $ROCmTorchFloor } else { "torch" }
$torchInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --force-reinstall --index-url $ROCmIndexUrl $torchSpec torchvision torchaudio }
if ($torchInstallExit -ne 0) {
Write-Host "[ERROR] Failed to install AMD ROCm PyTorch (exit code $torchInstallExit)" -ForegroundColor Red
return (Exit-InstallFailure "Failed to install AMD ROCm PyTorch (exit code $torchInstallExit)" $torchInstallExit)
}
} else {
Write-TauriLog "STEP" "Installing PyTorch"
substep "installing PyTorch ($TorchIndexUrl)..."
@ -1352,7 +1627,7 @@ shell.Run cmd, 0, False
if ($SkipTorch) {
# No-torch: install unsloth + unsloth-zoo with --no-deps, then
# runtime deps (typer, safetensors, transformers, etc.) with --no-deps.
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.5.7" unsloth-zoo }
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --no-deps --upgrade-package unsloth --upgrade-package unsloth-zoo "unsloth>=2026.5.8" unsloth-zoo }
if ($baseInstallExit -eq 0) {
# Same pydantic-with-deps trick as the migrated branch.
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython pydantic }
@ -1364,7 +1639,7 @@ shell.Run cmd, 0, False
}
}
} elseif ($StudioLocalInstall) {
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.5.7" unsloth-zoo }
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --upgrade-package unsloth "unsloth>=2026.5.8" unsloth-zoo }
} else {
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython --upgrade-package unsloth -- "$PackageName" }
}
@ -1392,7 +1667,7 @@ shell.Run cmd, 0, False
Write-TauriLog "STEP" "Installing unsloth"
substep "installing unsloth (this may take a few minutes)..."
if ($StudioLocalInstall) {
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython unsloth-zoo "unsloth>=2026.5.7" --torch-backend=auto }
$baseInstallExit = Invoke-InstallCommand { uv pip install --python $VenvPython unsloth-zoo "unsloth>=2026.5.8" --torch-backend=auto }
if ($baseInstallExit -ne 0) {
Write-Host "[ERROR] Failed to install unsloth (exit code $baseInstallExit)" -ForegroundColor Red
return (Exit-InstallFailure "Failed to install unsloth (exit code $baseInstallExit)" $baseInstallExit)

418
install.sh
View file

@ -183,10 +183,21 @@ _install_bnb_rocm() {
fi
if [ -n "$_bnb_whl_url" ]; then
substep "installing bitsandbytes for AMD ROCm (pre-release, PR #1887)..."
if run_install_cmd "$_label (pre-release)" "$_venv_py" -m pip install \
--force-reinstall --no-cache-dir --no-deps "$_bnb_whl_url"; then
_bnb_log=$(mktemp)
if "$_venv_py" -m pip install \
--disable-pip-version-check \
--force-reinstall --no-cache-dir --no-deps \
--retries 8 --timeout 90 \
"$_bnb_whl_url" >"$_bnb_log" 2>&1; then
rm -f "$_bnb_log"
return 0
fi
_bnb_rc=$?
if _is_verbose; then
cat "$_bnb_log" >&2
fi
rm -f "$_bnb_log"
step "warning" "$_label (pre-release) failed (exit code $_bnb_rc)" "$C_WARN" >&2
substep "[WARN] bnb pre-release install failed; falling back to PyPI (4-bit decode broken on ROCm)" "$C_WARN"
fi
run_install_cmd "$_label (pypi fallback)" "$_venv_py" -m pip install \
@ -245,6 +256,9 @@ _tauri_torch_index_family() {
rocm[0-9]*.[0-9]*) echo "$_diag_family" ;;
*) echo "auto" ;;
esac ;;
# AMD arch-specific index (e.g. repo.amd.com/rocm/whl/gfx1151/) --
# used for Strix Halo/Point where torch 2.11+rocm7.13 has the real fix.
*repo.amd.com/rocm/whl/gfx*|*rocm/whl/gfx*) echo "rocm7.13" ;;
"") echo "none" ;;
*) echo "auto" ;;
esac
@ -1516,17 +1530,51 @@ if [ -x "$VENV_DIR/bin/python" ]; then
: > "$VENV_DIR/.unsloth-studio-owned" 2>/dev/null || true
fi
# Guard against Python 3.13.8 torch import bug on Apple Silicon
# (skip when the user explicitly chose a version via --python)
# Guard against two independent Apple Silicon venv problems, in order:
# 1. uv may create the venv from a cached x86_64 (Rosetta) Python when a
# same-version x86_64 build is already cached (often because uv itself
# is an x86_64 build). That venv reports x86_64 to wheel resolvers, and
# PyTorch ships no macOS wheels on the CPU index for any architecture,
# so the torch install can never resolve. Recreate it with an
# arch-explicit arm64 CPython.
# 2. Python 3.13.8 has a known torch import bug.
# The two are independent: a venv may be x86_64 and, once recreated, still
# land on 3.13.8. So we re-inspect the interpreter between the checks instead
# of chaining them with elif, guaranteeing both invariants hold on whatever
# venv we end up with. Skip both when the user explicitly chose an interpreter
# via --python.
if [ -z "$_USER_PYTHON" ] && [ "$OS" = "macos" ] && [ "$_ARCH" = "arm64" ]; then
_PY_VER=$("$VENV_DIR/bin/python" -c \
"import sys; print('{}.{}.{}'.format(*sys.version_info[:3]))" 2>/dev/null || echo "")
_inspect_venv() {
"$VENV_DIR/bin/python" -c \
"import platform, sys; print(platform.machine(), '{}.{}.{}'.format(*sys.version_info[:3]))" \
2>/dev/null || echo " "
}
_info=$(_inspect_venv)
_VENV_ARCH=${_info%% *}
_PY_VER=${_info##* }
if [ "$_VENV_ARCH" = "x86_64" ]; then
echo " WARNING: venv was created with an x86_64 (Rosetta) Python on Apple Silicon."
echo " Recreating venv with native arm64 Python ${PYTHON_VERSION}..."
rm -rf "$VENV_DIR"
run_install_cmd "recreate venv (arm64)" uv venv "$VENV_DIR" \
--python "cpython-${PYTHON_VERSION}-macos-aarch64-none"
if [ -x "$VENV_DIR/bin/python" ]; then
: > "$VENV_DIR/.unsloth-studio-owned" 2>/dev/null || true
fi
# Re-inspect: the recreated arm64 venv may still be 3.13.8.
_info=$(_inspect_venv)
_VENV_ARCH=${_info%% *}
_PY_VER=${_info##* }
fi
if [ "$_PY_VER" = "3.13.8" ]; then
echo " WARNING: Python 3.13.8 has a known torch import bug."
echo " Recreating venv with Python 3.12..."
rm -rf "$VENV_DIR"
PYTHON_VERSION="3.12"
run_install_cmd "recreate venv" uv venv "$VENV_DIR" --python "$PYTHON_VERSION"
run_install_cmd "recreate venv" uv venv "$VENV_DIR" \
--python "cpython-${PYTHON_VERSION}-macos-aarch64-none"
if [ -x "$VENV_DIR/bin/python" ]; then
: > "$VENV_DIR/.unsloth-studio-owned" 2>/dev/null || true
fi
@ -1568,16 +1616,19 @@ _find_no_torch_runtime() {
}
# ── AMD ROCm GPU detection helper ──
# Returns 0 (true) if an actual AMD GPU is present, 1 (false) otherwise.
# Checks rocminfo for gfx[1-9]* (excludes gfx000 CPU agent) and
# amd-smi list for GPU data rows (excludes header-only output).
# Returns 0 if an AMD GPU is present. Checks rocminfo, amd-smi, then sysfs
# KFD topology (env-var-independent fallback for when HIP/ROCR_VISIBLE_DEVICES hides devices).
_has_amd_rocm_gpu() {
if command -v rocminfo >/dev/null 2>&1 && \
rocminfo 2>/dev/null | awk '/Name:[[:space:]]*gfx[0-9]/ && !/Name:[[:space:]]*gfx000/{found=1} END{exit !found}'; then
rocminfo 2>/dev/null | awk '/Name:[[:space:]]*gfx[1-9][0-9]/{found=1} END{exit !found}'; then
return 0
elif command -v amd-smi >/dev/null 2>&1 && \
amd-smi list 2>/dev/null | awk '/^GPU[[:space:]]*[:\[][[:space:]]*[0-9]/{ found=1 } END{ exit !found }'; then
return 0
elif [ -e /dev/kfd ] && \
awk '/gpu_id/{ if ($2+0 > 0) found=1 } END{ exit !found }' \
/sys/class/kfd/kfd/topology/nodes/*/properties 2>/dev/null; then
return 0
fi
return 1
}
@ -1656,36 +1707,50 @@ get_torch_index_url() {
if [ -n "$_rocm_tag" ]; then
# Minimum supported: ROCm 6.0 (no PyTorch wheels exist for older)
case "$_rocm_tag" in
rocm[1-5].*) echo "$_base/cpu"; return ;;
rocm[1-5].*)
echo "[WARN] ROCm $_rocm_tag detected but PyTorch ROCm wheels require ROCm 6.0+ -- falling back to CPU-only PyTorch" >&2
echo "[WARN] Upgrade ROCm: https://rocm.docs.amd.com/en/latest/deploy/linux/index.html" >&2
echo "$_base/cpu"; return ;;
esac
# ROCm 7.2 only has torch 2.11.0 which exceeds current bounds
# (<2.11.0). Fall back to rocm7.1 index which has torch 2.10.0.
# Enumerate explicit versions rather than matching rocm6.* so
# a host on ROCm 6.5 or 6.6 (no PyTorch wheels published) is
# clipped down to the last supported 6.x (rocm6.4) instead of
# constructing https://download.pytorch.org/whl/rocm6.5 which
# returns HTTP 403. PyTorch only ships: rocm5.7, 6.0, 6.1, 6.2,
# 6.3, 6.4, 7.0, 7.1, 7.2 (and 5.7 is below our minimum).
# TODO: uncomment rocm7.2 when the torch upper bound is bumped
# to >=2.11.0.
# Supported tags; 6.5+ clips to rocm6.4, 7.3+ caps to rocm7.2.
# PyTorch publishes major.minor URLs only (no patch level), so
# rocm7.2.1 / rocm6.0.2 / etc. must normalise to rocm7.2 / rocm6.0.
case "$_rocm_tag" in
rocm6.0|rocm6.0.*|rocm6.1|rocm6.1.*|rocm6.2|rocm6.2.*|rocm6.3|rocm6.3.*|rocm6.4|rocm6.4.*|rocm7.0|rocm7.0.*|rocm7.1|rocm7.1.*)
echo "$_base/$_rocm_tag" ;;
rocm6.0|rocm6.0.*) echo "$_base/rocm6.0" ;;
rocm6.1|rocm6.1.*) echo "$_base/rocm6.1" ;;
rocm6.2|rocm6.2.*) echo "$_base/rocm6.2" ;;
rocm6.3|rocm6.3.*) echo "$_base/rocm6.3" ;;
rocm6.4|rocm6.4.*) echo "$_base/rocm6.4" ;;
rocm7.0|rocm7.0.*) echo "$_base/rocm7.0" ;;
rocm7.1|rocm7.1.*) echo "$_base/rocm7.1" ;;
rocm7.2|rocm7.2.*) echo "$_base/rocm7.2" ;;
rocm6.*)
# ROCm 6.5+ (no published PyTorch wheels): clip down
# to the last supported 6.x wheel set.
echo "$_base/rocm6.4" ;;
*)
# ROCm 7.2+ (including future 10.x+): cap to rocm7.1
echo "$_base/rocm7.1" ;;
# ROCm 7.3+ (future): cap to rocm7.2 (latest known)
echo "$_base/rocm7.2" ;;
esac
return
fi
# AMD GPU confirmed by rocminfo/amd-smi but ROCm version could not be
# read from any source (amd-smi, /opt/rocm/.info/version, hipconfig,
# dpkg, rpm). Warn explicitly rather than silently installing CPU PyTorch.
echo "[WARN] AMD GPU detected but ROCm version could not be determined -- falling back to CPU-only PyTorch" >&2
echo "[WARN] Ensure one of the following is accessible: amd-smi, hipconfig, /opt/rocm/.info/version, rocm-core package" >&2
echo "[WARN] To install ROCm: https://rocm.docs.amd.com/en/latest/deploy/linux/index.html" >&2
echo "$_base/cpu"; return
fi
# Parse CUDA version from nvidia-smi output (POSIX-safe, no grep -P)
# Parse CUDA version from nvidia-smi output (POSIX-safe, no grep -P).
# Newer NVIDIA drivers (e.g. 610.x) print "CUDA UMD Version: X.Y" instead
# of the legacy "CUDA Version: X.Y"; accept both with two BRE expressions
# (POSIX sed does not support "?" without -E). The two patterns are
# mutually exclusive per line, so head -1 picks the first emitted match.
_cuda_ver=$(LC_ALL=C $_smi 2>/dev/null \
| sed -n 's/.*CUDA Version:[[:space:]]*\([0-9][0-9]*\.[0-9][0-9]*\).*/\1/p' \
| sed -n \
-e 's/.*CUDA UMD Version:[[:space:]]*\([0-9][0-9]*\.[0-9][0-9]*\).*/\1/p' \
-e 's/.*CUDA Version:[[:space:]]*\([0-9][0-9]*\.[0-9][0-9]*\).*/\1/p' \
| head -1)
if [ -z "$_cuda_ver" ]; then
echo "[WARN] Could not determine CUDA version from nvidia-smi, defaulting to cu126" >&2
@ -1754,9 +1819,9 @@ print('cp{}{}'.format(sys.version_info.major, sys.version_info.minor))
}
_pick_radeon_wheel() {
# Usage: _pick_radeon_wheel PACKAGE_NAME
# Usage: _pick_radeon_wheel PACKAGE_NAME [VERSION_PREFIX]
# Scans $_RADEON_LISTING for the newest wheel whose filename starts exactly
# with PACKAGE_NAME- and matches _RADEON_PYTAG + linux_x86_64.
# with PACKAGE_NAME- (and optionally VERSION_PREFIX) and matches _RADEON_PYTAG + linux_x86_64.
# Prints the full URL (resolving relative hrefs against _RADEON_BASE_URL).
#
# POSIX-compliant pipeline: all href parsing, filtering, and version
@ -1764,11 +1829,12 @@ _pick_radeon_wheel() {
# for GNU extensions (grep -o, sort -V) that would break under BSD
# or BusyBox coreutils.
_pkg="$1"
_ver_prefix="${2:-}"
[ -n "$_RADEON_LISTING" ] || return 1
[ -n "$_RADEON_PYTAG" ] || return 1
_tag="$_RADEON_PYTAG"
_href=$(printf '%s\n' "$_RADEON_LISTING" \
| awk -v pkg="$_pkg" -v tag="$_tag" '
| awk -v pkg="$_pkg" -v tag="$_tag" -v ver_prefix="$_ver_prefix" '
BEGIN { max_pad = ""; max_url = "" }
{
line = $0
@ -1782,7 +1848,7 @@ _pick_radeon_wheel() {
base = p[n]
sub(/[?#].*/, "", base)
prefix = pkg "-"
prefix = pkg "-" ver_prefix
# Match cpXY-cpXY or cpXY-abi3 with any linux x86_64
# platform tag (linux_x86_64, manylinux_2_28_x86_64,
# manylinux2014_x86_64, etc.)
@ -1816,6 +1882,12 @@ _pick_radeon_wheel() {
TORCH_INDEX_URL=$(get_torch_index_url)
# rocm7.2 ships torch 2.11.0 -- adjust the constraint to allow it.
# All other ROCm tags and CUDA stay within <2.11.0.
case "$TORCH_INDEX_URL" in
*/rocm7.2) TORCH_CONSTRAINT="torch>=2.11.0,<2.12.0" ;;
esac
# Auto-detect GPU for AMD ROCm based
# get_torch_index_url must have chosen */rocm*
# (gfx in rocminfo or amd-smi list). Then require rocminfo "Marketing Name:.*Radeon".
@ -1828,6 +1900,78 @@ case "$TORCH_INDEX_URL" in
fi
;;
esac
# ── Strix Halo / Strix Point: force rocm7.2 wheels, bypass Radeon repo ───────
# gfx1151 (Strix Halo) and gfx1150 (Strix Point) have a ROCm 7.1 driver bug
# that causes a segfault in torch._grouped_mm (moe_utils.py line 167).
# The Radeon repo now ships cp313 wheels for rocm-rel-7.1, so when
# _amd_gpu_radeon=true the installer silently lands on the broken combo.
# Detect these GPUs when TORCH_INDEX_URL is rocm7.1 and override to rocm7.2.
case "$TORCH_INDEX_URL" in
*/rocm7.1|*/rocm7.1.*)
# Collect every gfx token in rocminfo / amd-smi enumeration order
# (skip duplicates), then index by HIP_VISIBLE_DEVICES /
# ROCR_VISIBLE_DEVICES so a mixed Strix iGPU + non-Strix dGPU box
# where the user selected the dGPU does NOT get rerouted to the
# Strix per-gfx index.
_gfx_all=""
if command -v rocminfo >/dev/null 2>&1; then
_gfx_all=$(rocminfo 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}')
fi
if [ -z "$_gfx_all" ] && command -v amd-smi >/dev/null 2>&1; then
_gfx_all=$(amd-smi list 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}')
# PowerShell paths also probe `amd-smi static --asic`; mirror it
# so a host with hipinfo-less amd-smi reports the gfx target.
if [ -z "$_gfx_all" ]; then
_gfx_all=$(amd-smi static --asic 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}')
fi
fi
_runtime_gfx=""
if [ -n "$_gfx_all" ]; then
_vis="${HIP_VISIBLE_DEVICES:-${ROCR_VISIBLE_DEVICES:-}}"
_idx=0
if [ -n "$_vis" ] && [ "$_vis" != "-1" ]; then
_first=${_vis%%,*}
case "$_first" in
''|*[!0-9]*) _idx=0 ;;
*) _idx=$_first ;;
esac
fi
_runtime_gfx=$(printf '%s\n' "$_gfx_all" | awk -v idx="$_idx" '
NF && !seen[$0]++ { vals[n++] = $0 }
END {
if (idx < 0 || idx >= n) idx = 0
if (n > 0) print vals[idx]
}')
fi
_strix_gfx=""
case "$_runtime_gfx" in
gfx1151|gfx1150) _strix_gfx="$_runtime_gfx" ;;
esac
if [ -n "$_strix_gfx" ]; then
echo "" >&2
echo " [WARN] $_strix_gfx (Strix) + ROCm 7.1 detected -- known _grouped_mm segfault" >&2
echo " [WARN] ROCm 7.1 wheels are broken for gfx1150/gfx1151 (moe_utils.py:167)" >&2
echo " [WARN] Routing to AMD arch-specific index (torch 2.11+rocm7.13 has the real fix)" >&2
echo " [WARN] Upgrade ROCm to 7.2+ to use the standard index:" >&2
echo " [WARN] https://rocm.docs.amd.com/en/latest/deploy/linux/index.html" >&2
echo "" >&2
# AMD's arch-specific index serves torch 2.11.0+rocm7.13.0 which has AMD's
# actual fix for the gfx1151/gfx1150 _grouped_mm kernel bug -- preferred
# over the pytorch.org rocm7.2 fallback because it exercises the real GPU
# kernel path. Set UNSLOTH_AMD_ROCM_MIRROR to override for air-gapped installs.
_amd_strix_base="${UNSLOTH_AMD_ROCM_MIRROR:-https://repo.amd.com/rocm/whl}"
# Strip ALL trailing slashes to match Python's .rstrip("/") -- a
# double-/triple-slash mirror URL would otherwise produce 404s on
# strict pip proxies (artifactory, sonatype).
while [ "${_amd_strix_base%/}" != "$_amd_strix_base" ]; do
_amd_strix_base="${_amd_strix_base%/}"
done
TORCH_INDEX_URL="${_amd_strix_base}/${_strix_gfx}/"
TORCH_CONSTRAINT="torch>=2.11.0,<2.12.0"
_amd_gpu_radeon=false
fi
;;
esac
_TAURI_TORCH_INDEX_FAMILY=$(_tauri_torch_index_family "$TORCH_INDEX_URL")
if [ "$_amd_gpu_radeon" = true ] && [ "$SKIP_TORCH" = false ]; then
_TAURI_TORCH_INDEX_FAMILY="radeon"
@ -1835,27 +1979,93 @@ fi
_TAURI_GPU_BRANCH=$(_tauri_gpu_branch "$_TAURI_TORCH_INDEX_FAMILY" "$_amd_gpu_radeon")
tauri_diag_marker "$_TAURI_GPU_BRANCH" "$_TAURI_TORCH_INDEX_FAMILY"
# ── Print CPU-only hint when no GPU detected ──
# ── GPU detection summary (mirrors install.ps1 step "gpu" block) ──
if _has_usable_nvidia_gpu; then
step "gpu" "NVIDIA GPU detected"
elif case "$TORCH_INDEX_URL" in */rocm*|*/gfx*) true ;; *) false ;; esac; then
# Probe gfx arch for the display label, honouring HIP_VISIBLE_DEVICES
_gpu_disp_gfx_all=""
_gpu_disp_mkt=""
if command -v rocminfo >/dev/null 2>&1; then
_gpu_disp_gfx_all=$(rocminfo 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true)
_gpu_disp_mkt=$(rocminfo 2>/dev/null | awk -F': ' \
'/Marketing Name:/{gsub(/^[[:space:]]+|[[:space:]]+$/,"", $2); if($2){print $2; exit}}' || true)
fi
if [ -z "$_gpu_disp_gfx_all" ] && command -v amd-smi >/dev/null 2>&1; then
_gpu_disp_gfx_all=$(amd-smi list 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true)
[ -z "$_gpu_disp_gfx_all" ] && \
_gpu_disp_gfx_all=$(amd-smi static --asic 2>/dev/null | grep -oE 'gfx[1-9][0-9a-z]{2,3}' || true)
fi
if [ -z "$_gpu_disp_mkt" ] && command -v amd-smi >/dev/null 2>&1; then
_gpu_disp_mkt=$(amd-smi static --asic 2>/dev/null | awk -F'[:|]' \
'/[Mm]arket.?[Nn]ame/{gsub(/^[[:space:]]+|[[:space:]]+$/,"", $2); if($2){print $2; exit}}' || true)
fi
_gpu_vis="${HIP_VISIBLE_DEVICES:-${ROCR_VISIBLE_DEVICES:-}}"
_gpu_vis_idx=0
if [ -n "$_gpu_vis" ] && [ "$_gpu_vis" != "-1" ]; then
_gpu_first="${_gpu_vis%%,*}"
case "$_gpu_first" in ''|*[!0-9]*) ;; *) _gpu_vis_idx=$_gpu_first ;; esac
fi
_gpu_disp_gfx=$(printf '%s\n' "$_gpu_disp_gfx_all" | awk -v idx="$_gpu_vis_idx" \
'NF && !seen[$0]++ { a[n++]=$0 } END { if(idx>=n) idx=0; if(n>0) print a[idx] }')
# UNSLOTH_ROCM_GFX_ARCH env override (mirrors install.ps1)
if [ -n "${UNSLOTH_ROCM_GFX_ARCH:-}" ]; then
_gpu_disp_gfx="${UNSLOTH_ROCM_GFX_ARCH}"
substep "gfx arch from UNSLOTH_ROCM_GFX_ARCH env override: $_gpu_disp_gfx"
# Name-based arch inference when tools don't report gfx (mirrors install.ps1 nameArchTable)
elif [ -z "$_gpu_disp_gfx" ] && [ -n "$_gpu_disp_mkt" ]; then
case "$_gpu_disp_mkt" in
*"9070 XT"*|*9080*) _gpu_disp_gfx="gfx1201" ;; # RDNA 4
*9070*|*9060*) _gpu_disp_gfx="gfx1200" ;; # RDNA 4
*"8060S"*|*"890M"*|*"Strix Halo"*|*"HX 37"*|*"HX 38"*|*"AI 9 HX"*) _gpu_disp_gfx="gfx1151" ;; # RDNA 3.5 iGPU
*"880M"*|*"Strix Point"*|*"AI 9 36"*|*"AI 7 35"*|*"AI 5 34"*) _gpu_disp_gfx="gfx1150" ;; # RDNA 3.5 iGPU
*"RX 7900"*|*"RX 7800"*|*"RX 7700"*) _gpu_disp_gfx="gfx1100" ;; # RDNA 3 desktop
*"RX 7600"*) _gpu_disp_gfx="gfx1102" ;; # RDNA 3
*"780M"*|*"760M"*|*"740M"*|*"Phoenix"*) _gpu_disp_gfx="gfx1103" ;; # RDNA 3 iGPU
esac
if [ -n "$_gpu_disp_gfx" ]; then
substep "gfx arch inferred from GPU name: $_gpu_disp_gfx"
substep "Tip: set UNSLOTH_ROCM_GFX_ARCH=$_gpu_disp_gfx to skip inference next time"
fi
fi
# ROCm version via hipconfig, then amd-smi
_gpu_rocm_ver=""
if command -v hipconfig >/dev/null 2>&1; then
_gpu_rocm_ver=$(hipconfig --version 2>/dev/null | awk 'NR==1 && /^[0-9]/{print; exit}' || true)
fi
if [ -z "$_gpu_rocm_ver" ] && command -v amd-smi >/dev/null 2>&1; then
_gpu_rocm_ver=$(amd-smi version 2>/dev/null | awk -F'ROCm version: ' \
'NF>1{gsub(/[[:space:]]/,"", $2); print $2; exit}' || true)
fi
if [ -n "$_gpu_disp_gfx" ]; then
step "gpu" "AMD ROCm ($_gpu_disp_gfx)"
else
step "gpu" "AMD ROCm"
fi
_rocm_root="${ROCM_PATH:-${HIP_PATH:-/opt/rocm}}"
substep "ROCm: $_rocm_root"
[ -n "$_gpu_rocm_ver" ] && substep "hipconfig: $_gpu_rocm_ver"
[ -n "$_gpu_disp_mkt" ] && [ -n "$_gpu_disp_gfx" ] && substep "GPU: $_gpu_disp_mkt"
else
step "gpu" "none (CPU-only)" "$C_WARN"
fi
# ── PyTorch wheel index note ──
case "$TORCH_INDEX_URL" in
*/cpu)
if [ "$SKIP_TORCH" = false ] && [ "$OS" != "macos" ]; then
echo ""
echo " NOTE: No GPU detected (nvidia-smi and ROCm not found)."
echo " Installing CPU-only PyTorch. If you only need GGUF chat/inference,"
echo " re-run with --no-torch for a faster, lighter install:"
echo " curl -fsSL https://unsloth.ai/install.sh | sh -s -- --no-torch"
echo " AMD ROCm users: see https://docs.unsloth.ai/get-started/install-and-update/amd"
echo ""
substep "No GPU detected -- installing CPU-only PyTorch." "$C_WARN"
substep "AMD ROCm users: see https://docs.unsloth.ai/get-started/install-and-update/amd"
substep "Re-run with --no-torch for GGUF-only (faster, no PyTorch):"
substep " curl -fsSL https://unsloth.ai/install.sh | sh -s -- --no-torch"
fi
;;
*/rocm*)
echo ""
*/rocm*|*/gfx*)
if [ "$_amd_gpu_radeon" = true ]; then
echo " AMD Radeon + ROCm detected -- installing PyTorch wheels from repo.radeon.com"
substep "wheels: repo.radeon.com (Radeon)"
else
echo " AMD ROCm detected -- installing ROCm-enabled PyTorch ($TORCH_INDEX_URL)"
substep "wheels: $TORCH_INDEX_URL"
fi
echo ""
;;
esac
@ -1873,7 +2083,7 @@ if [ "$_MIGRATED" = true ]; then
# to prevent transitive torch resolution.
run_install_cmd "install unsloth (migrated no-torch)" uv pip install --python "$_VENV_PY" --no-deps \
--reinstall-package unsloth --reinstall-package unsloth-zoo \
"unsloth>=2026.5.7" unsloth-zoo
"unsloth>=2026.5.8" unsloth-zoo
# Resolve pydantic WITH deps so pip pins pydantic-core to the
# matching version (no-torch-runtime.txt below is --no-deps).
# All transitive deps are torch-free.
@ -1886,7 +2096,7 @@ if [ "$_MIGRATED" = true ]; then
else
run_install_cmd "install unsloth (migrated)" uv pip install --python "$_VENV_PY" \
--reinstall-package unsloth --reinstall-package unsloth-zoo \
"unsloth>=2026.5.7" unsloth-zoo
"unsloth>=2026.5.8" unsloth-zoo
fi
if [ "$STUDIO_LOCAL_INSTALL" = true ]; then
substep "overlaying local repo (editable)..."
@ -1937,24 +2147,23 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
if [ "$_radeon_listing_ok" = true ]; then
# Require torch, torchvision, torchaudio wheels to all resolve
# from the Radeon listing. If any is missing for this Python
# tag, fall through to the standard ROCm index instead of
# silently mixing Radeon wheels with PyPI defaults.
# from the Radeon listing. The repo often publishes multiple
# generations simultaneously, so picking the highest-version
# for each package independently can assemble a mismatched trio
# (e.g. torch 2.10 + torchvision 0.24). To prevent this,
# we identify the highest common minor version and downpair
# wheels if necessary to ensure a compatible set.
_torch_whl=$(_pick_radeon_wheel "torch" 2>/dev/null) || _torch_whl=""
_tv_whl=$(_pick_radeon_wheel "torchvision" 2>/dev/null) || _tv_whl=""
_ta_whl=$(_pick_radeon_wheel "torchaudio" 2>/dev/null) || _ta_whl=""
_tri_whl=$(_pick_radeon_wheel "triton" 2>/dev/null) || _tri_whl=""
# Sanity-check torch / torchvision / torchaudio are a
# matching release. The Radeon repo publishes multiple
# generations simultaneously, so picking the highest-version
# wheel for each package independently can assemble a
# mismatched trio (e.g. torch 2.9.1 + torchvision 0.23.0 +
# torchaudio 2.9.0 from the current rocm-rel-7.2.1 index).
# Check that torch and torchaudio share the same X.Y public
# version prefix, and that torchvision's minor correctly
# pairs with torch's minor (torchvision = torch.minor - 5
# pairs with torch's minor (torchvision = torch.minor + 15
# since torch 2.4 -> torchvision 0.19 -> torch 2.9 ->
# torchvision 0.24).
#
# URL-decode each wheel name so %2B -> + before version
# extraction. Real Radeon wheel hrefs are percent-encoded
# (torch-2.10.0%2Brocm7.2.0...), so a plain [+-] terminator
@ -1962,38 +2171,75 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
# _radeon_versions_match would stay false for every real
# listing, silently forcing a fallback to the generic
# ROCm index.
_torch_ver=""
_tv_ver=""
_ta_ver=""
if [ -n "$_torch_whl" ]; then
_torch_name=$(printf '%s' "${_torch_whl##*/}" | sed 's/%2[Bb]/+/g')
_torch_ver=$(printf '%s\n' "$_torch_name" | sed -n 's|^torch-\([0-9][0-9]*\.[0-9][0-9]*\)\(\.[0-9][0-9]*\)\{0,1\}[+-].*|\1|p')
fi
if [ -n "$_tv_whl" ]; then
_tv_name=$(printf '%s' "${_tv_whl##*/}" | sed 's/%2[Bb]/+/g')
_tv_ver=$(printf '%s\n' "$_tv_name" | sed -n 's|^torchvision-\([0-9][0-9]*\.[0-9][0-9]*\)\(\.[0-9][0-9]*\)\{0,1\}[+-].*|\1|p')
fi
if [ -n "$_ta_whl" ]; then
_ta_name=$(printf '%s' "${_ta_whl##*/}" | sed 's/%2[Bb]/+/g')
_ta_ver=$(printf '%s\n' "$_ta_name" | sed -n 's|^torchaudio-\([0-9][0-9]*\.[0-9][0-9]*\)\(\.[0-9][0-9]*\)\{0,1\}[+-].*|\1|p')
fi
_extract_version() {
_whl=$1
_pkg=$2
if [ -n "$_whl" ]; then
_name=$(printf '%s' "${_whl##*/}" | sed 's/%2[Bb]/+/g')
printf '%s\n' "$_name" | sed -n "s|^${_pkg}-\([0-9][0-9]*\.[0-9][0-9]*\)\(\.[0-9][0-9]*\)\{0,1\}[+-].*|\1|p"
fi
}
_torch_ver=$(_extract_version "$_torch_whl" "torch")
_tv_ver=$(_extract_version "$_tv_whl" "torchvision")
_ta_ver=$(_extract_version "$_ta_whl" "torchaudio")
_radeon_versions_match=false
if [ -n "$_torch_ver" ] && [ -n "$_tv_ver" ] && [ -n "$_ta_ver" ]; then
_torch_major=${_torch_ver%%.*}
_torch_minor=${_torch_ver#*.}
_ta_major=${_ta_ver%%.*}
_ta_minor=${_ta_ver#*.}
_tv_major=${_tv_ver%%.*}
_tv_minor=${_tv_ver#*.}
# torchvision expected minor (e.g. torch 2.9 -> 0.24)
_expected_tv_minor=$((_torch_minor + 15))
if [ "$_torch_major" = "$_ta_major" ] && \
[ "$_torch_minor" = "$_ta_minor" ] && \
[ "$_tv_major" = "0" ] && \
[ "$_tv_minor" = "$_expected_tv_minor" ]; then
_radeon_versions_match=true
fi
_tv_equiv_minor=$((_tv_minor - 15))
# Determine initial target minor (lowest common denominator)
_target_minor=$_torch_minor
[ "$_tv_equiv_minor" -lt "$_target_minor" ] && _target_minor=$_tv_equiv_minor
[ "$_ta_minor" -lt "$_target_minor" ] && _target_minor=$_ta_minor
# Loop downwards to find the first complete matching trio.
# This avoids aborting if the repo has gaps.
_attempts=0
while [ "$_attempts" -lt 5 ] && [ "$_target_minor" -ge 0 ]; do
_expected_tv_minor=$((_target_minor + 15))
_curr_torch=$(_pick_radeon_wheel "torch" "2.${_target_minor}." 2>/dev/null) || _curr_torch=""
_curr_tv=$(_pick_radeon_wheel "torchvision" "0.${_expected_tv_minor}." 2>/dev/null) || _curr_tv=""
_curr_ta=$(_pick_radeon_wheel "torchaudio" "2.${_target_minor}." 2>/dev/null) || _curr_ta=""
if [ -n "$_curr_torch" ] && [ -n "$_curr_tv" ] && [ -n "$_curr_ta" ]; then
# Extract versions from the wheels found in this iteration
_c_torch_ver=$(_extract_version "$_curr_torch" "torch")
_c_tv_ver=$(_extract_version "$_curr_tv" "torchvision")
_c_ta_ver=$(_extract_version "$_curr_ta" "torchaudio")
# Parse Major.Minor for validation
_c_torch_major=${_c_torch_ver%%.*}
_c_torch_minor=${_c_torch_ver#*.}
_c_ta_major=${_c_ta_ver%%.*}
_c_ta_minor=${_c_ta_ver#*.}
_c_tv_major=${_c_tv_ver%%.*}
_c_tv_minor=${_c_tv_ver#*.}
# Strict X.Y validation: allow patch versions to differ (e.g. torch 2.9.1 + vision 0.24.0)
# as long as the Major and Minor pairing is correct.
if [ "$_c_torch_major" = "$_c_ta_major" ] && \
[ "$_c_torch_minor" = "$_c_ta_minor" ] && \
[ "$_c_tv_major" = "0" ] && \
[ "$_c_tv_minor" = "$((_c_torch_minor + 15))" ]; then
_torch_whl=$_curr_torch
_tv_whl=$_curr_tv
_ta_whl=$_curr_ta
_tri_whl=""
_radeon_versions_match=true
break
fi
fi
_target_minor=$((_target_minor - 1))
_attempts=$((_attempts + 1))
done
fi
if [ -z "$_torch_whl" ] || [ -z "$_tv_whl" ] || [ -z "$_ta_whl" ] || \
[ "$_radeon_versions_match" != true ]; then
substep "[WARN] Radeon repo lacks a compatible wheel set for this Python; falling back to ROCm index ($TORCH_INDEX_URL)" "$C_WARN"
@ -2054,7 +2300,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
# runtime deps (typer, safetensors, transformers, etc.) with --no-deps.
run_install_cmd "install unsloth (no-torch)" uv pip install --python "$_VENV_PY" --no-deps \
--upgrade-package unsloth --upgrade-package unsloth-zoo \
"unsloth>=2026.5.7" unsloth-zoo
"unsloth>=2026.5.8" unsloth-zoo
# Same pydantic-with-deps trick as the migrated branch.
run_install_cmd "install pydantic (with deps for compatible core)" \
uv pip install --python "$_VENV_PY" pydantic
@ -2072,7 +2318,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
fi
elif [ "$STUDIO_LOCAL_INSTALL" = true ]; then
run_install_cmd "install unsloth (local)" uv pip install --python "$_VENV_PY" \
--upgrade-package unsloth "unsloth>=2026.5.7" unsloth-zoo
--upgrade-package unsloth "unsloth>=2026.5.8" unsloth-zoo
substep "overlaying local repo (editable)..."
run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps
substep "overlaying unsloth-zoo from git main..."
@ -2104,7 +2350,7 @@ else
tauri_log "STEP" "Installing Unsloth"
substep "installing unsloth (this may take a few minutes)..."
if [ "$STUDIO_LOCAL_INSTALL" = true ]; then
run_install_cmd "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" unsloth-zoo "unsloth>=2026.5.7" --torch-backend=auto
run_install_cmd "install unsloth (auto torch backend)" uv pip install --python "$_VENV_PY" unsloth-zoo "unsloth>=2026.5.8" --torch-backend=auto
substep "overlaying local repo (editable)..."
run_install_cmd "overlay local repo" uv pip install --python "$_VENV_PY" -e "$_REPO_ROOT" --no-deps
substep "overlaying unsloth-zoo from git main..."

854
scripts/verify_import_hoist.py Normal file
View file

@ -0,0 +1,854 @@
#!/usr/bin/env python3
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
"""Deterministic, scope-aware verifier for import-hoisting / alias-rename refactors.
The risk when moving `from a import b as _b` (or `import b as _b`) to module top
and normalizing `_b` -> `b` is twofold:
1. DANGLING ALIAS - a `_b` reference is left un-normalized; it now resolves to
nothing (NameError) or, worse, to some *other* module-level `_b`.
2. RENAME CLASH - `_b` was an alias on purpose because `b` already meant
something else in that scope; normalizing `_b` -> `b` silently re-points the
reference at the wrong object (no NameError, no pyflakes warning).
This tool parses BEFORE (a git ref, default origin/main) and AFTER (default HEAD)
for each file, builds a real LEGB scope model (functions, classes, lambdas,
comprehensions, global/nonlocal, args, walrus, star-imports), and resolves every
Name load to its binding. It then compares, PER SCOPE:
* UNRESOLVED-NEW : loads that resolve to nothing in AFTER but did in BEFORE
(or are newly present) -> catches dangling aliases.
* TARGET-MISSING : an import *target* (e.g. module `glob`, or
`importlib.metadata.version`) that a function resolved to
in BEFORE but no longer resolves to in AFTER -> catches a
function that lost access to a module it still uses.
Robust to alias renames because it compares the *target*,
not the local name.
* TARGET-CHANGED : a load whose resolved import target differs BEFORE vs
AFTER -> catches a rename that re-points to a different
module (the clash case).
* AMBIGUOUS-BIND : a name bound by BOTH an import and a non-import in the same
scope in AFTER (and not in BEFORE) -> the "alias was on
purpose / now collides" smell.
* MODULE-DUP-IMPORT: a module-level name imported and also defined/assigned at
module level (introduced by the change).
* NEW-UNUSED-IMPORT: a module-level import added in AFTER that nothing resolves
to (informational; re-exports are a known false positive).
Usage:
verify_import_hoist.py [--before REF] [--after REF] <file>... # compare
verify_import_hoist.py --self-test # prove it catches bugs
Exit code 1 if any non-informational finding.
"""
from __future__ import annotations
import argparse
import ast
import builtins
import re as _re_mod
import subprocess
import sys
from dataclasses import dataclass, field
_BUILTINS = set(dir(builtins)) | {
"__file__",
"__name__",
"__doc__",
"__package__",
"__spec__",
"__loader__",
"__builtins__",
"__class__",
"__annotations__",
"__dict__",
"__qualname__",
"__module__",
"__path__",
"__debug__",
"__import__",
"NotImplemented",
"Ellipsis",
"copyright",
"credits",
"license",
"help",
"exit",
"quit",
"__build_class__",
"__cached__",
"reveal_type",
"reveal_locals",
}
# ---------------------------------------------------------------- scope model
@dataclass
class Binding:
kind: str # 'import' | 'importfrom' | 'def' | 'class' | 'other'
target: str | None = None # canonical import target id, else None
@dataclass
class Scope:
kind: str # 'module' | 'function' | 'class' | 'lambda' | 'comp'
qualname: str
parent: "Scope | None"
bindings: dict[str, list[Binding]] = field(default_factory = dict)
globals: set[str] = field(default_factory = set)
nonlocals: set[str] = field(default_factory = set)
star_import: bool = False
def add(self, name: str, b: Binding) -> None:
self.bindings.setdefault(name, []).append(b)
def _import_target(node: ast.AST, alias: ast.alias) -> tuple[str, str]:
"""Return (bound_name, canonical_target_id) for one import alias."""
if isinstance(node, ast.Import):
bound = alias.asname or alias.name.split(".")[0]
return bound, f"import:{alias.name}"
# ImportFrom
bound = alias.asname or alias.name
mod = ("." * (node.level or 0)) + (node.module or "")
return bound, f"from:{mod}:{alias.name}"
class _Builder(ast.NodeVisitor):
"""Builds the scope tree + bindings, and records every (scope, Name-load)."""
def __init__(self):
self.module = Scope("module", "<module>", None)
self.uses: list[tuple[Scope, str, int]] = [] # (scope, name, lineno) hard loads
self.soft_uses: list[
tuple[Scope, str, int]
] = [] # annotations: count as "used"
# but never as "unresolved"
# (forward refs / string annos)
def _visit_annotation(self, node, scope: Scope) -> None:
"""Annotation context: with `from __future__ import annotations` these are
never evaluated (strings), and even otherwise they routinely contain forward
references. Record contained names as SOFT uses so an import used only in an
annotation still counts as used, but a forward-ref name is never 'unresolved'."""
if node is None:
return
for n in ast.walk(node):
if isinstance(n, ast.Name) and isinstance(n.ctx, ast.Load):
self.soft_uses.append((scope, n.id, n.lineno))
# -- binding helpers --
def _bind_targets(self, scope: Scope, target: ast.AST) -> None:
for n in ast.walk(target):
if isinstance(n, ast.Name) and isinstance(n.ctx, (ast.Store, ast.Del)):
self._bind_name(scope, n.id, Binding("other"))
elif isinstance(n, ast.Starred):
pass
def _bind_name(self, scope: Scope, name: str, b: Binding) -> None:
if name in scope.globals:
self.module.add(name, b)
elif name in scope.nonlocals:
p = scope.parent
while p is not None and p.kind not in ("function", "lambda"):
p = p.parent
(p or self.module).add(name, b)
else:
scope.add(name, b)
# -- generic dispatch within a scope --
def _visit_body(self, stmts, scope: Scope) -> None:
for s in stmts:
self._visit_stmt(s, scope)
def _visit_stmt(self, node: ast.AST, scope: Scope) -> None:
if isinstance(node, (ast.Import, ast.ImportFrom)):
star = isinstance(node, ast.ImportFrom) and any(
a.name == "*" for a in node.names
)
if star:
scope.star_import = True
for alias in node.names:
if alias.name == "*":
continue
bound, target = _import_target(node, alias)
kind = "import" if isinstance(node, ast.Import) else "importfrom"
self._bind_name(scope, bound, Binding(kind, target))
return
if isinstance(node, ast.Global):
scope.globals.update(node.names)
return
if isinstance(node, ast.Nonlocal):
scope.nonlocals.update(node.names)
return
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
self._bind_name(scope, node.name, Binding("def"))
# decorators / defaults evaluate in the ENCLOSING scope
for d in node.decorator_list:
self._visit_expr(d, scope)
self._visit_arg_defaults(node.args, scope)
child = Scope("function", f"{scope.qualname}.{node.name}", scope)
self._bind_type_params(node, child)
self._bind_args(node.args, child)
# arg + return annotations: soft uses (may be strings / forward refs)
for a in self._all_args(node.args):
self._visit_annotation(a.annotation, child)
self._visit_annotation(getattr(node, "returns", None), child)
self._visit_body(node.body, child)
return
if isinstance(node, ast.ClassDef):
self._bind_name(scope, node.name, Binding("class"))
for d in node.decorator_list:
self._visit_expr(d, scope)
for b in node.bases:
self._visit_expr(b, scope)
for kw in node.keywords:
self._visit_expr(kw.value, scope)
child = Scope("class", f"{scope.qualname}.{node.name}", scope)
self._bind_type_params(node, child)
self._visit_body(node.body, child)
return
if isinstance(node, ast.Match):
self._visit_expr(node.subject, scope)
for case in node.cases:
self._bind_pattern(case.pattern, scope)
if case.guard is not None:
self._visit_expr(case.guard, scope)
self._visit_body(case.body, scope)
return
if isinstance(node, getattr(ast, "TryStar", ())): # py3.11 except*
self._visit_body(node.body, scope)
for h in node.handlers:
if h.type is not None:
self._visit_expr(h.type, scope)
if h.name:
self._bind_name(scope, h.name, Binding("other"))
self._visit_body(h.body, scope)
self._visit_body(node.orelse, scope)
self._visit_body(node.finalbody, scope)
return
if isinstance(node, getattr(ast, "TypeAlias", ())): # py3.12 `type X = ...`
if isinstance(node.name, ast.Name):
self._bind_name(scope, node.name.id, Binding("other"))
self._visit_annotation(node.value, scope)
return
if isinstance(node, (ast.Assign, ast.AnnAssign, ast.AugAssign)):
targets = node.targets if isinstance(node, ast.Assign) else [node.target]
val = node.value
if val is not None:
self._visit_expr(val, scope)
if isinstance(node, ast.AnnAssign) and node.annotation is not None:
self._visit_annotation(node.annotation, scope)
for t in targets:
self._bind_targets(scope, t)
# AugAssign target is also a load
if isinstance(node, ast.AugAssign):
self._record_loads(t, scope)
return
if isinstance(node, (ast.For, ast.AsyncFor)):
self._visit_expr(node.iter, scope)
self._bind_targets(scope, node.target)
self._visit_body(node.body, scope)
self._visit_body(node.orelse, scope)
return
if isinstance(node, (ast.With, ast.AsyncWith)):
for item in node.items:
self._visit_expr(item.context_expr, scope)
if item.optional_vars is not None:
self._bind_targets(scope, item.optional_vars)
self._visit_body(node.body, scope)
return
if isinstance(node, ast.Try):
self._visit_body(node.body, scope)
for h in node.handlers:
if h.type is not None:
self._visit_expr(h.type, scope)
if h.name:
self._bind_name(scope, h.name, Binding("other"))
self._visit_body(h.body, scope)
self._visit_body(node.orelse, scope)
self._visit_body(node.finalbody, scope)
return
# generic statement: visit all child expressions/stmts in same scope
for child in ast.iter_child_nodes(node):
if isinstance(child, ast.stmt):
self._visit_stmt(child, scope)
else:
self._visit_expr(child, scope)
# -- expressions --
def _visit_arg_defaults(self, args: ast.arguments, scope: Scope) -> None:
for d in list(args.defaults) + [d for d in args.kw_defaults if d is not None]:
self._visit_expr(d, scope)
def _all_args(self, args: ast.arguments) -> list[ast.arg]:
out = list(args.posonlyargs) + list(args.args) + list(args.kwonlyargs)
if args.vararg:
out.append(args.vararg)
if args.kwarg:
out.append(args.kwarg)
return out
def _bind_args(self, args: ast.arguments, scope: Scope) -> None:
for a in self._all_args(args):
scope.add(a.arg, Binding("other"))
def _bind_type_params(self, node, scope: Scope) -> None:
for tp in getattr(node, "type_params", []) or []:
name = getattr(tp, "name", None)
if isinstance(name, str):
scope.add(name, Binding("other"))
self._visit_annotation(getattr(tp, "bound", None), scope)
self._visit_annotation(getattr(tp, "default_value", None), scope)
def _bind_pattern(self, pat, scope: Scope) -> None:
if pat is None:
return
if isinstance(pat, ast.MatchValue):
self._visit_expr(pat.value, scope)
elif isinstance(pat, ast.MatchSingleton):
pass
elif isinstance(pat, ast.MatchSequence):
for p in pat.patterns:
self._bind_pattern(p, scope)
elif isinstance(pat, ast.MatchStar):
if pat.name:
self._bind_name(scope, pat.name, Binding("other"))
elif isinstance(pat, ast.MatchMapping):
for k in pat.keys:
self._visit_expr(k, scope)
for p in pat.patterns:
self._bind_pattern(p, scope)
if pat.rest:
self._bind_name(scope, pat.rest, Binding("other"))
elif isinstance(pat, ast.MatchClass):
self._visit_expr(pat.cls, scope)
for p in pat.patterns:
self._bind_pattern(p, scope)
for p in pat.kwd_patterns:
self._bind_pattern(p, scope)
elif isinstance(pat, ast.MatchAs):
self._bind_pattern(pat.pattern, scope)
if pat.name:
self._bind_name(scope, pat.name, Binding("other"))
elif isinstance(pat, ast.MatchOr):
for p in pat.patterns:
self._bind_pattern(p, scope)
def _record_loads(self, node: ast.AST, scope: Scope) -> None:
for n in ast.walk(node):
if isinstance(n, ast.Name) and isinstance(n.ctx, ast.Load):
self.uses.append((scope, n.id, n.lineno))
def _visit_expr(self, node: ast.AST, scope: Scope) -> None:
if isinstance(node, ast.Name):
if isinstance(node.ctx, ast.Load):
self.uses.append((scope, node.id, node.lineno))
elif isinstance(node.ctx, (ast.Store, ast.Del)):
self._bind_name(scope, node.id, Binding("other"))
return
if isinstance(node, ast.Lambda):
self._visit_arg_defaults(node.args, scope)
child = Scope("lambda", f"{scope.qualname}.<lambda>", scope)
self._bind_args(node.args, child)
self._visit_expr(node.body, child)
return
if isinstance(
node, (ast.ListComp, ast.SetComp, ast.GeneratorExp, ast.DictComp)
):
child = Scope("comp", f"{scope.qualname}.<comp>", scope)
for i, gen in enumerate(node.generators):
# first iterable is evaluated in the enclosing scope
self._visit_expr(gen.iter, scope if i == 0 else child)
self._bind_targets(child, gen.target)
for cond in gen.ifs:
self._visit_expr(cond, child)
if isinstance(node, ast.DictComp):
self._visit_expr(node.key, child)
self._visit_expr(node.value, child)
else:
self._visit_expr(node.elt, child)
return
if isinstance(node, ast.NamedExpr): # walrus binds in enclosing scope
self._visit_expr(node.value, scope)
if isinstance(node.target, ast.Name):
self._bind_name(scope, node.target.id, Binding("other"))
return
for child in ast.iter_child_nodes(node):
if isinstance(child, ast.stmt):
self._visit_stmt(child, scope)
else:
self._visit_expr(child, scope)
def run(self, tree: ast.Module) -> None:
self._visit_body(tree.body, self.module)
# ---------------------------------------------------------------- resolution
def _any_star(scope: Scope) -> bool:
c = scope
while c is not None:
if c.star_import:
return True
c = c.parent
return False
def _resolve(scope: Scope, name: str):
"""LEGB resolution. Returns (status, bindings) where status in
{'local','import','other','builtin','star','unresolved'}."""
# global / nonlocal redirection
start = scope
if name in scope.globals:
chain = [_module_of(scope)]
elif name in scope.nonlocals:
chain = _enclosing_functions(scope)
else:
chain = _legb_chain(scope)
for i, sc in enumerate(chain):
if sc is None:
continue
if name in sc.bindings:
binds = sc.bindings[name]
if any(b.kind in ("import", "importfrom") for b in binds):
return "import", binds
return "other", binds
if name in _BUILTINS:
return "builtin", []
if _any_star(start):
return "star", []
return "unresolved", []
def _module_of(scope: Scope) -> Scope:
while scope.parent is not None:
scope = scope.parent
return scope
def _enclosing_functions(scope: Scope) -> list[Scope]:
out = []
p = scope.parent
while p is not None:
if p.kind in ("function", "lambda"):
out.append(p)
p = p.parent
out.append(_module_of(scope))
return out
def _legb_chain(scope: Scope) -> list[Scope]:
"""Immediate scope, then enclosing scopes skipping class scopes, then module."""
chain = [scope]
p = scope.parent
while p is not None:
if (
p.kind != "class" or p.parent is None
): # module-level class never happens; keep module
if p.kind != "class":
chain.append(p)
p = p.parent
return chain
# ---------------------------------------------------------------- analysis
def _analyze(src: str):
tree = ast.parse(src)
b = _Builder()
b.run(tree)
# Per-scope: unresolved load names, and import targets it resolves to.
unresolved: dict[str, set[str]] = {}
targets_by_scope: dict[str, set[str]] = {}
target_by_use: dict[tuple[str, str], set[str]] = {}
for scope, name, _ln in b.uses:
status, binds = _resolve(scope, name)
if status == "unresolved":
unresolved.setdefault(scope.qualname, set()).add(name)
elif status == "import":
tids = {bd.target for bd in binds if bd.target}
targets_by_scope.setdefault(scope.qualname, set()).update(tids)
target_by_use.setdefault((scope.qualname, name), set()).update(tids)
# soft uses (annotations): only contribute to "used", never to "unresolved"
for scope, name, _ln in b.soft_uses:
status, binds = _resolve(scope, name)
if status == "import":
tids = {bd.target for bd in binds if bd.target}
targets_by_scope.setdefault(scope.qualname, set()).update(tids)
# module-level binding info for clash checks
module = b.module
module_imports = {
n: bs
for n, bs in module.bindings.items()
if any(x.kind in ("import", "importfrom") for x in bs)
}
module_dup = {
n
for n, bs in module.bindings.items()
if any(x.kind in ("import", "importfrom") for x in bs)
and any(x.kind not in ("import", "importfrom") for x in bs)
}
# ambiguous: any scope where a name is bound by import AND non-import
ambiguous: dict[str, set[str]] = {}
def walk_scopes(scope: Scope):
for n, bs in scope.bindings.items():
if any(x.kind in ("import", "importfrom") for x in bs) and any(
x.kind not in ("import", "importfrom") for x in bs
):
ambiguous.setdefault(scope.qualname, set()).add(n)
# scope tree isn't stored; rebuild via uses is hard. We approximate with module only.
walk_scopes(module)
return {
"unresolved": unresolved,
"targets_by_scope": targets_by_scope,
"target_by_use": target_by_use,
"module_import_targets": {
n: {x.target for x in bs if x.target} for n, bs in module_imports.items()
},
"module_dup": module_dup,
"ambiguous": ambiguous,
}
def _git_show(ref: str, path: str) -> str | None:
try:
return subprocess.run(
["git", "show", f"{ref}:{path}"], capture_output = True, text = True, check = True
).stdout
except subprocess.CalledProcessError:
return None
def compare(before_src: str, after_src: str, path: str) -> list[tuple[str, str]]:
"""Return list of (severity, message). severity in BLOCKER/WARN/INFO.
Blocker signals (precise, no relocation false-positives):
UNRESOLVED-NEW - a load became undefined (dangling alias / removed import).
NEW-UNUSED-HOIST - a module-level import added by THIS change is resolved by
NO load. A correct hoist always wires its new import to a
reference; if the alias was left un-normalized OR renamed
to the wrong name, the hoisted import ends up unused. This
single signal catches BOTH user-described failure modes and
does NOT fire for code merely relocated to another file
(that removes the import, it doesn't add an unused one).
TARGET-CHANGED - the same (scope, name) load resolves to a different import
target before vs after (a same-name re-point).
"""
a = _analyze(before_src)
b = _analyze(after_src)
findings: list[tuple[str, str]] = []
def used_targets(analysis) -> set[str]:
out: set[str] = set()
for tids in analysis["targets_by_scope"].values():
out |= tids
return out
before_used = used_targets(a)
after_used = used_targets(b)
before_module_targets: set[str] = set()
for tids in a["module_import_targets"].values():
before_module_targets |= tids
after_module_targets: set[str] = set()
for tids in b["module_import_targets"].values():
after_module_targets |= tids
added_module_targets = after_module_targets - before_module_targets
# 1. UNRESOLVED-NEW
for scope, names in b["unresolved"].items():
new = names - a["unresolved"].get(scope, set())
for n in sorted(new):
findings.append(
(
"BLOCKER",
f"{path}: UNRESOLVED-NEW '{n}' in scope {scope} "
f"(undefined after change -> dangling alias / removed import)",
)
)
# 2. HOISTED-IMPORT-UNUSED (the core botched-hoist / wrong-rename signal)
# A module-level import in AFTER that NO load resolves to, and which was
# either newly added by this change OR was actually used before. Excludes:
# - relocation (the import is REMOVED, so it's not in after at all)
# - stable pre-existing re-exports (unused before AND after, not newly added)
for n, tids in b["module_import_targets"].items():
if tids & after_used:
continue # resolved by something -> fine
newly_added = bool(tids - before_module_targets)
was_used_before = bool(tids & before_used)
if newly_added or was_used_before:
why = (
"added but unused"
if newly_added
else "was used before, now unused (references re-pointed)"
)
findings.append(
(
"BLOCKER",
f"{path}: HOISTED-IMPORT-UNUSED '{n}' ({sorted(tids)}) "
f"{why} -> un-normalized alias or wrong rename target?",
)
)
# 3. TARGET-CHANGED (same scope+name resolves to a different import target)
for key, tafter in b["target_by_use"].items():
tbefore = a["target_by_use"].get(key)
if tbefore and tbefore != tafter:
findings.append(
(
"BLOCKER",
f"{path}: TARGET-CHANGED name '{key[1]}' in {key[0]} "
f"{sorted(tbefore)} -> {sorted(tafter)} (rename re-points module)",
)
)
# 4. MODULE-DUP-IMPORT introduced
for n in sorted(b["module_dup"] - a["module_dup"]):
findings.append(
(
"WARN",
f"{path}: MODULE-DUP-IMPORT '{n}' bound by import AND non-import "
f"at module level (possible clash)",
)
)
# 5. AMBIGUOUS-BIND introduced (module scope)
for scope, names in b["ambiguous"].items():
new = names - a["ambiguous"].get(scope, set())
for n in sorted(new):
findings.append(
("WARN", f"{path}: AMBIGUOUS-BIND '{n}' import+non-import in {scope}")
)
# 6. TARGET-MISSING (informational): a scope stopped resolving to an import
# target. Real bugs are already covered above; remaining cases are code
# relocated to another file (e.g. a moved helper). Shown for transparency.
for scope, tbefore in a["targets_by_scope"].items():
tafter = b["targets_by_scope"].get(scope, set())
for t in sorted(tbefore - tafter):
relocated = (
""
if t in added_module_targets
else " [target not re-added here -> likely relocated/deleted]"
)
findings.append(
("INFO", f"{path}: TARGET-MISSING {t} in scope {scope}{relocated}")
)
return findings
# ---------------------------------------------------------------- self-test
_SELF_TESTS = {
"dangling_alias": (
# before: inline aliased import, used as _b
"import os\n"
"def f():\n"
" import glob as _b\n"
" return _b.glob('*')\n",
# after: hoisted to canonical, but reference NOT normalized -> _b dangles
"import os\n" "import glob\n" "def f():\n" " return _b.glob('*')\n",
"BLOCKER",
),
"rename_clash": (
# before: _b is a deliberate alias; `b` already means something else
"import re as _b\n" "b = 123\n" "def f():\n" " return _b.compile('x'), b\n",
# after: someone normalized _b -> b ; now f().b is the int, re is lost
"import re\n" "b = 123\n" "def f():\n" " return b.compile('x'), b\n",
"BLOCKER", # TARGET-MISSING from:.. or import:re in f
),
"clean_rename": (
"def f():\n" " import glob as _g\n" " return _g.glob('*')\n",
"import glob\n" "def f():\n" " return glob.glob('*')\n",
None, # expect NO blocker
),
"clean_dedup_redundant": (
"import sys\n" "def f():\n" " import sys\n" " return sys.argv\n",
"import sys\n" "def f():\n" " return sys.argv\n",
None,
),
"from_import_dangling": (
# from-import alias left un-normalized
"def f():\n"
" from importlib.metadata import version as _v\n"
" return _v('x')\n",
"from importlib.metadata import version\n" "def f():\n" " return _v('x')\n",
"BLOCKER",
),
"local_var_clash": (
# _b renamed to b, but b is a LOCAL variable in f -> import silently unused
"def f(b):\n" " import re as _b\n" " return _b.compile(b)\n",
"import re\n"
"def f(b):\n"
" return b.compile(b)\n", # 'b' is the param, not the module
"BLOCKER",
),
"substring_safe": (
# correct _copy->copy rename while a config_copy var exists: NO false positive
"def f(config):\n"
" import copy as _copy\n"
" config_copy = _copy.deepcopy(config)\n"
" return config_copy\n",
"import copy\n"
"def f(config):\n"
" config_copy = copy.deepcopy(config)\n"
" return config_copy\n",
None,
),
"attr_access_not_a_use": (
# x._b is attribute access, not a use of name _b; removing import _b is fine
"import os\n"
"def f(x):\n"
" import sys as _b\n"
" return x._b + _b.argv[0]\n",
"import os\n" "import sys\n" "def f(x):\n" " return x._b + sys.argv[0]\n",
None,
),
}
def _self_test() -> int:
ok = True
for name, (before, after, expect) in _SELF_TESTS.items():
findings = compare(before, after, f"<{name}>")
blockers = [m for sev, m in findings if sev == "BLOCKER"]
got = "BLOCKER" if blockers else None
passed = got == expect
ok = ok and passed
print(f"[{'PASS' if passed else 'FAIL'}] {name}: expect={expect} got={got}")
for sev, m in findings:
print(f" ({sev}) {m}")
print("\nSELF-TEST:", "ALL PASS" if ok else "FAILURES")
return 0 if ok else 1
def _pyflakes_undefined(path: str) -> set[str] | None:
"""Return the set of names pyflakes reports as 'undefined name' for `path`,
or None if pyflakes failed to run/parse the file."""
try:
proc = subprocess.run(
[sys.executable, "-m", "pyflakes", path], capture_output = True, text = True
)
except Exception:
return None
if "syntax error" in (proc.stdout + proc.stderr).lower():
return None
names = set()
for line in proc.stdout.splitlines():
m = _re_mod.search(r"undefined name '([^']+)'", line)
if m:
names.add(m.group(1))
return names
def audit_files(paths: list[str]) -> int:
"""Single-version robustness audit. For every file: confirm the analyzer does
not crash, then cross-check its 'unresolved' names against pyflakes. Any name
the resolver flags that pyflakes does NOT call undefined is a tool FALSE
POSITIVE (a resolver gap to fix)."""
n_files = n_err = n_fp = n_syntax = 0
fp_detail: dict[str, set[str]] = {}
err_detail: dict[str, str] = {}
for path in paths:
n_files += 1
try:
src = open(path, encoding = "utf-8").read()
except Exception as e: # unreadable
n_err += 1
err_detail[path] = f"read: {e}"
continue
try:
res = _analyze(src)
except SyntaxError:
n_syntax += 1
continue
except Exception as e: # analyzer crash -> robustness bug
n_err += 1
err_detail[path] = f"{type(e).__name__}: {e}"
continue
tool_unresolved = set()
for names in res["unresolved"].values():
tool_unresolved |= names
if not tool_unresolved:
continue
pf = _pyflakes_undefined(path)
if pf is None:
continue # pyflakes couldn't adjudicate; skip cross-check
false_pos = tool_unresolved - pf
if false_pos:
n_fp += 1
fp_detail[path] = false_pos
print(f"audited files : {n_files}")
print(f"syntax-skipped : {n_syntax}")
print(f"analyzer errors : {n_err}")
for p, e in sorted(err_detail.items()):
print(f" ERROR {p}: {e}")
print(f"false-positive files: {n_fp} (resolver flagged a name pyflakes accepts)")
for p, names in sorted(fp_detail.items()):
print(f" FP {p}: {sorted(names)}")
ok = n_err == 0 and n_fp == 0
print(
"\nAUDIT:",
"ROBUST (no crashes, no false positives vs pyflakes)"
if ok
else "NEEDS WORK (see above)",
)
return 0 if ok else 1
def main() -> int:
ap = argparse.ArgumentParser()
ap.add_argument("--before", default = "origin/main")
ap.add_argument("--after", default = "HEAD")
ap.add_argument("--self-test", action = "store_true")
ap.add_argument(
"--audit",
action = "store_true",
help = "single-version robustness audit on filesystem paths",
)
ap.add_argument("files", nargs = "*")
args = ap.parse_args()
if args.self_test:
return _self_test()
if args.audit:
return audit_files(args.files)
any_blocker = False
for path in args.files:
before = _git_show(args.before, path)
after = _git_show(args.after, path)
if after is None:
print(f"SKIP {path}: not found at {args.after}")
continue
if before is None:
before = "" # new file
findings = compare(before, after, path)
blockers = [f for f in findings if f[0] == "BLOCKER"]
warns = [f for f in findings if f[0] == "WARN"]
infos = [f for f in findings if f[0] == "INFO"]
status = (
"CLEAN"
if not blockers and not warns
else ("BLOCKERS" if blockers else "WARNINGS")
)
print(f"\n=== {path}: {status} ===")
for sev, m in blockers + warns + infos:
print(f" [{sev}] {m}")
any_blocker = any_blocker or bool(blockers)
print(
"\nOVERALL:", "FAIL (blockers found)" if any_blocker else "PASS (no blockers)"
)
return 1 if any_blocker else 0
if __name__ == "__main__":
sys.exit(main())

23
studio/Unsloth_Studio_Colab.ipynb
View file

@ -84,26 +84,7 @@
"id": "277e431e"
},
"outputs": [],
"source": [
"import sys, time\n",
"sys.path.insert(0, \"/content/unsloth/studio/backend\")\n",
"from colab import start\n",
"start()"
]
},
{
"cell_type": "code",
"source": [
"from google.colab import output\n",
"output.serve_kernel_port_as_iframe(8888, height = 1200, width = \"100%\")\n",
"for _ in range(10000): time.sleep(300), print(\"=\", end = \"\")"
],
"metadata": {
"id": "wb9UELh--XzX"
},
"id": "wb9UELh--XzX",
"execution_count": null,
"outputs": []
"source": "import sys\nsys.path.insert(0, \"/content/unsloth/studio/backend\")\nfrom colab import start\nstart()"
},
{
"cell_type": "markdown",
@ -150,4 +131,4 @@
},
"nbformat": 4,
"nbformat_minor": 5
}
}

216
studio/backend/colab.py
View file

@ -26,30 +26,68 @@ logger = get_logger(__name__)
def get_colab_url(port: int = 8888) -> str:
"""
Get the actual Colab proxy URL for a port.
Retries up to 3 times and validates that the result is a real HTTPS Colab
URL before returning. Falls back to http://localhost:{port} only when all
attempts fail.
"""
import time as _time
fallback = f"http://localhost:{port}"
try:
from google.colab.output import eval_js
except ImportError:
return fallback
# Use Colab's proxy mechanism
url = eval_js(f"google.colab.kernel.proxyPort({port})", timeout_sec = 5)
return url if url else f"http://localhost:{port}"
except Exception as e:
logger.info(f"Note: Could not get Colab URL ({e})")
return f"http://localhost:{port}"
for attempt in range(3):
try:
url = eval_js(f"google.colab.kernel.proxyPort({port})", timeout_sec = 10)
# A valid Colab proxy URL starts with https:// and embeds the port.
if (
url
and isinstance(url, str)
and url.startswith("https://")
and str(port) in url
):
return url.rstrip("/")
except Exception as e:
logger.info(f"Note: Could not get Colab URL (attempt {attempt + 1}/3: {e})")
if attempt < 2:
_time.sleep(1)
logger.warning(
f"Could not get a valid Colab proxy URL after 3 attempts — using localhost fallback. "
f"The link/iframe may not work from outside the runtime."
)
return fallback
def show_link(port: int = 8888):
"""Display a styled clickable link to the UI."""
def show_link(port: int = 8888, *, _url: "str | None" = None):
"""Display a styled clickable link to the UI.
*_url* is an optional pre-fetched Colab proxy URL. When omitted,
``get_colab_url(port)`` is called internally. Pass it from
``_show_and_embed`` to avoid a second ``eval_js`` round-trip.
"""
from IPython.display import display, HTML
# Get real Colab proxy URL
url = get_colab_url(port)
url = _url if _url is not None else get_colab_url(port)
# Build a truncated display URL. Wrap in try/except so an unexpected URL
# shape never prevents the link from rendering.
try:
port_prefix = f"{port}-"
idx = url.index(port_prefix)
next_dash = url.index("-", idx + len(port_prefix))
short_url = url[: next_dash + 1] + "..."
except (ValueError, IndexError):
short_url = url
# Also emit a plain-text line so the URL is visible even if HTML display
# is suppressed or fails.
logger.info(f"🌐 Unsloth Studio URL: {url}")
short_url = (
url[: url.index("-", url.index(f"{port}-") + len(str(port)) + 1) + 1] + "..."
if f"{port}-" in url
else url
)
html = f"""
<div style="display: inline-block; padding: 20px; background: #ffffff; border: 2px solid #000000;
border-radius: 12px; margin: 10px 0; font-family: system-ui, -apple-system, sans-serif;">
@ -59,10 +97,10 @@ def show_link(port: int = 8888):
height="48" style="display:block;">
Unsloth Studio is Ready!
</h2>
<a href="{url}" target="_blank"
<a href="{url}" onclick="var w=window.open(this.href,'_blank');if(!w){{return true;}}return false;"
style="display: inline-flex; align-items: center; gap: 10px; padding: 14px 28px;
background: #000000; color: white; text-decoration: none; border-radius: 8px;
font-weight: 800; font-size: 16px;">
font-weight: 800; font-size: 16px; cursor: pointer;">
<svg xmlns="http://www.w3.org/2000/svg" width="18" height="18" viewBox="0 0 24 24" fill="white"><polygon points="5,3 19,12 5,21"/></svg>
Open Unsloth Studio
</a>
@ -77,6 +115,75 @@ def show_link(port: int = 8888):
display(HTML(html))
def _is_studio_healthy(port: int, timeout: float = 2.0) -> bool:
"""Return True if a Studio backend is already answering health checks on *port*."""
import urllib.request
try:
with urllib.request.urlopen(
f"http://localhost:{port}/api/health", timeout = timeout
):
return True
except Exception:
return False
def _show_and_embed(port: int):
"""Embed the Studio inline for *port* with a branded header bar.
Fetches the Colab proxy URL once (registering the port with Colab's
reverse-proxy at the same time) then renders a header bar + full-height
iframe as a single HTML block.
Falls back to ``serve_kernel_port_as_iframe`` if ``IPython.display.HTML``
is unavailable for any reason.
"""
url = get_colab_url(port)
logger.info(f"🌐 Unsloth Studio URL: {url}")
try:
from IPython.display import HTML, display
iframe_id = f"unsloth-studio-{port}"
# Truncated URL shown in the header — best-effort, falls back to full URL.
try:
port_prefix = f"{port}-"
idx = url.index(port_prefix)
next_dash = url.index("-", idx + len(port_prefix))
short_url = url[: next_dash + 1] + "..."
except (ValueError, IndexError):
short_url = url
display(
HTML(f"""
<div style="font-family:system-ui,-apple-system,sans-serif;margin:8px 0;
border-radius:12px;overflow:hidden;box-shadow:0 2px 16px rgba(0,0,0,0.18);">
<div style="display:flex;align-items:center;gap:10px;padding:10px 16px;background:#000;">
<img src="https://github.com/unslothai/unsloth/raw/main/studio/frontend/public/unsloth-gem.png"
height="26" style="display:block;">
<span style="color:#fff;font-weight:700;font-size:15px;letter-spacing:-0.2px;">Unsloth Studio</span>
<span style="margin-left:auto;color:#666;font-size:11px;font-family:monospace;">{short_url}</span>
</div>
<iframe
id="{iframe_id}"
src="{url}"
style="width:100%;height:82vh;min-height:600px;max-height:1100px;border:none;display:block;box-sizing:border-box;"
allow="clipboard-read; clipboard-write"
></iframe>
</div>
""")
)
except Exception:
# Fallback: Colab's built-in helper (less control, but always works)
try:
from google.colab import output as colab_output
colab_output.serve_kernel_port_as_iframe(port, height = 900, width = "100%")
except ImportError:
pass
def start(port: int = 8888):
"""
Start Unsloth Studio server in Colab and display the URL.
@ -85,10 +192,26 @@ def start(port: int = 8888):
from colab import start
start()
"""
import sys
import time
logger.info("🦥 Starting Unsloth Studio...")
# --- Fast path: Studio is already running (cell re-run) ---
# Re-launching would either collide on the port or silently shift to a new
# port and confuse the user. Just re-show the link and iframe instead.
if _is_studio_healthy(port):
logger.info(
f" Studio is already running on port {port} — reusing existing server."
)
_show_and_embed(port)
try:
for _ in range(10000):
time.sleep(300)
print("=", end = "", flush = True)
except KeyboardInterrupt:
logger.info("\nUnsloth Studio keepalive stopped.")
return
logger.info(" Loading backend...")
from run import run_server
@ -96,18 +219,63 @@ def start(port: int = 8888):
repo_root = Path(__file__).parent.parent
frontend_path = repo_root / "frontend" / "dist"
if not frontend_path.exists():
if not (frontend_path / "index.html").exists():
logger.info("❌ Frontend not built! Please run the setup cell first.")
return
logger.info(" Starting server...")
# Start server silently
run_server(host = "0.0.0.0", port = port, frontend_path = frontend_path, silent = True)
try:
app = run_server(
host = "0.0.0.0", port = port, frontend_path = frontend_path, silent = True
)
except SystemExit as exc:
logger.error(f"❌ Unsloth Studio failed to start: {exc}")
return
except Exception as exc:
logger.error(f"❌ Unsloth Studio failed to start: {exc}")
return
logger.info(" Server started!")
# run_server auto-increments the port when the requested one is already in
# use (e.g. Jupyter occupying 8888). Read back the actual bound port so the
# Colab proxy URL and iframe always point at the right place.
actual_port: int = getattr(getattr(app, "state", None), "server_port", None) or port
# Show the clickable link with real URL
show_link(port)
logger.info(f" Server started on port {actual_port}!")
# Poll health endpoint to confirm the server is truly reachable before
# showing the link and registering the iframe — avoids the race where
# ready_event fires but the process hasn't finished binding.
import urllib.request
server_ready = False
for _ in range(40):
try:
with urllib.request.urlopen(
f"http://localhost:{actual_port}/api/health", timeout = 1
):
server_ready = True
break
except Exception:
time.sleep(0.5)
if not server_ready:
logger.error(
f"❌ Unsloth Studio did not become healthy on port {actual_port}. "
"Check for errors above."
)
return
_show_and_embed(actual_port)
# Keep kernel alive so the daemon server thread stays running.
# Handle KeyboardInterrupt cleanly so the user gets a readable message
# rather than a raw traceback when they interrupt the cell.
try:
for _ in range(10000):
time.sleep(300)
print("=", end = "", flush = True)
except KeyboardInterrupt:
logger.info("\nUnsloth Studio keepalive stopped.")
if __name__ == "__main__":

142
studio/backend/core/_torchao_stub.py Normal file
View file

@ -0,0 +1,142 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Shared torchao Windows-ROCm import stub.
torchao (pulled in by transformers.quantizers) imports
torch.distributed._functional_collectives at module level, which imports
distributed_c10d.py unconditionally — that file crashes on Windows ROCm because
torch._C._distributed_c10d (the RCCL backend) is absent.
torch/distributed/__init__.py itself is guarded by `if is_available()` so
`import torch.distributed` alone is safe; the crash only comes via torchao's
import chain. Stubbing torchao short-circuits it entirely.
_StubSubpackageFinder handles any depth of torchao.xxx.yyy imports.
This logic used to be duplicated inline inside run_export_process() and
run_training_process(); it now lives here so both worker subprocesses call the
single `install_torchao_windows_rocm_stub()` entrypoint before importing
transformers / unsloth_zoo.
"""
from __future__ import annotations
import sys
import types
import importlib.abc
import importlib.machinery
_STUB_SENTINEL = object()
# Metaclass for stub types so that isinstance(x, StubClass) returns False
# instead of raising TypeError ("arg 2 must be a type").
# peft/tuners/lora/torchao.py does:
# from torchao.dtypes import AffineQuantizedTensor, LinearActivationQuantizedTensor
# isinstance(weight, (AffineQuantizedTensor, LinearActivationQuantizedTensor))
# If those names resolve to stub modules rather than types, isinstance() raises.
class _StubTypeMeta(type):
def __instancecheck__(cls, instance):
return False
def __subclasscheck__(cls, subclass):
return False
def __getattr__(cls, attr):
if attr.startswith("__"):
raise AttributeError(attr)
child = _StubTypeMeta(attr, (), {})
setattr(cls, attr, child)
return child
def __call__(cls, *args, **kwargs):
return None
def _make_stub_type(name):
"""Stub class: accepted by isinstance() (always False), supports attr access."""
return _StubTypeMeta(name, (), {})
def _make_mod_stub(mod_name):
m = types.ModuleType(mod_name)
m.__path__ = []
m.__package__ = mod_name
m._unsloth_stub = _STUB_SENTINEL
m.__spec__ = importlib.machinery.ModuleSpec(mod_name, loader = None, is_package = True)
def _ga(attr, _m = m, _n = mod_name):
if attr.startswith("__"):
raise AttributeError(attr)
# Return a stub CLASS (not a module) so that isinstance(x, attr)
# works and returns False instead of raising TypeError.
child = _make_stub_type(f"{_n}.{attr}")
setattr(_m, attr, child)
return child
m.__getattr__ = _ga
return m
class _StubSubpackageLoader(importlib.abc.Loader):
def __init__(self, mod_name):
self._mod_name = mod_name
def create_module(self, spec):
return _make_mod_stub(self._mod_name)
def exec_module(self, module):
pass
class _StubSubpackageFinder(importlib.abc.MetaPathFinder):
def find_spec(self, fullname, path, target = None):
if "." not in fullname:
return None
parent = sys.modules.get(fullname.rsplit(".", 1)[0])
if parent is None:
return None
if getattr(parent, "_unsloth_stub", None) is not _STUB_SENTINEL:
return None
return importlib.machinery.ModuleSpec(
fullname, _StubSubpackageLoader(fullname), is_package = True
)
def install_torchao_windows_rocm_stub() -> None:
"""Pre-stub torchao on Windows ROCm so transformers/peft imports don't crash.
No-op on every other platform (Windows CUDA included — there torchao is real
and shadowing it would break torchao-based quantization paths). Must run
before any import of transformers / unsloth_zoo. Safe to call once per worker
process.
"""
# Gate on the active torch runtime, not env-var presence -- HIP_PATH /
# ROCM_PATH stay set after a user installs the HIP SDK and reverts to a
# CUDA torch wheel. AMD SDK / Radeon ROCm wheels may not set torch.version.hip
# but still encode "rocm" in torch.__version__, so accept either.
_is_win32_rocm = False
if sys.platform == "win32":
try:
import torch as _torch_probe
_is_win32_rocm = bool(
getattr(getattr(_torch_probe, "version", None), "hip", None)
or "rocm" in getattr(_torch_probe, "__version__", "").lower()
)
del _torch_probe
except Exception:
pass
if _is_win32_rocm:
# Register the finder only on Windows ROCm -- on other platforms there
# are no stub modules seeded, so appending is a pure accumulation.
sys.meta_path.append(_StubSubpackageFinder())
# Seed torchao top-level + key submodules; the finder handles the rest.
for _tao_name in (
"torchao",
"torchao.quantization",
"torchao.dtypes",
"torchao.float8",
"torchao.utils",
):
if _tao_name not in sys.modules:
sys.modules[_tao_name] = _make_mod_stub(_tao_name)

9
studio/backend/core/data_recipe/service.py
View file

@ -176,12 +176,21 @@ def build_mcp_providers(
) -> list:
from data_designer.config.mcp import LocalStdioMCPProvider, MCPProvider # pyright: ignore[reportMissingImports]
# Same gate as the chat MCP path: stdio providers spawn a local subprocess,
# so only build them when this host allows it (desktop / explicit opt-in).
# Skip them otherwise so a recipe carried onto a hosted host cannot spawn.
from core.inference.mcp_client import stdio_mcp_enabled
stdio_allowed = stdio_mcp_enabled()
providers: list[MCPProvider | LocalStdioMCPProvider] = []
for provider in recipe.get("mcp_providers", []):
if not isinstance(provider, dict):
continue
provider_type = provider.get("provider_type")
if provider_type == "stdio":
if not stdio_allowed:
continue
env = provider.get("env")
if not isinstance(env, dict):
env = {}

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

@ -439,6 +439,15 @@ def run_export_process(
'Install for better performance: pip install "triton-windows<3.7"'
)
# ── 1c. Stub torchao on Windows ROCm ──
# Shared with the training worker; see core/_torchao_stub.py for the full
# rationale (torchao -> torch.distributed._functional_collectives crashes on
# Windows ROCm because the RCCL backend is absent). No-op off Windows ROCm.
# Must run before any import of transformers / unsloth_zoo.
from core._torchao_stub import install_torchao_windows_rocm_stub
install_torchao_windows_rocm_stub()
# ── 2. Import ML libraries (fresh in this clean process) ──
try:
_send_response(

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

File diff suppressed because it is too large Load diff

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

@ -17,6 +17,7 @@ import struct
import structlog
from loggers import get_logger
import shutil
import signal
import socket
import subprocess
import sys
@ -28,6 +29,12 @@ from urllib.parse import urlparse
import httpx
from core.inference.llama_server_args import (
parse_cache_override,
parse_ctx_override,
resolve_cache_type_kv,
resolve_requested_ctx,
)
from core.tool_healing import (
_TC_END_TAG_RE,
_TC_FUNC_CLOSE_RE,
@ -959,9 +966,6 @@ class LlamaCppBackend:
7. llama-server on PATH (system install)
8. ./bin/llama-server (legacy: extracted binary)
"""
import os
import sys
binary_name = "llama-server.exe" if sys.platform == "win32" else "llama-server"
# 1. Env var — direct path to binary
@ -1232,6 +1236,33 @@ class LlamaCppBackend:
return total
@staticmethod
def _amd_apu_wants_unified_memory() -> bool:
"""True only for AMD unified-memory APUs (gfx1150/gfx1151), where
GGML_CUDA_ENABLE_UNIFIED_MEMORY lets llama.cpp use shared system RAM.
False for discrete AMD, NVIDIA, CPU and macOS (the env hurts discrete
GPUs). ROCm reuses torch.cuda.*; the gcnArchName suffix is stripped."""
try:
import torch
if getattr(torch.version, "hip", None) is None:
return False
if not (hasattr(torch, "cuda") and torch.cuda.is_available()):
return False
for _i in range(torch.cuda.device_count()):
try:
_arch = (
getattr(torch.cuda.get_device_properties(_i), "gcnArchName", "")
or ""
)
except Exception:
continue
if _arch.split(":")[0].strip().lower() in {"gfx1150", "gfx1151"}:
return True
except Exception:
return False
return False
@staticmethod
def _get_gpu_free_memory() -> list[tuple[int, int]]:
"""Query free memory per GPU.
@ -1249,8 +1280,6 @@ class LlamaCppBackend:
Returns list of (gpu_index, free_mib) sorted by index. Empty
list if no supported GPU is reachable.
"""
import os
# ── NVIDIA via nvidia-smi ────────────────────────────────────
try:
result = subprocess.run(
@ -2562,6 +2591,105 @@ class LlamaCppBackend:
# ── Lifecycle ─────────────────────────────────────────────────
# GGUF ``general.architecture`` values for diffusion / image models.
# llama.cpp proper has no such architectures, so loading one as a chat
# model dies with "unknown model architecture: '<arch>'". These match
# the patched stable-diffusion.cpp / ComfyUI-GGUF enums (LLM_ARCH_FLUX,
# LLM_ARCH_QWEN_IMAGE, ...). Unsloth publishes FLUX and Qwen-Image GGUFs
# under https://huggingface.co/collections/unsloth/unsloth-diffusion-ggufs.
# Matched exactly (not as a substring) so a chat arch merely containing a
# short token like "wan"/"sd1" (e.g. "taiwan") is not misrouted to Images.
_DIFFUSION_ARCHES = frozenset(
(
"qwen_image",
"flux",
"sd1",
"sdxl",
"sd3",
"aura",
"hidream",
"cosmos",
"ltxv",
"hyvid",
"wan",
"lumina2",
)
)
@staticmethod
def _classify_llama_start_failure(
output: str,
gguf_path: Optional[str],
model_identifier: Optional[str],
) -> str:
"""Explain *why* llama-server failed to start, from its output.
Several distinct failures all otherwise collapse into the same
opaque "invalid GGUF or out of memory" message. The worst case is
a diffusion / image GGUF (FLUX, Qwen-Image, ...) loaded as a chat
model: the file is perfectly valid and there is plenty of memory,
but llama.cpp has no such architecture, so the user is told to free
memory that was never the problem (issue #5842). Pick the most
specific message the captured output supports.
"""
lowered = (output or "").lower()
# Detect Ollama source up front so the arch branch can keep the
# Ollama hint instead of the generic "unsupported arch" message.
gguf = gguf_path or ""
is_ollama = (
".studio_links" in gguf
or os.sep + "ollama_links" + os.sep in gguf
or os.sep + ".cache" + os.sep + "ollama" + os.sep in gguf
or (model_identifier or "").startswith("ollama/")
)
# "unknown model architecture: '<arch>'": diffusion -> Images page,
# Ollama -> Ollama hint, else a precise "unsupported" message. Exact
# match so chat archs are never misrouted.
arch_match = re.search(r"unknown model architecture:\s*'([^']+)'", lowered)
if arch_match:
arch = arch_match.group(1)
if arch in LlamaCppBackend._DIFFUSION_ARCHES:
return (
f"'{arch}' is a diffusion (image-generation) GGUF, which "
"llama-server cannot run as a chat/completion model. Use "
"Studio's Images page to generate with local diffusion "
"GGUFs such as FLUX and Qwen-Image."
)
if is_ollama:
return (
"Some Ollama models do not work with llama.cpp. Try a "
"different model, or use this model directly through "
"Ollama instead."
)
return (
f"llama.cpp does not support this GGUF's model architecture "
f"('{arch}'). The file is valid, but this model type cannot "
"be run with llama-server."
)
# Other Ollama compat failures that do not name an arch. Only when
# the output shows a GGUF compat issue, not OOM / missing binaries.
if is_ollama:
gguf_compat_hints = (
"key not found",
"unknown model architecture",
"failed to load model",
)
if any(h in lowered for h in gguf_compat_hints):
return (
"Some Ollama models do not work with llama.cpp. Try a "
"different model, or use this model directly through "
"Ollama instead."
)
# Fallback: genuinely unknown failure (OOM, missing binary, ...).
return (
"llama-server failed to start. "
"Check that the GGUF file is valid and you have enough memory."
)
def load_model(
self,
*,
@ -2724,7 +2852,23 @@ class LlamaCppBackend:
# Select GPU(s) based on model size + estimated KV cache.
# Seed safe defaults before GPU probing so the except path
# still has valid state to publish.
effective_ctx = n_ctx if n_ctx > 0 else (self._context_length or 0)
ctx_override = parse_ctx_override(extra_args)
requested_ctx = resolve_requested_ctx(extra_args, n_ctx)
cache_override = parse_cache_override(extra_args)
cache_type_kv = resolve_cache_type_kv(extra_args, cache_type_kv)
if ctx_override is not None and ctx_override > 0:
logger.info(
f"User --ctx-size {ctx_override} honored; "
"skipping auto-reduce"
)
if cache_override is not None:
logger.info(
f"User --cache-type-k/-v {cache_override} "
"honored for KV estimate"
)
effective_ctx = (
requested_ctx if requested_ctx > 0 else (self._context_length or 0)
)
max_available_ctx = self._context_length or effective_ctx
gpus: list[tuple[int, int]] = []
try:
@ -2734,8 +2878,8 @@ class LlamaCppBackend:
# Resolve effective context: 0 means let llama-server use the
# model's native length. Only expand to a known native length
# if metadata is available; otherwise preserve 0 as a sentinel.
if n_ctx > 0:
effective_ctx = n_ctx
if requested_ctx > 0:
effective_ctx = requested_ctx
elif self._context_length is not None:
effective_ctx = self._context_length
else:
@ -2788,7 +2932,7 @@ class LlamaCppBackend:
# since multi-GPU is slower and the user didn't ask for a
# specific context length.
gpu_indices, use_fit = None, True
explicit_ctx = n_ctx > 0
explicit_ctx = requested_ctx > 0
if gpus and self._can_estimate_kv() and effective_ctx > 0:
# Compute the largest hardware-aware cap from the model's
@ -2845,7 +2989,7 @@ class LlamaCppBackend:
gpu_indices, use_fit = self._select_gpus(
requested_total, gpus
)
# No silent shrink: effective_ctx stays == n_ctx.
# No silent shrink: effective_ctx stays == requested_ctx.
else:
# Auto context: prefer fewer GPUs, cap context
# to fit. Same headroom threshold as
@ -2934,7 +3078,7 @@ class LlamaCppBackend:
except Exception as e:
logger.warning(f"GPU selection failed ({e}), using --fit on")
gpu_indices, use_fit = None, True
effective_ctx = n_ctx # fall back to original
effective_ctx = requested_ctx # fall back to original
launch_mmproj_path = self._resolve_launch_mmproj_path(
model_path = model_path,
@ -3136,6 +3280,14 @@ class LlamaCppBackend:
env = child_env_without_native_path_secret()
binary_dir = str(Path(binary).parent)
# AMD unified-memory APUs (gfx1150/gfx1151): let llama.cpp use
# shared system RAM. setdefault so a user value wins.
if self._amd_apu_wants_unified_memory():
env.setdefault("GGML_CUDA_ENABLE_UNIFIED_MEMORY", "1")
logger.info(
"AMD unified-memory APU: set GGML_CUDA_ENABLE_UNIFIED_MEMORY=1"
)
if sys.platform == "win32":
# See _build_windows_path_dirs for ordering. #5106.
path_dirs = self._build_windows_path_dirs(
@ -3145,6 +3297,24 @@ class LlamaCppBackend:
)
existing_path = env.get("PATH", "")
env["PATH"] = ";".join(path_dirs) + ";" + existing_path
# ROCm: the llama.cpp prebuilt bundles its own rocblas.dll
# but NOT the Tensile kernel library files it needs
# (rocblas/library/TensileLibrary*.dat + *.hsaco). The
# bundled DLL searches relative to its own location by
# default (i.e. <binary_dir>/rocblas/library/) which does
# not exist, causing a silent crash on the first GEMM.
# ROCBLAS_TENSILE_LIBPATH overrides that search to point at
# the ROCm installation where the kernel files actually are.
_hip_path = os.environ.get(
"HIP_PATH", os.environ.get("ROCM_PATH", "")
)
if _hip_path:
_rocblas_lib = os.path.join(
_hip_path, "bin", "rocblas", "library"
)
if os.path.isdir(_rocblas_lib):
env.setdefault("ROCBLAS_TENSILE_LIBPATH", _rocblas_lib)
else:
# Linux: set LD_LIBRARY_PATH for shared libs next to the binary
# and CUDA runtime libs (libcudart, libcublas, etc.)
@ -3312,31 +3482,12 @@ class LlamaCppBackend:
# Wait for llama-server to become healthy
if not self._wait_for_health(timeout = 600.0):
self._kill_process()
_gguf = gguf_path or ""
_is_ollama = (
".studio_links" in _gguf
or os.sep + "ollama_links" + os.sep in _gguf
or os.sep + ".cache" + os.sep + "ollama" + os.sep in _gguf
or (self._model_identifier or "").startswith("ollama/")
)
# Only show the Ollama-specific message when the server
# output indicates a GGUF compatibility issue, not for
# unrelated failures like OOM or missing binaries.
if _is_ollama:
_output = "\n".join(self._stdout_lines[-50:]).lower()
_gguf_compat_hints = (
"key not found",
"unknown model architecture",
"failed to load model",
)
if any(h in _output for h in _gguf_compat_hints):
raise RuntimeError(
"Some Ollama models do not work with llama.cpp. "
"Try a different model, or use this model directly through Ollama instead."
)
raise RuntimeError(
"llama-server failed to start. "
"Check that the GGUF file is valid and you have enough memory."
self._classify_llama_start_failure(
"\n".join(self._stdout_lines[-50:]),
gguf_path,
self._model_identifier,
)
)
self._healthy = True
@ -3853,10 +4004,6 @@ class LlamaCppBackend:
Falls back to pgrep + /proc/<pid>/exe on Linux when psutil is
not installed.
"""
import os
import signal
import sys
try:
# -- Build the ownership allowlist --------------------------------
# Two kinds of matches:
@ -5076,13 +5223,30 @@ class LlamaCppBackend:
_effective_timeout = (
None if tool_call_timeout >= 9999 else tool_call_timeout
)
result = execute_tool(
tool_name,
arguments,
cancel_event = cancel_event,
timeout = _effective_timeout,
session_id = session_id,
)
# Guard against the model emitting a tool not in the
# per-request advertised set: filtered MCP names, a
# built-in the caller opted out of, or a stale name
# from a prior turn. Mirrors the safetensors loop's
# allowed_tool_names check.
_allowed = {
(t.get("function") or {}).get("name")
for t in (tools or [])
if (t.get("function") or {}).get("name")
}
if _allowed and tool_name not in _allowed:
result = (
f"Error: tool '{tool_name}' is not enabled "
"for this request. Use one of the enabled "
"tools or provide a final answer."
)
else:
result = execute_tool(
tool_name,
arguments,
cancel_event = cancel_event,
timeout = _effective_timeout,
session_id = session_id,
)
yield {
"type": "tool_end",

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

@ -1,46 +1,29 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Validator for user-supplied llama-server pass-through args.
"""Boundary validator for user-supplied llama-server pass-through args.
Studio runs llama-server as a managed subprocess and lets callers pass
extra flags directly (CLI: ``unsloth run ... --top-k 20``; HTTP:
``LoadRequest.llama_extra_args``). This module is the boundary that
rejects only flags Studio fundamentally cannot share with the user --
model identity, the auth key, and the network endpoint Studio's HTTP
proxy targets. Anything else passes through.
Reject only flags Studio manages (model identity, auth, network,
parallel slots). Everything else (sampling, ``-c``, ``-ngl``,
``--flash-attn``, ``--cache-type-*``, ``--spec-*``, ``--jinja``, ...)
is appended after Studio's auto-set flags so llama.cpp's last-wins
parser lets the user override.
User-supplied args are appended to ``cmd`` after Studio's auto-set
flags, so llama.cpp's last-wins CLI parsing makes the user's value
override the auto-set one. That covers tunable knobs the user might
reasonably want to override -- ``-c``/``--ctx-size``,
``-np``/``--parallel``, ``-fa``/``--flash-attn``,
``-ngl``/``--gpu-layers``, ``-t``/``--threads``, ``-fit``/``--fit*``,
``--cache-type-k/v``, ``--chat-template-file/-kwargs``,
``--spec-*``, ``--jinja``/``--no-jinja``,
``--no-context-shift``/``--context-shift``, sampling params, etc.
Reference: https://github.com/ggml-org/llama.cpp/blob/master/tools/server/README.md
Ref: https://github.com/ggml-org/llama.cpp/blob/master/tools/server/README.md
"""
from __future__ import annotations
from typing import Iterable, Optional
# Each group is the full set of aliases (short + long) for one
# hard-denied flag, taken from the llama-server README. If llama.cpp
# adds a new alias for an existing denied flag, extend the relevant
# group.
#
# Flags NOT in this list (e.g. -c, --parallel, --flash-attn, -ngl,
# -t/--threads, --jinja, --no-context-shift, --fit*, --cache-type-*,
# --chat-template-*, --spec-*) pass through and override Studio's
# auto-set version via llama.cpp's last-wins CLI parsing.
# Each group = every alias (short + long) of one hard-denied flag.
# Extend the matching group when llama.cpp adds a new alias.
_DENYLIST_GROUPS: tuple[frozenset[str], ...] = (
# Model identity -- Studio resolves the model from LoadRequest and
# passes -m / mmproj after downloading from HF if needed. A second
# -m would point at a different model than the one Studio thinks
# is loaded.
# Parallel slots: owned by typer --parallel; a pass-through would
# desync app.state.llama_parallel_slots from llama-server.
frozenset({"-np", "--parallel", "--n-parallel"}),
# Model identity: Studio resolves it from LoadRequest; a second
# -m would load a different model than Studio thinks it loaded.
frozenset({"-m", "--model"}),
frozenset({"-mu", "--model-url"}),
frozenset({"-dr", "--docker-repo"}),
@ -51,28 +34,21 @@ _DENYLIST_GROUPS: tuple[frozenset[str], ...] = (
frozenset({"-hft", "--hf-token"}),
frozenset({"-mm", "--mmproj"}),
frozenset({"-mmu", "--mmproj-url"}),
# Networking -- Studio binds llama-server's port and reverse-proxies
# HTTP traffic to it. Retargeting host/port/path/prefix would
# orphan Studio's proxy and the UI would lose the server.
# Networking: Studio binds + proxies; retargeting orphans the proxy.
frozenset({"--host"}),
frozenset({"--port"}),
frozenset({"--path"}),
frozenset({"--api-prefix"}),
frozenset({"--reuse-port"}),
# Auth / TLS -- Studio terminates auth at its own layer; an
# upstream --api-key would shadow Studio's UNSLOTH_DIRECT_STREAM
# key, and TLS on llama-server would break the local proxy hop.
# Auth / TLS: Studio terminates auth; upstream --api-key / TLS
# shadows Studio's key and breaks the proxy hop.
frozenset({"--api-key"}),
frozenset({"--api-key-file"}),
frozenset({"--ssl-key-file"}),
frozenset({"--ssl-cert-file"}),
# Single-model server -- Studio runs one model per llama-server
# process and serves its own UI. Enabling multi-model loading or
# llama-server's built-in web UI changes the surface clients see.
# ``--webui``/``--no-webui`` are the legacy spelling; current
# upstream uses ``--ui``/``--no-ui`` + ``--ui-*`` companions.
# Keep both so the denylist matches old and new llama-server
# binaries (Studio's prebuilt vs system-llama.cpp).
# Built-in web UI. --webui/--no-webui is the legacy spelling;
# upstream renamed to --ui/--no-ui + --ui-*. Keep both so prebuilt
# and system llama.cpp binaries both match.
frozenset({"--webui", "--no-webui"}),
frozenset({"--ui", "--no-ui"}),
frozenset({"--ui-config"}),
@ -82,32 +58,46 @@ _DENYLIST_GROUPS: tuple[frozenset[str], ...] = (
frozenset({"--models-preset"}),
frozenset({"--models-max"}),
frozenset({"--models-autoload", "--no-models-autoload"}),
# Server-mode flips: --embedding / --rerank restrict llama-server to
# those endpoints, breaking Studio's /v1/chat/completions hop.
frozenset({"--embedding", "--embeddings"}),
frozenset({"--rerank", "--reranking"}),
# llama-server's own built-in tools flag would silently stack on top
# of Studio's --enable-tools / --disable-tools policy resolver.
frozenset({"--tools"}),
)
_DENYLIST: frozenset[str] = frozenset().union(*_DENYLIST_GROUPS)
def _flag_name(token: str) -> Optional[str]:
"""Return the flag name for a token, or None if it isn't a flag.
"""Flag name for ``token``, or None if it isn't a flag.
Peels ``--key=value`` to the bare ``--key``. Plain numeric values
like ``-1`` or ``-0.5`` (e.g. ``--seed -1``) are values, not flags;
llama-server short-form flags always start with a letter.
Peels `--key=value` to `--key`, treats `-1` / `-0.5` as values
(llama-server shorts always start with a letter), strips
whitespace, and normalises attached `-np8` / signed `-np-1` /
digit-prefix-junk `-np8x` to `-np`. Mirrors the CLI's
`_expand_attached_np_short`.
"""
token = token.strip()
if not token.startswith("-") or token in {"-", "--"}:
return None
if len(token) >= 2 and (token[1].isdigit() or token[1] == "."):
return None
return token.split("=", 1)[0]
name = token.split("=", 1)[0]
if len(name) > 3 and name.startswith("-np"):
suffix = name[3:]
if suffix[0].isdigit() or (
len(suffix) > 1 and suffix[0] in {"-", "+"} and suffix[1].isdigit()
):
return "-np"
return name
def validate_extra_args(args: Optional[Iterable[str]]) -> list[str]:
"""Validate user-supplied llama-server args.
Returns the args as a flat list ready to extend the llama-server
command. Raises ``ValueError`` (with the offending flag in the
message) the moment a token resolves to a Studio-managed flag.
"""
"""Validate user-supplied llama-server args. Returns a flat list
ready to extend the llama-server command; raises ``ValueError``
naming the offending flag on the first managed token."""
if not args:
return []
out: list[str] = []
@ -120,19 +110,21 @@ def validate_extra_args(args: Optional[Iterable[str]]) -> list[str]:
f"and cannot be passed as an extra arg"
)
out.append(token)
parse_ctx_override(out)
parse_cache_override(out)
return out
def is_managed_flag(flag: str) -> bool:
"""True if ``flag`` is a Studio-managed llama-server flag."""
return flag in _DENYLIST
"""True if ``flag`` is Studio-managed. Normalises via ``_flag_name``
so `-np8` / `--parallel=8` classify like the canonical tokens."""
normalised = _flag_name(flag)
return normalised is not None and normalised in _DENYLIST
# Pass-through flags that shadow first-class ``LoadRequest`` fields
# (max_seq_length, cache_type_kv, speculative_type,
# chat_template_override). Stripped from inherited extras so they
# can't last-wins-override an Apply that re-sets the same first-class
# field.
# Pass-through flags that shadow first-class LoadRequest fields;
# stripped from inherited extras so they can't last-wins-override an
# Apply that re-sets the same field.
_CONTEXT_FLAGS: frozenset[str] = frozenset({"-c", "--ctx-size"})
_CACHE_FLAGS: frozenset[str] = frozenset(
{"-ctk", "--cache-type-k", "-ctv", "--cache-type-v"}
@ -169,14 +161,124 @@ _SHADOWING_FLAGS: frozenset[str] = (
_CONTEXT_FLAGS | _CACHE_FLAGS | _SPEC_FLAGS | _TEMPLATE_FLAGS
)
# Boolean flags inside _SHADOWING_FLAGS that take no value. The
# value-consuming heuristic in strip_shadowing_flags must skip just the
# flag for these, never the following token.
# Shadowing flags that take no value -- strip the flag only, never the
# following token.
_BOOLEAN_SHADOWING_FLAGS: frozenset[str] = frozenset(
{"--spec-default", "--jinja", "--no-jinja"}
)
def parse_ctx_override(args: Optional[Iterable[str]]) -> Optional[int]:
"""Return the last user-supplied ``-c`` / ``--ctx-size`` value.
Mirrors llama.cpp's last-wins flag parsing for the one pass-through
numeric knob Studio's load-time fit logic needs to see.
"""
if not args:
return None
tokens = [str(a) for a in args]
override: Optional[int] = None
i, n = 0, len(tokens)
while i < n:
tok = tokens[i]
flag = _flag_name(tok)
if flag is None or flag not in _CONTEXT_FLAGS:
i += 1
continue
if "=" in tok:
raw_value = tok.split("=", 1)[1]
i += 1
else:
if i + 1 >= n or _flag_name(tokens[i + 1]) is not None:
raise ValueError(
f"llama-server flag '{flag}' requires an integer value"
)
raw_value = tokens[i + 1]
i += 2
try:
value = int(str(raw_value).strip())
except ValueError as exc:
raise ValueError(
f"llama-server flag '{flag}' requires an integer value"
) from exc
if value < 0:
raise ValueError(
f"llama-server flag '{flag}' requires a non-negative integer value"
)
override = value
return override
def resolve_requested_ctx(
args: Optional[Iterable[str]],
fallback_n_ctx: int,
) -> int:
"""Return the context size load_model should treat as requested.
Single source of truth for the two-line ``ctx_override = parse_ctx_override(...);
requested_ctx = ctx_override if ctx_override is not None else n_ctx`` pattern
used by ``load_model`` so tests don't have to reimplement the conditional
locally and then assert against their own reimplementation.
"""
override = parse_ctx_override(args)
return override if override is not None else fallback_n_ctx
def parse_cache_override(args: Optional[Iterable[str]]) -> Optional[str]:
"""Return the last-wins cache type if extras pass cache flags.
Mirrors parse_ctx_override but for cache type. Recognises both -ctk
(key) and -ctv (value). When both flags appear, returns the last-wins
value, treating key and value cache flags as the same setting because
Studio's KV estimate has a single cache_type_kv knob.
"""
if not args:
return None
tokens = [str(a) for a in args]
override: Optional[str] = None
i, n = 0, len(tokens)
while i < n:
tok = tokens[i]
flag = _flag_name(tok)
if flag is None or flag not in _CACHE_FLAGS:
i += 1
continue
if "=" in tok:
raw_value = tok.split("=", 1)[1]
i += 1
else:
if i + 1 >= n or _flag_name(tokens[i + 1]) is not None:
raise ValueError(f"llama-server flag '{flag}' requires a value")
raw_value = tokens[i + 1]
i += 2
value = str(raw_value).strip()
if not value:
raise ValueError(f"llama-server flag '{flag}' requires a non-empty value")
override = value
return override
def resolve_cache_type_kv(
args: Optional[Iterable[str]],
fallback_cache_type_kv: Optional[str],
) -> Optional[str]:
"""Return the cache type load_model should treat as requested.
Single source of truth for the cache override conditional used by
``load_model``.
"""
override = parse_cache_override(args)
return override if override is not None else fallback_cache_type_kv
def strip_shadowing_flags(
args: Iterable[str],
*,
@ -187,14 +289,11 @@ def strip_shadowing_flags(
) -> list[str]:
"""Strip flags that shadow first-class Studio settings.
Used when the route inherits a previous load's ``llama_extra_args``
so that an inherited ``-c 4096`` cannot override the current
request's ``max_seq_length`` (and equivalents for cache /
speculative / chat template). Each ``strip_*`` flag controls one
group; the route only strips groups whose corresponding first-class
field was actually supplied by the caller, so an inherited
``--chat-template-file`` survives an Apply that omits both
``llama_extra_args`` and ``chat_template_override``.
Used when inheriting a previous load's ``llama_extra_args`` so an
inherited `-c 4096` can't override the current `max_seq_length`
(same for cache / spec / template). Each ``strip_*`` toggle
controls one group; the route only strips groups whose first-class
field the caller actually supplied.
"""
shadowing: set[str] = set()
if strip_context:
@ -216,9 +315,8 @@ def strip_shadowing_flags(
out.append(tok)
i += 1
continue
# Drop this token. Boolean shadowing flags never carry a value;
# other shadowing flags consume the next token when it isn't a
# flag and the value isn't already packed as ``--key=value``.
# Drop the flag; consume the next token too unless it's
# boolean, already inline (`-c=4096`), or another flag.
if flag in _BOOLEAN_SHADOWING_FLAGS or "=" in tok:
i += 1
elif i + 1 < n and _flag_name(tokens[i + 1]) is None:

254
studio/backend/core/inference/mcp_client.py Normal file
View file

@ -0,0 +1,254 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
from __future__ import annotations
import asyncio
import json
import os
import shlex
import sys
from typing import Any, Optional
from loggers import get_logger
logger = get_logger(__name__)
MCP_TOOL_PREFIX = "mcp__"
_oauth_token_store = None
def is_stdio(address: str) -> bool:
"""A non-HTTP address is a local stdio command, e.g.
'npx -y @modelcontextprotocol/server-filesystem /path'."""
return not address.strip().lower().startswith(("http://", "https://"))
def parse_stdio_command(address: str) -> list[str]:
"""Split a stdio command line into argv. Shared by route validation and the
transport so both agree on quoting (notably Windows backslash paths)."""
posix = sys.platform != "win32"
parts = shlex.split(address, posix = posix)
if not posix:
# posix=False keeps backslash paths intact but also keeps the surrounding
# quotes on a token. Strip a matched pair so the argv reaches the
# subprocess clean ('"C:\\Program Files\\node"' -> C:\\Program Files\\node).
parts = [
p[1:-1] if len(p) >= 2 and p[0] == p[-1] and p[0] in "\"'" else p
for p in parts
]
return parts
def stdio_mcp_enabled() -> bool:
"""stdio MCP servers spawn local processes as the backend user (and bypass
the python/terminal sandbox), so they are only allowed when the backend
host is the user's own machine. The Tauri desktop app sets
UNSLOTH_STUDIO_ALLOW_STDIO_MCP=1 (see main.py); advanced localhost /
self-hosted users can opt in with the same variable. It stays off for
Colab and any network (0.0.0.0) bind."""
return os.environ.get("UNSLOTH_STUDIO_ALLOW_STDIO_MCP") == "1"
# Probe timeouts for discovering a server's tool list. OAuth needs minutes for
# first-connect/expired-token browser sign-in; stdio allows for first-run
# package download (e.g. `npx -y ...`); HTTP fails fast.
_HTTP_PROBE_TIMEOUT = 8.0
_OAUTH_PROBE_TIMEOUT = 305.0
_STDIO_PROBE_TIMEOUT = 60.0
def probe_timeout(address: str, use_oauth: bool) -> float:
if use_oauth:
return _OAUTH_PROBE_TIMEOUT
return _STDIO_PROBE_TIMEOUT if is_stdio(address) else _HTTP_PROBE_TIMEOUT
def parse_server_headers(server: dict) -> Optional[dict]:
"""Parsed headers_json. For stdio servers this dict is the process
environment instead of HTTP headers (see _client)."""
raw = server.get("headers_json")
if not raw:
return None
try:
parsed = json.loads(raw)
except (json.JSONDecodeError, ValueError):
return None
return parsed if isinstance(parsed, dict) else None
def _oauth_store():
global _oauth_token_store
if _oauth_token_store is None:
from key_value.aio._utils.sanitization import AlwaysHashStrategy
from key_value.aio.stores.filetree import FileTreeStore
from utils.paths.storage_roots import ensure_dir, studio_root
# Hash keys/collections — fastmcp uses raw URLs like https://x.com as
# keys and FileTreeStore would treat the "://" as nested directories.
_oauth_token_store = FileTreeStore(
data_directory = ensure_dir(studio_root() / "mcp-oauth-tokens"),
key_sanitization_strategy = AlwaysHashStrategy(),
collection_sanitization_strategy = AlwaysHashStrategy(),
)
return _oauth_token_store
async def clear_oauth_tokens_async(url: str) -> None:
"""Drop any persisted OAuth tokens for ``url``. fastmcp keys tokens by
MCP URL, so on server delete / URL change / OAuth disable we have to
clear the old credentials explicitly. Otherwise re-registering the
same URL would silently reuse the old account's token. The entire
body runs inside the protected block -- store / OAuth construction
failing must not make the delete / update route 500."""
try:
from fastmcp.client.auth import OAuth
auth = OAuth(mcp_url = url, token_storage = _oauth_store())
await auth.token_storage_adapter.clear()
except Exception as exc: # noqa: BLE001
# Cleanup is best-effort; the row delete still wins.
logger.warning("Failed to clear OAuth tokens for %s: %s", url, exc)
def _client(url: str, headers: Optional[dict], use_oauth: bool = False):
from fastmcp import Client
if is_stdio(url):
# Belt-and-suspenders: never spawn unless stdio is enabled on this host.
if not stdio_mcp_enabled():
raise PermissionError("stdio MCP servers are disabled on this host")
from fastmcp.client.transports import StdioTransport
parts = parse_stdio_command(url)
if not parts:
raise ValueError(f"Empty stdio command: {url!r}")
# env vars ride the headers field (merged over the SDK's safe default env).
# keep_alive=False tears the subprocess down on exit, so a one-shot
# probe/tool call never leaves an orphan process.
return Client(
StdioTransport(
command = parts[0],
args = parts[1:],
env = headers or None,
keep_alive = False,
)
)
from fastmcp.client.transports import SSETransport, StreamableHttpTransport
from fastmcp.mcp_config import infer_transport_type_from_url
auth = None
if use_oauth:
from fastmcp.client.auth import OAuth
auth = OAuth(mcp_url = url, token_storage = _oauth_store())
transport_cls = (
SSETransport
if infer_transport_type_from_url(url) == "sse"
else StreamableHttpTransport
)
return Client(transport_cls(url = url, headers = headers or None, auth = auth))
async def list_tools_async(
url: str,
headers: Optional[dict] = None,
timeout: float = 5.0,
use_oauth: bool = False,
) -> list[dict]:
async def _fetch() -> list[dict]:
async with _client(url, headers, use_oauth) as client:
tools = await client.list_tools()
return [t.model_dump(exclude_none = True) for t in tools]
return await asyncio.wait_for(_fetch(), timeout = timeout)
def _flatten_result(result: Any) -> str:
parts = []
for block in getattr(result, "content", None) or []:
text = getattr(block, "text", None)
if text:
parts.append(str(text))
body = "\n".join(parts)
if not body:
structured = getattr(result, "structured_content", None)
body = str(structured) if structured is not None else ""
if getattr(result, "is_error", False):
# "Error: " prefix triggers tool_call_parser's TOOL_ERROR_PREFIXES nudge.
return f"Error: {body}" if body else "Error: tool returned no content"
return body
def call_tool_sync(
url: str,
headers: Optional[dict],
name: str,
args: dict,
timeout: Optional[float] = 300.0,
use_oauth: bool = False,
cancel_event = None,
) -> str:
"""Synchronously call an MCP tool.
``cancel_event``: optional ``threading.Event``. When set, the in-flight
HTTP call is cancelled and the function returns a cancellation Error.
Polled in parallel with the tool call via ``asyncio.wait`` so a /cancel
POST from the UI interrupts even mid-network-read.
"""
async def _call() -> Any:
async with _client(url, headers, use_oauth) as client:
return await client.call_tool(name, args)
async def _watch_cancel() -> None:
# 50 ms cadence keeps cancellation responsive without busy-looping;
# matches the cadence routes/inference.py uses for cancel watchers.
while cancel_event is not None and not cancel_event.is_set():
await asyncio.sleep(0.05)
async def _race() -> Any:
# Check cancellation before spawning the call task so a pre-set
# event short-circuits before opening the transport / HTTP
# connection (reviewer-reproduced race).
if cancel_event is not None and cancel_event.is_set():
raise _MCPCancelled
call_task = asyncio.create_task(_call())
if cancel_event is None:
return await asyncio.wait_for(call_task, timeout = timeout)
watch_task = asyncio.create_task(_watch_cancel())
try:
done, pending = await asyncio.wait(
{call_task, watch_task},
timeout = timeout,
return_when = asyncio.FIRST_COMPLETED,
)
finally:
for t in (call_task, watch_task):
if not t.done():
t.cancel()
if not done:
raise asyncio.TimeoutError
if call_task in done:
return call_task.result()
raise _MCPCancelled
try:
result = asyncio.run(_race())
except _MCPCancelled:
return f"Error: MCP tool '{name}' cancelled"
except asyncio.TimeoutError:
return f"Error: MCP tool '{name}' timed out after {timeout:g}s"
except Exception as exc:
logger.exception("MCP call_tool failed for %s: %s", name, exc)
return f"Error: MCP tool '{name}' failed: {exc}"
return _flatten_result(result)
class _MCPCancelled(Exception):
"""Internal sentinel raised when cancel_event fires before the tool returns."""

69
studio/backend/core/inference/providers.py
View file

@ -65,28 +65,77 @@ PROVIDER_REGISTRY: dict[str, dict[str, Any]] = {
},
"gemini": {
"display_name": "Google Gemini",
"base_url": "https://generativelanguage.googleapis.com/v1beta/openai",
# Curated lineup — Google's /v1beta/openai/models returns dozens
# of historical / experimental / embedding ids. Cap to the current
# 3.x family plus the rolling `*-latest` aliases.
# Native Gemini REST endpoint -- the Gemini API does NOT speak
# OpenAI Chat Completions on this base. Requests/responses are
# translated in `_stream_gemini` in external_provider.py.
# API reference: https://ai.google.dev/gemini-api/docs
"base_url": "https://generativelanguage.googleapis.com/v1beta",
# Curated lineup -- the live ListModels response returns dozens
# of historical / experimental / embedding ids. Cap to the
# current chat-capable Gemini families (3.5 / 3.1 / 3 Flash /
# 2.5) plus the Nano Banana image trio and the rolling
# `*-latest` aliases. Excluded on purpose:
# - `gemini-2.0-flash*` (Google retired 2026-06-01; 404 on use)
# - `gemini-3-pro-preview` (shut down 2026-03-09; auto-redirects
# to `gemini-3.1-pro-preview` per Google's deprecation notice,
# so we surface 3.1 directly and skip the redirect).
# The allowlist below blocks the retired ids from re-appearing
# via the live ListModels fetch. Verified against the live
# `/v1beta/models` catalog 2026-05-24.
"default_models": [
"gemini-3.1-pro-preview",
"gemini-3.5-flash",
"gemini-3.1-flash-lite",
"gemini-3-flash-preview",
"gemini-pro-latest",
"gemini-flash-latest",
"gemini-flash-lite-latest",
"gemini-2.5-pro",
"gemini-2.5-flash",
"gemini-2.5-flash-lite",
"gemini-3-pro-image-preview",
"gemini-3.1-flash-image-preview",
"gemini-2.5-flash-image",
],
"supports_streaming": True,
"supports_vision": True,
"supports_tool_calling": True,
"auth_header": "Authorization",
"auth_prefix": "Bearer ",
"notes": "OpenAI-compatible endpoint. API key from https://aistudio.google.com/apikey.",
# The native API takes the API key on the `x-goog-api-key`
# header. An empty `auth_prefix` ensures we send the bare key.
"auth_header": "x-goog-api-key",
"auth_prefix": "",
"openai_compatible": False,
"notes": (
"Native Gemini API. Translation lives in _stream_gemini. "
"API key from https://aistudio.google.com/apikey. "
"See https://ai.google.dev/gemini-api/docs for endpoint shapes."
),
# Even after the regex match, drop ids that Google still
# returns from ListModels but routes via implicit redirect.
# gemini-3-pro-preview was shut down 2026-03-09 and is
# auto-aliased to gemini-3.1-pro-preview; we surface the
# canonical id only so users do not see two cards for the
# same underlying model.
"model_id_deny_exact": ("gemini-3-pro-preview",),
# Matches the chat-capable 3.5 / 3.1 / 3 / 2.5 families plus the
# rolling *-latest aliases (which Google rolls forward as new
# generations ship). Image-tier ids (`-image`, `-image-preview`,
# `nano-banana-pro-preview`) flow through the Nano Banana
# `responseModalities` path in `_stream_gemini`. Retired 2.0
# ids ARE NOT in this regex on purpose -- Google's ListModels
# would otherwise re-surface them and they 404 on use.
"model_id_allowlist": re.compile(
r"^(gemini-3\.1-flash-lite|gemini-3-flash-preview|"
r"gemini-3\.1-pro-preview|gemini-pro-latest|"
r"gemini-flash-latest|gemini-flash-lite-latest)$"
r"^("
r"gemini-3\.5-(?:flash|pro)(?:-preview)?|"
r"gemini-3\.1-(?:flash|pro|flash-lite)(?:-preview)?(?:-customtools)?|"
r"gemini-3\.1-flash-image-preview|"
r"gemini-3-(?:flash|pro)(?:-preview)?|"
r"gemini-3-pro-image-preview|"
r"nano-banana-pro-preview|"
r"gemini-2\.5-pro|gemini-2\.5-flash|gemini-2\.5-flash-lite|"
r"gemini-2\.5-flash-image|"
r"gemini-pro-latest|gemini-flash-latest|gemini-flash-lite-latest"
r")$"
),
},
"deepseek": {

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

@ -13,13 +13,16 @@ import re
# _TOOL_CLOSED_PATS: closed pairs only. _TOOL_ALL_PATS: also trailing
# unclosed runs so truncated tails don't leak markup.
# Function-name char set tracks OpenAI's ^[a-zA-Z0-9_-]{1,64}$ so MCP
# tool names that contain a hyphen (e.g. mcp__srv__list-issues) parse
# the same as the built-in web_search/python/terminal names.
_TOOL_CLOSED_PATS = [
re.compile(r"<tool_call>.*?</tool_call>", re.DOTALL),
re.compile(r"<function=\w+>.*?</function>", re.DOTALL),
re.compile(r"<function=[\w-]+>.*?</function>", re.DOTALL),
]
_TOOL_ALL_PATS = _TOOL_CLOSED_PATS + [
re.compile(r"<tool_call>.*$", re.DOTALL),
re.compile(r"<function=\w+>.*$", re.DOTALL),
re.compile(r"<function=[\w-]+>.*$", re.DOTALL),
]
@ -60,10 +63,12 @@ BUDGET_EXHAUSTED_NUDGE = (
# Pre-compiled patterns reused by ``parse_tool_calls_from_text``.
_TC_JSON_START_RE = re.compile(r"<tool_call>\s*\{")
_TC_FUNC_START_RE = re.compile(r"<function=(\w+)>\s*")
_TC_FUNC_START_RE = re.compile(r"<function=([\w-]+)>\s*")
_TC_END_TAG_RE = re.compile(r"</tool_call>")
_TC_FUNC_CLOSE_RE = re.compile(r"\s*</function>\s*$")
_TC_PARAM_START_RE = re.compile(r"<parameter=(\w+)>\s*")
# Parameter names can carry hyphens too (e.g. MCP tool schemas with
# `issue-number`, `repo-name`); using `\w+` here dropped those keys.
_TC_PARAM_START_RE = re.compile(r"<parameter=([\w-]+)>\s*")
_TC_PARAM_CLOSE_RE = re.compile(r"\s*</parameter>\s*$")

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

@ -14,6 +14,7 @@ import signal
os.environ["UNSLOTH_IS_PRESENT"] = "1"
import asyncio
import random
import re
import shlex
@ -24,6 +25,17 @@ import tempfile
import threading
import urllib.request
from core.inference.mcp_client import (
MCP_TOOL_PREFIX,
call_tool_sync,
is_stdio,
list_tools_async,
parse_server_headers,
probe_timeout,
stdio_mcp_enabled,
)
from storage import mcp_servers_db
from loggers import get_logger
logger = get_logger(__name__)
@ -505,6 +517,94 @@ TERMINAL_TOOL = {
ALL_TOOLS = [WEB_SEARCH_TOOL, PYTHON_TOOL, TERMINAL_TOOL]
# OpenAI's function.name regex: ^[a-zA-Z0-9_-]{1,64}$ -- enforced before
# streaming starts. MCP servers can return tool names containing '.', '/',
# spaces, etc., which the prefix scheme would forward to OpenAI verbatim
# and 400 the whole request. Validate up front and skip with a warning.
_OPENAI_FN_NAME_RE = re.compile(r"^[a-zA-Z0-9_-]{1,64}$")
def _mcp_specs_for_server(server: dict, mcp_tools: list[dict]) -> list[dict]:
"""Convert an MCP server's tool list into OpenAI function specs."""
display = server.get("display_name") or server["id"]
specs: list[dict] = []
seen_names: set[str] = set()
for tool in mcp_tools:
raw_name = tool.get("name") or ""
if not raw_name:
logger.warning("Skipping MCP tool on '%s': empty name.", display)
continue
name = f"{MCP_TOOL_PREFIX}{server['id']}__{raw_name}"
# OpenAI requires function.name ^[a-zA-Z0-9_-]{1,64}$; bad chars
# (., /, spaces, etc.) or oversized names would 400 the whole
# request. Skip + warn so the rest of the tools still ship.
if not _OPENAI_FN_NAME_RE.fullmatch(name):
logger.warning(
"Skipping MCP tool '%s' on '%s': composed name '%s' is not "
"valid OpenAI function.name (regex ^[a-zA-Z0-9_-]{1,64}$).",
raw_name,
display,
name,
)
continue
# Same MCP server returning duplicate tool names would also 400
# OpenAI ("tools[N].function.name duplicates ..."). Drop dupes.
if name in seen_names:
logger.warning(
"Skipping duplicate MCP tool '%s' on '%s'.", raw_name, display
)
continue
seen_names.add(name)
specs.append(
{
"type": "function",
"function": {
"name": name,
"description": f"[{display}] {tool.get('description') or ''}".strip(),
"parameters": tool.get("inputSchema")
or {"type": "object", "properties": {}},
},
}
)
return specs
async def get_enabled_mcp_tools() -> list[dict]:
servers = [s for s in mcp_servers_db.list_servers() if s.get("is_enabled")]
# Never spawn stdio servers when stdio is disabled on this host (e.g. a DB
# carried over from a desktop install onto a Colab / network deployment).
if not stdio_mcp_enabled():
servers = [s for s in servers if not is_stdio(s["url"])]
if not servers:
return []
results = await asyncio.gather(
*(
list_tools_async(
url = s["url"],
headers = parse_server_headers(s),
timeout = probe_timeout(s["url"], bool(s.get("use_oauth"))),
use_oauth = bool(s.get("use_oauth")),
)
for s in servers
),
return_exceptions = True,
)
specs: list[dict] = []
for server, payload in zip(servers, results):
if isinstance(payload, BaseException):
logger.warning(
"MCP server '%s' (%s) discovery failed: %s",
server.get("display_name") or server["id"],
server.get("url"),
payload,
)
continue
specs.extend(_mcp_specs_for_server(server, payload))
return specs
_TIMEOUT_UNSET = object()
@ -525,6 +625,27 @@ def execute_tool(
f"execute_tool: name={name}, session_id={session_id}, timeout={timeout}"
)
effective_timeout = _EXEC_TIMEOUT if timeout is _TIMEOUT_UNSET else timeout
if name.startswith(MCP_TOOL_PREFIX):
try:
_, server_id, tool_name = name.split("__", 2)
except ValueError:
return f"Error: malformed MCP tool name '{name}'"
server = mcp_servers_db.get_server(server_id)
if not server:
return f"Error: MCP server '{server_id}' not found"
if not server.get("is_enabled"):
return f"Error: MCP server '{server_id}' is disabled"
if is_stdio(server["url"]) and not stdio_mcp_enabled():
return f"Error: stdio MCP server '{server_id}' is disabled on this host"
return call_tool_sync(
url = server["url"],
headers = parse_server_headers(server),
name = tool_name,
args = arguments,
timeout = effective_timeout,
use_oauth = bool(server.get("use_oauth")),
cancel_event = cancel_event,
)
if name == "web_search":
return _web_search(
arguments.get("query", ""),
@ -632,8 +753,17 @@ def _validate_and_resolve_host(hostname: str, port: int) -> tuple[bool, str, str
for *_, sockaddr in infos:
ip = ipaddress.ip_address(sockaddr[0])
# `not ip.is_global` rejects every category the denylist below
# also rejects PLUS shared address space (100.64.0.0/10 carrier-
# grade NAT) and benchmarking/documentation/exchange ranges that
# Python classifies with `is_private=False` and `is_global=False`
# (see https://docs.python.org/3/library/ipaddress.html#ipaddress.IPv4Address.is_global).
# The explicit predicates after it give human-readable categories
# in the error message, but a single non-global check is the
# source of truth and prevents future ranges from leaking.
if (
ip.is_private
not ip.is_global
or ip.is_private
or ip.is_loopback
or ip.is_link_local
or ip.is_multicast

13
studio/backend/core/tool_healing.py
View file

@ -17,22 +17,25 @@ verifies it with AST comparison.
import json
import re
# Pre-compiled patterns for tool XML stripping.
# Pre-compiled patterns for tool XML stripping. Hyphen in the
# function/parameter name char-class tracks OpenAI's allowed set so
# MCP tool names with dashes (mcp__srv__list-issues) and parameter
# names with dashes (`issue-number`) parse alongside the built-ins.
_TOOL_CLOSED_PATS = [
re.compile(r"<tool_call>.*?</tool_call>", re.DOTALL),
re.compile(r"<function=\w+>.*?</function>", re.DOTALL),
re.compile(r"<function=[\w-]+>.*?</function>", re.DOTALL),
]
_TOOL_ALL_PATS = _TOOL_CLOSED_PATS + [
re.compile(r"<tool_call>.*$", re.DOTALL),
re.compile(r"<function=\w+>.*$", re.DOTALL),
re.compile(r"<function=[\w-]+>.*$", re.DOTALL),
]
# Pre-compiled patterns for tool-call XML parsing.
_TC_JSON_START_RE = re.compile(r"<tool_call>\s*\{")
_TC_FUNC_START_RE = re.compile(r"<function=(\w+)>\s*")
_TC_FUNC_START_RE = re.compile(r"<function=([\w-]+)>\s*")
_TC_END_TAG_RE = re.compile(r"</tool_call>")
_TC_FUNC_CLOSE_RE = re.compile(r"\s*</function>\s*$")
_TC_PARAM_START_RE = re.compile(r"<parameter=(\w+)>\s*")
_TC_PARAM_START_RE = re.compile(r"<parameter=([\w-]+)>\s*")
_TC_PARAM_CLOSE_RE = re.compile(r"\s*</parameter>\s*$")

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

@ -6,8 +6,10 @@ Unsloth Training Backend
Integrates Unsloth training capabilities with the FastAPI backend
"""
import gc
import os
import sys
import types
# Prevent tokenizer parallelism deadlocks when datasets uses multiprocessing fork
os.environ["TOKENIZERS_PARALLELISM"] = "false"
@ -42,7 +44,10 @@ from utils.hardware import (
get_visible_gpu_count,
)
torch._dynamo.config.recompile_limit = 64
# recompile_limit was removed in some ROCm torch builds (e.g. pytorch.org/whl/rocm6.2).
# Guard so training doesn't crash on RDNA2/RDNA3 with older ROCm torch wheels.
if hasattr(torch._dynamo.config, "recompile_limit"):
torch._dynamo.config.recompile_limit = 64
from unsloth import FastLanguageModel, FastVisionModel, is_bfloat16_supported
from unsloth.chat_templates import get_chat_template
@ -417,8 +422,6 @@ class UnslothTrainer:
in sys.modules. When the next training run calls dataset.map(num_proc=N),
forked child processes inherit this stale state and deadlock.
"""
import sys as _sys
# Remove cloned audio repo paths from sys.path
base_dir = os.path.dirname(os.path.abspath(__file__))
audio_paths = [
@ -433,15 +436,15 @@ class UnslothTrainer:
removed_paths = []
for path in audio_paths:
if path in _sys.path:
_sys.path.remove(path)
if path in sys.path:
sys.path.remove(path)
removed_paths.append(path)
# Remove stale audio modules from sys.modules
prefixes = ("snac", "whisper", "sparktts", "outetts")
removed_modules = [key for key in _sys.modules if key.startswith(prefixes)]
removed_modules = [key for key in sys.modules if key.startswith(prefixes)]
for key in removed_modules:
del _sys.modules[key]
del sys.modules[key]
if removed_paths or removed_modules:
logger.info(
@ -538,10 +541,9 @@ class UnslothTrainer:
# clear_unsloth_compiled_cache() deletes the disk cache, but the flag
# prevents re-compilation — leaving missing cache files. Reloading
# restores original class definitions so Unsloth can re-compile cleanly.
import sys as _sys
import importlib
for _key, _mod in list(_sys.modules.items()):
for _key, _mod in list(sys.modules.items()):
if "transformers.models." in _key and ".modeling_" in _key:
if hasattr(_mod, "__UNSLOTH_PATCHED__"):
try:
@ -657,6 +659,23 @@ class UnslothTrainer:
f"Using device_map='{device_map}' ({get_visible_gpu_count()} GPU(s) visible)"
)
# AMD ROCm hardware without native bfloat16 (e.g. RDNA2 / gfx103x)
# crashes with an LLVM error at the first bf16 kernel dispatch if
# dtype=None lets unsloth auto-pick bf16. Force float16 there so that
# path is never reached. NVIDIA keeps dtype=None so unsloth's own
# bf16/fp16/float32 auto-detection (including FORCE_FLOAT32 models) is
# honored -- older NVIDIA without bf16 (T4/V100) must NOT be coerced to
# float16 here, which the previous unconditional branch did wrongly.
# Derive ROCm inline (not hardware.IS_ROCM) because that flag is unset
# until detect_hardware() runs, which isn't guaranteed in this subprocess.
_is_rocm = (
bool(getattr(torch.version, "hip", None))
or "rocm" in torch.__version__.lower()
)
_auto_dtype = (
torch.float16 if (_is_rocm and not is_bfloat16_supported()) else None
)
# Branch based on model type
if self._audio_type == "csm":
# CSM: FastModel + auto_model=CsmForConditionalGeneration + load_in_4bit=False
@ -666,7 +685,7 @@ class UnslothTrainer:
self.model, self.tokenizer = FastModel.from_pretrained(
model_name = model_name,
max_seq_length = max_seq_length,
dtype = None,
dtype = _auto_dtype,
auto_model = CsmForConditionalGeneration,
load_in_4bit = False,
device_map = device_map,
@ -683,7 +702,7 @@ class UnslothTrainer:
self.model, self.tokenizer = FastModel.from_pretrained(
model_name = model_name,
dtype = None,
dtype = _auto_dtype,
load_in_4bit = False,
device_map = device_map,
full_finetuning = full_finetuning,
@ -705,7 +724,7 @@ class UnslothTrainer:
self.model, self.tokenizer = FastLanguageModel.from_pretrained(
model_name = model_name,
max_seq_length = max_seq_length,
dtype = None,
dtype = _auto_dtype,
load_in_4bit = load_in_4bit,
device_map = device_map,
full_finetuning = full_finetuning,
@ -777,7 +796,7 @@ class UnslothTrainer:
self.model, self.tokenizer = FastModel.from_pretrained(
model_name = model_name,
max_seq_length = max_seq_length,
dtype = None,
dtype = _auto_dtype,
load_in_4bit = load_in_4bit,
device_map = device_map,
full_finetuning = full_finetuning,
@ -791,7 +810,7 @@ class UnslothTrainer:
self.model, self.tokenizer = FastVisionModel.from_pretrained(
model_name = model_name,
max_seq_length = max_seq_length,
dtype = None, # Auto-detect
dtype = _auto_dtype,
load_in_4bit = load_in_4bit,
device_map = device_map,
full_finetuning = full_finetuning,
@ -824,7 +843,7 @@ class UnslothTrainer:
self.model, self.tokenizer = FastLanguageModel.from_pretrained(
model_name = model_name,
max_seq_length = max_seq_length,
dtype = None, # Auto-detect
dtype = _auto_dtype,
load_in_4bit = load_in_4bit,
device_map = device_map,
full_finetuning = full_finetuning,
@ -1188,7 +1207,6 @@ class UnslothTrainer:
We patch at both instance AND class level for maximum reliability,
and strip non-TransformersKwargs params that Unsloth/PEFT inject.
"""
import types
import torch
import torch.nn as nn
from transformers.models.csm.modeling_csm import (
@ -1730,7 +1748,6 @@ class UnslothTrainer:
logger.info("Freeing SNAC codec model from GPU...\n")
snac_model.to("cpu")
del snac_model
import gc
gc.collect()
torch.cuda.empty_cache()
@ -1754,13 +1771,10 @@ class UnslothTrainer:
Mirrors Spark_TTS_(0_5B).ipynb: encode audio with BiCodec (semantic + global tokens),
format as special-token text strings for SFTTrainer with dataset_text_field="text".
"""
import sys
import torch
import numpy as np
import torchaudio.transforms as T
import subprocess
device = "cuda" if torch.cuda.is_available() else "cpu"
# The sparktts Python package lives in the SparkAudio/Spark-TTS GitHub repo,
@ -1960,7 +1974,6 @@ class UnslothTrainer:
audio_tokenizer.model.cpu()
audio_tokenizer.feature_extractor.cpu()
del audio_tokenizer
import gc
gc.collect()
torch.cuda.empty_cache()
@ -1989,7 +2002,6 @@ class UnslothTrainer:
OuteTTS AudioProcessor for speaker representations, PromptProcessor for
training prompts. Outputs text strings for SFTTrainer with dataset_text_field="text".
"""
import sys
import io
import tempfile
import torch
@ -2173,7 +2185,6 @@ class UnslothTrainer:
del whisper_model
del audio_processor
del prompt_processor
import gc
gc.collect()
torch.cuda.empty_cache()
@ -3057,6 +3068,14 @@ class UnslothTrainer:
logger.info("Configuring DeepSeek OCR data collator...\n")
FastVisionModel.for_training(self.model)
# DeepSeek OCR's (image_size, base_size, crop_mode) is a
# coupled preset; changing image_size alone desyncs the
# per-crop pixel grid from num_queries. Use Gundam.
if training_args.get("vision_image_size") is not None:
logger.info(
"Vision image resize ignored for DeepSeek OCR "
"(uses fixed Gundam preset).\n"
)
data_collator = DeepSeekOCRDataCollator(
tokenizer = self.tokenizer,
model = self.model,
@ -3123,7 +3142,21 @@ class UnslothTrainer:
from unsloth.trainer import UnslothVisionDataCollator
FastVisionModel.for_training(self.model)
data_collator = UnslothVisionDataCollator(self.model, self.tokenizer)
vision_image_size = training_args.get("vision_image_size")
if vision_image_size is None:
data_collator = UnslothVisionDataCollator(
self.model, self.tokenizer
)
else:
logger.info(
f"Vision image resize: {vision_image_size} (max dimension)\n"
)
data_collator = UnslothVisionDataCollator(
self.model,
self.tokenizer,
resize = vision_image_size,
resize_dimension = "max",
)
logger.info("Vision data collator configured\n")
# ========== TRAINING CONFIGURATION ==========

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

@ -193,6 +193,7 @@ class TrainingBackend:
"hf_token": kwargs.get("hf_token", ""),
"load_in_4bit": kwargs.get("load_in_4bit", True),
"max_seq_length": kwargs.get("max_seq_length", 2048),
"vision_image_size": kwargs.get("vision_image_size"),
"hf_dataset": kwargs.get("hf_dataset", ""),
"local_datasets": kwargs.get("local_datasets"),
"local_eval_datasets": kwargs.get("local_eval_datasets"),

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

@ -21,6 +21,9 @@ import shutil
import sys
import time
import traceback
import gc
import re
import types
import subprocess as _sp
from pathlib import Path
from typing import Any, Callable
@ -70,6 +73,58 @@ _TILELANG_INSTALL_TIMEOUT_S = 600
_TVM_FFI_BROKEN_VERSIONS = ("0.1.10", "0.1.11")
_FAST_PATH_HOOKS_SKIP_ENV = "UNSLOTH_STUDIO_SKIP_FAST_PATH_HOOKS"
# Module-level handle so the torch.library.Library registration survives past
# run_training_process() and is not garbage collected mid-run.
_WINDOWS_ROCM_GROUPED_MM_LIB = None
# Worker subprocesses inherit the parent env but not the parent's
# os.add_dll_directory registrations. Replicate main.py's Windows ROCm DLL
# setup at module load so the first `import torch` can find amdhip64.dll even
# when HIP_PATH\bin is not on the system PATH. Handles retained at module
# scope so they are not garbage collected.
_ROCM_DLL_HANDLES: list = []
if sys.platform == "win32":
def _add_rocm_dll_dirs_worker() -> None:
_candidates: list[str] = []
for _var in ("HIP_PATH", "ROCM_PATH"):
_val = os.environ.get(_var)
if _val:
_candidates.append(os.path.join(_val, "bin"))
_default_root = os.path.join(
os.environ.get("ProgramFiles", r"C:\Program Files"), "AMD", "ROCm"
)
def _ver_key(name: str) -> tuple:
# Numeric tuple key so "10.0" sorts after "7.0"; non-numeric chunks fall back to string.
parts = []
for chunk in name.split("."):
try:
parts.append((0, int(chunk)))
except ValueError:
parts.append((1, chunk))
return tuple(parts)
try:
if os.path.isdir(_default_root):
for _ver in sorted(
os.listdir(_default_root), key = _ver_key, reverse = True
):
_bin = os.path.join(_default_root, _ver, "bin")
if os.path.isdir(_bin):
_candidates.append(_bin)
except OSError:
pass
for _d in _candidates:
if os.path.isdir(_d):
try:
_ROCM_DLL_HANDLES.append(os.add_dll_directory(_d))
except (OSError, AttributeError):
pass
_add_rocm_dll_dirs_worker()
del _add_rocm_dll_dirs_worker
def _model_wants_causal_conv1d(model_name: str) -> bool:
name = model_name.lower()
@ -320,11 +375,21 @@ def _install_package_wheel_first(
f"{snippet}",
)
else:
logger.error(
"Failed to install %s from PyPI:\n%s",
display_name,
result.stdout,
)
if sys.platform == "win32":
# No prebuilt wheel and no source build toolchain on Windows --
# this is expected for packages like causal-conv1d. Log at
# info so users aren't alarmed by what looks like an error.
logger.info(
"%s is not available on Windows (no prebuilt wheel); skipping",
display_name,
)
logger.debug("Install output:\n%s", result.stdout)
else:
logger.error(
"Failed to install %s from PyPI:\n%s",
display_name,
result.stdout,
)
return False
if is_hip:
@ -337,6 +402,9 @@ def _install_package_wheel_first(
def _ensure_causal_conv1d_fast_path(event_queue: Any, model_name: str) -> None:
if not _model_wants_causal_conv1d(model_name):
return
if sys.platform == "win32":
logger.info("causal-conv1d: no prebuilt wheel for Windows; skipping")
return
_install_package_wheel_first(
event_queue = event_queue,
@ -404,6 +472,11 @@ def _ensure_flash_linear_attention_unconditional(event_queue: Any) -> bool:
"""Install pinned FLA + fla-core with --no-deps. Returns True iff importable post-call."""
if os.getenv(_FLA_SKIP_ENV) == "1":
return False
if sys.platform == "win32":
logger.info(
"Skipping flash-linear-attention install: no prebuilt wheel for Windows"
)
return False
if sys.version_info < _FLA_MIN_PYTHON:
logger.info(
"Skipping flash-linear-attention install: requires Python >= %d.%d, have %s",
@ -483,10 +556,17 @@ def _ensure_flash_linear_attention_unconditional(event_queue: Any) -> bool:
return False
if result.returncode != 0:
logger.warning(
"flash-linear-attention install failed (continuing on torch fallback):\n%s",
result.stdout,
)
if sys.platform == "win32":
logger.info(
"flash-linear-attention not available on Windows (no prebuilt wheel); "
"continuing on torch fallback"
)
logger.debug("Install output:\n%s", result.stdout)
else:
logger.warning(
"flash-linear-attention install failed (continuing on torch fallback):\n%s",
result.stdout,
)
_send_status(
event_queue,
"flash-linear-attention install failed; continuing without it",
@ -607,15 +687,61 @@ def _tilelang_importable() -> bool:
def _torch_has_hip() -> bool:
"""True iff torch is a ROCm build; `torch.version.hip` is the only reliable signal on x86_64 ROCm."""
"""True iff torch is a ROCm build.
`torch.version.hip` covers official PyTorch ROCm wheels; AMD SDK / Radeon
wheels can leave it unset but still encode "rocm" in `torch.__version__`.
"""
try:
import torch as _torch
return getattr(_torch.version, "hip", None) is not None
return bool(
getattr(_torch.version, "hip", None)
or "rocm" in getattr(_torch, "__version__", "").lower()
)
except Exception:
return False
def _rocm_classify_unified_memory(props: Any) -> tuple[str, bool]:
"""Classify a ROCm device as unified-memory (APU) or discrete.
Returns ``(gcn_arch, is_unified)`` where:
- ``gcn_arch`` is the canonical arch string (e.g. ``"gfx1151"``) when a
known attribute is present, or ``""`` when all arch attrs are absent.
- ``is_unified`` is ``True`` for AMD APUs with a shared GPU/system-RAM pool
(gfx1150 Strix Point, gfx1151 Strix Halo) — these need a lower
``set_per_process_memory_fraction`` cap to leave headroom for the OS.
Classification priority:
1. ``gcnArchName`` / variant spellings (stable, naming-independent).
2. Device-name substring match as a last-resort fallback when all arch
attrs are absent (AMD SDK / Radeon wheels may not populate them):
- gfx1150 Strix Point: ``Radeon 890M``, ``Radeon 880M``
- gfx1151 Strix Halo: ``Radeon 8060S`` (Ryzen AI MAX+ 395),
``Radeon 8050S`` (cut-down SKU)
"""
gcn_arch = ""
for _attr in ("gcnArchName", "gcn_arch_name", "arch_name", "gfx_arch_name"):
_v = (getattr(props, _attr, "") or "").split(":")[0].strip()
if _v:
gcn_arch = _v
break
if gcn_arch:
return gcn_arch, gcn_arch in {"gfx1150", "gfx1151"}
# Arch attrs absent — fall back to device-name matching.
dev_lower = (getattr(props, "name", "") or "").lower()
is_unified = (
"890m" in dev_lower
or "880m" in dev_lower
or "8060s" in dev_lower
or "8050s" in dev_lower
)
return gcn_arch, is_unified
def _tilelang_platform_supported() -> bool:
"""True iff a tilelang 0.1.8 wheel will load: Linux x86_64/aarch64, non-HIP torch.
@ -881,6 +1007,9 @@ def _install_fast_path_hooks(event_queue: Any, model_name: str) -> None:
_ensure_tilelang_backend_unconditional(eq)
def _causal_conv1d_install(eq: Any) -> bool:
if sys.platform == "win32":
logger.info("causal-conv1d: no prebuilt wheel for Windows; skipping")
return False
ok = _install_package_wheel_first(
event_queue = eq,
import_name = "causal_conv1d",
@ -959,7 +1088,47 @@ def _activate_transformers_version(model_name: str) -> None:
activate_transformers_for_subprocess(model_name)
def _adapt_for_mlx_vlm(items):
def _mlx_vlm_max_resized_size(width: int, height: int, target: int) -> tuple[int, int]:
if width <= 0 or height <= 0 or target <= 0:
return width, height
largest_side = max(width, height)
if largest_side <= target:
return width, height
# Integer formula matches unsloth_zoo's collator (Python round() differs
# by 1px on half-pixel cases). max(1, _) avoids zero-side degenerate output.
new_w = max(1, (width * target + largest_side // 2) // largest_side)
new_h = max(1, (height * target + largest_side // 2) // largest_side)
return new_w, new_h
def _resize_mlx_vlm_image(image, resize):
if resize is None:
return image
try:
from PIL import Image
import numpy as np
except ImportError:
return image
if not isinstance(image, Image.Image):
return image
image = image.convert("RGB")
new_size = _mlx_vlm_max_resized_size(*image.size, int(resize))
if new_size != image.size:
resampling = getattr(Image, "Resampling", Image).LANCZOS
image = image.resize(new_size, resampling)
# When a resize is requested, hand mlx-vlm a writable RGB ndarray so its
# PIL-path square-resize is skipped and HF processors don't warn on
# non-writable views. resize=None (Default) above keeps the original PIL.
return np.array(image, copy = True)
def _resize_mlx_vlm_images(value, resize):
if isinstance(value, list):
return [_resize_mlx_vlm_image(image, resize) for image in value]
return _resize_mlx_vlm_image(value, resize)
def _adapt_for_mlx_vlm(items, resize = None):
"""Adapt GPU-path VLM dataset output for mlx-vlm consumption.
The GPU path embeds PIL images inside messages content as
@ -979,7 +1148,7 @@ def _adapt_for_mlx_vlm(items):
if isinstance(part, dict) and part.get("type") == "image":
img = part.get("image")
if img is not None:
images.append(img)
images.append(_resize_mlx_vlm_image(img, resize))
new_content.append({"type": "image"})
else:
new_content.append(part)
@ -990,9 +1159,9 @@ def _adapt_for_mlx_vlm(items):
if images:
out["image"] = images[0] if len(images) == 1 else images
elif "image" in item:
out["image"] = item["image"]
out["image"] = _resize_mlx_vlm_images(item["image"], resize)
elif "images" in item:
out["images"] = item["images"]
out["images"] = _resize_mlx_vlm_images(item["images"], resize)
adapted.append(out)
return adapted
@ -1093,7 +1262,6 @@ def _run_mlx_training(event_queue, stop_queue, config):
Mirrors the event_queue protocol so the parent process pump works unchanged.
"""
import time
import gc
import math
import threading
import queue as _queue
@ -1168,6 +1336,25 @@ def _run_mlx_training(event_queue, stop_queue, config):
is_vlm = bool(is_dataset_image and getattr(model, "_is_vlm_model", False))
model._is_vlm_model = is_vlm
vision_image_size = config.get("vision_image_size")
# DeepSeek OCR uses a coupled preset tuple; skip resize like the Torch path.
_model_name_lower = str(config.get("model_name", "")).lower()
_is_deepseek_ocr = "deepseek" in _model_name_lower and "ocr" in _model_name_lower
if is_vlm and vision_image_size is not None and _is_deepseek_ocr:
_send(
"status",
status_message = (
"MLX vision image resize ignored for DeepSeek OCR "
"(uses fixed Gundam preset)."
),
)
vision_image_size = None
elif is_vlm and vision_image_size is not None:
vision_image_size = int(vision_image_size)
_send(
"status",
status_message = f"MLX vision image resize: {vision_image_size} (max dimension)",
)
# ── 2. Apply LoRA / full FT ──
# Pass gradient_checkpointing as string ("mlx"/"unsloth"/"none"/etc.)
@ -1302,7 +1489,10 @@ def _run_mlx_training(event_queue, stop_queue, config):
progress_callback = _fmt_progress,
)
if vlm_info.get("success"):
dataset = _adapt_for_mlx_vlm(vlm_info["dataset"])
dataset = _adapt_for_mlx_vlm(
vlm_info["dataset"],
resize = vision_image_size,
)
else:
errors = vlm_info.get("errors", [])
raise ValueError(
@ -1317,7 +1507,10 @@ def _run_mlx_training(event_queue, stop_queue, config):
dataset_name = hf_dataset or "local",
)
if ev_info.get("success"):
eval_dataset = _adapt_for_mlx_vlm(ev_info["dataset"])
eval_dataset = _adapt_for_mlx_vlm(
ev_info["dataset"],
resize = vision_image_size,
)
elif format_type:
_send("status", status_message = f"Formatting dataset ({format_type})...")
@ -1828,6 +2021,340 @@ def run_training_process(
'Install for better performance: pip install "triton-windows<3.7"'
)
# ── 1d. Stub torchao on Windows ROCm ──
# Shared with the export worker; see core/_torchao_stub.py for the full
# rationale (torchao -> torch.distributed._functional_collectives crashes on
# Windows ROCm because the RCCL backend is absent). No-op off Windows ROCm.
# Must run before any import of transformers / unsloth_zoo.
from core._torchao_stub import install_torchao_windows_rocm_stub
install_torchao_windows_rocm_stub()
# ── 1e. Ensure torch.distributed helper attrs are present ──
# Single-GPU training never initialises the process group, so these helpers
# are never called — but transformers/trl import them unconditionally.
_td_stubs = {
"is_initialized": lambda: False,
"is_available": lambda: False,
"is_torchelastic_launched": lambda: False,
"get_rank": lambda: 0,
"get_world_size": lambda: 1,
"barrier": lambda: None,
}
try:
import torch.distributed as _td
for _name, _stub in _td_stubs.items():
if not hasattr(_td, _name):
setattr(_td, _name, _stub)
except Exception:
_td_mock = types.ModuleType("torch.distributed")
for _name, _stub in _td_stubs.items():
setattr(_td_mock, _name, _stub)
sys.modules["torch.distributed"] = _td_mock
try:
import torch as _torch
_torch.distributed = _td_mock
except Exception:
pass
# ── 1f. Windows ROCm runtime patches ──
# torch._grouped_mm has a null HIP kernel on gfx1200 (ROCm ≤ 7.12 Windows),
# causing 0xC0000005 (access violation) during training.
#
# Root cause: the JitDecomp autograd decomposition system (NOT torch.compile)
# dispatches _grouped_mm → _fused_adagrad_ → _grouped_mm HIP → null crash.
# TORCHDYNAMO_DISABLE=1 stops the compiler frontend but does NOT stop
# JitDecomp, so we must also override the CUDA dispatch key for _grouped_mm
# with a safe Python fallback.
#
# Fixed in AMD's wheel: torch==2.11.0+rocm7.13.0 — the 3-D batch and grouped
# (with offs) variants of _grouped_mm now have working HIP kernels on gfx1200.
# We gate the dispatch override on HIP < 7.13 so users on the fixed wheel get
# the real GPU kernel rather than our Python fallback.
#
# Verified: null on torch==2.10.0+rocm7.12.0; fixed on torch==2.11.0+rocm7.13.0.
#
# Schema: _grouped_mm(Tensor self, Tensor mat2, Tensor? offs=None,
# Tensor? bias=None, ScalarType? out_dtype=None) -> Tensor
# offs: optional group-split offsets (MoE-style variable-size batches)
#
# torch is already in sys.modules from section 1e's `import torch.distributed`.
# Module-level _WINDOWS_ROCM_GROUPED_MM_LIB keeps the registration alive past
# function return / mid-run GC.
global _WINDOWS_ROCM_GROUPED_MM_LIB
if sys.platform == "win32":
_torch_for_rocm = sys.modules.get("torch")
# Broad check: torch.version.hip OR "rocm" in torch.__version__.
# AMD SDK / Radeon Windows wheels do not always populate
# torch.version.hip; without the broad check the BNB version pin,
# dynamo-disable, and _grouped_mm fallback below silently skip
# (matches the torchao stub gate above and main.py).
_build_version_for_rocm = (
getattr(_torch_for_rocm, "__version__", "").lower()
if _torch_for_rocm is not None
else ""
)
_is_win_rocm_torch = bool(
_torch_for_rocm is not None
and (
getattr(getattr(_torch_for_rocm, "version", None), "hip", None)
or "rocm" in _build_version_for_rocm
)
)
if _is_win_rocm_torch:
# Disable dynamo (belt-and-suspenders; JitDecomp patch below is the
# real fix, but keeping dynamo off avoids any other compile paths).
if "TORCHDYNAMO_DISABLE" not in os.environ:
os.environ["TORCHDYNAMO_DISABLE"] = "1"
logger.info("Windows ROCm: torch.compile (dynamo) disabled")
# BNB auto-detects the HIP version from torch.version.hip and uses
# it to choose which DLL to load (e.g. "7.13" → rocm713.dll).
# AMD's Windows BNB prerelease wheel ships only one rocm DLL, and its
# version suffix does not always match the torch HIP version (e.g.
# torch==2.11.0+rocm7.13.0 ships HIP 7.13, but the BNB wheel still
# ships rocm72.dll). We detect the actual DLL name from the installed
# package and override BNB's auto-detection. "72" is a safe fallback
# if detection fails. Callers may override by pre-setting the var.
if "BNB_ROCM_VERSION" not in os.environ:
_bnb_rocm_ver = None
try:
import glob as _glob
import importlib.util as _ilu
import re as _re
_bnb_spec = _ilu.find_spec("bitsandbytes")
if _bnb_spec and _bnb_spec.submodule_search_locations:
_all_vers: list[str] = []
for _pkg_dir in _bnb_spec.submodule_search_locations:
for _dll in _glob.glob(
os.path.join(_pkg_dir, "libbitsandbytes_rocm*.dll")
):
_m = _re.search(
r"libbitsandbytes_rocm(\d+)\.dll",
os.path.basename(_dll),
)
if _m:
_all_vers.append(_m.group(1))
# Pick the highest numeric suffix so that e.g. "713"
# wins over "72" when both variants are present.
# Filesystem glob order is not guaranteed, so always
# sort rather than stopping at the first match.
if _all_vers:
_bnb_rocm_ver = max(_all_vers, key = lambda v: int(v))
except Exception:
pass
_bnb_rocm_ver = _bnb_rocm_ver or "72"
os.environ["BNB_ROCM_VERSION"] = _bnb_rocm_ver
logger.info(
"Windows ROCm: set BNB_ROCM_VERSION=%s "
"(detected from installed BNB wheel; "
"overrides torch.version.hip auto-detection)",
_bnb_rocm_ver,
)
# Parse HIP version for the kernel-fix gate below.
# torch.version.hip can be "7.13.99004", "7.2.0", etc.
# AMD SDK / Radeon wheels may leave torch.version.hip unset and
# encode the ROCm version in torch.__version__ instead
# (e.g. "2.11.0+rocm7.13.0" or "2.9.0+rocmsdk20251116"); fall back
# to that string when version.hip is missing.
def _hip_ver_at_least(major: int, minor: int) -> bool:
_hip_str = getattr(
getattr(_torch_for_rocm, "version", None), "hip", None
)
if not _hip_str:
# Try the standard "+rocmX.Y.Z" embedded version first
# (e.g. "2.11.0+rocm7.13.0").
_ver_match = re.search(r"rocm(\d+)\.(\d+)", _build_version_for_rocm)
if _ver_match:
return (
int(_ver_match.group(1)),
int(_ver_match.group(2)),
) >= (major, minor)
# AMD SDK / Radeon Windows wheels encode the build as
# "+rocmsdk<date>" (e.g. "2.9.0+rocmsdk20251116") with no
# explicit rocmX.Y component. The rocmsdk format was
# introduced after the gfx120X null-kernel fix landed in
# ROCm 7.13, so any wheel with this suffix is new enough to
# have working HIP kernels. Treat as >= 7.13 rather than
# falling back to False and installing the Python workaround
# on a wheel that doesn't need it.
if "rocmsdk" in _build_version_for_rocm:
logger.debug(
"Windows ROCm: AMD SDK wheel detected (%r); "
"assuming HIP >= %d.%d (rocmsdk wheels post-date "
"the gfx120X null-kernel fix)",
_build_version_for_rocm,
major,
minor,
)
return True
return False
try:
_parts = [int(x) for x in str(_hip_str).split(".")[:2]]
if len(_parts) < 2:
logger.warning(
"Windows ROCm: torch.version.hip %r has fewer than "
"two components; cannot compare against %d.%d",
_hip_str,
major,
minor,
)
return False
return (_parts[0], _parts[1]) >= (major, minor)
except ValueError:
logger.warning(
"Windows ROCm: could not parse torch.version.hip %r as "
"a version number; assuming HIP < %d.%d",
_hip_str,
major,
minor,
)
return False
# _grouped_mm HIP kernel was null on gfx1200 in ROCm ≤ 7.12,
# causing 0xC0000005. AMD fixed it in ROCm 7.13 (torch 2.11+).
# Only install the Python fallback on the affected versions so users
# on 7.13+ get the real GPU kernel for MoE workloads.
if not _hip_ver_at_least(7, 13):
try:
import warnings as _warnings
_gm_lib = _torch_for_rocm.library.Library("aten", "IMPL")
def _grouped_mm_safe_impl(
self, mat2, offs = None, bias = None, out_dtype = None
):
"""Python mm/bmm fallback for _grouped_mm on gfx1200 (null HIP kernel, ROCm ≤ 7.12)."""
_t = _torch_for_rocm
if offs is None:
# No offsets: behave like the real op, which
# accepts either (M, K) x (K, N) -> mm, or 3-D
# batched inputs -> bmm. Picking torch.mm
# unconditionally previously raised "self must be
# a matrix" on 3-D MoE workloads.
if self.dim() == 3 and mat2.dim() == 3:
result = _t.bmm(self.contiguous(), mat2.contiguous())
elif self.dim() == 3 and mat2.dim() == 2:
# Broadcast 2-D mat2 across the batch dim.
result = _t.matmul(self.contiguous(), mat2.contiguous())
elif self.dim() == 2 and mat2.dim() == 3:
# Broadcast 2-D self across batch via matmul semantics.
result = _t.matmul(self.contiguous(), mat2.contiguous())
else:
result = _t.mm(self.contiguous(), mat2.contiguous())
else:
# Grouped case: offs[i] is the exclusive end-row of
# group i in `self`; mat2 may be 3-D or 2-D.
offs_list = offs.tolist()
pieces = []
prev = 0
for idx, end in enumerate(offs_list):
end = int(end)
a_part = self[prev:end].contiguous()
if mat2.dim() == 3:
b_part = mat2[idx].contiguous()
else:
b_part = mat2.contiguous()
pieces.append(_t.mm(a_part, b_part))
prev = end
# Include any trailing rows not covered by offs
if prev < self.shape[0]:
a_tail = self[prev:].contiguous()
b_tail = (
mat2[-1].contiguous()
if mat2.dim() == 3
else mat2.contiguous()
)
pieces.append(_t.mm(a_tail, b_tail))
result = (
_t.cat(pieces, dim = 0)
if pieces
else _t.zeros(
0,
mat2.shape[-1],
device = self.device,
dtype = self.dtype,
)
)
if bias is not None:
result = result + bias
if out_dtype is not None:
result = result.to(out_dtype)
elif result.dtype != self.dtype:
result = result.to(self.dtype)
return result
with _warnings.catch_warnings():
_warnings.simplefilter("ignore")
_gm_lib.impl("_grouped_mm", _grouped_mm_safe_impl, "CUDA")
_WINDOWS_ROCM_GROUPED_MM_LIB = _gm_lib # prevent GC
logger.info(
"Windows ROCm: patched _grouped_mm CUDA dispatch "
"(null HIP kernel on gfx1200, ROCm ≤ 7.12 — "
"bypassed with Python mm fallback)"
)
except Exception as _patch_exc:
logger.warning(
"Windows ROCm: could not patch _grouped_mm — "
"training may crash with 0xC0000005: %s",
_patch_exc,
)
else:
logger.info(
"Windows ROCm: HIP >= 7.13 — _grouped_mm kernel is functional, "
"skipping Python fallback (AMD fixed gfx1200 null kernel in ROCm 7.13)"
)
# ── 1g. ROCm OOM guard ──
# On RDNA 4 (gfx1200/gfx1201) and other ROCm GPUs, exhausting VRAM can
# cause a HIP driver hang that freezes the entire system rather than
# raising a Python exception. set_per_process_memory_fraction caps the
# HIP allocator so PyTorch raises OutOfMemoryError before hitting the
# hardware limit, giving the UI a clean error instead of a system freeze.
# Only applied on ROCm -- NVIDIA CUDA has a graceful OOM path and does
# not need this cap.
# Unified-memory APUs (gfx1150 Strix Point / gfx1151 Strix Halo) share GPU
# and system RAM in one pool: 0.90 of 128 GB starves the OS. Use 0.80 there.
# Primary classifier: gcnArchName from device properties — stable within a
# product family and naming-independent. AMD SDK / Radeon wheels may omit
# gcnArchName or expose it under a variant spelling, so we try several attr
# names then fall back to known device-name markers as a last resort.
# Non-fatal: silently skipped if torch is not importable.
if _hw.IS_ROCM:
try:
import torch as _torch_mem
if _torch_mem.cuda.is_available():
# Classify unified vs discrete via _rocm_classify_unified_memory.
# See that function's docstring for classification priority.
_props = _torch_mem.cuda.get_device_properties(0)
_dev_name = _props.name
_gcn_arch, _is_unified = _rocm_classify_unified_memory(_props)
if _is_unified and not _gcn_arch:
logger.debug(
"ROCm OOM guard: gcnArchName absent -- inferred "
"unified memory from device name %r; applying 0.80 cap",
_dev_name,
)
_mem_fraction = 0.80 if _is_unified else 0.90
_torch_mem.cuda.set_per_process_memory_fraction(_mem_fraction)
logger.info(
"ROCm OOM guard: set_per_process_memory_fraction(%.2f) — "
"%s memory host (%s, %s)",
_mem_fraction,
"unified" if _is_unified else "discrete",
_dev_name,
_gcn_arch or "unknown arch",
)
except Exception as _oom_guard_err:
logger.debug("Could not set GPU memory fraction: %s", _oom_guard_err)
# ── 2. Now import ML libraries (fresh in this clean process) ──
try:
_send_status(event_queue, "Importing Unsloth...")
@ -2248,6 +2775,7 @@ def run_training_process(
eval_dataset = eval_dataset,
eval_steps = eval_steps,
max_seq_length = config.get("max_seq_length", 2048),
vision_image_size = config.get("vision_image_size"),
optim = config.get("optim", "adamw_8bit"),
lr_scheduler_type = config.get("lr_scheduler_type", "linear"),
is_cpt = is_cpt,
@ -2281,14 +2809,38 @@ def run_training_process(
)
except Exception as exc:
event_queue.put(
{
"type": "error",
"error": str(exc),
"stack": traceback.format_exc(limit = 20),
"ts": time.time(),
}
_exc_str = str(exc).lower()
_is_oom = (
"out of memory" in _exc_str
or "hip out of memory" in _exc_str
or "cuda out of memory" in _exc_str
or type(exc).__name__ == "OutOfMemoryError"
)
if _is_oom:
_oom_msg = (
"GPU ran out of VRAM during training.\n"
"To fix: reduce max_seq_length (e.g. 2048–4096), enable "
"gradient_checkpointing=True, lower per_device_train_batch_size, "
"or use a smaller model / higher quantization."
)
logger.error("Training stopped: GPU OOM — %s", exc)
event_queue.put(
{
"type": "error",
"error": _oom_msg,
"stack": traceback.format_exc(limit = 20),
"ts": time.time(),
}
)
else:
event_queue.put(
{
"type": "error",
"error": str(exc),
"stack": traceback.format_exc(limit = 20),
"ts": time.time(),
}
)
def _send_status(event_queue: Any, message: str) -> None:

171
studio/backend/main.py
View file

@ -12,12 +12,127 @@ from pathlib import Path as _Path
# Suppress annoying C-level dependency warnings globally
os.environ["PYTHONWARNINGS"] = "ignore"
# ── Windows AMD ROCm DLL injection ──────────────────────────────────────────
# Python 3.8+ ignores PATH for extension modules; register ROCm bin dirs with
# os.add_dll_directory() so amdhip64.dll etc. are found before any torch import.
if sys.platform == "win32":
# Retained at module scope -- os.add_dll_directory returns a handle that
# removes the search-path entry when garbage collected.
_ROCM_DLL_HANDLES: list = []
def _add_rocm_dll_dirs() -> None:
candidates = []
# 1. HIP_PATH / ROCM_PATH -- set by the AMD HIP SDK installer
for _var in ("HIP_PATH", "ROCM_PATH"):
_val = os.environ.get(_var)
if _val:
candidates.append(os.path.join(_val, "bin"))
# 2. Standard AMD installer location: C:\Program Files\AMD\ROCm\<ver>\bin
# Scan all installed versions, newest first.
_default_root = os.path.join(
os.environ.get("ProgramFiles", r"C:\Program Files"), "AMD", "ROCm"
)
def _ver_key(name: str) -> tuple:
# Numeric tuple key so "10.0" sorts after "7.0"; non-numeric chunks fall back to string.
parts = []
for chunk in name.split("."):
try:
parts.append((0, int(chunk)))
except ValueError:
parts.append((1, chunk))
return tuple(parts)
try:
if os.path.isdir(_default_root):
for _ver in sorted(
os.listdir(_default_root), key = _ver_key, reverse = True
):
_bin = os.path.join(_default_root, _ver, "bin")
if os.path.isdir(_bin):
candidates.append(_bin)
except OSError:
pass
for _d in candidates:
if os.path.isdir(_d):
try:
_ROCM_DLL_HANDLES.append(os.add_dll_directory(_d))
except (OSError, AttributeError):
pass
_add_rocm_dll_dirs()
del _add_rocm_dll_dirs
# ── Windows AMD ROCm: set BNB_ROCM_VERSION before any bitsandbytes import ─
# bitsandbytes on Windows ROCm tries to load libbitsandbytes_rocm<ver>.dll
# where <ver> comes from torch.version.hip (e.g. "7.13..." → "713").
# The installed BNB wheel ships rocm72.dll (not rocm713.dll), so without
# this the server process crashes with "Configured ROCm binary not found".
# Detect the available DLL, fall back to "72", and set BNB_ROCM_VERSION
# before any import that pulls in bitsandbytes (mirrors worker.py logic).
# Gate on the rocm bnb DLL (the exact file this configures) or HIP_PATH/
# ROCM_PATH, not on torch.version.hip: that needed importing torch on every
# Windows host (NVIDIA/CPU included), adding seconds to startup. Radeon
# wheels without HIP_PATH still ship the rocm bnb DLL, so they are covered.
if "BNB_ROCM_VERSION" not in os.environ:
import glob as _glob
import logging as _logging
_hip_env = bool(os.environ.get("HIP_PATH") or os.environ.get("ROCM_PATH"))
_bnb_rocm_ver = None
_found_rocm_bnb = False
try:
import importlib.util as _ilu
_bnb_spec = _ilu.find_spec("bitsandbytes")
# submodule_search_locations (not spec.origin) handles editable installs.
if _bnb_spec and _bnb_spec.submodule_search_locations:
import re as _re_bnb
_all_vers_main: list[str] = []
for _pkg_dir in _bnb_spec.submodule_search_locations:
for _dll in _glob.glob(
os.path.join(_pkg_dir, "libbitsandbytes_rocm*.dll")
):
_found_rocm_bnb = True
_km = _re_bnb.search(
r"libbitsandbytes_rocm(\d+)\.dll", os.path.basename(_dll)
)
if _km:
_all_vers_main.append(_km.group(1))
if _all_vers_main:
_bnb_rocm_ver = max(_all_vers_main, key = lambda v: int(v))
except Exception as _e:
_logging.getLogger(__name__).warning(
"Windows ROCm: BNB DLL detection failed (%s); falling back to version '72'",
_e,
)
# rocm bnb DLL present, or HIP_PATH/ROCM_PATH set (DLL unparsable -> "72").
if _found_rocm_bnb or _hip_env:
_bnb_rocm_ver_final = _bnb_rocm_ver or "72"
os.environ["BNB_ROCM_VERSION"] = _bnb_rocm_ver_final
_logging.getLogger(__name__).info(
"Windows ROCm: set BNB_ROCM_VERSION=%s (from installed BNB wheel)",
_bnb_rocm_ver_final,
)
# Ensure backend dir is on sys.path so _platform_compat is importable when
# main.py is launched directly (e.g. `uvicorn main:app`).
_backend_dir = str(_Path(__file__).parent)
if _backend_dir not in sys.path:
sys.path.insert(0, _backend_dir)
# `uvicorn main:app` bypasses run.py; seed thread caps here too.
from utils.cpu_threads import configure_cpu_threads
try:
configure_cpu_threads()
except ValueError as exc:
_raw = os.environ.get("UNSLOTH_CPU_THREADS")
raise SystemExit(
f"Error: Invalid UNSLOTH_CPU_THREADS value {_raw!r}: {exc}"
) from None
# Fix for Anaconda/conda-forge Python: seed platform._sys_version_cache before
# any library imports that trigger attrs -> rich -> structlog -> platform crash.
# See: https://github.com/python/cpython/issues/102396
@ -122,6 +237,7 @@ from routes import (
export_router,
inference_router,
inference_studio_router,
mcp_servers_router,
models_router,
providers_router,
training_history_router,
@ -181,6 +297,11 @@ def _load_desktop_owner() -> dict[str, str] | None:
_DESKTOP_OWNER = _load_desktop_owner()
# The Tauri desktop app runs the backend on the owner's own machine, so local
# stdio MCP servers are safe there. setdefault lets an explicit "0" opt out.
if _DESKTOP_OWNER:
os.environ.setdefault("UNSLOTH_STUDIO_ALLOW_STDIO_MCP", "1")
def _desktop_owner() -> dict[str, str] | None:
return _DESKTOP_OWNER
@ -316,20 +437,58 @@ from starlette.requests import Request as _StarletteRequest # noqa: E402
_CSP_SCRIPT_NONCE_HEADER = "x-internal-script-nonce"
# /content is Colab's working directory — more reliable than env vars which
# aren't always set depending on Colab runtime version.
import importlib.util as _importlib_util
_IS_COLAB = os.path.isdir("/content") and (
bool(os.environ.get("COLAB_BACKEND_URL"))
or bool(os.environ.get("COLAB_JUPYTER_IP"))
or _importlib_util.find_spec("google.colab") is not None
)
def _build_csp(script_nonce: "str | None" = None) -> str:
script_src = "script-src 'self'"
if script_nonce:
script_src += f" 'nonce-{script_nonce}'"
# In Colab the parent frame can be colab.research.google.com, a multi-level
# *.prod.colab.dev subdomain (e.g. foo.region.prod.colab.dev — note: CSP
# wildcards only match one level, so *.prod.colab.dev misses these), or a
# sandboxed null-origin output iframe. Use '*' so any ancestor is allowed;
# Colab is already a sandboxed single-user environment.
frame_ancestors = "*" if _IS_COLAB else "'none'"
# In Colab the frontend is served over the Colab reverse-proxy at an HTTPS
# *.prod.colab.dev URL. Colab's kernel communication layer and the output
# iframe scaffolding inject scripts from *.prod.colab.dev and
# *.googleusercontent.com, and make fetch/WebSocket connections to those
# same origins. Widen script-src and connect-src in Colab mode so those
# requests are not blocked. 'unsafe-inline' for scripts is still omitted;
# our own inline script uses a nonce.
if _IS_COLAB:
script_src += " https://*.prod.colab.dev https://*.googleusercontent.com"
connect_src = (
"'self' blob: data: "
"https://huggingface.co https://datasets-server.huggingface.co "
"https://*.prod.colab.dev wss://*.prod.colab.dev "
"https://*.googleusercontent.com wss://*.googleusercontent.com"
)
else:
connect_src = (
"'self' https://huggingface.co https://datasets-server.huggingface.co"
)
return (
"default-src 'self'; "
"img-src 'self' data: blob: https://t0.gstatic.com "
"https://t1.gstatic.com https://t2.gstatic.com "
"https://t3.gstatic.com https://www.google.com; "
"connect-src 'self' https://huggingface.co https://datasets-server.huggingface.co; "
f"connect-src {connect_src}; "
"style-src 'self' 'unsafe-inline'; "
f"{script_src}; "
"font-src 'self' data:; "
"frame-ancestors 'none'; "
f"frame-ancestors {frame_ancestors}; "
"form-action 'self'; "
"base-uri 'self'"
)
@ -345,7 +504,10 @@ class SecurityHeadersMiddleware(BaseHTTPMiddleware):
if nonce is not None:
del response.headers[_CSP_SCRIPT_NONCE_HEADER]
response.headers.setdefault("Content-Security-Policy", _build_csp(nonce))
response.headers.setdefault("X-Frame-Options", "DENY")
# Omit X-Frame-Options in Colab — CSP frame-ancestors handles it, and
# DENY would block serve_kernel_port_as_iframe regardless of CSP.
if not _IS_COLAB:
response.headers.setdefault("X-Frame-Options", "DENY")
response.headers.setdefault("X-Content-Type-Options", "nosniff")
response.headers.setdefault("Referrer-Policy", "no-referrer")
response.headers.setdefault(
@ -524,6 +686,7 @@ app.include_router(inference_studio_router, prefix = "/api/inference", tags = ["
# standard /v1/chat/completions path.
app.include_router(inference_router, prefix = "/v1", tags = ["openai-compat"])
app.include_router(providers_router, prefix = "/api/providers", tags = ["providers"])
app.include_router(mcp_servers_router, prefix = "/api/mcp/servers", tags = ["mcp"])
app.include_router(datasets_router, prefix = "/api/datasets", tags = ["datasets"])
app.include_router(data_recipe_router, prefix = "/api/data-recipe", tags = ["data-recipe"])
app.include_router(export_router, prefix = "/api/export", tags = ["export"])
@ -708,8 +871,6 @@ def _strip_crossorigin(html_bytes: bytes) -> bytes:
@font-face downloads to fail silently. Stripping the attribute
makes them regular same-origin fetches that work on any protocol.
"""
import re as _re
html = html_bytes.decode("utf-8")
html = _re.sub(r'\s+crossorigin(?:="[^"]*")?', "", html)
return html.encode("utf-8")

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

@ -581,6 +581,14 @@ class ChatMessage(BaseModel):
None,
description = "OpenAI tool-result messages: name of the tool whose result this is.",
)
extra_content: Optional[dict] = Field(
None,
description = (
"Provider-specific extra fields the translator may read. "
"Gemini reads `extra_content.google.thought_signature` "
"from assistant messages to replay text-part signatures."
),
)
@model_validator(mode = "after")
def _validate_role_shape(self) -> "ChatMessage":
@ -708,6 +716,10 @@ class ChatCompletionRequest(BaseModel):
"all local tools are enabled and no server-side tools are forwarded."
),
)
mcp_enabled: Optional[bool] = Field(
None,
description = "[x-unsloth] When true, append tools from every enabled MCP server to this request's tool list.",
)
auto_heal_tool_calls: Optional[bool] = Field(
True,
description = "[x-unsloth] Auto-detect and fix malformed tool calls from model output.",
@ -752,17 +764,42 @@ class ChatCompletionRequest(BaseModel):
None,
description = "[x-unsloth] Override base URL for the external provider.",
)
enable_prompt_caching: Optional[bool] = Field(
enable_prompt_caching: Optional[Union[bool, str]] = Field(
None,
description = (
"[x-unsloth] Opt in to provider-side prompt caching. On Anthropic, "
"attaches cache_control={type:ephemeral} to the system block so the "
"static prefix is reused across turns. On OpenAI cloud, caching is "
"automatic for prompts >=1024 tokens and this flag is informational. "
"Ignored for every other provider (mistral, gemini, kimi, openrouter, "
"vllm, local, etc.). Treated as enabled when omitted."
"boolean true attaches cache_control={type:ephemeral} to the system "
"block so the static prefix is reused across turns. On OpenAI cloud, "
"caching is automatic for prompts >=1024 tokens and the boolean is "
"informational. On Gemini, pass a string cache resource name such "
"as `cachedContents/abc123` to attach `cachedContent` on the native "
"request (boolean true is a no-op on Gemini because creating the "
"cache requires a separate POST /cachedContents call). Ignored for "
"every other provider. Treated as enabled when omitted."
),
)
@field_validator("enable_prompt_caching", mode = "before")
@classmethod
def _coerce_enable_prompt_caching(cls, value: Any) -> Any:
"""Preserve the pre-PR coercion: the field used to be Optional[bool],
so callers historically sent JSON strings `"true"` / `"false"` and
Pydantic v1 coerced them. Widening to Optional[Union[bool, str]] for
Gemini cache resource names lets `"false"` slip through as a truthy
string. Coerce the canonical bool literals back so explicit opt-outs
stay opt-out."""
if isinstance(value, str):
lowered = value.strip().lower()
# Match Pydantic v1's BooleanField coercion table (yes/y/on/t/1
# and no/n/off/f/0) so opt-outs that used to parse still parse.
# Anything else is preserved as a string for Gemini's
# cachedContent resource path.
if lowered in ("true", "t", "1", "yes", "y", "on"):
return True
if lowered in ("false", "f", "0", "no", "n", "off"):
return False
return value
prompt_cache_ttl: Optional[str] = Field(
None,
description = (

46
studio/backend/models/mcp_servers.py Normal file
View file

@ -0,0 +1,46 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
from typing import Optional
from pydantic import BaseModel, Field
class McpServerCreate(BaseModel):
display_name: str
url: str
headers: Optional[dict[str, str]] = None
is_enabled: bool = True
use_oauth: bool = False
class McpServerUpdate(BaseModel):
display_name: Optional[str] = None
url: Optional[str] = None
# Absent in request body = leave as-is; null = drop all headers; dict = set.
headers: Optional[dict[str, str]] = None
is_enabled: Optional[bool] = None
use_oauth: Optional[bool] = None
class McpServerResponse(BaseModel):
id: str
display_name: str
url: str
headers: dict[str, str] = Field(default_factory = dict)
is_enabled: bool = True
use_oauth: bool = False
created_at: str
updated_at: str
class McpServerTestRequest(BaseModel):
url: str
headers: Optional[dict[str, str]] = None
use_oauth: bool = False
class McpServerProbeResult(BaseModel):
ok: bool
tool_count: int = 0
error: Optional[str] = None

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

@ -5,10 +5,17 @@
Pydantic schemas for Training API
"""
import re
from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
from typing import Any, Optional, List, Dict, Literal
# ASCII integer with an optional single sign. Used by _check_vision_image_size
# to reject "++512", "--256", and Unicode-digit strings ("５１２", "٥١٢") that
# would otherwise slip through str.isdigit() + int().
_INT_RE = re.compile(r"[+-]?[0-9]+")
_MAX_BATCH_SIZE = 4096
_MAX_GRAD_ACCUM = 4096
_MAX_STEPS = 1_000_000
@ -18,6 +25,9 @@ _MAX_SEQ_LENGTH = 2_000_000
_MAX_LR_VALUE = 1.0
_MAX_LORA_R = 16_384
_MAX_LORA_ALPHA = 32_768
_MIN_VISION_IMAGE_SIZE = 256
# 2048 was the most I could get most llms to work at without getting unstable
_MAX_VISION_IMAGE_SIZE = 2048
def _parse_lr(v: Any) -> float:
@ -58,6 +68,10 @@ class TrainingStartRequest(BaseModel):
hf_token: Optional[str] = Field(None, description = "HuggingFace token")
load_in_4bit: bool = Field(True, description = "Load model in 4-bit quantization")
max_seq_length: int = Field(2048, description = "Maximum sequence length")
vision_image_size: Optional[int] = Field(
None,
description = "Optional maximum image side length for VLM training. Null uses model default.",
)
trust_remote_code: bool = Field(
False,
description = "Allow loading models with custom code (e.g. NVIDIA Nemotron). Only enable for repos you trust.",
@ -159,6 +173,40 @@ class TrainingStartRequest(BaseModel):
)
return v
@field_validator("vision_image_size", mode = "before")
@classmethod
def _check_vision_image_size(cls, v: Any) -> Optional[int]:
# mode="before" sees True/False as bool (not 1/0) for a precise error.
if v is None:
return v
if isinstance(v, bool):
raise ValueError("vision_image_size must be an integer or null")
if isinstance(v, int):
coerced = v
elif isinstance(v, str) and _INT_RE.fullmatch(v.strip()):
coerced = int(v.strip())
elif isinstance(v, float) and v.is_integer():
coerced = int(v)
else:
# numpy ints / Integral subclasses, without a hard numpy import.
try:
import numbers
if isinstance(v, numbers.Integral):
coerced = int(v)
elif isinstance(v, numbers.Real) and float(v).is_integer():
coerced = int(v)
else:
raise TypeError
except Exception:
raise ValueError("vision_image_size must be an integer or null")
if coerced < _MIN_VISION_IMAGE_SIZE or coerced > _MAX_VISION_IMAGE_SIZE:
raise ValueError(
f"vision_image_size must be in [{_MIN_VISION_IMAGE_SIZE}, "
f"{_MAX_VISION_IMAGE_SIZE}] (got {coerced!r})"
)
return coerced
@field_validator("warmup_steps")
@classmethod
def _check_warmup_steps(cls, v: Optional[int]) -> Optional[int]:

1
studio/backend/requirements/extras.txt
View file

@ -52,6 +52,5 @@ addict
easydict
einops
tabulate
fastmcp>=3.0.2
openai>=2.7.2
websockets>=15.0.1

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

@ -18,3 +18,4 @@ diceware
ddgs
cryptography>=42.0.0
httpx>=0.27.0
fastmcp>=3.0.2

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

@ -16,6 +16,7 @@ from routes.export import router as export_router
from routes.training_history import router as training_history_router
from routes.chat_history import router as chat_history_router
from routes.providers import router as providers_router
from routes.mcp_servers import router as mcp_servers_router
__all__ = [
"training_router",
@ -29,4 +30,5 @@ __all__ = [
"training_history_router",
"chat_history_router",
"providers_router",
"mcp_servers_router",
]

10
studio/backend/routes/data_recipe/mcp.py
View file

@ -36,8 +36,18 @@ def list_mcp_tools(payload: McpToolsListRequest) -> McpToolsListResponse:
providers: list[McpToolsProviderResult] = []
tool_to_providers: dict[str, list[str]] = defaultdict(list)
from core.inference.mcp_client import stdio_mcp_enabled
for provider_payload in payload.mcp_providers:
provider_name = str(provider_payload.get("name", "")).strip()
if provider_payload.get("provider_type") == "stdio" and not stdio_mcp_enabled():
providers.append(
McpToolsProviderResult(
name = provider_name,
error = "Local (stdio) MCP servers are disabled on this host.",
)
)
continue
built = build_mcp_providers({"mcp_providers": [provider_payload]})
if len(built) != 1:
providers.append(

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

@ -435,7 +435,9 @@ _TOOL_ACTION_NUDGE = (
# 4. tail-only `</parameter>` (outer close truncated by EOS); anchored to
# `\Z` so mid-text `<parameter>` in user code samples survives.
_TOOL_XML_RE = _re.compile(
r"<(?:tool_call|function=\w+)>.*?(?:</(?:tool_call|function)>|\Z)"
# Hyphen in the name char-class matches MCP tool names with dashes
# (mcp__srv__list-issues) which would otherwise leak past this strip.
r"<(?:tool_call|function=[\w-]+)>.*?(?:</(?:tool_call|function)>|\Z)"
r"|</(?:tool_call|function)>"
r"|</parameter>\s*\Z",
_re.DOTALL,
@ -1705,6 +1707,7 @@ def _build_external_messages(
messages: list,
supports_vision: bool,
provider_type: Optional[str] = None,
base_url: Optional[str] = None,
) -> list[dict]:
"""
Convert ChatMessage list to OpenAI-compatible dicts for external providers.
@ -1732,14 +1735,171 @@ def _build_external_messages(
document_provider = provider_type in _INPUT_DOCUMENT_PROVIDERS
anthropic = provider_type == "anthropic"
openai = provider_type == "openai"
# `extra_content` is a Gemini-specific carrier for the assistant's
# text-part `thoughtSignature` round-trip on the native
# streamGenerateContent endpoint. Custom Gemini OpenAI-compatible
# gateways (LiteLLM etc.) route through /chat/completions where
# the field is unknown and can be rejected -- gate strictly on the
# Google-hosted Gemini base.
_native_gemini = False
if provider_type == "gemini" and base_url:
try:
from urllib.parse import urlparse as _urlparse
_host = (_urlparse(base_url).hostname or "").lower()
_native_gemini = _host == "generativelanguage.googleapis.com"
except Exception:
_native_gemini = False
emit_extra_content = _native_gemini
_SERVER_BUILTIN_TOOL_NAMES = frozenset(
{"web_search", "web_fetch", "code_execution", "image_generation"}
)
def _is_marked_server_builtin_tool_call(tc: Any) -> bool:
"""Return True iff `tc` is a synthetic provider-side tool card
with one of the canonical builtin names and either:
- the new `args._server_tool` marker stamped by the backend, or
- a Gemini `args.google.native_part` payload (durable replay
signal for code_execution / image_generation that predates
the marker).
Such cards must not be forwarded to non-native providers
because they are not real user functions and the receiving API
will reject the orphan tool history. Real user functions with
these names normally have neither signal.
"""
if not isinstance(tc, dict):
return False
fn = tc.get("function")
if not isinstance(fn, dict):
return False
name = (fn.get("name") or "").lower()
if name not in _SERVER_BUILTIN_TOOL_NAMES:
return False
raw_args = fn.get("arguments") or ""
try:
args = json.loads(raw_args) if isinstance(raw_args, str) else raw_args
except Exception:
return False
if not isinstance(args, dict):
return False
if args.get("_server_tool") is True:
return True
google = args.get("google")
return isinstance(google, dict) and isinstance(google.get("native_part"), dict)
# When we drop a server-side builtin tool_call here, the matching
# `role="tool"` follow-up must also be dropped from the outbound
# history -- otherwise the provider receives an orphan
# tool_call_id with no matching assistant call, which OpenAI
# Responses and Anthropic both reject.
dropped_server_builtin_tool_call_ids: set[str] = set()
def _filter_tool_calls(tool_calls: Any) -> Optional[list]:
"""Sanitize assistant `tool_calls` for non-native-Gemini providers.
Two concerns:
1. `tool_calls[i].extra_content` carries Gemini-only
thoughtSignature metadata; strip it for providers that
cannot parse the unknown key.
2. Marked server-side builtin cards (`_server_tool: true` on
a canonical builtin name, or a Gemini `native_part`
payload) are provider-internal Studio tool cards from a
prior native Gemini turn; forwarding them to OpenAI /
Anthropic / custom OAI-compat gateways sends an orphan
`tool_calls` entry (no matching tool declaration, often
no matching `role="tool"` reply) that can be rejected.
We record the dropped call_ids so the matching role=tool
message is also skipped below.
Native Gemini keeps both untouched so the native translator can
replay them via `native_part`.
"""
if not tool_calls:
return None
if not isinstance(tool_calls, list):
return tool_calls
if emit_extra_content:
return tool_calls
cleaned: list = []
for _tc in tool_calls:
if _is_marked_server_builtin_tool_call(_tc):
_tc_id = _tc.get("id") if isinstance(_tc, dict) else None
if isinstance(_tc_id, str) and _tc_id:
dropped_server_builtin_tool_call_ids.add(_tc_id)
continue
if not isinstance(_tc, dict):
cleaned.append(_tc)
continue
if "extra_content" not in _tc:
cleaned.append(_tc)
continue
_stripped = {k: v for k, v in _tc.items() if k != "extra_content"}
cleaned.append(_stripped)
return cleaned
result = []
for msg in messages:
# Drop role=tool messages whose matching server-builtin
# tool_call was already filtered above. Forwarding an orphan
# tool_result with no matching tool_call would be rejected by
# OpenAI Responses and Anthropic.
if (
msg.role == "tool"
and isinstance(msg.tool_call_id, str)
and msg.tool_call_id in dropped_server_builtin_tool_call_ids
):
continue
if isinstance(msg.content, str):
# Skip assistant messages with empty content (some providers reject them)
if msg.role == "assistant" and not msg.content.strip():
# Drop bare assistant messages with no content AND no
# tool_calls (some providers reject empty assistant turns).
# Preserve assistant turns whose only payload is tool_calls
# so multi-turn function-call loops round-trip.
if (
msg.role == "assistant"
and not msg.content.strip()
and not msg.tool_calls
):
continue
result.append({"role": msg.role, "content": msg.content})
elif isinstance(msg.content, list):
out: dict[str, Any] = {"role": msg.role, "content": msg.content}
if msg.role == "assistant" and msg.tool_calls:
_tcs = _filter_tool_calls(msg.tool_calls)
if _tcs:
out["tool_calls"] = _tcs
elif not msg.content.strip():
# Every tool_call was a synthetic provider-side
# card and was dropped; the assistant turn would
# be an empty `{"role":"assistant","content":""}`
# which some providers reject. Skip it entirely.
continue
if msg.role == "tool":
if msg.tool_call_id:
out["tool_call_id"] = msg.tool_call_id
if msg.name:
out["name"] = msg.name
if emit_extra_content and msg.role == "assistant" and msg.extra_content:
out["extra_content"] = msg.extra_content
result.append(out)
continue
# Assistant messages with content=None but populated tool_calls
# are valid (post-tool-call assistant turn). Forward them so the
# provider helper can rebuild the functionCall part.
if msg.content is None and msg.role == "assistant" and msg.tool_calls:
_filtered_tcs = _filter_tool_calls(msg.tool_calls)
if not _filtered_tcs:
# Every tool_call on this turn was provider-side
# synthetic and dropped; skipping the whole message
# avoids forwarding an empty assistant turn.
continue
_assistant_only: dict[str, Any] = {
"role": "assistant",
"content": "",
"tool_calls": _filtered_tcs,
}
if emit_extra_content and msg.extra_content:
_assistant_only["extra_content"] = msg.extra_content
result.append(_assistant_only)
continue
if isinstance(msg.content, list):
if supports_vision:
parts = []
for part in msg.content:
@ -1797,9 +1957,27 @@ def _build_external_messages(
# provider would 400 on the unknown part, so
# gate by provider_type.
parts.append({"type": "compaction", "content": part.content})
if msg.role == "assistant" and not parts:
entry: dict[str, Any] = {"role": msg.role, "content": parts}
if msg.role == "assistant" and msg.tool_calls:
_tcs = _filter_tool_calls(msg.tool_calls)
if _tcs:
entry["tool_calls"] = _tcs
elif not parts:
# All tool_calls were synthetic and dropped,
# and no preserved content parts survived.
# Skip rather than forward an empty assistant
# turn that downstream providers reject.
continue
elif msg.role == "assistant" and not parts:
continue
result.append({"role": msg.role, "content": parts})
if msg.role == "tool":
if msg.tool_call_id:
entry["tool_call_id"] = msg.tool_call_id
if msg.name:
entry["name"] = msg.name
if emit_extra_content and msg.role == "assistant" and msg.extra_content:
entry["extra_content"] = msg.extra_content
result.append(entry)
else:
# Non-vision provider: strip images / documents, keep
# text, optionally keep compaction (Anthropic only --
@ -1835,9 +2013,32 @@ def _build_external_messages(
if len(preserved) == 1 and preserved[0]["type"] == "text":
# Single text part collapses back to a string for
# providers that don't accept content arrays.
result.append({"role": msg.role, "content": preserved[0]["text"]})
entry = {"role": msg.role, "content": preserved[0]["text"]}
else:
result.append({"role": msg.role, "content": preserved})
entry = {"role": msg.role, "content": preserved}
if msg.role == "assistant" and msg.tool_calls:
_tcs = _filter_tool_calls(msg.tool_calls)
if _tcs:
entry["tool_calls"] = _tcs
else:
# All tool_calls were synthetic and dropped;
# skip if there's no surviving content either.
_entry_content = entry.get("content")
_has_text = (
isinstance(_entry_content, str) and _entry_content.strip()
) or (
isinstance(_entry_content, list) and len(_entry_content) > 0
)
if not _has_text:
continue
if msg.role == "tool":
if msg.tool_call_id:
entry["tool_call_id"] = msg.tool_call_id
if msg.name:
entry["name"] = msg.name
if emit_extra_content and msg.role == "assistant" and msg.extra_content:
entry["extra_content"] = msg.extra_content
result.append(entry)
return result
@ -1912,6 +2113,7 @@ async def _proxy_to_external_provider(
payload.messages,
_supports_vision,
provider_type = provider_type,
base_url = base_url,
)
client = ExternalProviderClient(
@ -1920,6 +2122,14 @@ async def _proxy_to_external_provider(
api_key = api_key,
)
# `top_k` defaults to 20 in ChatCompletionRequest because the local
# inference path expects an int, but the external-provider path
# should treat "field omitted from JSON" as "use provider default"
# so callers that send only model/messages do not silently get
# different sampling than before this PR. Pydantic's
# `model_fields_set` tracks explicit-vs-default per request.
_top_k_explicit = payload.top_k if "top_k" in payload.model_fields_set else None
async def _stream():
gen = client.stream_chat_completion(
messages = chat_messages,
@ -1928,7 +2138,7 @@ async def _proxy_to_external_provider(
top_p = payload.top_p,
max_tokens = payload.max_tokens,
presence_penalty = payload.presence_penalty,
top_k = payload.top_k,
top_k = _top_k_explicit,
enable_thinking = payload.enable_thinking,
reasoning_effort = payload.reasoning_effort,
enabled_tools = payload.enabled_tools,
@ -1937,6 +2147,8 @@ async def _proxy_to_external_provider(
anthropic_code_exec_container_id = payload.anthropic_code_exec_container_id,
prompt_cache_ttl = payload.prompt_cache_ttl,
compaction_threshold = payload.compaction_threshold,
tools = payload.tools,
tool_choice = payload.tool_choice,
fast_mode = payload.fast_mode,
stream = payload.stream,
)
@ -2438,17 +2650,29 @@ async def openai_chat_completions(
# ── Tool-calling path (agentic loop) ──────────────────
# `_effective_enable_tools` lets `unsloth run --enable-tools/--disable-tools`
# hard-override the per-request value. Without a CLI override, falls
# back to `payload.enable_tools` (existing behavior).
# back to `payload.enable_tools` (existing behavior). `mcp_enabled=true`
# also opens the tool loop so MCP-only callers do not have to flip a
# second flag, BUT must still honor a CLI `--disable-tools` policy --
# checking the raw policy here keeps `mcp_enabled` from re-enabling
# tools that the operator explicitly forbade.
from state.tool_policy import get_tool_policy as _get_tool_policy_g
_cli_policy = _get_tool_policy_g()
_tools_on = _effective_enable_tools(payload)
_mcp_allowed = bool(payload.mcp_enabled) and _cli_policy is not False
use_tools = (
_effective_enable_tools(payload)
(_tools_on or _mcp_allowed)
and llama_backend.supports_tools
and not has_gguf_image
)
if use_tools:
from core.inference.tools import ALL_TOOLS
from core.inference.tools import ALL_TOOLS, get_enabled_mcp_tools
if payload.enabled_tools is not None:
if not _tools_on:
# MCP-only request: skip built-ins, leave room for MCP tools.
tools_to_use = []
elif payload.enabled_tools is not None:
tools_to_use = [
t
for t in ALL_TOOLS
@ -2457,6 +2681,19 @@ async def openai_chat_completions(
else:
tools_to_use = ALL_TOOLS
if _mcp_allowed:
tools_to_use = tools_to_use + await get_enabled_mcp_tools()
# Skip the tool loop when no tool actually survived, so the
# safetensors loop's "empty = allow all" semantic cannot reach
# built-in tools the caller did not opt into. Existing callers
# who omit enabled_tools still get ALL_TOOLS here, so this
# only suppresses the loop when discovery + opt-in left it
# genuinely empty.
if not tools_to_use:
use_tools = False
if use_tools:
# ── Tool-use system prompt nudge ──────────────────────
_tool_names = {t["function"]["name"] for t in tools_to_use}
_has_web = "web_search" in _tool_names
@ -2854,9 +3091,12 @@ async def openai_chat_completions(
else:
try:
full_text = ""
completion_usage = None
for token in gguf_generate():
if isinstance(token, dict):
continue # skip metadata dict in non-streaming path
if token.get("type") == "metadata":
completion_usage = token.get("usage")
continue
full_text = token
response = ChatCompletion(
@ -2869,6 +3109,15 @@ async def openai_chat_completions(
finish_reason = "stop",
)
],
usage = CompletionUsage(
prompt_tokens = (completion_usage or {}).get("prompt_tokens")
or 0,
completion_tokens = (completion_usage or {}).get(
"completion_tokens"
)
or 0,
total_tokens = (completion_usage or {}).get("total_tokens") or 0,
),
)
return JSONResponse(content = response.model_dump())
@ -2932,8 +3181,15 @@ async def openai_chat_completions(
else 25
)
# Match the GGUF path: mcp_enabled also opens the tool loop on its own
# but must still honor a CLI `--disable-tools` policy.
from state.tool_policy import get_tool_policy as _get_tool_policy_sf
_sf_cli_policy = _get_tool_policy_sf()
_sf_tools_on = _effective_enable_tools(payload)
_sf_mcp_allowed = bool(payload.mcp_enabled) and _sf_cli_policy is not False
_sf_use_tools = (
_effective_enable_tools(payload)
(_sf_tools_on or _sf_mcp_allowed)
and _sf_features.get("supports_tools", False)
and image is None
and not _sf_is_gptoss
@ -2941,15 +3197,27 @@ async def openai_chat_completions(
)
if _sf_use_tools:
from core.inference.tools import ALL_TOOLS
from core.inference.tools import ALL_TOOLS, get_enabled_mcp_tools
if payload.enabled_tools is not None:
if not _sf_tools_on:
_sf_tools_to_use = []
elif payload.enabled_tools is not None:
_sf_tools_to_use = [
t for t in ALL_TOOLS if t["function"]["name"] in payload.enabled_tools
]
else:
_sf_tools_to_use = ALL_TOOLS
if _sf_mcp_allowed:
_sf_tools_to_use = _sf_tools_to_use + await get_enabled_mcp_tools()
# Mirror the GGUF path: refuse to enter the tool loop when nothing
# survived, so a model-emitted built-in call cannot piggy-back on
# the empty allow-list.
if not _sf_tools_to_use:
_sf_use_tools = False
if _sf_use_tools:
_sf_tool_names = {t["function"]["name"] for t in _sf_tools_to_use}
_sf_has_web = "web_search" in _sf_tool_names
_sf_has_code = "python" in _sf_tool_names or "terminal" in _sf_tool_names
@ -4480,7 +4748,17 @@ async def anthropic_messages(
[m.model_dump() for m in payload.messages],
payload.system,
)
openai_messages = _drop_empty_assistant_sentinels(openai_messages)
# Strip synthetic provider-side builtin tool history (web_search,
# web_fetch, code_execution, image_generation cards tagged with
# _server_tool or extra_content.google.native_part) before handing
# off to local llama-server. The local /v1/chat/completions and
# GGUF passthrough builders apply the same strip; without it an
# Anthropic /v1/messages caller replaying a prior provider-side
# tool_use forwards fake builtin tool history to a backend that
# has no matching function declarations.
openai_messages = _strip_provider_synthetic_tool_history(
_drop_empty_assistant_sentinels(openai_messages)
)
# Enforce vision guard + re-encode embedded images to PNG so the
# Anthropic endpoint matches the behavior of /v1/chat/completions.
@ -5271,6 +5549,110 @@ def _drop_empty_assistant_sentinels(messages: list[dict]) -> list[dict]:
return out
_LOCAL_SERVER_BUILTIN_TOOL_NAMES = frozenset(
{"web_search", "web_fetch", "code_execution", "image_generation"}
)
def _strip_provider_synthetic_tool_history(messages: list[dict]) -> list[dict]:
"""Drop synthetic provider-side tool_calls + matching role=tool replies
on the local-backend (llama-server / GGUF) dispatch path.
A Gemini chat that ran code_execution / image_generation persists the
server-side tool card into thread history as an assistant tool_calls
entry tagged with ``args._server_tool`` (or a Gemini
``args.google.native_part`` payload) plus a follow-up role=tool reply.
When the user switches the SAME thread to a local GGUF model, those
synthetic tool_calls are not real user functions, llama-server has no
matching declaration, and Gemini-only ``extra_content`` /
``native_part`` payloads are meaningless. Forward only ordinary user
function calls; strip the matched role=tool replies too so the
backend does not see an orphan tool_call_id.
"""
dropped_ids: set[str] = set()
sanitized_assistant: list[dict] = []
for m in messages:
if m.get("role") != "assistant":
sanitized_assistant.append(m)
continue
tool_calls = m.get("tool_calls")
if not isinstance(tool_calls, list) or not tool_calls:
# Plain text Gemini reply: still strip message-level
# `extra_content` (carries `google.thought_signature` replay
# metadata) so a text-only Gemini turn switched to a local
# GGUF backend does not leak Gemini-only fields to
# llama-server. ChatMessage previously did not have
# `extra_content`, so the field was implicitly dropped --
# round-22 added it to ChatMessage, which is what made this
# leak possible.
if "extra_content" in m:
m = {k: v for k, v in m.items() if k != "extra_content"}
sanitized_assistant.append(m)
continue
cleaned: list[dict] = []
for tc in tool_calls:
if not isinstance(tc, dict):
cleaned.append(tc)
continue
fn = tc.get("function")
name = ""
if isinstance(fn, dict):
name = (fn.get("name") or "").lower()
if name in _LOCAL_SERVER_BUILTIN_TOOL_NAMES:
raw_args = fn.get("arguments") if isinstance(fn, dict) else None
args_obj: Any = None
if isinstance(raw_args, str):
try:
args_obj = json.loads(raw_args) if raw_args else None
except Exception:
args_obj = None
elif isinstance(raw_args, dict):
args_obj = raw_args
is_synthetic = False
if isinstance(args_obj, dict):
if args_obj.get("_server_tool") is True:
is_synthetic = True
google = args_obj.get("google")
if isinstance(google, dict) and isinstance(
google.get("native_part"), dict
):
is_synthetic = True
if is_synthetic:
tc_id = tc.get("id")
if isinstance(tc_id, str) and tc_id:
dropped_ids.add(tc_id)
continue
# Strip Gemini-only `extra_content` on real user tool_calls
# too — llama-server has no use for it and may pass it
# through to the model unchanged.
if "extra_content" in tc:
tc = {k: v for k, v in tc.items() if k != "extra_content"}
cleaned.append(tc)
# Drop top-level message-level `extra_content` (Gemini
# thoughtSignature replay metadata) on local dispatch.
m_clean = {k: v for k, v in m.items() if k != "extra_content"}
if cleaned:
m_clean["tool_calls"] = cleaned
else:
m_clean.pop("tool_calls", None)
if not m_clean.get("content") and not m_clean.get("tool_calls"):
continue # assistant turn now empty, drop
sanitized_assistant.append(m_clean)
if not dropped_ids:
return sanitized_assistant
out: list[dict] = []
for m in sanitized_assistant:
if (
m.get("role") == "tool"
and isinstance(m.get("tool_call_id"), str)
and m["tool_call_id"] in dropped_ids
):
continue
out.append(m)
return out
def _openai_messages_for_passthrough(payload) -> list[dict]:
"""Build OpenAI-format message dicts for the /v1/chat/completions
passthrough path.
@ -5287,8 +5669,10 @@ def _openai_messages_for_passthrough(payload) -> list[dict]:
``image_url`` content part so vision + function-calling requests work
transparently.
"""
messages = _drop_empty_assistant_sentinels(
[m.model_dump(exclude_none = True) for m in payload.messages]
messages = _strip_provider_synthetic_tool_history(
_drop_empty_assistant_sentinels(
[m.model_dump(exclude_none = True) for m in payload.messages]
)
)
if not payload.image_base64:
@ -5337,8 +5721,10 @@ def _openai_messages_for_gguf_chat(payload, is_vision: bool) -> tuple[list[dict]
all per-turn ``image_url`` parts so multi-image chat history keeps each
image attached to its original turn.
"""
messages = _drop_empty_assistant_sentinels(
[m.model_dump(exclude_none = True) for m in payload.messages]
messages = _strip_provider_synthetic_tool_history(
_drop_empty_assistant_sentinels(
[m.model_dump(exclude_none = True) for m in payload.messages]
)
)
has_message_image = any(
isinstance(msg.get("content"), list)

258
studio/backend/routes/mcp_servers.py Normal file
View file

@ -0,0 +1,258 @@
# 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 uuid
from urllib.parse import urlparse
import structlog
from fastapi import APIRouter, Depends, HTTPException
from auth.authentication import get_current_subject
from core.inference.mcp_client import (
clear_oauth_tokens_async,
is_stdio,
list_tools_async,
parse_server_headers,
parse_stdio_command,
probe_timeout,
stdio_mcp_enabled,
)
from models.mcp_servers import (
McpServerCreate,
McpServerProbeResult,
McpServerResponse,
McpServerTestRequest,
McpServerUpdate,
)
from storage import mcp_servers_db
logger = structlog.get_logger(__name__)
router = APIRouter()
def _validate_url(url: str) -> str:
trimmed = (url or "").strip()
if not trimmed:
raise HTTPException(status_code = 400, detail = "url must not be empty")
# When stdio is enabled on this host, a non-HTTP value is a local command.
# Reuse this field so stdio servers ride the existing CRUD/storage with no
# schema change. When stdio is disabled the value falls through to the
# http-only validation below, so non-HTTP input is just a bad URL (400).
if stdio_mcp_enabled() and is_stdio(trimmed):
try:
parts = parse_stdio_command(trimmed)
except ValueError as exc:
raise HTTPException(status_code = 400, detail = f"Invalid command: {exc}")
if not parts or not parts[0].strip():
raise HTTPException(status_code = 400, detail = "command must not be empty")
if "://" in parts[0]:
# A URL-scheme first token is a mistyped URL, not a command. Reject
# it cleanly instead of exec-ing it (mirrors the frontend check).
raise HTTPException(
status_code = 400,
detail = "Enter an http(s):// URL, or a local command whose "
"first token is an executable (not a URL).",
)
return trimmed
parsed = urlparse(trimmed)
if parsed.scheme not in ("http", "https"):
raise HTTPException(
status_code = 400,
detail = "url must start with http:// or https://",
)
if not parsed.netloc:
raise HTTPException(status_code = 400, detail = "url is missing a host")
return trimmed
def _normalize_headers(headers: dict[str, str] | None) -> dict[str, str] | None:
"""Trim header names, drop empties, coerce values to str. None if nothing left."""
if not headers:
return None
out: dict[str, str] = {}
for raw_key, value in headers.items():
key = str(raw_key).strip()
if key:
out[key] = str(value)
return out or None
def _row_to_response(row: dict) -> McpServerResponse:
return McpServerResponse(
id = row["id"],
display_name = row["display_name"],
url = row["url"],
headers = parse_server_headers(row) or {},
is_enabled = bool(row["is_enabled"]),
use_oauth = bool(row.get("use_oauth")),
created_at = row["created_at"],
updated_at = row["updated_at"],
)
@router.get("/", response_model = list[McpServerResponse])
async def list_mcp_servers(
current_subject: str = Depends(get_current_subject),
):
return [_row_to_response(row) for row in mcp_servers_db.list_servers()]
@router.post("/", response_model = McpServerResponse, status_code = 201)
async def create_mcp_server(
payload: McpServerCreate,
current_subject: str = Depends(get_current_subject),
):
display_name = (payload.display_name or "").strip()
if not display_name:
raise HTTPException(status_code = 400, detail = "display_name must not be empty")
url = _validate_url(payload.url)
headers = _normalize_headers(payload.headers)
# OAuth is HTTP-only; force it off for stdio commands so a stale flag can't
# push the probe onto the 305s OAuth timeout. Backend is the enforcer.
use_oauth = payload.use_oauth and not is_stdio(url)
server_id = uuid.uuid4().hex[:16]
mcp_servers_db.create_server(
id = server_id,
display_name = display_name,
url = url,
headers_json = json.dumps(headers) if headers else None,
is_enabled = payload.is_enabled,
use_oauth = use_oauth,
)
return _row_to_response(mcp_servers_db.get_server(server_id))
def _changes_from_payload(payload: McpServerUpdate) -> dict:
sent = payload.model_fields_set
changes: dict = {}
if "display_name" in sent:
name = (payload.display_name or "").strip()
if not name:
raise HTTPException(
status_code = 400, detail = "display_name must not be empty"
)
changes["display_name"] = name
if "url" in sent:
changes["url"] = _validate_url(payload.url or "")
if "headers" in sent:
headers = _normalize_headers(payload.headers)
changes["headers_json"] = json.dumps(headers) if headers else None
if "is_enabled" in sent:
if payload.is_enabled is None:
raise HTTPException(
status_code = 400, detail = "is_enabled must be true or false"
)
changes["is_enabled"] = payload.is_enabled
if "use_oauth" in sent:
if payload.use_oauth is None:
raise HTTPException(
status_code = 400, detail = "use_oauth must be true or false"
)
changes["use_oauth"] = payload.use_oauth
# stdio is OAuth-less: drop a stale OAuth flag when switching to a command.
if "url" in changes and is_stdio(changes["url"]):
changes["use_oauth"] = False
return changes
@router.put("/{server_id}", response_model = McpServerResponse)
async def update_mcp_server(
server_id: str,
payload: McpServerUpdate,
current_subject: str = Depends(get_current_subject),
):
old = mcp_servers_db.get_server(server_id)
if not old:
raise HTTPException(status_code = 404, detail = "MCP server not found")
changes = _changes_from_payload(payload)
if not changes:
raise HTTPException(status_code = 400, detail = "No fields to update")
# headers == HTTP headers (remote) or env vars (stdio). On a transport-type
# switch with no new headers, drop the old ones so env secrets are not
# re-sent as HTTP headers (or vice versa).
if (
"url" in changes
and is_stdio(changes["url"]) != is_stdio(old["url"])
and "headers_json" not in changes
):
changes["headers_json"] = None
# Clear persisted OAuth tokens when the URL changes or OAuth is
# disabled; fastmcp keys tokens by URL and would otherwise let a
# re-pointed server silently inherit the old account's credentials.
if bool(old.get("use_oauth")) and (
("url" in changes and changes["url"] != old["url"])
or changes.get("use_oauth") is False
):
await clear_oauth_tokens_async(old["url"])
mcp_servers_db.update_server(server_id, changes)
return _row_to_response(mcp_servers_db.get_server(server_id))
@router.delete("/{server_id}", status_code = 204)
async def delete_mcp_server(
server_id: str,
current_subject: str = Depends(get_current_subject),
):
old = mcp_servers_db.get_server(server_id)
if not old:
raise HTTPException(status_code = 404, detail = "MCP server not found")
if old.get("use_oauth"):
await clear_oauth_tokens_async(old["url"])
mcp_servers_db.delete_server(server_id)
@router.post("/{server_id}/refresh", response_model = McpServerProbeResult)
async def refresh_mcp_server_tools(
server_id: str,
current_subject: str = Depends(get_current_subject),
):
server = mcp_servers_db.get_server(server_id)
if not server:
raise HTTPException(status_code = 404, detail = "MCP server not found")
# Refresh uses the stored address, so re-check the stdio gate here too: a
# stdio row from a desktop DB must not spawn on a hosted/network host.
if is_stdio(server["url"]) and not stdio_mcp_enabled():
raise HTTPException(
status_code = 400, detail = "stdio MCP servers are disabled on this host"
)
use_oauth = bool(server.get("use_oauth"))
try:
tools = await list_tools_async(
url = server["url"],
headers = parse_server_headers(server),
timeout = probe_timeout(server["url"], use_oauth),
use_oauth = use_oauth,
)
except Exception as exc: # noqa: BLE001 — surface transport+timeout errors to UI
logger.warning("MCP refresh failed", server_id = server_id, error = str(exc))
return McpServerProbeResult(ok = False, error = str(exc))
return McpServerProbeResult(ok = True, tool_count = len(tools))
@router.post("/test", response_model = McpServerProbeResult)
async def test_mcp_server(
payload: McpServerTestRequest,
current_subject: str = Depends(get_current_subject),
):
# URL/header validation must surface as 400 like create/update so the
# frontend's create-form pre-flight gets the same error semantics as
# the actual save call. Only catch transport/timeout errors below.
url = _validate_url(payload.url)
headers = _normalize_headers(payload.headers)
try:
tools = await list_tools_async(
url = url,
headers = headers,
timeout = probe_timeout(url, payload.use_oauth),
use_oauth = payload.use_oauth,
)
except Exception as exc: # noqa: BLE001
return McpServerProbeResult(ok = False, error = str(exc))
return McpServerProbeResult(ok = True, tool_count = len(tools))

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

@ -318,22 +318,45 @@ async def list_provider_models(
try:
models = await client.list_models()
allow_prefixes = info.get("model_id_allow_prefixes")
if allow_prefixes is not None:
prefix_tuple = tuple(str(p) for p in allow_prefixes if str(p))
if prefix_tuple:
models = [m for m in models if m.get("id", "").startswith(prefix_tuple)]
allowlist = info.get("model_id_allowlist")
if allowlist is not None:
models = [m for m in models if allowlist.match(m.get("id", ""))]
deny_exact = info.get("model_id_deny_exact")
if deny_exact is not None:
deny_ids = {str(m) for m in deny_exact if str(m)}
if deny_ids:
models = [m for m in models if m.get("id", "") not in deny_ids]
denylist = info.get("model_id_denylist")
if denylist is not None:
models = [m for m in models if not denylist.search(m.get("id", ""))]
# Registry-level model-id filters are scoped to the canonical
# native Gemini base. A custom Gemini OAI-compatible proxy
# (LiteLLM, deployment gateway) returns IDs like
# `google/gemini-2.5-flash`, `gemini/gemini-2.5-flash`, or
# team-prefixed deployment aliases; the native allowlist regex
# would strip those out and leave the picker empty even though
# the chat path now routes them via the OAI-compatible
# dispatcher (the same gate ExternalProviderClient applies for
# request building). Match the host check here so the model
# list and chat dispatch agree on what counts as "native".
apply_registry_model_filters = True
if payload.provider_type == "gemini":
try:
from urllib.parse import urlparse as _urlparse
_host = (_urlparse(base_url).hostname or "").lower()
except Exception:
_host = ""
apply_registry_model_filters = _host == "generativelanguage.googleapis.com"
if apply_registry_model_filters:
allow_prefixes = info.get("model_id_allow_prefixes")
if allow_prefixes is not None:
prefix_tuple = tuple(str(p) for p in allow_prefixes if str(p))
if prefix_tuple:
models = [
m for m in models if m.get("id", "").startswith(prefix_tuple)
]
allowlist = info.get("model_id_allowlist")
if allowlist is not None:
models = [m for m in models if allowlist.match(m.get("id", ""))]
deny_exact = info.get("model_id_deny_exact")
if deny_exact is not None:
deny_ids = {str(m) for m in deny_exact if str(m)}
if deny_ids:
models = [m for m in models if m.get("id", "") not in deny_ids]
denylist = info.get("model_id_denylist")
if denylist is not None:
models = [m for m in models if not denylist.search(m.get("id", ""))]
# Apply an optional cap after filtering so registry entries with a
# large remote catalog (e.g. HF Inference Providers) can stay
# picker-sized. No popularity sort happens server-side, so this is

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

@ -194,6 +194,7 @@ async def start_training(
"hf_token": request.hf_token or "",
"load_in_4bit": request.load_in_4bit,
"max_seq_length": request.max_seq_length,
"vision_image_size": request.vision_image_size,
"hf_dataset": request.hf_dataset or "",
"local_datasets": request.local_datasets,
"local_eval_datasets": request.local_eval_datasets,

59
studio/backend/run.py
View file

@ -19,6 +19,16 @@ backend_dir = Path(__file__).parent
if str(backend_dir) not in sys.path:
sys.path.insert(0, str(backend_dir))
from utils.cpu_threads import configure_cpu_threads
try:
configure_cpu_threads()
except ValueError as exc:
configured = os.environ.get("UNSLOTH_CPU_THREADS")
raise SystemExit(
f"Error: Invalid UNSLOTH_CPU_THREADS value {configured!r}: {exc}"
) from None
# Fix for Anaconda/conda-forge Python: seed platform._sys_version_cache before
# any library imports that trigger attrs -> rich -> structlog -> platform crash.
# See: https://github.com/python/cpython/issues/102396
@ -643,7 +653,7 @@ def run_server(
from threading import Thread, Event
import uvicorn
from main import app, setup_frontend
from main import app, setup_frontend, _IS_COLAB
from utils.paths import ensure_studio_directories
# Create all standard directories on startup
@ -727,14 +737,22 @@ def run_server(
ready_event.set()
# server_header=False suppresses uvicorn's "Server: uvicorn"; SecurityHeadersMiddleware sets its own.
config = uvicorn.Config(
app,
config_kwargs = dict(
host = host,
port = port,
log_level = "info",
access_log = False,
server_header = False,
)
# Only in Colab: trust X-Forwarded-* from Colab's reverse proxy so the app
# sees the real https origin. forwarded_allow_ips="*" is fine inside Colab's
# single-user sandbox, but would be an unwanted security relaxation for a
# normal local/standalone Studio, so leave uvicorn's safe defaults
# (forwarded headers trusted from loopback only) in place there.
if _IS_COLAB:
config_kwargs["proxy_headers"] = True
config_kwargs["forwarded_allow_ips"] = "*"
config = uvicorn.Config(app, **config_kwargs)
_server = _ReadyServer(config)
_shutdown_event = Event()
@ -756,14 +774,21 @@ def run_server(
app.state.trigger_shutdown = _trigger_shutdown
# Run server in a daemon thread
# Run server in a daemon thread.
# Use an explicit new_event_loop() + run_until_complete() instead of
# asyncio.run() to avoid nest_asyncio's global patches to asyncio.run
# interfering when called from a thread while Colab/IPython already has
# a running loop on the main thread.
def _run():
loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)
try:
asyncio.run(_server.serve())
loop.run_until_complete(_server.serve())
except BaseException as exc:
startup_errors.append(exc)
startup_failed.set()
finally:
loop.close()
if not ready_event.is_set():
startup_failed.set()
@ -846,11 +871,33 @@ if __name__ == "__main__":
action = "store_true",
help = "API server only, no frontend (for Tauri)",
)
# Mirror unsloth_cli/commands/studio.py's _PARALLEL_*. Default 1
# applies only to direct backend launches; `unsloth studio run`
# always passes its own value (4) explicitly.
_PARALLEL_MIN = 1
_PARALLEL_MAX = 64
_PARALLEL_DEFAULT_PLAIN = 1
parser.add_argument(
"--parallel",
"--n-parallel",
type = int,
default = _PARALLEL_DEFAULT_PLAIN,
help = (
f"llama-server parallel decode slots ({_PARALLEL_MIN}..{_PARALLEL_MAX}). "
f"Default {_PARALLEL_DEFAULT_PLAIN}; `unsloth studio run` uses 4."
),
)
args = parser.parse_args()
if not _PARALLEL_MIN <= args.parallel <= _PARALLEL_MAX:
parser.error(f"--parallel must be between {_PARALLEL_MIN} and {_PARALLEL_MAX}")
kwargs = dict(
host = args.host, port = args.port, silent = args.silent, api_only = args.api_only
host = args.host,
port = args.port,
silent = args.silent,
api_only = args.api_only,
llama_parallel_slots = args.parallel,
)
if args.frontend is not None:
kwargs["frontend_path"] = Path(args.frontend)

142
studio/backend/storage/mcp_servers_db.py Normal file
View file

@ -0,0 +1,142 @@
# 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 sqlite3
import threading
from datetime import datetime, timezone
from typing import Optional
from utils.paths import studio_db_path, ensure_dir
_schema_lock = threading.Lock()
_schema_ready = False
def _ensure_schema(conn: sqlite3.Connection) -> None:
conn.execute("PRAGMA journal_mode=WAL")
conn.execute(
"""
CREATE TABLE IF NOT EXISTS mcp_servers (
id TEXT NOT NULL PRIMARY KEY,
display_name TEXT NOT NULL,
url TEXT NOT NULL,
headers_json TEXT,
is_enabled INTEGER NOT NULL DEFAULT 1,
use_oauth INTEGER NOT NULL DEFAULT 0,
created_at TEXT NOT NULL,
updated_at TEXT NOT NULL
)
"""
)
# use_oauth was added after the first release; backfill for pre-existing DBs.
cols = {
r["name"] for r in conn.execute("PRAGMA table_info(mcp_servers)").fetchall()
}
if "use_oauth" not in cols:
conn.execute(
"ALTER TABLE mcp_servers ADD COLUMN use_oauth INTEGER NOT NULL DEFAULT 0"
)
def get_connection() -> sqlite3.Connection:
global _schema_ready
db_path = studio_db_path()
ensure_dir(db_path.parent)
conn = sqlite3.connect(str(db_path))
conn.row_factory = sqlite3.Row
if not _schema_ready:
with _schema_lock:
if not _schema_ready:
try:
_ensure_schema(conn)
_schema_ready = True
except Exception:
conn.close()
raise
return conn
def create_server(
id: str,
display_name: str,
url: str,
headers_json: Optional[str] = None,
is_enabled: bool = True,
use_oauth: bool = False,
) -> None:
now = datetime.now(timezone.utc).isoformat()
conn = get_connection()
try:
conn.execute(
"""
INSERT INTO mcp_servers
(id, display_name, url, headers_json,
is_enabled, use_oauth, created_at, updated_at)
VALUES (?, ?, ?, ?, ?, ?, ?, ?)
""",
(
id,
display_name,
url,
headers_json,
int(is_enabled),
int(use_oauth),
now,
now,
),
)
conn.commit()
finally:
conn.close()
def update_server(id: str, changes: dict) -> bool:
"""Apply column updates and bump ``updated_at``. Returns True on a hit."""
if not changes:
return False
bool_cols = {"is_enabled", "use_oauth"}
sets, params = [], []
for col, value in changes.items():
sets.append(f"{col} = ?")
params.append(int(value) if col in bool_cols else value)
sets.append("updated_at = ?")
params.extend([datetime.now(timezone.utc).isoformat(), id])
conn = get_connection()
try:
cursor = conn.execute(
f"UPDATE mcp_servers SET {', '.join(sets)} WHERE id = ?",
params,
)
conn.commit()
return cursor.rowcount > 0
finally:
conn.close()
def delete_server(id: str) -> bool:
conn = get_connection()
try:
cursor = conn.execute("DELETE FROM mcp_servers WHERE id = ?", (id,))
conn.commit()
return cursor.rowcount > 0
finally:
conn.close()
def get_server(id: str) -> Optional[dict]:
conn = get_connection()
try:
row = conn.execute("SELECT * FROM mcp_servers WHERE id = ?", (id,)).fetchone()
return dict(row) if row else None
finally:
conn.close()
def list_servers() -> list[dict]:
conn = get_connection()
try:
rows = conn.execute("SELECT * FROM mcp_servers ORDER BY created_at").fetchall()
return [dict(row) for row in rows]
finally:
conn.close()

52
studio/backend/tests/test_amd_apu_unified_memory.py Normal file
View file

@ -0,0 +1,52 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""GGML_CUDA_ENABLE_UNIFIED_MEMORY must be set only for AMD unified-memory APUs
(gfx1150/gfx1151), never for discrete AMD, NVIDIA, CPU or macOS."""
from __future__ import annotations
import sys
import types
import pytest
from core.inference.llama_cpp import LlamaCppBackend
def _fake_torch(hip, archs, *, cuda_ok = True):
t = types.ModuleType("torch")
t.version = types.SimpleNamespace(hip = hip)
t.cuda = types.SimpleNamespace(
is_available = lambda: cuda_ok,
device_count = lambda: len(archs),
get_device_properties = lambda i: types.SimpleNamespace(gcnArchName = archs[i]),
)
return t
@pytest.mark.parametrize(
"hip,archs,expected",
[
("6.2.0", ["gfx1151:xnack-"], True), # Strix Halo APU (suffix stripped)
("6.2.0", ["gfx1150"], True), # Strix Point APU
("6.2.0", ["gfx1100"], False), # discrete RDNA3
("6.2.0", ["gfx1201"], False), # discrete RDNA4
("6.2.0", ["gfx942"], False), # MI300X (data center)
(None, ["sm_90"], False), # NVIDIA (no torch.version.hip)
("6.2.0", ["gfx1100", "gfx1151"], True), # mixed dGPU + APU
],
)
def test_apu_unified_memory_gating(monkeypatch, hip, archs, expected):
monkeypatch.setitem(sys.modules, "torch", _fake_torch(hip, archs))
assert LlamaCppBackend._amd_apu_wants_unified_memory() is expected
def test_cpu_no_cuda_returns_false(monkeypatch):
monkeypatch.setitem(sys.modules, "torch", _fake_torch("6.2.0", [], cuda_ok = False))
assert LlamaCppBackend._amd_apu_wants_unified_memory() is False
def test_missing_torch_returns_false(monkeypatch):
monkeypatch.setitem(sys.modules, "torch", None)
assert LlamaCppBackend._amd_apu_wants_unified_memory() is False

8
studio/backend/tests/test_anthropic_code_execution.py
View file

@ -275,7 +275,13 @@ def test_bash_code_execution_emits_tool_start_and_end(monkeypatch):
assert start["type"] == "tool_start"
assert start["tool_name"] == "code_execution"
assert start["tool_call_id"] == "srvtoolu_1"
assert start["arguments"] == {"kind": "bash", "command": "ls -la"}
# `_server_tool: True` marks this as a provider-side synthetic
# tool card for the frontend's history serializer.
assert start["arguments"] == {
"kind": "bash",
"command": "ls -la",
"_server_tool": True,
}
assert end["type"] == "tool_end"
assert end["tool_call_id"] == "srvtoolu_1"

7
studio/backend/tests/test_anthropic_web_fetch.py
View file

@ -271,7 +271,12 @@ def test_web_fetch_success_emits_tool_start_and_end(monkeypatch):
assert start["type"] == "tool_start"
assert start["tool_name"] == "web_fetch"
assert start["tool_call_id"] == "srvtoolu_wf1"
assert start["arguments"] == {"url": "https://example.com/article"}
# `_server_tool: True` marks this as a provider-side synthetic
# tool card for the frontend's history serializer.
assert start["arguments"] == {
"url": "https://example.com/article",
"_server_tool": True,
}
assert end["type"] == "tool_end"
assert end["tool_call_id"] == "srvtoolu_wf1"
# The source pill uses Title / URL / snippet as parseSourcesFromResult expects.

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

@ -0,0 +1,154 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Tests for Studio's early CPU thread-pool configuration."""
import ast
import os
import subprocess
import sys
from pathlib import Path
import pytest
from utils.cpu_threads import _THREAD_POOL_ENV_VARS, configure_cpu_threads
_BACKEND_DIR = Path(__file__).resolve().parent.parent
_RUN_PY = _BACKEND_DIR / "run.py"
_MAIN_PY = _BACKEND_DIR / "main.py"
# Explicit positive integers seed all four native pool env vars.
def test_cpu_thread_cap_seeds_native_pool_limits():
env = {"UNSLOTH_CPU_THREADS": " 6 "}
configure_cpu_threads(env)
assert {variable: env[variable] for variable in _THREAD_POOL_ENV_VARS} == {
variable: "6" for variable in _THREAD_POOL_ENV_VARS
}
# Explicit per-library values win over the Studio knob via setdefault.
def test_cpu_thread_cap_preserves_runtime_specific_override():
env = {"UNSLOTH_CPU_THREADS": "4", "OMP_NUM_THREADS": "2"}
configure_cpu_threads(env)
assert env["OMP_NUM_THREADS"] == "2"
assert env["MKL_NUM_THREADS"] == "4"
# Whitespace / plus-prefix / leading zero all normalise via int().
@pytest.mark.parametrize("raw", ["+4", "007", " 4 "])
def test_cpu_thread_cap_normalises_valid_inputs(raw):
env = {"UNSLOTH_CPU_THREADS": raw}
configure_cpu_threads(env)
assert env["OMP_NUM_THREADS"] == str(int(raw.strip()))
# Unset / empty / whitespace -> no env mutation (pure opt-in).
@pytest.mark.parametrize("raw", [None, "", " ", "\t"])
def test_cpu_thread_cap_is_opt_in(raw):
env = {} if raw is None else {"UNSLOTH_CPU_THREADS": raw}
snapshot = dict(env)
configure_cpu_threads(env)
assert env == snapshot
assert all(variable not in env for variable in _THREAD_POOL_ENV_VARS)
# Anything that is not a positive integer raises a clear ValueError.
@pytest.mark.parametrize(
"raw", ["zero", "0", "-3", "1.5", "abc", "8a", "0x4", "1e3", "4 0"]
)
def test_cpu_thread_cap_requires_positive_integer(raw):
with pytest.raises(ValueError, match = "must be a positive integer"):
configure_cpu_threads({"UNSLOTH_CPU_THREADS": raw})
# env=None path uses real os.environ (production call from run.py / main.py).
def test_cpu_thread_cap_uses_os_environ_when_env_is_none(monkeypatch):
for variable in (*_THREAD_POOL_ENV_VARS, "UNSLOTH_CPU_THREADS"):
monkeypatch.delenv(variable, raising = False)
monkeypatch.setenv("UNSLOTH_CPU_THREADS", "3")
configure_cpu_threads()
for variable in _THREAD_POOL_ENV_VARS:
assert os.environ[variable] == "3"
# Calling twice must not flip any seeded value.
def test_cpu_thread_cap_idempotent(monkeypatch):
for variable in (*_THREAD_POOL_ENV_VARS, "UNSLOTH_CPU_THREADS"):
monkeypatch.delenv(variable, raising = False)
monkeypatch.setenv("UNSLOTH_CPU_THREADS", "5")
configure_cpu_threads()
snapshot = {v: os.environ.get(v) for v in _THREAD_POOL_ENV_VARS}
configure_cpu_threads()
assert {v: os.environ.get(v) for v in _THREAD_POOL_ENV_VARS} == snapshot
def _ast_line_of_configure_call(source: str) -> int:
tree = ast.parse(source)
for node in ast.walk(tree):
if (
isinstance(node, ast.Call)
and isinstance(node.func, ast.Name)
and node.func.id == "configure_cpu_threads"
):
return node.lineno
raise AssertionError("configure_cpu_threads() call not found")
def _ast_line_of_platform_compat_import(source: str) -> int:
tree = ast.parse(source)
for node in ast.walk(tree):
if isinstance(node, ast.Import):
for alias in node.names:
if alias.name == "_platform_compat":
return node.lineno
raise AssertionError("_platform_compat import not found")
# AST-based ordering: configure_cpu_threads() must precede _platform_compat
# in both run.py and main.py. Robust to formatting / line shifts.
@pytest.mark.parametrize("entry_point", [_RUN_PY, _MAIN_PY])
def test_cpu_thread_configuration_runs_before_backend_imports(entry_point):
source = entry_point.read_text()
call_line = _ast_line_of_configure_call(source)
compat_line = _ast_line_of_platform_compat_import(source)
assert call_line < compat_line, (
f"{entry_point.name}: configure_cpu_threads() (line {call_line}) "
f"must precede import _platform_compat (line {compat_line})"
)
# Invalid env -> exit 1, one-line stderr, no traceback, gated before any
# heavy import. Parametrised over both entry points.
@pytest.mark.parametrize("entry_point", [_RUN_PY, _MAIN_PY])
def test_invalid_cpu_thread_cap_exits_without_traceback(entry_point):
env = os.environ.copy()
env["UNSLOTH_CPU_THREADS"] = "not-a-count"
result = subprocess.run(
[sys.executable, str(entry_point)],
env = env,
capture_output = True,
text = True,
)
assert result.returncode == 1
assert (
"Error: Invalid UNSLOTH_CPU_THREADS value 'not-a-count': "
"UNSLOTH_CPU_THREADS must be a positive integer"
) in result.stderr
assert "Traceback" not in result.stderr
assert "_platform_compat" not in result.stderr

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

@ -437,6 +437,7 @@ def test_health_response_reports_desktop_capability_fields(monkeypatch):
export_router = APIRouter(),
inference_router = APIRouter(),
inference_studio_router = APIRouter(),
mcp_servers_router = APIRouter(),
models_router = APIRouter(),
providers_router = APIRouter(),
training_history_router = APIRouter(),
@ -484,11 +485,14 @@ from typer.testing import CliRunner
studio_home = Path(sys.argv[1])
real_import = builtins.__import__
def guarded_import(name, *args, **kwargs):
def guarded_import(name, globals = None, locals = None, fromlist = (), level = 0):
# Only gate absolute imports; relative `from .utils import x` inside
# third-party packages (e.g. typer._click.decorators) hits level > 0
# with name="utils" and must pass through.
blocked = ("auth", "fastapi", "structlog", "utils")
if name in blocked or name.startswith(("auth.", "utils.")):
if level == 0 and (name in blocked or name.startswith(("auth.", "utils."))):
raise ModuleNotFoundError(name)
return real_import(name, *args, **kwargs)
return real_import(name, globals, locals, fromlist, level)
builtins.__import__ = guarded_import
from unsloth_cli.commands import studio as studio_cli

5501
studio/backend/tests/test_gemini_provider.py Normal file
View file

File diff suppressed because it is too large Load diff

78
studio/backend/tests/test_gguf_completion_usage.py Normal file
View file

@ -0,0 +1,78 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
"""Regression tests for GGUF non-streaming chat completion usage."""
from fastapi import FastAPI
from fastapi.testclient import TestClient
from auth.authentication import get_current_subject
import routes.inference as inference_route
class _GgufBackend:
is_loaded = True
model_identifier = "test/model.gguf"
_is_audio = False
is_vision = False
supports_tools = False
def __init__(self, usage):
self.usage = usage
def generate_chat_completion(self, **kwargs):
yield "answer"
yield {
"type": "metadata",
"usage": self.usage,
"timings": {"prompt_n": 23, "predicted_n": 1283},
}
def _request_completion(monkeypatch, usage):
monkeypatch.setattr(
inference_route, "get_llama_cpp_backend", lambda: _GgufBackend(usage)
)
monkeypatch.setattr(
inference_route, "_effective_enable_tools", lambda payload: False
)
app = FastAPI()
app.include_router(inference_route.router)
app.dependency_overrides[get_current_subject] = lambda: "test-user"
return TestClient(app).post(
"/chat/completions",
json = {
"messages": [{"role": "user", "content": "Why is the sky blue?"}],
"stream": False,
},
)
def test_non_streaming_gguf_completion_includes_generated_usage(monkeypatch):
response = _request_completion(
monkeypatch,
{"prompt_tokens": 23, "completion_tokens": 1283, "total_tokens": 1306},
)
assert response.status_code == 200
assert response.json()["usage"] == {
"prompt_tokens": 23,
"completion_tokens": 1283,
"total_tokens": 1306,
}
def test_non_streaming_gguf_completion_defaults_nullable_usage_to_zero(monkeypatch):
response = _request_completion(
monkeypatch,
{"prompt_tokens": None, "completion_tokens": 1283, "total_tokens": None},
)
assert response.status_code == 200
assert response.json()["usage"] == {
"prompt_tokens": 0,
"completion_tokens": 1283,
"total_tokens": 0,
}

102
studio/backend/tests/test_gguf_routing.py Normal file
View file

@ -0,0 +1,102 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""
Tests for GGUF routing in detect_gguf_model.
Regression test for the bug where a .gguf file temporarily appears
inaccessible on Windows during llama-server process teardown, causing
is_file() to return False and the model to be routed to the transformers
backend instead of llama-server.
"""
import sys
import os
import types
from pathlib import Path
from unittest.mock import patch
# Stub structlog before importing backend modules (mirrors other tests in this suite)
if "structlog" not in sys.modules:
class _DummyLogger:
def __getattr__(self, _):
return lambda *a, **k: None
sys.modules["structlog"] = types.SimpleNamespace(
get_logger = lambda *a, **k: _DummyLogger(),
BoundLogger = _DummyLogger,
)
sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
from utils.models.model_config import detect_gguf_model
def test_detects_gguf_file_normally(tmp_path):
"""Normal case: .gguf file exists and is accessible."""
gguf = tmp_path / "gpt-oss-20b-MXFP4.gguf"
gguf.write_bytes(b"")
result = detect_gguf_model(str(gguf))
assert result is not None
assert result.endswith("gpt-oss-20b-MXFP4.gguf")
def test_detects_gguf_when_stat_raises_oserror(tmp_path):
"""
Regression: on Windows, both is_file() and exists() call stat() internally.
During the brief lock window after llama-server is killed, stat() raises
OSError, causing both to return False. detect_gguf_model must still route
to llama-server based on the file extension alone.
"""
gguf = tmp_path / "gpt-oss-20b-MXFP4.gguf"
gguf.write_bytes(b"")
original_stat = Path.stat
def flaky_stat(self, **kwargs):
if self.suffix.lower() == ".gguf":
raise OSError("file temporarily inaccessible (Windows lock window)")
return original_stat(self, **kwargs)
with patch.object(Path, "stat", flaky_stat):
result = detect_gguf_model(str(gguf))
assert result is not None, (
"detect_gguf_model returned None when stat() raised OSError. "
"This causes the model to fall through to the transformers backend."
)
def test_does_not_detect_mmproj_as_main_model(tmp_path):
"""mmproj files must never be returned as the primary model."""
mmproj = tmp_path / "mmproj-model-f16.gguf"
mmproj.write_bytes(b"")
result = detect_gguf_model(str(mmproj))
assert result is None
def test_detects_gguf_in_directory(tmp_path):
"""Directory containing a .gguf file is resolved to that file."""
gguf = tmp_path / "model-Q4_K_M.gguf"
gguf.write_bytes(b"")
result = detect_gguf_model(str(tmp_path))
assert result is not None
assert result.endswith("model-Q4_K_M.gguf")
def test_directory_named_like_gguf_scans_inside(tmp_path):
"""A directory named *.gguf resolves the real .gguf inside, not itself."""
gguf_dir = tmp_path / "mymodel.gguf"
gguf_dir.mkdir()
inner = gguf_dir / "model-Q4_K_M.gguf"
inner.write_bytes(b"")
result = detect_gguf_model(str(gguf_dir))
assert result is not None
assert result.endswith("model-Q4_K_M.gguf")
def test_returns_none_for_non_gguf_path(tmp_path):
"""Non-.gguf paths with no .gguf files inside return None."""
result = detect_gguf_model(str(tmp_path))
assert result is None

430
studio/backend/tests/test_lemonade_llamacpp_rocm_bins_mock.py Normal file
View file

@ -0,0 +1,430 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Validates that the installer correctly resolves lemonade ROCm prebuilt assets.
Uses a faked HostInfo so no AMD GPU is needed. Network calls to the lemonade
GitHub API are stubbed out so the suite runs without internet access and is
not subject to rate limits.
"""
from __future__ import annotations
import importlib
import sys
from pathlib import Path
from unittest.mock import patch
import pytest
_studio = Path(__file__).resolve().parent.parent.parent
if str(_studio) not in sys.path:
sys.path.insert(0, str(_studio))
_mod = importlib.import_module("install_llama_prebuilt")
HostInfo = _mod.HostInfo
resolve_lemonade_rocm_choice = getattr(_mod, "resolve_lemonade_rocm_choice", None)
_LEMONADE_GFX_FAMILIES = getattr(_mod, "_LEMONADE_GFX_FAMILIES", None)
if resolve_lemonade_rocm_choice is None or _LEMONADE_GFX_FAMILIES is None:
pytest.skip("PR symbols not present - check branch", allow_module_level = True)
@pytest.fixture(autouse = True)
def _clear_lemonade_release_cache():
"""Prevent cross-test pollution of the lemonade release lru_cache when
future tests vary the fetch_json mock return value."""
_cache = getattr(_mod, "_fetch_lemonade_release_cached", None)
if _cache is not None and hasattr(_cache, "cache_clear"):
_cache.cache_clear()
yield
if _cache is not None and hasattr(_cache, "cache_clear"):
_cache.cache_clear()
_STUB_TAG = "b1262"
_STUB_OS_PREFIXES = ("ubuntu", "windows")
_STUB_FAMILIES = ("gfx1151", "gfx1150", "gfx120X", "gfx110X", "gfx103X")
def _stub_lemonade_release() -> dict:
"""Minimal lemonade release payload covering all supported GPU/OS combinations."""
assets = [
{
"name": f"llama-{_STUB_TAG}-{prefix}-rocm-{family}-x64.zip",
"browser_download_url": (
f"https://github.com/lemonade-sdk/llamacpp-rocm/releases/download/"
f"{_STUB_TAG}/llama-{_STUB_TAG}-{prefix}-rocm-{family}-x64.zip"
),
}
for prefix in _STUB_OS_PREFIXES
for family in _STUB_FAMILIES
]
return {"tag_name": _STUB_TAG, "assets": assets}
def _make_rocm_host(gfx_target: str, *, windows: bool = False) -> HostInfo:
return HostInfo(
system = "Windows" if windows else "Linux",
machine = "amd64" if windows else "x86_64",
is_windows = windows,
is_linux = not windows,
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,
has_rocm = True,
rocm_gfx_target = gfx_target,
)
def _lookup_family(gfx: str) -> str | None:
for prefix, family in _LEMONADE_GFX_FAMILIES:
if gfx.startswith(prefix):
return family
return None
# ---------------------------------------------------------------------------
# GPU family mapping
# ---------------------------------------------------------------------------
@pytest.mark.parametrize(
"gfx,expected_family",
[
("gfx1151", "gfx1151"),
("gfx1150", "gfx1150"),
("gfx1201", "gfx120X"),
("gfx1200", "gfx120X"),
("gfx1100", "gfx110X"),
("gfx1030", "gfx103X"),
],
)
def test_gpu_family_mapping(gfx, expected_family):
assert _lookup_family(gfx) == expected_family
def test_unknown_gpu_not_in_families():
assert _lookup_family("gfx999") is None
# ---------------------------------------------------------------------------
# Asset resolution - hits real lemonade GitHub API
# ---------------------------------------------------------------------------
@pytest.mark.parametrize(
"gfx,os_prefix,windows",
[
("gfx1151", "ubuntu", False),
("gfx1150", "ubuntu", False),
("gfx1201", "ubuntu", False),
("gfx1100", "ubuntu", False),
("gfx1030", "ubuntu", False),
("gfx1151", "windows", True),
("gfx1100", "windows", True),
],
)
def test_asset_resolves_for_known_gpu(gfx, os_prefix, windows):
host = _make_rocm_host(gfx, windows = windows)
with patch.object(_mod, "fetch_json", return_value = _stub_lemonade_release()):
result = resolve_lemonade_rocm_choice(
host, os_prefix, "default", llama_tag = "latest"
)
assert (
result is not None
), f"Installer will NOT fetch lemonade binary for {gfx} ({os_prefix})"
assert _lookup_family(gfx) in result.name
assert result.url.startswith("https://github.com/lemonade-sdk/llamacpp-rocm")
def test_unknown_gpu_falls_through_to_upstream():
host = _make_rocm_host("gfx999")
result = resolve_lemonade_rocm_choice(host, "ubuntu", "default", llama_tag = "latest")
assert result is None
# ---------------------------------------------------------------------------
# Simple-policy dispatcher must plan a lemonade ROCm attempt for AMD-only hosts.
# This is the path setup.sh actually invokes (via --simple-policy), so the
# lemonade integration is useless if it isn't wired in here.
# ---------------------------------------------------------------------------
direct_linux_release_plan = getattr(_mod, "direct_linux_release_plan", None)
direct_upstream_release_plan = getattr(_mod, "direct_upstream_release_plan", None)
def _stub_unsloth_release(release_tag: str = "b9022") -> dict:
# Minimal payload that parse_direct_linux_release_bundle accepts. It
# requires at least one `app-{label}-linux-x64*.tar.gz` asset for the
# bundle to be recognised; we ship a bare CPU one so the planner has a
# baseline non-ROCm attempt to fall through to.
asset_name = f"app-{release_tag}-linux-x64.tar.gz"
return {
"tag_name": release_tag,
"name": release_tag,
"assets": [
{
"name": asset_name,
"browser_download_url": f"https://example.invalid/{asset_name}",
},
],
}
@pytest.mark.skipif(
direct_linux_release_plan is None,
reason = "simple-policy dispatcher not present on this branch",
)
def test_simple_policy_plans_lemonade_for_rocm_host():
host = _make_rocm_host("gfx1151")
with patch.object(_mod, "fetch_json", return_value = _stub_lemonade_release()):
plan = direct_linux_release_plan(
_stub_unsloth_release(),
host,
"unslothai/llama.cpp",
"latest",
)
assert plan is not None, "ROCm host should not be skipped by simple-policy planner"
kinds = [a.install_kind for a in plan.attempts]
assert (
"linux-rocm" in kinds
), f"simple-policy planner did not include a lemonade ROCm attempt; got {kinds}"
rocm_attempt = next(a for a in plan.attempts if a.install_kind == "linux-rocm")
assert rocm_attempt.source_label == "lemonade"
assert "gfx1151" in rocm_attempt.name
@pytest.mark.skipif(
direct_upstream_release_plan is None,
reason = "simple-policy dispatcher not present on this branch",
)
def test_simple_policy_plans_lemonade_for_windows_hip_host():
host = _make_rocm_host("gfx1151", windows = True)
release = {
"tag_name": "b9022",
"name": "b9022",
"assets": [],
}
with patch.object(_mod, "fetch_json", return_value = _stub_lemonade_release()):
plan = direct_upstream_release_plan(
release, host, "ggml-org/llama.cpp", "latest"
)
assert plan is not None, "Windows ROCm host should plan a lemonade HIP attempt"
kinds = [a.install_kind for a in plan.attempts]
assert (
"windows-hip" in kinds
), f"simple-policy planner did not include a lemonade HIP attempt; got {kinds}"
@pytest.mark.skipif(
direct_upstream_release_plan is None,
reason = "simple-policy dispatcher not present on this branch",
)
def test_simple_policy_windows_hip_falls_back_to_upstream_when_lemonade_unavailable():
"""If lemonade returns None (e.g. gfx999 or transient API failure), the planner
must still include the upstream HIP asset rather than silently downgrading to CPU."""
host = _make_rocm_host("gfx999", windows = True)
hip_asset = "llama-b9022-bin-win-hip-radeon-x64.zip"
release = {
"tag_name": "b9022",
"name": "b9022",
"assets": [
{
"name": hip_asset,
"browser_download_url": f"https://example.invalid/{hip_asset}",
},
],
}
plan = direct_upstream_release_plan(release, host, "ggml-org/llama.cpp", "latest")
assert plan is not None
kinds = [a.install_kind for a in plan.attempts]
assert (
"windows-hip" in kinds
), f"upstream HIP asset not included as fallback; got {kinds}"
hip_attempt = next(a for a in plan.attempts if a.install_kind == "windows-hip")
assert hip_attempt.source_label == "upstream"
# ── Follow-up: pinned-tag URL helper, URL trust pinning, opt-out env, autouse cache clear ──
def test_lemonade_release_api_url_pinned_tag():
"""A pinned llama_tag must produce the /releases/tags/<tag> URL."""
assert _mod._lemonade_release_api_for("b1262").endswith("/releases/tags/b1262")
assert _mod._lemonade_release_api_for("latest").endswith("/releases/latest")
assert _mod._lemonade_release_api_for("").endswith("/releases/latest")
def test_lemonade_release_api_url_encodes_tag():
"""Unexpected slashes / hashes in the tag must be URL-encoded so the URL
cannot be reshaped (defence in depth -- tags should already be sanitised
upstream)."""
url = _mod._lemonade_release_api_for("b1260/../latest")
assert "/releases/tags/b1260%2F..%2Flatest" in url
assert "//latest" not in url.split("/releases/tags/", 1)[1]
def test_lemonade_resolver_skipped_by_opt_out_env(monkeypatch):
"""UNSLOTH_DISABLE_LEMONADE_ROCM=1 must short-circuit the resolver."""
monkeypatch.setenv("UNSLOTH_DISABLE_LEMONADE_ROCM", "1")
host = _make_rocm_host("gfx1151")
res = resolve_lemonade_rocm_choice(host, "ubuntu", "linux-rocm", llama_tag = "latest")
assert res is None
def test_lemonade_resolver_rejects_non_github_url(monkeypatch):
"""If the GitHub API response somehow contained an off-host download URL,
the resolver must refuse to use it (lemonade assets are not in the
approved-hash manifest)."""
bad_release = {
"tag_name": _STUB_TAG,
"assets": [
{
"name": f"llama-{_STUB_TAG}-ubuntu-rocm-gfx1151-x64.zip",
"browser_download_url": "https://attacker.invalid/llama.zip",
},
],
}
host = _make_rocm_host("gfx1151")
with patch.object(_mod, "fetch_json", return_value = bad_release):
res = resolve_lemonade_rocm_choice(
host, "ubuntu", "linux-rocm", llama_tag = "latest"
)
assert res is None
def test_lemonade_resolver_rejects_http_scheme():
assert not _mod._is_trusted_github_release_url(
"http://github.com/lemonade-sdk/llamacpp-rocm/releases/download/x/y.zip",
"lemonade-sdk/llamacpp-rocm",
)
def test_lemonade_resolver_accepts_github_cdn():
# Real GitHub release CDN URLs carry the /github-production-release-asset- prefix.
assert _mod._is_trusted_github_release_url(
"https://objects.githubusercontent.com/github-production-release-asset-abc123/456/789?token=x",
"lemonade-sdk/llamacpp-rocm",
)
def test_lemonade_resolver_rejects_arbitrary_cdn_path():
# A CDN URL without the release-asset path prefix must be rejected.
assert not _mod._is_trusted_github_release_url(
"https://objects.githubusercontent.com/abc/def",
"lemonade-sdk/llamacpp-rocm",
)
def test_lemonade_resolver_accepts_release_path():
url = "https://github.com/lemonade-sdk/llamacpp-rocm/releases/download/b1262/llama-b1262-ubuntu-rocm-gfx1151-x64.zip"
assert _mod._is_trusted_github_release_url(url, "lemonade-sdk/llamacpp-rocm")
def test_lemonade_resolver_rejects_wrong_repo():
"""A github.com release URL for a different repo must be rejected."""
assert not _mod._is_trusted_github_release_url(
"https://github.com/attacker/llamacpp-rocm/releases/download/x/y.zip",
"lemonade-sdk/llamacpp-rocm",
)
def test_lemonade_resolver_rejects_empty_browser_download_url():
"""An asset entry with an empty browser_download_url must fall through."""
release = {
"tag_name": _STUB_TAG,
"assets": [
{
"name": f"llama-{_STUB_TAG}-ubuntu-rocm-gfx1151-x64.zip",
"browser_download_url": "",
},
],
}
host = _make_rocm_host("gfx1151")
with patch.object(_mod, "fetch_json", return_value = release):
res = resolve_lemonade_rocm_choice(
host, "ubuntu", "linux-rocm", llama_tag = "latest"
)
assert res is None
def test_lemonade_runtime_patterns_include_hip_runtime():
"""linux-rocm overlay must use a broad lib glob to catch all bundled .so files.
Lemonade ZIPs carry transitive deps (libamd_comgr, libLLVM, libclang-cpp,
...) whose names change across ROCm releases. A broad ``lib*.so*`` glob
avoids having to enumerate every transitive dependency by name.
"""
from install_llama_prebuilt import runtime_patterns_for_choice, AssetChoice
choice = AssetChoice(
repo = "lemonade-sdk/llamacpp-rocm",
tag = "b1262",
name = "llama-b1262-ubuntu-rocm-gfx1151-x64.zip",
url = "https://github.com/lemonade-sdk/llamacpp-rocm/releases/download/b1262/x.zip",
source_label = "lemonade",
install_kind = "linux-rocm",
)
pats = runtime_patterns_for_choice(choice)
# The broad glob must be present so every .so in the lemonade bundle
# (including transitive deps added in future ROCm releases) gets overlaid.
assert "lib*.so*" in pats, f"'lib*.so*' missing from linux-rocm patterns: {pats}"
_pick_rocm_gfx_target = getattr(_mod, "_pick_rocm_gfx_target", None)
@pytest.mark.skipif(
_pick_rocm_gfx_target is None,
reason = "_pick_rocm_gfx_target not present on this branch",
)
def test_pick_rocm_gfx_target_honors_cuda_visible_devices(monkeypatch):
"""AMD HIP honours CUDA_VISIBLE_DEVICES identically to HIP_VISIBLE_DEVICES;
on a gfx1151 + gfx1100 mixed host, CUDA_VISIBLE_DEVICES=1 must select gfx1100."""
# Two GPUs; rocminfo reports each token twice (as in the real tool output).
probe_out = "gfx1151\ngfx1151\ngfx1100\ngfx1100"
monkeypatch.delenv("HIP_VISIBLE_DEVICES", raising = False)
monkeypatch.delenv("ROCR_VISIBLE_DEVICES", raising = False)
monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "1")
assert _pick_rocm_gfx_target(probe_out) == "gfx1100"
@pytest.mark.skipif(
_pick_rocm_gfx_target is None,
reason = "_pick_rocm_gfx_target not present on this branch",
)
def test_pick_rocm_gfx_target_cuda_visible_devices_minus_one_returns_none(monkeypatch):
"""CUDA_VISIBLE_DEVICES=-1 means no GPU visible; resolver must return None."""
probe_out = "gfx1151\ngfx1100"
monkeypatch.delenv("HIP_VISIBLE_DEVICES", raising = False)
monkeypatch.delenv("ROCR_VISIBLE_DEVICES", raising = False)
monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "-1")
assert _pick_rocm_gfx_target(probe_out) is None
@pytest.mark.skipif(
_pick_rocm_gfx_target is None,
reason = "_pick_rocm_gfx_target not present on this branch",
)
def test_pick_rocm_gfx_target_same_arch_multi_gpu(monkeypatch):
"""Regression: [gfx1100, gfx1100, gfx1151] with HIP_VISIBLE_DEVICES=2 must
return gfx1151, not fall back to GPU 0 due to dict.fromkeys collapsing the
two gfx1100 entries into one and making index 2 out of range."""
# Simulate rocminfo output for 3 GPUs (2x gfx1100 dGPU + 1x gfx1151 APU).
# Each GPU gets its own Agent section with a few token mentions.
probe_out = (
"***\nAgent 1\n***\n gfx1100 some info\n gfx1100\n"
"***\nAgent 2\n***\n gfx1100 some info\n gfx1100\n"
"***\nAgent 3\n***\n gfx1151 some info\n gfx1151\n"
)
monkeypatch.delenv("ROCR_VISIBLE_DEVICES", raising = False)
monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising = False)
monkeypatch.setenv("HIP_VISIBLE_DEVICES", "2")
assert _pick_rocm_gfx_target(probe_out) == "gfx1151"

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

@ -84,6 +84,7 @@ _httpx_stub.Client = type(
sys.modules.setdefault("httpx", _httpx_stub)
from core.inference.llama_cpp import LlamaCppBackend
from core.inference.llama_server_args import parse_ctx_override, resolve_requested_ctx
# ---------------------------------------------------------------------------
@ -131,6 +132,7 @@ def _drive(
native_ctx = 131072,
kv_per_token_bytes = 325_000,
can_estimate_kv = True,
extra_args = None,
):
"""Drive the post-metadata portion of load_model with stubbed inputs.
@ -148,11 +150,16 @@ def _drive(
inst._can_estimate_kv = lambda: can_estimate_kv
context_length = inst._context_length
# Use the production helper instead of reimplementing the conditional
# locally; reimplementing makes the test pass for the test's own logic
# rather than production's, and silent drift won't be caught.
ctx_override = parse_ctx_override(extra_args)
requested_ctx = resolve_requested_ctx(extra_args, n_ctx)
effective_ctx = n_ctx if n_ctx > 0 else (context_length or 0)
effective_ctx = requested_ctx if requested_ctx > 0 else (context_length or 0)
max_available_ctx = context_length or effective_ctx
if n_ctx > 0:
effective_ctx = n_ctx
if requested_ctx > 0:
effective_ctx = requested_ctx
elif context_length is not None:
effective_ctx = context_length
else:
@ -161,7 +168,7 @@ def _drive(
max_available_ctx = context_length or effective_ctx
gpu_indices, use_fit = None, True
explicit_ctx = n_ctx > 0
explicit_ctx = requested_ctx > 0
if gpus and inst._can_estimate_kv() and effective_ctx > 0:
native_ctx_for_cap = context_length or effective_ctx
@ -236,6 +243,7 @@ def _drive(
"gpu_indices": gpu_indices,
"max_available_ctx": max_available_ctx,
"original_ctx": original_ctx,
"ctx_override": ctx_override,
}
@ -349,6 +357,48 @@ class TestExplicitCtxRespectsUser:
assert plan["c_arg"] == 2048
# ---------------------------------------------------------------------------
# Pass-through --ctx-size participates in context fit (#5676).
# ---------------------------------------------------------------------------
class TestExtraArgsCtxOverride:
def test_ctx_size_extra_honored_over_auto(self):
plan = _drive(
n_ctx = 0,
model_gib = 131,
gpus = [(0, 97_000)],
native_ctx = 196608,
extra_args = ["--ctx-size", "128000"],
)
assert plan["ctx_override"] == 128000
assert plan["original_ctx"] == 128000
assert plan["c_arg"] == 128000
assert plan["use_fit"] is True
def test_ctx_size_short_alias_honored_over_auto(self):
plan = _drive(
n_ctx = 0,
model_gib = 131,
gpus = [(0, 97_000)],
native_ctx = 196608,
extra_args = ["-c", "128000"],
)
assert plan["c_arg"] == 128000
assert plan["use_fit"] is True
def test_ctx_size_extra_wins_over_first_class_field(self):
plan = _drive(
n_ctx = 4096,
model_gib = 8,
gpus = [(0, 24_000)],
native_ctx = 131072,
extra_args = ["--ctx-size", "128000"],
)
assert plan["original_ctx"] == 128000
assert plan["c_arg"] == 128000
# ---------------------------------------------------------------------------
# Non-regression: fittable + auto still auto-picks largest fitting ctx
# ---------------------------------------------------------------------------

144
studio/backend/tests/test_llama_cpp_start_failure_classification.py Normal file
View file

@ -0,0 +1,144 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Tests for LlamaCppBackend._classify_llama_start_failure.
When llama-server exits before becoming healthy, load_model turns its
captured stdout/stderr into a user-facing reason. A diffusion / image
GGUF (FLUX, Qwen-Image, ...) is a valid file with plenty of memory, so
the generic "invalid file or out of memory" message is actively
misleading (issue #5842). These tests pin the classification.
"""
from __future__ import annotations
import sys
import types as _types
from pathlib import Path
import pytest
_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
if _BACKEND_DIR not in sys.path:
sys.path.insert(0, _BACKEND_DIR)
# Match the stubbing pattern in sibling tests so the module imports in a
# lightweight env without fastapi.
_loggers_stub = _types.ModuleType("loggers")
_loggers_stub.get_logger = lambda name: __import__("logging").getLogger(name)
sys.modules.setdefault("loggers", _loggers_stub)
# Give the structlog stub a real get_logger: a bare ModuleType poisons
# sys.modules for later tests that call structlog.get_logger at import time.
_structlog_stub = _types.ModuleType("structlog")
_structlog_stub.get_logger = lambda *a, **k: __import__("logging").getLogger(
"structlog"
)
sys.modules.setdefault("structlog", _structlog_stub)
if not hasattr(sys.modules["structlog"], "get_logger"):
sys.modules["structlog"].get_logger = _structlog_stub.get_logger
from core.inference.llama_cpp import LlamaCppBackend # noqa: E402
_classify = LlamaCppBackend._classify_llama_start_failure
# Real llama-server failure lines (lower-cased downstream anyway).
_QWEN_IMAGE_OUT = (
"load_model: loading model 'qwen-image-edit-2511-Q4_K_M.gguf'\n"
"llama_model_load: error loading model: unknown model architecture: 'qwen_image'\n"
"llama_model_load_from_file_impl: failed to load model"
)
_OOM_OUT = (
"ggml_backend_cuda_buffer_type_alloc_buffer: allocating 12000.00 MiB on "
"device 0: cudaMalloc failed: out of memory"
)
class TestDiffusionArchitectures:
def test_qwen_image_routes_to_images_page(self):
msg = _classify(_QWEN_IMAGE_OUT, "/models/qwen-image.gguf", "local/qwen-image")
assert "diffusion" in msg.lower()
assert "Images page" in msg
assert "qwen_image" in msg
# Must NOT keep blaming memory / file validity.
assert "out of memory" not in msg.lower()
assert "enough memory" not in msg.lower()
# Parametrize over the production set so new arches are auto-covered.
@pytest.mark.parametrize("arch", sorted(LlamaCppBackend._DIFFUSION_ARCHES))
def test_every_diffusion_arch_is_recognised(self, arch):
out = f"error loading model: unknown model architecture: '{arch}'"
msg = _classify(out, f"/models/{arch}.gguf", f"local/{arch}")
assert "diffusion" in msg.lower()
assert "Images page" in msg
assert arch in msg
class TestUnsupportedNonDiffusionArchitecture:
def test_unknown_llm_arch_says_unsupported_not_oom(self):
out = "error loading model: unknown model architecture: 'some_new_llm'"
msg = _classify(out, "/models/x.gguf", "local/x")
assert "some_new_llm" in msg
assert "architecture" in msg.lower()
# Specific, not the misleading memory message.
assert "enough memory" not in msg.lower()
assert "diffusion" not in msg.lower()
# Exact match: a chat arch merely containing a diffusion token (wan,
# sd1, flux, ...) must not be routed to the Images page.
@pytest.mark.parametrize(
"arch",
[
"taiwan", # contains "wan"
"swan_llm", # contains "wan"
"fluxion", # contains "flux"
"sd1234", # contains "sd1"
"sd3_chat", # contains "sd3"
"aura2_text", # contains "aura"
"cosmos_reason", # contains "cosmos"
"qwen_image_text", # contains "qwen_image"
],
)
def test_arch_containing_diffusion_token_is_not_misrouted(self, arch):
out = f"error loading model: unknown model architecture: '{arch}'"
msg = _classify(out, f"/models/{arch}.gguf", f"local/{arch}")
assert arch in msg
assert "does not support" in msg.lower()
assert "diffusion" not in msg.lower()
assert "Images page" not in msg
class TestOllamaAndFallback:
_OLLAMA_GGUF = (
f"/home/u/.ollama{__import__('os').sep}ollama_links"
f"{__import__('os').sep}m.gguf"
)
def test_ollama_compat_message_still_works(self):
out = "llama_model_load: error loading model: key not found"
msg = _classify(out, self._OLLAMA_GGUF, "ollama/llama3")
assert "Ollama" in msg
def test_ollama_unknown_arch_keeps_ollama_guidance(self):
# Ollama + non-diffusion unknown arch keeps the Ollama hint, not the
# generic llama.cpp "unsupported" message.
out = "error loading model: unknown model architecture: 'some_new_llm'"
msg = _classify(out, self._OLLAMA_GGUF, "ollama/some-new")
assert "Ollama" in msg
assert "directly through Ollama" in msg
assert "does not support" not in msg.lower()
def test_ollama_diffusion_arch_still_routes_to_images(self):
# Diffusion routing wins over the Ollama hint.
out = "error loading model: unknown model architecture: 'flux'"
msg = _classify(out, self._OLLAMA_GGUF, "ollama/flux")
assert "diffusion" in msg.lower()
assert "Images page" in msg
def test_generic_oom_keeps_memory_message(self):
msg = _classify(_OOM_OUT, "/models/big.gguf", "local/big")
assert "enough memory" in msg.lower()
assert "diffusion" not in msg.lower()
def test_empty_output_is_safe(self):
msg = _classify("", None, None)
assert "llama-server failed to start" in msg

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

@ -3,21 +3,38 @@
"""Unit tests for the llama-server pass-through args validator.
The validator is the security boundary between user-supplied CLI / HTTP
input and the llama-server subprocess command. These tests pin the
denylist behavior so the boundary doesn't quietly regress when new
managed flags are added.
The validator is the boundary between user CLI/HTTP input and the
llama-server subprocess. These tests pin denylist behaviour so it
doesn't quietly regress when new managed flags are added.
"""
from __future__ import annotations
import importlib.util
import re
from pathlib import Path
import pytest
from core.inference.llama_server_args import (
is_managed_flag,
strip_shadowing_flags,
validate_extra_args,
# Load llama_server_args.py directly so this test doesn't drag in the
# full backend chain (fastapi / structlog / loggers / utils.hardware)
# via core/inference/__init__.py. The validator is intentionally
# dependency-free and unit-tests should reflect that.
_LSA_PATH = (
Path(__file__).resolve().parent.parent
/ "core"
/ "inference"
/ "llama_server_args.py"
)
_spec = importlib.util.spec_from_file_location("_lsa_test_only", _LSA_PATH)
_lsa = importlib.util.module_from_spec(_spec)
_spec.loader.exec_module(_lsa)
is_managed_flag = _lsa.is_managed_flag
parse_cache_override = _lsa.parse_cache_override
parse_ctx_override = _lsa.parse_ctx_override
resolve_cache_type_kv = _lsa.resolve_cache_type_kv
strip_shadowing_flags = _lsa.strip_shadowing_flags
validate_extra_args = _lsa.validate_extra_args
# ── Pass-through (allowed) ───────────────────────────────────────────
@ -60,13 +77,12 @@ from core.inference.llama_server_args import (
# Reasoning controls
["--reasoning-format", "deepseek"],
["-rea", "auto"],
# Soft-managed flags the user may want to override on the CLI;
# llama.cpp's last-wins parsing means these win over Studio's
# auto-set version.
# Soft-managed: user-supplied flags last-wins-override Studio's
# auto-set version. --parallel / -np / --n-parallel are NOT
# here -- they're hard-denied (KV-cache + slot count would
# desync). Use `unsloth studio run --parallel N` instead.
["-c", "131072"],
["--ctx-size", "8192"],
["--parallel", "1"],
["-np", "8"],
["--flash-attn", "off"],
["-fa", "on"],
["--no-context-shift"],
@ -99,8 +115,7 @@ def test_value_with_equals_form_passes_through():
def test_non_flag_token_passes_through():
# A bare positional value (not preceded by a flag) is preserved
# verbatim. llama-server may reject it, but that's not our job.
# Bare positionals are passed through; llama-server can reject them.
assert validate_extra_args(["foo"]) == ["foo"]
@ -110,18 +125,33 @@ def test_non_flag_token_passes_through():
@pytest.mark.parametrize(
"denied",
[
# Model identity
# Parallel slots -- owned by the typer --parallel flag.
"-np",
"--parallel",
"--n-parallel",
# Model identity (every alias; bumping llama.cpp must keep
# every form rejected, not just the long).
"-m",
"--model",
"-mu",
"--model-url",
"-dr",
"--docker-repo",
"-hf",
"-hfr",
"--hf-repo",
"-hff",
"--hf-file",
"-hfv",
"-hfrv",
"--hf-repo-v",
"-hffv",
"--hf-file-v",
"-hft",
"--hf-token",
"-mm",
"--mmproj",
"-mmu",
"--mmproj-url",
# Networking (Studio binds + proxies)
"--host",
@ -134,11 +164,28 @@ def test_non_flag_token_passes_through():
"--api-key-file",
"--ssl-key-file",
"--ssl-cert-file",
# Single-model server
# Single-model server (legacy --webui + current --ui group)
"--webui",
"--no-webui",
"--ui",
"--no-ui",
"--ui-config",
"--ui-config-file",
"--ui-mcp-proxy",
"--no-ui-mcp-proxy",
"--models-dir",
"--models-preset",
"--models-max",
"--models-autoload",
"--no-models-autoload",
# Server-mode flips: --embedding / --rerank would restrict
# llama-server to those endpoints and break Studio's chat hop.
"--embedding",
"--embeddings",
"--rerank",
"--reranking",
# llama-server's own --tools clashes with Studio's tool policy.
"--tools",
],
)
def test_denylist_rejects_all_aliases(denied):
@ -146,14 +193,65 @@ def test_denylist_rejects_all_aliases(denied):
validate_extra_args([denied, "value"])
@pytest.mark.parametrize(
"args,offending",
[
# Pass-through --parallel would last-wins-override the real
# slot count while Studio's KV-cache fit + llama_parallel_slots
# stay at the typer value -- plan vs. process disagree.
(["--parallel", "8"], "--parallel"),
(["--parallel=8"], "--parallel"),
(["--n-parallel", "16"], "--n-parallel"),
(["--n-parallel=16"], "--n-parallel"),
(["-np", "32"], "-np"),
# Attached short form: Click clusters it CLI-side; HTTP /load
# with `["-np8"]` must still resolve to managed.
(["-np8"], "-np"),
(["-np64"], "-np"),
# Out-of-range values that would bypass the typer 1..64 guard.
(["--parallel", "999"], "--parallel"),
(["-np", "0"], "-np"),
(["-np999"], "-np"),
# Signed attached forms; `-np-1` must not slip past.
(["-np-1"], "-np"),
(["-np+1"], "-np"),
],
)
def test_parallel_flags_are_managed(args, offending):
with pytest.raises(ValueError, match = re.escape(offending)):
validate_extra_args(args)
def test_denylist_rejects_equals_form():
with pytest.raises(ValueError, match = "--port"):
validate_extra_args(["--port=9000"])
@pytest.mark.parametrize(
"padded",
[" --parallel", "--parallel ", "\t--parallel", " -np", "-np \n", "-np\t"],
)
def test_denylist_rejects_whitespace_padded_forms(padded):
# `_flag_name` trims whitespace before lookup; otherwise a trailing
# space could slip a managed flag past the boundary.
with pytest.raises(ValueError, match = "parallel|np"):
validate_extra_args([padded, "8"])
@pytest.mark.parametrize(
"attached",
["-np8x", "-np-1foo", "-np+1bar", "-np9zzz"],
)
def test_denylist_rejects_np_with_digit_prefix_and_junk(attached):
# Backend `_flag_name` must classify the same forms the CLI
# rewriter expands, else HTTP /load could smuggle `-np8x` through.
with pytest.raises(ValueError, match = "np"):
validate_extra_args([attached])
def test_denylist_rejects_short_form_when_long_is_denied():
# -m is the short form of the hard-denied --model; rejecting only
# the long form would leave a trivial bypass.
# `-m` is the short form of --model; rejecting only the long
# form would leave a trivial bypass.
with pytest.raises(ValueError, match = "-m"):
validate_extra_args(["-m", "/some/other/path.gguf"])
@ -165,9 +263,7 @@ def test_denylist_message_names_offending_flag():
def test_first_denied_flag_short_circuits():
# Validation stops at the first denied flag; later denied flags
# in the same call don't matter for behaviour, but the message
# should name the first one we hit.
# Validation stops at the first denied flag; the message names it.
with pytest.raises(ValueError, match = "--port"):
validate_extra_args(["--port", "1", "--host", "x"])
@ -177,8 +273,7 @@ def test_first_denied_flag_short_circuits():
@pytest.mark.parametrize("value", ["-1", "-0.5", "-42", "-.5"])
def test_negative_number_value_is_not_flag(value):
# ``--seed -1`` is a value, not a flag. Validator must not try
# to look up "-1" in the denylist.
# `--seed -1`: the -1 is a value, not a flag.
assert validate_extra_args(["--seed", value]) == ["--seed", value]
@ -190,6 +285,15 @@ def test_is_managed_flag_true_for_denied():
assert is_managed_flag("--api-key") is True
assert is_managed_flag("-m") is True
assert is_managed_flag("--model") is True
# Parallel slots owned by the typer --parallel flag.
assert is_managed_flag("--parallel") is True
assert is_managed_flag("--n-parallel") is True
assert is_managed_flag("-np") is True
# Normalised forms must classify like the canonical token so
# is_managed_flag filtering stays in sync with validate_extra_args.
assert is_managed_flag("-np8") is True
assert is_managed_flag("--parallel=8") is True
assert is_managed_flag("--port=9000") is True
def test_is_managed_flag_false_for_pass_through():
@ -199,7 +303,6 @@ def test_is_managed_flag_false_for_pass_through():
# Soft-managed flags pass through (last-wins override)
assert is_managed_flag("-c") is False
assert is_managed_flag("--ctx-size") is False
assert is_managed_flag("--parallel") is False
assert is_managed_flag("--flash-attn") is False
assert is_managed_flag("-ngl") is False
assert is_managed_flag("--threads") is False
@ -231,8 +334,8 @@ def test_strip_shadowing_flags_keeps_context_when_not_requested():
def test_strip_shadowing_flags_keeps_chat_template_when_template_disabled():
# Caller did not supply chat_template_override; the inherited
# --chat-template-file must survive the strip.
# No chat_template_override supplied; inherited
# --chat-template-file must survive.
out = strip_shadowing_flags(
["--chat-template-file", "/tmp/custom.jinja", "--top-k", "20"],
strip_context = True,
@ -282,7 +385,7 @@ def test_strip_shadowing_flags_keeps_spec_when_spec_disabled():
def test_strip_shadowing_flags_drops_mtp_flags_when_requested():
# MTP / draft-mtp flags must be stripped when speculative_type is re-applied.
# MTP / draft-mtp flags must drop when speculative_type re-applies.
out = strip_shadowing_flags(
[
"--spec-type",
@ -310,9 +413,88 @@ def test_is_managed_flag_false_for_mtp_pass_through():
assert is_managed_flag("--spec-ngram-mod-n-max") is False
# ── parse_ctx_override ───────────────────────────────────────────────
@pytest.mark.parametrize(
"args,expected",
[
(None, None),
([], None),
(["--top-k", "20"], None),
(["--ctx-size", "128000"], 128000),
(["--ctx-size=128000"], 128000),
(["-c", "128000"], 128000),
(["-c=128000"], 128000),
(["-c", "4096", "--ctx-size", "128000"], 128000),
],
)
def test_parse_ctx_override(args, expected):
assert parse_ctx_override(args) == expected
@pytest.mark.parametrize(
"args",
[
["--ctx-size"],
["--ctx-size", "--top-k"],
["--ctx-size", "abc"],
["--ctx-size=abc"],
["-c", "-1"],
],
)
def test_parse_ctx_override_rejects_malformed_values(args):
with pytest.raises(ValueError, match = "ctx-size|'-c'"):
parse_ctx_override(args)
def test_validate_extra_args_rejects_malformed_ctx_override():
with pytest.raises(ValueError, match = "ctx-size"):
validate_extra_args(["--ctx-size", "abc"])
# ── parse_cache_override ─────────────────────────────────────────────
@pytest.mark.parametrize(
"args,expected",
[
(None, None),
([], None),
(["--top-k", "20"], None),
(["--cache-type-k", "q8_0"], "q8_0"),
(["-ctk", "q4_0"], "q4_0"),
(["-ctv", "q4_0"], "q4_0"),
(["--cache-type-k=q4_0"], "q4_0"),
(["-ctk", "f16", "-ctk", "q8_0"], "q8_0"),
],
)
def test_parse_cache_override(args, expected):
assert parse_cache_override(args) == expected
@pytest.mark.parametrize(
"args",
[
["-ctk"],
["-ctk", "-c", "4096"],
],
)
def test_parse_cache_override_rejects_malformed_values(args):
with pytest.raises(ValueError, match = "cache-type|'-ctk'"):
parse_cache_override(args)
def test_resolve_cache_type_kv_uses_override_when_present():
assert resolve_cache_type_kv(["--cache-type-k", "q8_0"], "f16") == "q8_0"
def test_resolve_cache_type_kv_uses_fallback_without_override():
assert resolve_cache_type_kv(["--top-k", "20"], "f16") == "f16"
def test_strip_shadowing_flags_boolean_does_not_consume_next_token():
# --spec-default is a boolean shadowing flag; the value-skipping
# heuristic must skip just the flag, not the following positional.
# `--spec-default` is boolean; drop just the flag, keep the next token.
out = strip_shadowing_flags(["--spec-default", "ngram-mod"], strip_spec = True)
assert out == ["ngram-mod"]
@ -343,8 +525,8 @@ def test_strip_shadowing_flags_handles_empty_input():
def test_strip_shadowing_flags_defaults_strip_everything():
# The route's already-loaded comparator calls strip_shadowing_flags
# with no kwargs to detect ANY shadowing flag in stored extras.
# The route's already-loaded comparator calls with no kwargs to
# detect ANY shadowing flag in stored extras.
out = strip_shadowing_flags(
["-c", "4096", "--cache-type-k", "q8_0", "--spec-default", "--jinja"]
)

24
studio/backend/tests/test_log_filter_no_truncation.py
View file

@ -2,27 +2,11 @@
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""
Regression tests for studio.backend.loggers.handlers.filter_sensitive_data.
Regression tests for loggers.handlers.filter_sensitive_data.
Context: filter_sensitive_data was originally written with a base64-detection
heuristic that truncated any string >100 chars containing ',' or '/' down to
20 chars + '...'. The block was dormant until PR #5246 wired the processor
into the structlog chain to redact native-path leases. Once active, the
heuristic ate normal log lines emitted by llama_cpp_backend (GGUF size
summary, mmproj selection, the full llama-server command line) and any
exception traceback that happened to contain a file path.
These tests pin two properties:
1. Long, comma- or slash-bearing log messages flow through filter_sensitive_data
unchanged. The exact strings exercised match the call sites at
studio/backend/core/inference/llama_cpp.py:2117, :2283, and :2312 that
were truncated in the original bug report.
2. PR #5246's native-path lease redaction still fires for both the inline
``native_path_lease=...`` regex form and the ``nativePathLease`` dict-key
form. This guards against future regressions that strip redaction along
with the truncation block.
Pins two properties: (1) long strings with commas/slashes pass through
unchanged (the base64-truncation heuristic from PR #5246 was too aggressive),
and (2) native-path lease redaction still fires for both inline and dict-key forms.
"""
from loggers.handlers import filter_sensitive_data

632
studio/backend/tests/test_mcp_servers.py Normal file
View file

@ -0,0 +1,632 @@
# 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 pytest
from fastapi import HTTPException
from storage import mcp_servers_db
def _reset_db(tmp_path, monkeypatch):
monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path))
monkeypatch.setattr(mcp_servers_db, "_schema_ready", False)
# ── storage: mcp_servers_db ─────────────────────────────────────────
def test_create_and_get_server(tmp_path, monkeypatch):
_reset_db(tmp_path, monkeypatch)
mcp_servers_db.create_server(
id = "srv1",
display_name = "GitHub",
url = "https://example.com/mcp",
headers_json = '{"Authorization": "Bearer x"}',
is_enabled = True,
use_oauth = False,
)
row = mcp_servers_db.get_server("srv1")
assert row["id"] == "srv1"
assert row["display_name"] == "GitHub"
assert row["url"] == "https://example.com/mcp"
assert row["headers_json"] == '{"Authorization": "Bearer x"}'
assert row["is_enabled"] == 1
assert row["use_oauth"] == 0
def test_list_servers_ordered_by_created_at(tmp_path, monkeypatch):
_reset_db(tmp_path, monkeypatch)
mcp_servers_db.create_server(id = "a", display_name = "A", url = "https://a/m")
mcp_servers_db.create_server(id = "b", display_name = "B", url = "https://b/m")
rows = mcp_servers_db.list_servers()
assert [r["id"] for r in rows] == ["a", "b"]
def test_update_server_coerces_bools(tmp_path, monkeypatch):
_reset_db(tmp_path, monkeypatch)
mcp_servers_db.create_server(id = "srv1", display_name = "A", url = "https://a/m")
assert mcp_servers_db.update_server(
"srv1", {"is_enabled": False, "use_oauth": True}
)
row = mcp_servers_db.get_server("srv1")
assert row["is_enabled"] == 0
assert row["use_oauth"] == 1
def test_update_server_empty_changes_returns_false(tmp_path, monkeypatch):
_reset_db(tmp_path, monkeypatch)
mcp_servers_db.create_server(id = "srv1", display_name = "A", url = "https://a/m")
assert mcp_servers_db.update_server("srv1", {}) is False
def test_delete_server_roundtrip(tmp_path, monkeypatch):
_reset_db(tmp_path, monkeypatch)
mcp_servers_db.create_server(id = "srv1", display_name = "A", url = "https://a/m")
assert mcp_servers_db.delete_server("srv1") is True
assert mcp_servers_db.delete_server("srv1") is False
assert mcp_servers_db.get_server("srv1") is None
# ── routes/mcp_servers: pure helpers ────────────────────────────────
def test_validate_url_accepts_http_and_https():
from routes.mcp_servers import _validate_url
assert _validate_url("http://example.com/mcp") == "http://example.com/mcp"
assert _validate_url("https://example.com/mcp") == "https://example.com/mcp"
assert _validate_url(" https://example.com/mcp ") == "https://example.com/mcp"
@pytest.mark.parametrize("bad", ["", " ", "ftp://x", "http://", "noscheme.com"])
def test_validate_url_rejects_bad(bad):
from routes.mcp_servers import _validate_url
with pytest.raises(HTTPException) as exc:
_validate_url(bad)
assert exc.value.status_code == 400
def test_normalize_headers():
from routes.mcp_servers import _normalize_headers
assert _normalize_headers({" Auth ": "Bearer x", "": "ignored"}) == {
"Auth": "Bearer x"
}
assert _normalize_headers({"X": 42}) == {"X": "42"}
assert _normalize_headers({}) is None
assert _normalize_headers(None) is None
assert _normalize_headers({" ": "x"}) is None
def test_changes_from_payload_tristate_headers():
from routes.mcp_servers import _changes_from_payload
from models.mcp_servers import McpServerUpdate
# omitted → key absent
assert "headers_json" not in _changes_from_payload(
McpServerUpdate(display_name = "x")
)
# null → stored as None (clear all headers)
assert _changes_from_payload(McpServerUpdate(headers = None))["headers_json"] is None
# dict → serialised JSON
assert (
_changes_from_payload(McpServerUpdate(headers = {"a": "1"}))["headers_json"]
== '{"a": "1"}'
)
# ── core/inference/tools: MCP wiring ────────────────────────────────
def test_mcp_specs_skip_oversized_names():
from core.inference.tools import _mcp_specs_for_server
server = {"id": "s" * 30, "display_name": "S"}
tools = [
{"name": "ok", "description": "fine"},
{"name": "x" * 40, "description": "too long"},
]
specs = _mcp_specs_for_server(server, tools)
assert len(specs) == 1
assert specs[0]["function"]["name"].endswith("__ok")
assert len(specs[0]["function"]["name"]) <= 64
def test_execute_tool_malformed_mcp_name():
from core.inference.tools import execute_tool
out = execute_tool("mcp__no_double_underscore", {})
assert out.startswith("Error: malformed MCP tool name")
def test_execute_tool_unknown_server(tmp_path, monkeypatch):
_reset_db(tmp_path, monkeypatch)
from core.inference.tools import execute_tool
assert (
execute_tool("mcp__missing__do_thing", {})
== "Error: MCP server 'missing' not found"
)
def test_execute_tool_disabled_server(tmp_path, monkeypatch):
_reset_db(tmp_path, monkeypatch)
mcp_servers_db.create_server(
id = "srv1",
display_name = "A",
url = "https://a/m",
is_enabled = False,
)
from core.inference.tools import execute_tool
assert (
execute_tool("mcp__srv1__do_thing", {})
== "Error: MCP server 'srv1' is disabled"
)
def test_mcp_specs_skip_invalid_openai_function_names():
"""OpenAI requires function.name ^[a-zA-Z0-9_-]{1,64}$; tools whose
names contain '.', '/', spaces, etc. would 400 the whole request."""
from core.inference.tools import _mcp_specs_for_server
server = {"id": "srv", "display_name": "S"}
tools = [
{"name": "ok"},
{"name": "with.dot"},
{"name": "weird/slash"},
{"name": "has space"},
{"name": "good-dash_ok"},
]
specs = _mcp_specs_for_server(server, tools)
names = {s["function"]["name"] for s in specs}
assert {"mcp__srv__ok", "mcp__srv__good-dash_ok"} == names
def test_mcp_specs_skip_empty_tool_name():
from core.inference.tools import _mcp_specs_for_server
server = {"id": "srv", "display_name": "S"}
specs = _mcp_specs_for_server(server, [{"name": "", "description": "x"}])
assert specs == []
def test_mcp_specs_drops_duplicate_names():
"""Same tool name twice from one MCP server -> OpenAI rejects the
request as 'duplicates'. Drop the duplicate before forwarding."""
from core.inference.tools import _mcp_specs_for_server
server = {"id": "srv", "display_name": "S"}
tools = [{"name": "echo"}, {"name": "echo"}]
specs = _mcp_specs_for_server(server, tools)
assert len(specs) == 1
def test_call_tool_sync_respects_pre_set_cancel_event(monkeypatch):
"""cancel_event already set before the call -> immediate Error: cancelled
without making a network round-trip."""
import threading
from core.inference import mcp_client
# Stub _client so the test doesn't need a real MCP server.
class _StubClient:
async def __aenter__(self):
return self
async def __aexit__(self, *args):
return False
async def call_tool(self, name, args):
import asyncio as _asyncio
await _asyncio.sleep(30) # never finishes within the test
monkeypatch.setattr(mcp_client, "_client", lambda *a, **kw: _StubClient())
cancel = threading.Event()
cancel.set()
out = mcp_client.call_tool_sync(
url = "https://example/mcp",
headers = None,
name = "slow",
args = {},
timeout = 30.0,
cancel_event = cancel,
)
assert "cancelled" in out.lower()
def test_clear_oauth_tokens_async_no_op_safe(tmp_path, monkeypatch):
"""clear_oauth_tokens_async on a URL with no stored token must not raise --
the delete + update handlers call it best-effort regardless of prior state."""
import asyncio
monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path))
from core.inference import mcp_client
monkeypatch.setattr(mcp_client, "_oauth_token_store", None)
asyncio.run(mcp_client.clear_oauth_tokens_async("https://example.com/mcp"))
def test_delete_server_calls_oauth_cleanup_when_oauth_was_on(tmp_path, monkeypatch):
"""delete_mcp_server route helper should invoke clear_oauth_tokens_async
when the deleted row had use_oauth=true."""
import asyncio
_reset_db(tmp_path, monkeypatch)
from core.inference import mcp_client
monkeypatch.setattr(mcp_client, "_oauth_token_store", None)
mcp_servers_db.create_server(
id = "oauth1",
display_name = "GH",
url = "https://gh-mcp.example/mcp",
is_enabled = True,
use_oauth = True,
)
calls: list[str] = []
async def fake_clear(url):
calls.append(url)
monkeypatch.setattr(mcp_client, "clear_oauth_tokens_async", fake_clear)
# Re-import the route's binding through the module so the patch is seen.
import routes.mcp_servers as routes_mcp
monkeypatch.setattr(routes_mcp, "clear_oauth_tokens_async", fake_clear)
asyncio.run(routes_mcp.delete_mcp_server("oauth1", current_subject = "u"))
assert calls == ["https://gh-mcp.example/mcp"]
assert mcp_servers_db.get_server("oauth1") is None
def test_delete_server_skips_oauth_cleanup_when_oauth_off(tmp_path, monkeypatch):
"""No OAuth token cleanup when the deleted server never had OAuth."""
import asyncio
_reset_db(tmp_path, monkeypatch)
from core.inference import mcp_client
import routes.mcp_servers as routes_mcp
monkeypatch.setattr(mcp_client, "_oauth_token_store", None)
mcp_servers_db.create_server(
id = "noauth",
display_name = "Plain",
url = "https://plain/mcp",
is_enabled = True,
use_oauth = False,
)
calls: list[str] = []
async def fake_clear(url):
calls.append(url)
monkeypatch.setattr(routes_mcp, "clear_oauth_tokens_async", fake_clear)
asyncio.run(routes_mcp.delete_mcp_server("noauth", current_subject = "u"))
assert calls == []
def test_update_server_clears_oauth_on_url_change(tmp_path, monkeypatch):
"""Changing the URL on an OAuth server must drop the old URL's tokens
so the new URL doesn't silently inherit credentials."""
import asyncio
_reset_db(tmp_path, monkeypatch)
from core.inference import mcp_client
from models.mcp_servers import McpServerUpdate
import routes.mcp_servers as routes_mcp
monkeypatch.setattr(mcp_client, "_oauth_token_store", None)
mcp_servers_db.create_server(
id = "s1",
display_name = "A",
url = "https://old/mcp",
is_enabled = True,
use_oauth = True,
)
calls: list[str] = []
async def fake_clear(url):
calls.append(url)
monkeypatch.setattr(routes_mcp, "clear_oauth_tokens_async", fake_clear)
asyncio.run(
routes_mcp.update_mcp_server(
"s1",
McpServerUpdate(url = "https://new/mcp"),
current_subject = "u",
)
)
assert calls == ["https://old/mcp"]
row = mcp_servers_db.get_server("s1")
assert row["url"] == "https://new/mcp"
def test_update_server_clears_oauth_when_oauth_disabled(tmp_path, monkeypatch):
"""Flipping use_oauth false must drop the old URL's tokens."""
import asyncio
_reset_db(tmp_path, monkeypatch)
from core.inference import mcp_client
from models.mcp_servers import McpServerUpdate
import routes.mcp_servers as routes_mcp
monkeypatch.setattr(mcp_client, "_oauth_token_store", None)
mcp_servers_db.create_server(
id = "s1",
display_name = "A",
url = "https://u/mcp",
is_enabled = True,
use_oauth = True,
)
calls: list[str] = []
async def fake_clear(url):
calls.append(url)
monkeypatch.setattr(routes_mcp, "clear_oauth_tokens_async", fake_clear)
asyncio.run(
routes_mcp.update_mcp_server(
"s1",
McpServerUpdate(use_oauth = False),
current_subject = "u",
)
)
assert calls == ["https://u/mcp"]
def test_changes_from_payload_rejects_null_is_enabled():
"""Explicit null for is_enabled used to hit int(None) -> TypeError 500."""
from routes.mcp_servers import _changes_from_payload
from models.mcp_servers import McpServerUpdate
with pytest.raises(HTTPException) as exc:
_changes_from_payload(McpServerUpdate(is_enabled = None))
assert exc.value.status_code == 400
def test_changes_from_payload_rejects_null_use_oauth():
"""Explicit null for use_oauth used to hit int(None) -> TypeError 500."""
from routes.mcp_servers import _changes_from_payload
from models.mcp_servers import McpServerUpdate
with pytest.raises(HTTPException) as exc:
_changes_from_payload(McpServerUpdate(use_oauth = None))
assert exc.value.status_code == 400
def test_test_endpoint_surfaces_url_validation_as_400(tmp_path, monkeypatch):
"""POST /api/mcp/servers/test must 400 on invalid URL like create/update;
previously the same input returned 200 with {"ok": false}."""
import asyncio
_reset_db(tmp_path, monkeypatch)
from routes.mcp_servers import test_mcp_server
from models.mcp_servers import McpServerTestRequest
with pytest.raises(HTTPException) as exc:
asyncio.run(
test_mcp_server(
McpServerTestRequest(url = "ftp://nope"),
current_subject = "u",
)
)
assert exc.value.status_code == 400
def test_tool_xml_parser_handles_hyphenated_parameter_names():
"""MCP tool schemas commonly use hyphenated property names like
`issue-number` / `repo-name`; the XML parser's `<parameter=\\w+>` regex
dropped those keys. Verify hyphenated parameter names round-trip."""
from core.inference.tool_call_parser import parse_tool_calls_from_text
import json as _json
calls = parse_tool_calls_from_text(
"<function=mcp__srv__create-issue>"
"<parameter=issue-title>Bug report</parameter>"
"<parameter=repo-name>octocat/hello</parameter>"
"</function>"
)
assert len(calls) == 1
args = _json.loads(calls[0]["function"]["arguments"])
assert args == {"issue-title": "Bug report", "repo-name": "octocat/hello"}
def test_tool_healing_strip_handles_hyphenated_function_names():
"""GGUF's core/tool_healing.py has its own copy of the XML strip
regex; the round-4 fix to the shared parser missed this file."""
from core.tool_healing import strip_tool_call_markup
out = strip_tool_call_markup(
"before <function=mcp__srv__list-issues>"
"<parameter=q>x</parameter></function> after"
)
assert out == "before after"
def test_gguf_allow_list_blocks_unadvertised_tool(monkeypatch):
"""When the model emits a tool call not in the per-request tool list
the GGUF agentic loop must refuse to dispatch -- mirroring the
safetensors path. Previously execute_tool ran the call regardless."""
from core.inference import tools as tools_mod
captured: list[str] = []
def fake_execute(name, args, **kw):
captured.append(name)
return "executed"
monkeypatch.setattr(tools_mod, "execute_tool", fake_execute)
# Re-create the allow-list check inline so we can unit-test the
# behavior without spinning up llama-server.
def _gate(tools_advertised, called_name, args):
allowed = {
(t.get("function") or {}).get("name")
for t in (tools_advertised or [])
if (t.get("function") or {}).get("name")
}
if allowed and called_name not in allowed:
return "Error: tool '" + called_name + "' is not enabled"
return fake_execute(called_name, args)
# Built-in not in advertised list -> blocked.
out = _gate(
[{"function": {"name": "mcp__srv__echo"}}],
"terminal",
{"command": "echo x"},
)
assert "not enabled" in out
assert captured == []
# Tool in advertised list -> runs.
out = _gate(
[{"function": {"name": "mcp__srv__echo"}}],
"mcp__srv__echo",
{"text": "hi"},
)
assert out == "executed"
assert captured == ["mcp__srv__echo"]
def test_call_tool_sync_short_circuits_on_pre_set_cancel(monkeypatch):
"""cancel_event set BEFORE call_tool_sync runs -> no HTTP request
is made. Previously the call task was created before the cancel
check, opening a transport that the watcher then had to cancel."""
from core.inference import mcp_client
opened: list[str] = []
class _StubClient:
async def __aenter__(self):
opened.append("opened")
return self
async def __aexit__(self, *args):
return False
async def call_tool(self, name, args):
return "ran"
monkeypatch.setattr(mcp_client, "_client", lambda *a, **kw: _StubClient())
import threading
ev = threading.Event()
ev.set()
out = mcp_client.call_tool_sync(
url = "https://example/mcp",
headers = None,
name = "x",
args = {},
timeout = 5.0,
cancel_event = ev,
)
assert "cancelled" in out.lower()
# The client must NOT have been opened.
assert opened == []
def test_clear_oauth_tokens_swallows_constructor_errors(tmp_path, monkeypatch):
"""clear_oauth_tokens_async is best-effort; an OAuth constructor
failure (e.g. missing fastmcp.client.auth) must not bubble out into
a 500 from the delete / update routes."""
import asyncio
from core.inference import mcp_client
monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path))
monkeypatch.setattr(mcp_client, "_oauth_token_store", None)
# Patch the OAuth import path to raise so the entire body fails.
class _BoomOAuth:
def __init__(self, *a, **kw):
raise RuntimeError("simulated")
import sys as _sys
fake_mod = type(_sys)("fastmcp.client.auth")
fake_mod.OAuth = _BoomOAuth
monkeypatch.setitem(_sys.modules, "fastmcp.client.auth", fake_mod)
# Must not raise.
asyncio.run(mcp_client.clear_oauth_tokens_async("https://x/mcp"))
def test_tool_xml_parser_handles_hyphenated_function_names():
"""MCP tool names are advertised as `mcp__srv__list-issues` (the regex
fix allows '-'); the XML tool-call parser must parse them too,
otherwise the model can call the tool but Studio cannot dispatch."""
from core.inference.tool_call_parser import parse_tool_calls_from_text
calls = parse_tool_calls_from_text(
"<function=mcp__srv__list-issues>"
"<parameter=repo>octocat/hello</parameter>"
"</function>"
)
assert len(calls) == 1
assert calls[0]["function"]["name"] == "mcp__srv__list-issues"
import json as _json
args = _json.loads(calls[0]["function"]["arguments"])
assert args == {"repo": "octocat/hello"}
def test_tool_xml_strip_handles_hyphenated_function_names():
"""routes/inference.py:_TOOL_XML_RE must strip a `<function=name-with-dash>`
block; otherwise hyphenated MCP tool-call XML leaks into chat history."""
import re as _re
from pathlib import Path
src = (Path(__file__).resolve().parent.parent / "routes/inference.py").read_text()
m = _re.search(r"_TOOL_XML_RE = _re\.compile\((.*?)\n\)", src, _re.DOTALL)
assert m, "could not extract _TOOL_XML_RE"
ns: dict = {"_re": _re}
exec(f"_TOOL_XML_RE = _re.compile({m.group(1)})", ns)
rx = ns["_TOOL_XML_RE"]
stripped = rx.sub(
"",
"before <function=mcp__srv__list-issues>"
"<parameter=q>x</parameter></function> after",
)
assert stripped == "before after"
def test_safetensors_agentic_empty_allowlist_still_means_allow_all():
"""Document existing contract: at the safetensors_agentic layer,
tools=[] is still treated as "no constraint" (so existing callers
work unchanged). The real fix for the MCP-only-no-discovery case
lives at the route level in inference.py, which refuses to enter
use_tools when the resolved tool list is empty."""
import threading
from core.inference.safetensors_agentic import run_safetensors_tool_loop
calls: list[str] = []
def fake_execute(name, args, **kw):
calls.append(name)
return "ran"
iteration = {"n": 0}
def fake_single_turn(messages):
iteration["n"] += 1
if iteration["n"] == 1:
txt = '<tool_call>{"name":"python","arguments":{"code":"1"}}</tool_call>'
buf = ""
for ch in txt:
buf += ch
yield buf
else:
yield "done"
list(
run_safetensors_tool_loop(
single_turn = fake_single_turn,
messages = [{"role": "user", "content": "x"}],
tools = [],
execute_tool = fake_execute,
cancel_event = threading.Event(),
max_tool_iterations = 1,
)
)
# Empty allow-list = run anything (preserved contract).
assert calls == [("python", {"code": "1"})] or len(calls) >= 1

236
studio/backend/tests/test_mcp_stdio_improvements.py Normal file
View file

@ -0,0 +1,236 @@
"""Tests for the proposed PR #5863 improvements.
Covers: _client() self-gating + keep_alive, OAuth normalised off for stdio
(create + update), env/header dropped on a transport-type switch, and the
backend rejecting a command whose first token is a URL scheme.
Run from studio/backend: python -m pytest tests/test_mcp_stdio_improvements.py -q
"""
import asyncio
import pytest
from fastapi import HTTPException
from core.inference import mcp_client
from storage import mcp_servers_db
def _reset_db(tmp_path, monkeypatch):
monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path))
monkeypatch.setattr(mcp_servers_db, "_schema_ready", False)
def _enable(monkeypatch):
monkeypatch.setenv("UNSLOTH_STUDIO_ALLOW_STDIO_MCP", "1")
def _disable(monkeypatch):
monkeypatch.delenv("UNSLOTH_STUDIO_ALLOW_STDIO_MCP", raising = False)
# ── P1: _client() self-gates the stdio sink ─────────────────────────
def test_client_refuses_stdio_when_disabled(monkeypatch):
_disable(monkeypatch)
with pytest.raises(PermissionError):
mcp_client._client("npx -y server /tmp", None)
def test_client_builds_stdio_when_enabled_without_spawning(monkeypatch):
_enable(monkeypatch)
# Constructing the Client must not spawn the subprocess (spawn happens on
# __aenter__); we only assert it builds.
client = mcp_client._client("npx -y server /tmp", {"K": "v"})
assert client is not None
def test_client_http_unaffected_by_gate(monkeypatch):
_disable(monkeypatch)
assert mcp_client._client("https://example.com/mcp", None) is not None
# ── P3: OAuth normalised off for stdio (create + update) ────────────
def test_create_forces_oauth_off_for_stdio(tmp_path, monkeypatch):
import routes.mcp_servers as routes_mcp
from models.mcp_servers import McpServerCreate
_reset_db(tmp_path, monkeypatch)
_enable(monkeypatch)
resp = asyncio.run(
routes_mcp.create_mcp_server(
McpServerCreate(
display_name = "FS", url = "npx -y server /tmp", use_oauth = True
),
current_subject = "u",
)
)
assert resp.use_oauth is False
assert mcp_servers_db.get_server(resp.id)["use_oauth"] == 0
def test_create_keeps_oauth_for_http(tmp_path, monkeypatch):
import routes.mcp_servers as routes_mcp
from models.mcp_servers import McpServerCreate
_reset_db(tmp_path, monkeypatch)
_enable(monkeypatch)
resp = asyncio.run(
routes_mcp.create_mcp_server(
McpServerCreate(display_name = "GH", url = "https://gh/mcp", use_oauth = True),
current_subject = "u",
)
)
assert resp.use_oauth is True
def test_update_url_to_stdio_clears_oauth(tmp_path, monkeypatch):
import routes.mcp_servers as routes_mcp
from models.mcp_servers import McpServerUpdate
_reset_db(tmp_path, monkeypatch)
_enable(monkeypatch)
monkeypatch.setattr(mcp_client, "_oauth_token_store", None)
monkeypatch.setattr(
routes_mcp, "clear_oauth_tokens_async", lambda *a, **k: asyncio.sleep(0)
)
mcp_servers_db.create_server(
id = "s1", display_name = "A", url = "https://a/mcp", use_oauth = True
)
resp = asyncio.run(
routes_mcp.update_mcp_server(
"s1", McpServerUpdate(url = "npx -y server /tmp"), current_subject = "u"
)
)
assert resp.use_oauth is False
# ── P4: env/headers dropped on a transport-type switch ──────────────
def test_switch_stdio_to_http_drops_env(tmp_path, monkeypatch):
import routes.mcp_servers as routes_mcp
from models.mcp_servers import McpServerUpdate
_reset_db(tmp_path, monkeypatch)
_enable(monkeypatch)
mcp_servers_db.create_server(
id = "s1",
display_name = "A",
url = "npx server",
headers_json = '{"API_KEY": "secret"}',
)
resp = asyncio.run(
routes_mcp.update_mcp_server(
"s1", McpServerUpdate(url = "https://remote/mcp"), current_subject = "u"
)
)
# the stdio env must NOT survive as HTTP headers on the remote endpoint
assert resp.headers == {}
assert mcp_servers_db.get_server("s1")["headers_json"] is None
def test_switch_keeps_explicitly_supplied_headers(tmp_path, monkeypatch):
import routes.mcp_servers as routes_mcp
from models.mcp_servers import McpServerUpdate
_reset_db(tmp_path, monkeypatch)
_enable(monkeypatch)
mcp_servers_db.create_server(
id = "s1",
display_name = "A",
url = "npx server",
headers_json = '{"API_KEY": "secret"}',
)
resp = asyncio.run(
routes_mcp.update_mcp_server(
"s1",
McpServerUpdate(
url = "https://remote/mcp", headers = {"Authorization": "Bearer new"}
),
current_subject = "u",
)
)
assert resp.headers == {"Authorization": "Bearer new"}
def test_same_transport_edit_keeps_headers(tmp_path, monkeypatch):
import routes.mcp_servers as routes_mcp
from models.mcp_servers import McpServerUpdate
_reset_db(tmp_path, monkeypatch)
_enable(monkeypatch)
mcp_servers_db.create_server(
id = "s1",
display_name = "A",
url = "npx server",
headers_json = '{"API_KEY": "secret"}',
)
# editing only the display name (still stdio) must not wipe env vars
resp = asyncio.run(
routes_mcp.update_mcp_server(
"s1", McpServerUpdate(display_name = "B"), current_subject = "u"
)
)
assert resp.headers == {"API_KEY": "secret"}
# ── P5: reject a command whose first token is a URL scheme ───────────
def test_validate_url_rejects_url_scheme_command_when_enabled(monkeypatch):
from routes.mcp_servers import _validate_url
_enable(monkeypatch)
for bad in ["ftp://host/x", "file:///etc/passwd", "ws://h/y"]:
with pytest.raises(HTTPException) as exc:
_validate_url(bad)
assert exc.value.status_code == 400
def test_validate_url_allows_url_in_argument(monkeypatch):
from routes.mcp_servers import _validate_url
_enable(monkeypatch)
# :// inside an ARGUMENT (not the first token) is still a valid command
assert _validate_url("npx server --url https://x/mcp") == (
"npx server --url https://x/mcp"
)
# ── P6: Data Recipe stdio path obeys the same host gate ─────────────
# build_mcp_providers needs the data_designer plugin, which is only installed in
# the Studio test job; skip there rather than fail the core matrix.
_STDIO_RECIPE = {
"mcp_providers": [
{
"provider_type": "stdio",
"name": "fs",
"command": "npx",
"args": ["-y", "@modelcontextprotocol/server-filesystem", "/tmp"],
"env": {},
}
]
}
def test_data_recipe_skips_stdio_when_disabled(monkeypatch):
pytest.importorskip("data_designer")
_disable(monkeypatch)
from core.data_recipe.service import build_mcp_providers
# gate off -> the stdio provider is dropped (no subprocess can be spawned)
assert build_mcp_providers(_STDIO_RECIPE) == []
def test_data_recipe_builds_stdio_when_enabled(monkeypatch):
pytest.importorskip("data_designer")
_enable(monkeypatch)
from core.data_recipe.service import build_mcp_providers
built = build_mcp_providers(_STDIO_RECIPE)
assert len(built) == 1 # constructed (not spawned) only when enabled

367
studio/backend/tests/test_mcp_stdio_pr5863.py Normal file
View file

@ -0,0 +1,367 @@
"""Verification tests for PR #5863 (stdio MCP server support).
Covers the pure helpers (is_stdio / parse_stdio_command / stdio_mcp_enabled /
probe_timeout), the route-level _validate_url gate, and - most importantly -
that the UNSLOTH_STUDIO_ALLOW_STDIO_MCP gate blocks the stdio transport at all
five enforcement points (create, update, test, refresh, discovery, execute)
when disabled, and reaches it when enabled. The transport (_client) is stubbed
so no real subprocess is spawned; a recorder asserts whether it was reached.
Run from studio/backend: python -m pytest tests/test_mcp_stdio_pr5863.py -q
"""
import sys
import pytest
from fastapi import HTTPException
from core.inference import mcp_client
from storage import mcp_servers_db
def _reset_db(tmp_path, monkeypatch):
monkeypatch.setenv("UNSLOTH_STUDIO_HOME", str(tmp_path))
monkeypatch.setattr(mcp_servers_db, "_schema_ready", False)
def _enable(monkeypatch):
monkeypatch.setenv("UNSLOTH_STUDIO_ALLOW_STDIO_MCP", "1")
def _disable(monkeypatch):
monkeypatch.delenv("UNSLOTH_STUDIO_ALLOW_STDIO_MCP", raising = False)
# ── transport stub + recorder ───────────────────────────────────────
class _FakeTool:
def __init__(self, name):
self._name = name
def model_dump(self, exclude_none = True):
return {"name": self._name, "description": f"{self._name} tool"}
class _Block:
def __init__(self, text):
self.type = "text"
self.text = text
class _FakeResult:
is_error = False
def __init__(self, text):
self.content = [_Block(text)]
class _RecordingClient:
"""Stands in for fastmcp.Client; records that the transport was opened."""
def __init__(self, url, headers, use_oauth, recorder):
recorder.append({"url": url, "headers": headers, "use_oauth": use_oauth})
async def __aenter__(self):
return self
async def __aexit__(self, *a):
return False
async def list_tools(self):
return [_FakeTool("list_directory"), _FakeTool("write_file")]
async def call_tool(self, name, args):
return _FakeResult(f"called {name}")
@pytest.fixture
def transport(monkeypatch):
"""Patch mcp_client._client with a recorder. Returns the recorder list;
empty == the stdio transport was never reached."""
recorder = []
monkeypatch.setattr(
mcp_client,
"_client",
lambda url, headers, use_oauth = False: _RecordingClient(
url, headers, use_oauth, recorder
),
)
return recorder
# ── 1. is_stdio ─────────────────────────────────────────────────────
@pytest.mark.parametrize(
"addr",
[
"http://localhost:8000/mcp",
"https://example.com/mcp",
" https://example.com/mcp ",
"HTTPS://EXAMPLE.COM/mcp",
],
)
def test_is_stdio_false_for_http(addr):
assert mcp_client.is_stdio(addr) is False
@pytest.mark.parametrize(
"addr",
[
"npx -y @modelcontextprotocol/server-filesystem /tmp",
"python -m some.module",
"uvx some-server --flag",
"/usr/local/bin/my-server",
],
)
def test_is_stdio_true_for_commands(addr):
assert mcp_client.is_stdio(addr) is True
# ── 2. parse_stdio_command ──────────────────────────────────────────
def test_parse_basic_argv():
assert mcp_client.parse_stdio_command(
"npx -y @modelcontextprotocol/server-filesystem /tmp"
) == ["npx", "-y", "@modelcontextprotocol/server-filesystem", "/tmp"]
def test_parse_keeps_url_argument_as_one_command():
# gemini "high": a :// inside an ARGUMENT must not break the command.
assert mcp_client.parse_stdio_command(
"npx server --endpoint https://example.com/mcp"
) == ["npx", "server", "--endpoint", "https://example.com/mcp"]
def test_parse_quoted_arg():
assert mcp_client.parse_stdio_command('python -m mod --name "a b"') == [
"python",
"-m",
"mod",
"--name",
"a b",
]
def test_parse_empty_returns_empty_list():
assert mcp_client.parse_stdio_command(" ") == []
def test_parse_unclosed_quote_raises_valueerror():
with pytest.raises(ValueError):
mcp_client.parse_stdio_command('npx "unclosed')
def test_parse_windows_strips_wrapping_quotes(monkeypatch):
# gemini "medium": posix=False keeps backslash paths but also the wrapping
# quotes; the PR strips a matched pair so argv[0] reaches the OS clean.
monkeypatch.setattr(sys, "platform", "win32")
parts = mcp_client.parse_stdio_command(
r'"C:\Program Files\node\node.exe" server.js'
)
assert parts[0] == r"C:\Program Files\node\node.exe"
assert parts[1] == "server.js"
# ── 3. stdio_mcp_enabled ────────────────────────────────────────────
@pytest.mark.parametrize("val", ["0", "false", "true", "", " 1 ", "yes", "2"])
def test_stdio_disabled_for_non_exact_one(monkeypatch, val):
monkeypatch.setenv("UNSLOTH_STUDIO_ALLOW_STDIO_MCP", val)
assert mcp_client.stdio_mcp_enabled() is False
def test_stdio_enabled_only_for_exact_one(monkeypatch):
_disable(monkeypatch)
assert mcp_client.stdio_mcp_enabled() is False
monkeypatch.setenv("UNSLOTH_STUDIO_ALLOW_STDIO_MCP", "1")
assert mcp_client.stdio_mcp_enabled() is True
# ── 4. probe_timeout ────────────────────────────────────────────────
def test_probe_timeout_matrix():
assert mcp_client.probe_timeout("https://x/mcp", False) == 8.0
assert mcp_client.probe_timeout("https://x/mcp", True) == 305.0
assert mcp_client.probe_timeout("npx server", False) == 60.0
# oauth wins regardless of address kind (documented behaviour)
assert mcp_client.probe_timeout("npx server", True) == 305.0
# ── 5. _validate_url gate ───────────────────────────────────────────
def test_validate_url_gate_off_rejects_stdio(monkeypatch):
_disable(monkeypatch)
from routes.mcp_servers import _validate_url
assert _validate_url("https://example.com/mcp") == "https://example.com/mcp"
for bad in ["npx server", "python -m mod", "ftp://host"]:
with pytest.raises(HTTPException) as exc:
_validate_url(bad)
assert exc.value.status_code == 400
def test_validate_url_gate_on_accepts_stdio(monkeypatch):
_enable(monkeypatch)
from routes.mcp_servers import _validate_url
assert _validate_url("npx -y server /tmp") == "npx -y server /tmp"
# http still works when stdio is on
assert _validate_url("https://x/mcp") == "https://x/mcp"
# url-bearing argument accepted as a command
assert _validate_url("npx server --url https://x/mcp") == (
"npx server --url https://x/mcp"
)
# empty / unparseable still rejected
for bad in [" ", '"unclosed']:
with pytest.raises(HTTPException) as exc:
_validate_url(bad)
assert exc.value.status_code == 400
# ── 6. gate enforcement at every spawn path (mocked transport) ──────
def test_create_route_gate(tmp_path, monkeypatch, transport):
import asyncio
from models.mcp_servers import McpServerCreate
import routes.mcp_servers as routes_mcp
_reset_db(tmp_path, monkeypatch)
payload = McpServerCreate(display_name = "FS", url = "npx -y server /tmp")
_disable(monkeypatch)
with pytest.raises(HTTPException) as exc:
asyncio.run(routes_mcp.create_mcp_server(payload, current_subject = "u"))
assert exc.value.status_code == 400
_enable(monkeypatch)
resp = asyncio.run(routes_mcp.create_mcp_server(payload, current_subject = "u"))
assert resp.url == "npx -y server /tmp"
def test_update_http_to_stdio_blocked_when_off(tmp_path, monkeypatch):
import asyncio
from models.mcp_servers import McpServerUpdate
import routes.mcp_servers as routes_mcp
_reset_db(tmp_path, monkeypatch)
_disable(monkeypatch)
mcp_servers_db.create_server(id = "s1", display_name = "A", url = "https://a/mcp")
# editing url -> stdio command must 400 (http->stdio edit bypass closed)
with pytest.raises(HTTPException) as exc:
asyncio.run(
routes_mcp.update_mcp_server(
"s1", McpServerUpdate(url = "npx server"), current_subject = "u"
)
)
assert exc.value.status_code == 400
def test_test_route_gate(tmp_path, monkeypatch, transport):
import asyncio
from models.mcp_servers import McpServerTestRequest
import routes.mcp_servers as routes_mcp
_reset_db(tmp_path, monkeypatch)
req = McpServerTestRequest(url = "npx -y server /tmp")
_disable(monkeypatch)
with pytest.raises(HTTPException) as exc:
asyncio.run(routes_mcp.test_mcp_server(req, current_subject = "u"))
assert exc.value.status_code == 400
assert transport == [] # transport never opened
_enable(monkeypatch)
res = asyncio.run(routes_mcp.test_mcp_server(req, current_subject = "u"))
assert res.ok and res.tool_count == 2
assert len(transport) == 1
def test_refresh_route_gate(tmp_path, monkeypatch, transport):
import asyncio
import routes.mcp_servers as routes_mcp
_reset_db(tmp_path, monkeypatch)
# a stdio row as if carried over from a desktop DB
mcp_servers_db.create_server(id = "stdio1", display_name = "FS", url = "npx server")
_disable(monkeypatch)
with pytest.raises(HTTPException) as exc:
asyncio.run(routes_mcp.refresh_mcp_server_tools("stdio1", current_subject = "u"))
assert exc.value.status_code == 400
assert transport == []
_enable(monkeypatch)
res = asyncio.run(
routes_mcp.refresh_mcp_server_tools("stdio1", current_subject = "u")
)
assert res.ok and res.tool_count == 2
assert len(transport) == 1
def test_discovery_gate(tmp_path, monkeypatch, transport):
import asyncio
from core.inference.tools import get_enabled_mcp_tools
_reset_db(tmp_path, monkeypatch)
mcp_servers_db.create_server(
id = "stdio1", display_name = "FS", url = "npx server", is_enabled = True
)
_disable(monkeypatch)
assert asyncio.run(get_enabled_mcp_tools()) == []
assert transport == [] # filtered out before any probe
_enable(monkeypatch)
specs = asyncio.run(get_enabled_mcp_tools())
assert len(specs) == 2
assert len(transport) == 1
def test_execute_gate(tmp_path, monkeypatch, transport):
from core.inference.tools import execute_tool
_reset_db(tmp_path, monkeypatch)
mcp_servers_db.create_server(
id = "stdio1", display_name = "FS", url = "npx server", is_enabled = True
)
_disable(monkeypatch)
out = execute_tool("mcp__stdio1__list_directory", {"path": "/tmp"})
assert "disabled on this host" in out
assert transport == []
_enable(monkeypatch)
out = execute_tool("mcp__stdio1__list_directory", {"path": "/tmp"})
assert out == "called list_directory"
assert len(transport) == 1
# ── 7. env vars ride headers_json as the subprocess env ─────────────
def test_stdio_env_passed_through(tmp_path, monkeypatch, transport):
from core.inference.tools import execute_tool
_reset_db(tmp_path, monkeypatch)
_enable(monkeypatch)
mcp_servers_db.create_server(
id = "stdio1",
display_name = "FS",
url = "npx server",
headers_json = '{"API_KEY": "sk-test"}',
is_enabled = True,
)
execute_tool("mcp__stdio1__list_directory", {})
assert transport[-1]["headers"] == {"API_KEY": "sk-test"}

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

@ -66,6 +66,7 @@ def _load_worker_module():
_worker = _load_worker_module()
_normalize_mlx_studio_optimizer = _worker._normalize_mlx_studio_optimizer
_normalize_mlx_studio_scheduler = _worker._normalize_mlx_studio_scheduler
_mlx_vlm_max_resized_size = _worker._mlx_vlm_max_resized_size
def test_mlx_studio_optimizer_aliases_are_explicit():
@ -82,3 +83,14 @@ def test_mlx_studio_rejects_unknown_optimizer():
def test_mlx_studio_rejects_unknown_scheduler():
with pytest.raises(ValueError, match = "Unsupported LR scheduler for MLX training"):
_normalize_mlx_studio_scheduler("linear_typo")
def test_mlx_vlm_resize_uses_max_dimension_like_torch_trainer():
assert _mlx_vlm_max_resized_size(1000, 500, 512) == (512, 256)
assert _mlx_vlm_max_resized_size(500, 1000, 512) == (256, 512)
assert _mlx_vlm_max_resized_size(1000, 1000, 512) == (512, 512)
assert _mlx_vlm_max_resized_size(256, 128, 1536) == (256, 128)
assert _mlx_vlm_max_resized_size(512, 256, 512) == (512, 256)
# Half-pixel cases must match the Torch collator (not banker's round).
assert _mlx_vlm_max_resized_size(333, 1000, 500) == (167, 500)
assert _mlx_vlm_max_resized_size(1000, 333, 500) == (500, 167)

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

@ -269,7 +269,15 @@ def test_shell_call_emits_tool_start_and_end(monkeypatch):
assert len(ends) == 1
assert starts[0]["tool_name"] == "code_execution"
assert starts[0]["tool_call_id"] == "scall_1"
assert starts[0]["arguments"] == {"kind": "bash", "command": "ls -la"}
# `_server_tool: True` is the synthetic-builtin marker the
# backend stamps onto every provider-side tool_start so the
# frontend serializer can distinguish hosted tools from
# user-declared functions on history replay.
assert starts[0]["arguments"] == {
"kind": "bash",
"command": "ls -la",
"_server_tool": True,
}
assert ends[0]["tool_call_id"] == "scall_1"
assert "total 24" in ends[0]["result"]

3
studio/backend/tests/test_openai_image_generation.py
View file

@ -207,9 +207,12 @@ def test_image_generation_done_emits_tool_event_chunks(monkeypatch):
ends = [e for e in image_events if e.get("type") == "tool_end"]
assert len(starts) == 1, image_events
assert len(ends) == 1, image_events
# `_server_tool: True` marks this as a provider-side synthetic
# tool card on the frontend's history serializer.
assert starts[0]["arguments"] == {
"kind": "image",
"prompt": "A photorealistic cat sitting",
"_server_tool": True,
"openai_image_generation_call_id": "img_abc",
}
assert ends[0]["image_b64"] == "AAAA"

248
studio/backend/tests/test_openai_responses_translation.py
View file

@ -215,6 +215,254 @@ def test_responses_sse_translates_to_chat_completions_chunks(monkeypatch):
assert payloads[-1] == "[DONE]"
def test_responses_function_call_output_translates_to_delta_tool_calls(monkeypatch):
"""Round 12: caller-supplied function tools forwarded into /v1/responses
must have their `function_call` output items translated back into Chat
Completions delta.tool_calls, and the terminal chunk must emit
finish_reason="tool_calls" so the frontend's accumulator runs the
function instead of seeing finish_reason="stop"."""
def handler(request: httpx.Request) -> httpx.Response:
events = [
{"type": "response.created"},
{
"type": "response.output_item.done",
"item": {
"type": "function_call",
"id": "fc_abc",
"call_id": "call_xyz",
"name": "get_weather",
"arguments": '{"city":"SF"}',
},
},
{"type": "response.completed", "response": {}},
]
return httpx.Response(
200,
content = _responses_sse(events),
headers = {"content-type": "text/event-stream"},
)
_mock_http_client(monkeypatch, handler)
async def run():
client = _make_client()
lines = await _collect(
client._stream_openai_responses(
messages = [{"role": "user", "content": "weather?"}],
model = "gpt-5.5",
temperature = 0.7,
top_p = 0.95,
max_tokens = None,
enable_thinking = None,
reasoning_effort = None,
tools = [
{
"type": "function",
"function": {
"name": "get_weather",
"parameters": {
"type": "object",
"properties": {"city": {"type": "string"}},
},
},
}
],
)
)
await client.close()
return lines
lines = _drive(run())
payloads = [
json.loads(line[len("data:") :].strip())
for line in lines
if line.startswith("data:") and line[len("data:") :].strip() != "[DONE]"
]
tool_call_deltas = [
p
for p in payloads
if isinstance(p, dict)
and p.get("choices")
and p["choices"][0].get("delta", {}).get("tool_calls")
]
assert tool_call_deltas, payloads
tc = tool_call_deltas[0]["choices"][0]["delta"]["tool_calls"][0]
assert tc["id"] == "call_xyz"
assert tc["function"]["name"] == "get_weather"
assert tc["function"]["arguments"] == '{"city":"SF"}'
# Final chunk reports tool_calls instead of stop.
terminal = next(
p
for p in payloads
if isinstance(p, dict)
and p.get("choices")
and p["choices"][0].get("finish_reason") in ("stop", "tool_calls")
)
assert terminal["choices"][0]["finish_reason"] == "tool_calls", payloads
def test_responses_parallel_function_calls_get_distinct_indices(monkeypatch):
"""Round 13: parallel function_call items must land on distinct
delta.tool_calls[].index slots so index-keyed clients don't
collapse the second call into the first."""
def handler(request: httpx.Request) -> httpx.Response:
events = [
{"type": "response.created"},
{
"type": "response.output_item.done",
"item": {
"type": "function_call",
"id": "fc_a",
"call_id": "call_a",
"name": "lookup_a",
"arguments": "{}",
},
},
{
"type": "response.output_item.done",
"item": {
"type": "function_call",
"id": "fc_b",
"call_id": "call_b",
"name": "lookup_b",
"arguments": "{}",
},
},
{"type": "response.completed", "response": {}},
]
return httpx.Response(
200,
content = _responses_sse(events),
headers = {"content-type": "text/event-stream"},
)
_mock_http_client(monkeypatch, handler)
async def run():
client = _make_client()
lines = await _collect(
client._stream_openai_responses(
messages = [{"role": "user", "content": "x"}],
model = "gpt-5.5",
temperature = 0.7,
top_p = 0.95,
max_tokens = None,
enable_thinking = None,
reasoning_effort = None,
tools = [
{
"type": "function",
"function": {
"name": "lookup_a",
"parameters": {"type": "object"},
},
},
{
"type": "function",
"function": {
"name": "lookup_b",
"parameters": {"type": "object"},
},
},
],
)
)
await client.close()
return lines
lines = _drive(run())
indices: list[int] = []
for raw in lines:
if not raw.startswith("data:"):
continue
payload = raw[len("data:") :].strip()
if payload == "[DONE]":
continue
try:
obj = json.loads(payload)
except Exception:
continue
delta = (obj.get("choices") or [{}])[0].get("delta") or {}
for tc in delta.get("tool_calls") or []:
indices.append(tc.get("index"))
assert indices == [0, 1], indices
def test_responses_follow_up_tool_result_uses_function_call_output_items(monkeypatch):
"""Round 13: a second turn after a Responses function call must
serialize the tool_calls history and tool result as Responses
`function_call` / `function_call_output` input items, not as
Chat Completions role="tool" content."""
captured: dict = {}
def handler(request: httpx.Request) -> httpx.Response:
captured["body"] = json.loads(request.content.decode("utf-8"))
return httpx.Response(
200,
content = _responses_sse(
[
{"type": "response.created"},
{"type": "response.completed", "response": {}},
]
),
headers = {"content-type": "text/event-stream"},
)
_mock_http_client(monkeypatch, handler)
async def run():
client = _make_client()
await _collect(
client._stream_openai_responses(
messages = [
{"role": "user", "content": "weather?"},
{
"role": "assistant",
"content": "",
"tool_calls": [
{
"id": "call_xyz",
"type": "function",
"function": {
"name": "get_weather",
"arguments": '{"city":"SF"}',
},
}
],
},
{
"role": "tool",
"tool_call_id": "call_xyz",
"content": "sunny",
},
{"role": "user", "content": "thanks"},
],
model = "gpt-5.5",
temperature = 0.7,
top_p = 0.95,
max_tokens = None,
enable_thinking = None,
reasoning_effort = None,
)
)
await client.close()
_drive(run())
items = captured["body"]["input"]
types = [it.get("type") or it.get("role") for it in items]
assert "function_call" in types, items
assert "function_call_output" in types, items
fc = next(it for it in items if it.get("type") == "function_call")
assert fc["call_id"] == "call_xyz"
assert fc["name"] == "get_weather"
assert fc["arguments"] == '{"city":"SF"}'
fco = next(it for it in items if it.get("type") == "function_call_output")
assert fco["call_id"] == "call_xyz"
assert fco["output"] == "sunny"
def test_responses_response_incomplete_maps_to_length_finish_reason(monkeypatch):
def handler(request: httpx.Request) -> httpx.Response:
events = [

176
studio/backend/tests/test_rocm_oom_guard.py Normal file
View file

@ -0,0 +1,176 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Unit tests for _rocm_classify_unified_memory (ROCm OOM-guard classifier).
Covers the three classification paths:
Path 1 – canonical gcnArchName attribute present.
Path 2 – gcnArchName absent, alternate-spelling attribute present.
Path 3 – ALL arch attrs absent; falls back to device-name substring match.
Regression for: Strix Halo (gfx1151) misclassified as discrete on AMD SDK /
Radeon wheels that populate props.name = "Radeon 8060S Graphics" but do NOT
set any gcnArchName attribute. Without the 8060s/8050s name patterns the
fallback returned is_unified=False, applying the 0.90 fraction instead of
0.80 and leaving only ~12.8 GiB OS headroom on a 128 GiB unified-memory pool.
"""
from __future__ import annotations
from types import SimpleNamespace
import pytest
from core.training.worker import _rocm_classify_unified_memory
# ── helpers ──────────────────────────────────────────────────────────────────
def _props(**kwargs) -> SimpleNamespace:
"""Build a fake device-properties object with the given attributes."""
return SimpleNamespace(**kwargs)
# ── Path 1: canonical gcnArchName ────────────────────────────────────────────
class TestCanonicalGcnArchName:
"""gcnArchName is present and populated."""
@pytest.mark.parametrize(
"arch, expected_unified",
[
("gfx1150", True), # Strix Point
("gfx1151", True), # Strix Halo
("gfx1100", False), # Navi 31 (RX 7900 XTX) — discrete
("gfx906", False), # MI50 — discrete server GPU
("gfx1201", False), # RX 9070 XT — discrete
],
)
def test_canonical_attr(self, arch: str, expected_unified: bool) -> None:
props = _props(gcnArchName = arch, name = "irrelevant")
gcn, is_unified = _rocm_classify_unified_memory(props)
assert gcn == arch
assert is_unified is expected_unified
def test_arch_with_colon_suffix_stripped(self) -> None:
"""gcnArchName can carry xnack/sramecc suffix; only the base is kept."""
props = _props(gcnArchName = "gfx1151:xnack-", name = "irrelevant")
gcn, is_unified = _rocm_classify_unified_memory(props)
assert gcn == "gfx1151"
assert is_unified is True
def test_canonical_attr_wins_over_name(self) -> None:
"""Arch attr takes priority; device name should be ignored."""
# Discrete arch, but name looks like a unified SKU — arch must win.
props = _props(gcnArchName = "gfx1100", name = "Radeon 890M")
gcn, is_unified = _rocm_classify_unified_memory(props)
assert gcn == "gfx1100"
assert is_unified is False
# ── Path 2: alternate-spelling fallback ──────────────────────────────────────
class TestAlternateSpellingFallback:
"""gcnArchName is missing but an alternate attr spelling is present."""
@pytest.mark.parametrize(
"attr_name",
["gcn_arch_name", "arch_name", "gfx_arch_name"],
)
def test_alternate_attr_unified(self, attr_name: str) -> None:
props = _props(**{attr_name: "gfx1151"}, name = "Radeon 8060S Graphics")
gcn, is_unified = _rocm_classify_unified_memory(props)
assert gcn == "gfx1151"
assert is_unified is True
@pytest.mark.parametrize(
"attr_name",
["gcn_arch_name", "arch_name", "gfx_arch_name"],
)
def test_alternate_attr_discrete(self, attr_name: str) -> None:
props = _props(**{attr_name: "gfx1201"}, name = "Radeon RX 9070 XT")
gcn, is_unified = _rocm_classify_unified_memory(props)
assert gcn == "gfx1201"
assert is_unified is False
def test_first_non_empty_attr_wins(self) -> None:
"""When multiple alternate attrs are present the first non-empty one wins."""
props = _props(gcn_arch_name = "gfx1151", arch_name = "gfx1100", name = "irrelevant")
gcn, is_unified = _rocm_classify_unified_memory(props)
assert gcn == "gfx1151"
assert is_unified is True
# ── Path 3: device-name fallback ─────────────────────────────────────────────
class TestDeviceNameFallback:
"""ALL arch attrs absent — classifier must rely solely on device name."""
# --- unified-memory devices that MUST be detected ---
@pytest.mark.parametrize(
"device_name",
[
# gfx1150 Strix Point
"Radeon 890M",
"AMD Radeon 890M Graphics",
"RADEON 890M", # case-insensitive
"Radeon 880M",
"AMD Radeon 880M Graphics",
# gfx1151 Strix Halo — the regression case from the review
"Radeon 8060S Graphics", # Ryzen AI MAX+ 395 (as returned by torch)
"AMD Radeon 8060S",
"Radeon 8050S Graphics", # cut-down Strix Halo SKU
"AMD Radeon 8050S",
# case variants
"RADEON 8060S GRAPHICS",
"radeon 8050s",
],
)
def test_unified_memory_detected(self, device_name: str) -> None:
props = _props(name = device_name)
gcn, is_unified = _rocm_classify_unified_memory(props)
assert gcn == "", f"expected empty gcn_arch, got {gcn!r}"
assert (
is_unified is True
), f"device {device_name!r} should be classified as unified-memory"
# --- discrete devices that must NOT be mis-classified ---
@pytest.mark.parametrize(
"device_name",
[
"Radeon RX 9070 XT",
"AMD Radeon RX 7900 XTX",
"Radeon RX 6900 XT",
"Radeon Pro W7900",
"AMD Instinct MI300X",
# Names that contain superficially similar substrings but are discrete
"Radeon RX 580",
"Radeon VII",
],
)
def test_discrete_not_misclassified(self, device_name: str) -> None:
props = _props(name = device_name)
gcn, is_unified = _rocm_classify_unified_memory(props)
assert gcn == ""
assert (
is_unified is False
), f"discrete device {device_name!r} should NOT be classified as unified-memory"
def test_empty_name_returns_false(self) -> None:
"""Completely absent name must not crash and must default to discrete."""
props = _props() # no 'name' attr at all
gcn, is_unified = _rocm_classify_unified_memory(props)
assert gcn == ""
assert is_unified is False
def test_none_name_returns_false(self) -> None:
props = _props(name = None)
gcn, is_unified = _rocm_classify_unified_memory(props)
assert gcn == ""
assert is_unified is False

48
studio/backend/tests/test_studio_train_validation.py
View file

@ -18,6 +18,8 @@ from models.training import (
_MAX_LORA_ALPHA,
_MAX_LORA_R,
_MAX_SEQ_LENGTH,
_MAX_VISION_IMAGE_SIZE,
_MIN_VISION_IMAGE_SIZE,
)
@ -62,6 +64,52 @@ class TestBatchSizeCap:
_check_field("batch_size", 0)
class TestVisionImageSizeCap:
def test_none_accepts_model_default(self):
_check_field("vision_image_size", None)
@pytest.mark.parametrize(
"value",
[_MIN_VISION_IMAGE_SIZE, 640, 1000, _MAX_VISION_IMAGE_SIZE],
)
def test_in_range_accepts(self, value):
_check_field("vision_image_size", value)
assert _MIN_VISION_IMAGE_SIZE == 256
assert _MAX_VISION_IMAGE_SIZE == 2048
@pytest.mark.parametrize(
"value",
[_MIN_VISION_IMAGE_SIZE - 1, _MAX_VISION_IMAGE_SIZE + 1, 640.5, True],
)
def test_invalid_rejects(self, value):
with pytest.raises(ValidationError):
_check_field("vision_image_size", value)
@pytest.mark.parametrize("value", [True, False])
def test_bool_error_says_integer_not_range(self, value):
# Regression guard: bools must say "integer or null", not "in [256, 2048]".
with pytest.raises(ValidationError) as exc:
_check_field("vision_image_size", value)
assert "integer or null" in str(exc.value)
@pytest.mark.parametrize("value", ["++512", "--256", "+-+512", "+", "-"])
def test_multi_sign_string_says_integer_not_raw(self, value):
# Regression guard: multi-sign strings must not leak int()'s raw
# "invalid literal" message; precise contract is "integer or null".
with pytest.raises(ValidationError) as exc:
_check_field("vision_image_size", value)
assert "integer or null" in str(exc.value)
assert "invalid literal" not in str(exc.value)
@pytest.mark.parametrize("value", ["５１２", "٥١٢", "१०२४"])
def test_unicode_digit_string_rejected(self, value):
# Full-width / Arabic-Indic / Devanagari digits must be rejected so the
# value reaching the backend equals the ASCII the user typed.
with pytest.raises(ValidationError) as exc:
_check_field("vision_image_size", value)
assert "integer or null" in str(exc.value)
class TestLoraRCap:
def test_at_cap_accepts(self):
_check_field("lora_r", _MAX_LORA_R)

39
studio/backend/utils/cpu_threads.py Normal file
View file

@ -0,0 +1,39 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Early CPU thread-pool configuration for Studio processes."""
import os
from typing import MutableMapping, Optional
_THREAD_POOL_ENV_VARS = (
"OMP_NUM_THREADS",
"MKL_NUM_THREADS",
"OPENBLAS_NUM_THREADS",
"NUMEXPR_NUM_THREADS",
)
def configure_cpu_threads(env: Optional[MutableMapping[str, str]] = None) -> None:
"""Apply ``UNSLOTH_CPU_THREADS`` to native CPU pools when configured.
This must run before importing libraries that initialize an OpenMP or
BLAS thread pool. Library-specific variables are left untouched so users
can override a single runtime independently.
"""
environ = os.environ if env is None else env
configured = environ.get("UNSLOTH_CPU_THREADS", "").strip()
if not configured:
return
try:
thread_count = int(configured)
except ValueError as exc:
raise ValueError("UNSLOTH_CPU_THREADS must be a positive integer") from exc
if thread_count < 1:
raise ValueError("UNSLOTH_CPU_THREADS must be a positive integer")
value = str(thread_count)
for variable in _THREAD_POOL_ENV_VARS:
environ.setdefault(variable, value)

62
studio/backend/utils/hardware/amd.py
View file

@ -11,18 +11,35 @@ nvidia.py counterparts.
import json
import math
import os
import platform
import re
import subprocess
import sys
from typing import Any, Optional
from loggers import get_logger
from utils.native_path_leases import child_env_without_native_path_secret
from utils.subprocess_compat import windows_hidden_subprocess_kwargs
logger = get_logger(__name__)
# amd-smi on Windows must initialise the full ROCm runtime on first call, which
# can take 15-25 s on cold hardware. Linux is consistently < 2 s.
_AMD_SMI_DEFAULT_TIMEOUT = 30 if platform.system() == "Windows" else 10
def _run_amd_smi(*args: str, timeout: int = 5) -> Optional[Any]:
# Circuit breaker: stop calling amd-smi after this many consecutive failures.
# On Windows, each failed call spawns a process that may show a UAC/DiskPart
# elevation prompt. Once we know amd-smi doesn't work we stop polling it.
_AMD_SMI_FAILURE_LIMIT = 3
_amd_smi_consecutive_failures = 0
_amd_smi_disabled = False
def _run_amd_smi(*args: str, timeout: int = _AMD_SMI_DEFAULT_TIMEOUT) -> Optional[Any]:
"""Run amd-smi with the given arguments and return parsed JSON, or None."""
global _amd_smi_consecutive_failures, _amd_smi_disabled
if _amd_smi_disabled:
return None
try:
result = subprocess.run(
["amd-smi", *args, "--json"],
@ -30,13 +47,40 @@ def _run_amd_smi(*args: str, timeout: int = 5) -> Optional[Any]:
text = True,
timeout = timeout,
env = child_env_without_native_path_secret(),
**windows_hidden_subprocess_kwargs(),
)
except (OSError, subprocess.TimeoutExpired) as e:
logger.warning("amd-smi query failed: %s", e)
if isinstance(e, FileNotFoundError):
# amd-smi ships with Adrenalin, not the HIP SDK -- absence is
# expected on HIP SDK-only Windows setups. Log at debug only.
logger.debug("amd-smi not found (not in PATH): %s", e)
else:
logger.warning("amd-smi query failed: %s", e)
_amd_smi_consecutive_failures += 1
if _amd_smi_consecutive_failures >= _AMD_SMI_FAILURE_LIMIT:
logger.info(
"amd-smi not available (not installed; expected on HIP SDK-only systems); "
"GPU VRAM polling disabled"
)
_amd_smi_disabled = True
return None
if result.returncode != 0 or not result.stdout.strip():
if result.returncode != 0:
logger.warning("amd-smi returned code %d", result.returncode)
_amd_smi_consecutive_failures += 1
if _amd_smi_consecutive_failures >= _AMD_SMI_FAILURE_LIMIT:
logger.info(
"amd-smi not available (not installed; expected on HIP SDK-only systems); "
"GPU VRAM polling disabled"
)
_amd_smi_disabled = True
return None
if not result.stdout.strip():
# amd-smi exited successfully but produced no output (e.g. no GPUs
# visible on this query, or a version that emits nothing for --json).
# This is not a tool failure, so don't count against the circuit breaker.
logger.debug("amd-smi exited 0 but returned no output")
return None
_amd_smi_consecutive_failures = 0 # reset on success
try:
return json.loads(result.stdout)
except json.JSONDecodeError:
@ -352,7 +396,7 @@ def get_visible_gpu_utilization(
)
parsed_id = _parse_numeric(raw_id)
if parsed_id is None:
logger.debug(
logger.warning(
"amd-smi GPU id %r could not be parsed; falling back to "
"enumeration index %d",
raw_id,
@ -360,7 +404,15 @@ def get_visible_gpu_utilization(
)
idx = fallback_idx
else:
idx = int(parsed_id)
rounded = round(parsed_id)
if rounded != parsed_id:
logger.warning(
"amd-smi GPU id %r parsed as non-integer %r; truncating to %d",
raw_id,
parsed_id,
rounded,
)
idx = int(rounded)
if idx not in visible_set:
continue
metrics = _extract_gpu_metrics(gpu_data)

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

@ -16,8 +16,16 @@ Usage:
...
"""
import copy
import gc
import glob
import os
import platform
import re
import subprocess
import sys
import types
from importlib.metadata import PackageNotFoundError, version as pkg_version
import structlog
from loggers import get_logger
from enum import Enum
@ -120,11 +128,13 @@ def detect_hardware() -> DeviceType:
# Distinguish AMD ROCm (HIP) from NVIDIA CUDA for display purposes.
# DeviceType stays CUDA since torch.cuda.* works on ROCm via HIP.
if getattr(torch.version, "hip", None) is not None:
# AMD's repo.radeon.com SDK wheels (e.g. 2.9.0+rocmsdk20251116) do
# not set torch.version.hip, so fall back to checking __version__.
_hip_ver = getattr(torch.version, "hip", None)
if _hip_ver is not None or "rocm" in torch.__version__.lower():
IS_ROCM = True
print(
f"Hardware detected: ROCm (HIP {torch.version.hip}) -- {device_name}"
)
_hip_label = _hip_ver or torch.__version__
print(f"Hardware detected: ROCm (HIP {_hip_label}) -- {device_name}")
else:
print(f"Hardware detected: CUDA -- {device_name}")
return DEVICE
@ -176,8 +186,6 @@ def clear_gpu_cache():
Clear GPU memory cache for the current device.
Safe to call on any platform — no-ops gracefully.
"""
import gc
gc.collect()
device = get_device()
@ -359,8 +367,6 @@ def get_package_versions() -> Dict[str, Optional[str]]:
Returns dict with keys: unsloth, torch, transformers, cuda.
Missing packages yield None.
"""
from importlib.metadata import version as pkg_version, PackageNotFoundError
packages = ("unsloth", "torch", "transformers")
versions: Dict[str, Optional[str]] = {}
@ -466,7 +472,7 @@ def _smi_query(func_name: str, *args, **kwargs) -> Optional[Dict[str, Any]]:
try:
func = getattr(_backend, func_name)
result = func(*args, **kwargs)
if result.get("available"):
if isinstance(result, dict) and result.get("available"):
return result
except Exception as e:
logger.warning("%s %s query failed: %s", backend_name, func_name, e)
@ -479,9 +485,6 @@ def _read_apple_gpu_stats() -> Dict[str, Any]:
Returns dict with utilization_pct, vram_used_bytes (system-wide GPU memory).
Returns empty dict on failure.
"""
import subprocess
import re
try:
result = subprocess.run(
["ioreg", "-r", "-c", "AGXAccelerator"],
@ -506,6 +509,133 @@ def _read_apple_gpu_stats() -> Dict[str, Any]:
}
def _rocm_linux_sysfs_gpu_busy_pct() -> Optional[float]:
"""Query AMD GPU compute utilization via Linux DRM sysfs gpu_busy_percent."""
if platform.system() != "Linux":
return None
try:
files = glob.glob("/sys/class/drm/card*/device/gpu_busy_percent")
if not files:
return None
values = [int(open(f).read().strip()) for f in files]
return round(sum(values) / len(values), 1)
except Exception:
return None
def _rocm_linux_sysfs_temp_c() -> Optional[float]:
"""Query AMD GPU edge temperature via Linux DRM hwmon sysfs (temp1_input, millidegrees C)."""
if platform.system() != "Linux":
return None
try:
files = glob.glob("/sys/class/drm/card*/device/hwmon/hwmon*/temp1_input")
if not files:
return None
temps = [int(open(f).read().strip()) / 1000.0 for f in files]
return round(max(temps), 1)
except Exception:
return None
def _rocm_linux_sysfs_power_w() -> Optional[float]:
"""Query AMD GPU average power draw via Linux DRM hwmon sysfs (microwatts)."""
if platform.system() != "Linux":
return None
try:
for pattern in (
"/sys/class/drm/card*/device/hwmon/hwmon*/power1_average",
"/sys/class/drm/card*/device/hwmon/hwmon*/power1_input",
):
files = glob.glob(pattern)
if files:
watts = sum(int(open(f).read().strip()) / 1_000_000.0 for f in files)
return round(watts, 1)
return None
except Exception:
return None
def _rocm_windows_perf_counter_gpu_util_pct() -> Optional[float]:
"""Query AMD GPU compute utilization via Windows Performance Counters (3D engine nodes)."""
if platform.system() != "Windows":
return None
try:
ps = (
"$s=(Get-Counter '\\GPU Engine(*engtype_3D*)\\Utilization Percentage'"
" -ErrorAction SilentlyContinue).CounterSamples;"
"if($s){[math]::Min(($s|Measure-Object CookedValue -Sum).Sum,100)}else{-1}"
)
r = subprocess.run(
["powershell", "-NoProfile", "-NonInteractive", "-Command", ps],
capture_output = True,
text = True,
timeout = 5,
)
if r.returncode != 0 or not r.stdout.strip():
return None
val = float(r.stdout.strip())
return round(val, 1) if val >= 0 else None
except Exception:
return None
def _rocm_linux_sysfs_vram_gb() -> tuple[Optional[float], Optional[float]]:
"""Query system-wide AMD GPU VRAM via Linux DRM sysfs.
Reads /sys/class/drm/card*/device/mem_info_vram_* which the kernel
updates in real-time across all processes. No tools required.
Returns (used_gb, total_gb) or (None, None) on failure.
"""
if platform.system() != "Linux":
return None, None
try:
used_files = glob.glob("/sys/class/drm/card*/device/mem_info_vram_used")
total_files = glob.glob("/sys/class/drm/card*/device/mem_info_vram_total")
if not used_files or not total_files:
return None, None
used_bytes = sum(int(open(f).read().strip()) for f in used_files)
total_bytes = sum(int(open(f).read().strip()) for f in total_files)
if total_bytes == 0:
return None, None
return round(used_bytes / (1024**3), 2), round(total_bytes / (1024**3), 2)
except Exception:
return None, None
def _rocm_windows_perf_counter_vram_gb() -> tuple[Optional[float], Optional[float]]:
"""Query system-wide dedicated GPU VRAM via Windows Performance Counters.
Uses the same data source as Task Manager so it reflects cross-process
usage accurately. Works for any GPU vendor without amd-smi or nvidia-smi.
Returns (used_gb, total_gb) or (None, None) on failure.
"""
if platform.system() != "Windows":
return None, None
try:
ps = (
"$s=(Get-Counter '\\GPU Adapter Memory(*)\\Dedicated Usage'"
" -ErrorAction SilentlyContinue).CounterSamples;"
"if($s){($s|Measure-Object CookedValue -Sum).Sum}else{-1}"
)
r = subprocess.run(
["powershell", "-NoProfile", "-NonInteractive", "-Command", ps],
capture_output = True,
text = True,
timeout = 5,
)
if r.returncode != 0 or not r.stdout.strip():
return None, None
used_bytes = float(r.stdout.strip())
if used_bytes < 0:
return None, None
import torch as _torch
total_bytes = _torch.cuda.get_device_properties(0).total_memory
return round(used_bytes / (1024**3), 2), round(total_bytes / (1024**3), 2)
except Exception:
return None, None
def get_gpu_utilization() -> Dict[str, Any]:
"""Return a live snapshot of device utilization information."""
device = get_device()
@ -514,7 +644,78 @@ def get_gpu_utilization() -> Dict[str, Any]:
result = _smi_query("get_primary_gpu_utilization")
if result is not None:
result["backend"] = _backend_label(device)
if IS_ROCM:
# Fix unified-memory VRAM on AMD iGPUs (Strix Halo etc.)
_reconcile_primary_rocm_unified_memory(
result, _get_parent_visible_gpu_spec()
)
return result
# SMI tool unavailable or returned no usable data. On Windows, query
# the Performance Counter API (same source as Task Manager) for
# system-wide dedicated VRAM — covers cross-process usage that
# torch.cuda.mem_get_info cannot see from the Studio server process.
if IS_ROCM and platform.system() == "Windows":
_win_used, _win_total = _rocm_windows_perf_counter_vram_gb()
if _win_used is not None and _win_total is not None:
_win_util = _rocm_windows_perf_counter_gpu_util_pct()
return {
"available": True,
"backend": _backend_label(device),
"gpu_utilization_pct": _win_util,
"temperature_c": None,
"vram_used_gb": _win_used,
"vram_total_gb": _win_total,
"vram_utilization_pct": round((_win_used / _win_total) * 100, 1)
if _win_total > 0
else None,
"power_draw_w": None,
"power_limit_w": None,
"power_utilization_pct": None,
}
# Linux: DRM sysfs gives system-wide VRAM across all processes, no tools needed.
if IS_ROCM and platform.system() == "Linux":
_linux_used, _linux_total = _rocm_linux_sysfs_vram_gb()
if _linux_used is not None and _linux_total is not None:
_linux_util = _rocm_linux_sysfs_gpu_busy_pct()
_linux_temp = _rocm_linux_sysfs_temp_c()
_linux_power = _rocm_linux_sysfs_power_w()
return {
"available": True,
"backend": _backend_label(device),
"gpu_utilization_pct": _linux_util,
"temperature_c": _linux_temp,
"vram_used_gb": _linux_used,
"vram_total_gb": _linux_total,
"vram_utilization_pct": round((_linux_used / _linux_total) * 100, 1)
if _linux_total > 0
else None,
"power_draw_w": _linux_power,
"power_limit_w": None,
"power_utilization_pct": None,
}
# Last resort: torch mem_get_info (process-local).
_visible_spec = _get_parent_visible_gpu_spec()
_numeric_ids = _visible_spec.get("numeric_ids") or [0]
_primary_idx = [_numeric_ids[0]] if _numeric_ids else [0]
_torch_devices = _torch_get_per_device_info(_primary_idx)
if _torch_devices:
_td = _torch_devices[0]
_total = _td["total_gb"]
_used = _td["used_gb"]
return {
"available": True,
"backend": _backend_label(device),
"gpu_utilization_pct": None,
"temperature_c": None,
"vram_used_gb": _used,
"vram_total_gb": _total,
"vram_utilization_pct": round((_used / _total) * 100, 1)
if _total > 0
else None,
"power_draw_w": None,
"power_limit_w": None,
"power_utilization_pct": None,
}
# MLX path: single _read_apple_gpu_stats() call carries both VRAM-used
# bytes and GPU utilization %. psutil for unified-memory total is cheap.
@ -578,6 +779,77 @@ def get_gpu_utilization() -> Dict[str, Any]:
return {"available": False, "backend": _backend_label(device)}
def _apply_unified_memory_correction(
device_metrics: Dict[str, Any], torch_info: Dict[str, Any]
) -> None:
"""Per-device reconciliation: when torch reports a larger memory total
than amd-smi, overwrite the smi VRAM fields in place.
Used by both the multi-device and primary-device reconciliation helpers
so the two endpoints stay in sync on AMD iGPUs with unified memory.
"""
torch_total_gb = torch_info["total_gb"]
smi_total_gb = device_metrics.get("vram_total_gb") or 0.0
if torch_total_gb > smi_total_gb:
torch_used_gb = torch_info["used_gb"]
device_metrics["vram_total_gb"] = torch_total_gb
device_metrics["vram_used_gb"] = torch_used_gb
device_metrics["vram_utilization_pct"] = (
round((torch_used_gb / torch_total_gb) * 100, 1)
if torch_total_gb > 0
else None
)
logger.debug(
"ROCm unified memory: replaced amd-smi VRAM (%.2f GB) with "
"torch mem_get_info total (%.2f GB) for device %s",
smi_total_gb,
torch_total_gb,
torch_info.get("index"),
)
def _reconcile_rocm_unified_memory(
utilization: Dict[str, Any], device_indices: list[int]
) -> None:
"""Fix amd-smi VRAM for ROCm unified-memory GPUs (e.g. Strix Halo).
amd-smi reports only the dedicated slice (~512 MB); torch sees the full
GTT pool (~128 GB). When torch total > smi total, overwrite per-device
VRAM fields so GPU selection uses the real available memory.
"""
torch_devices = _torch_get_per_device_info(device_indices)
if not torch_devices:
return
torch_by_index = {td["index"]: td for td in torch_devices}
for dev in utilization.get("devices", []):
td = torch_by_index.get(dev.get("index"))
if td is None:
continue
_apply_unified_memory_correction(dev, td)
def _reconcile_primary_rocm_unified_memory(
utilization: Dict[str, Any], parent_visible_spec: Dict[str, Any]
) -> None:
"""Same fix as _reconcile_rocm_unified_memory for the flat primary-GPU dict."""
numeric_ids = parent_visible_spec.get("numeric_ids")
if numeric_ids is None:
# No visibility env var set: torch ordinal 0 is the primary device.
primary_idx = [0]
elif len(numeric_ids) == 0:
# Empty mask (HIP_VISIBLE_DEVICES="" or "-1"): no GPU is visible to
# this process. Querying torch device 0 would raise a RuntimeError or
# return stale/wrong data, so bail out rather than writing bad values
# into the utilization dict.
return
else:
primary_idx = [int(numeric_ids[0])]
torch_devices = _torch_get_per_device_info(primary_idx)
if not torch_devices:
return
_apply_unified_memory_correction(utilization, torch_devices[0])
def get_visible_gpu_utilization() -> Dict[str, Any]:
device = get_device()
@ -590,6 +862,10 @@ def get_visible_gpu_utilization() -> Dict[str, Any]:
)
if result is not None:
result["backend"] = _backend_label(device)
numeric_ids = parent_visible_spec.get("numeric_ids")
if IS_ROCM and numeric_ids is not None:
# Fix unified-memory VRAM on AMD iGPUs (Strix Halo etc.)
_reconcile_rocm_unified_memory(result, numeric_ids)
return result
# Torch-based fallback for CUDA (nvidia-smi unavailable, AMD ROCm) and XPU (Intel)
@ -689,7 +965,15 @@ def _get_parent_visible_gpu_spec() -> Dict[str, Any]:
# Use explicit None checks (not `or`) so empty string "" is honoured
# as "no visible GPUs" rather than falling through to CUDA_VISIBLE_DEVICES.
cuda_visible = None
if IS_ROCM:
# Prefer ROCm masks only on a ROCm host, or when no CUDA mask is set, so a
# stale HIP_VISIBLE_DEVICES on an NVIDIA host can't override CUDA_VISIBLE_DEVICES.
_is_rocm_spec = IS_ROCM or (
"CUDA_VISIBLE_DEVICES" not in os.environ
and (
"HIP_VISIBLE_DEVICES" in os.environ or "ROCR_VISIBLE_DEVICES" in os.environ
)
)
if _is_rocm_spec:
hip_vis = os.environ.get("HIP_VISIBLE_DEVICES")
rocr_vis = os.environ.get("ROCR_VISIBLE_DEVICES")
if hip_vis is not None:
@ -865,7 +1149,57 @@ def _load_config_for_gpu_estimate(model_name: str, hf_token: Optional[str] = Non
def _determine_attention_impl_for_gpu_estimate(config) -> str:
import copy as _copy
# torch.distributed is incomplete on Windows ROCm — torch._C is a C
# extension (not a package), so Python cannot import the submodule
# torch._C._distributed_c10d that torch.distributed depends on.
# Inject an empty stub into sys.modules BEFORE importing torch.distributed
# so the import succeeds, then patch the missing process-group helpers.
if sys.platform == "win32" and IS_ROCM:
# Dummy class for any name torch.distributed tries to import from these stubs
class _Dummy:
pass
for _c10d_name in (
"torch._C._distributed_c10d",
"torch._C._distributed_autograd",
"torch._C._distributed_rpc",
):
if _c10d_name not in sys.modules:
_stub = types.ModuleType(_c10d_name)
# torch.distributed imports these names from _distributed_c10d;
# provide no-op dummies so the import doesn't raise AttributeError.
for _sym in (
"FakeProcessGroup",
"ProcessGroup",
"Work",
"Store",
"PrefixStore",
"FileStore",
"TCPStore",
"HashStore",
"Reducer",
"Logger",
"DistributedDebugLevel",
"GradBucket",
"BuiltinCommHookType",
):
setattr(_stub, _sym, _Dummy)
sys.modules[_c10d_name] = _stub
try:
import torch.distributed as _td
for _attr, _stub in (
("is_initialized", lambda: False),
("is_available", lambda: False),
("get_rank", lambda: 0),
("get_world_size", lambda: 1),
("is_torchelastic_launched", lambda: False),
):
if not hasattr(_td, _attr):
setattr(_td, _attr, _stub)
except ImportError:
pass
from unsloth.models._utils import resolve_attention_implementation
from transformers import AutoModel, AutoModelForCausalLM
@ -875,7 +1209,7 @@ def _determine_attention_impl_for_gpu_estimate(config) -> str:
# `sub_configs` and propagates to nested text_config / sub-configs, so a
# shallow copy still mutates those shared inner objects on the cached
# config returned by _load_config_for_gpu_estimate. Deepcopy isolates them.
config_copy = _copy.deepcopy(config)
config_copy = copy.deepcopy(config)
model_class = None
for auto_model in (AutoModelForCausalLM, AutoModel):
@ -1062,7 +1396,10 @@ def estimate_required_model_memory_gb(
_determine_attention_impl_for_gpu_estimate(config)
)
except Exception as e:
logger.warning(
# Log at debug: on Windows ROCm the torch.distributed stub does
# not implement Store, so this fires on every estimate call.
# It is expected and non-actionable -- eager is the safe fallback.
logger.debug(
"Could not resolve attention implementation for '%s': %s",
estimate_model,
e,
@ -1552,14 +1889,35 @@ def apply_gpu_ids(gpu_ids) -> None:
# parent process already set a ROCm visibility variable -- that
# way a downstream ROCm process inherits the narrowed mask even
# before Studio's hardware detection has classified the host.
# Final fallback: probe torch.version.hip so AMD workers without
# HIP_VISIBLE_DEVICES still get the correct ROCm visibility mask.
_inherits_rocm_visibility = (
"HIP_VISIBLE_DEVICES" in os.environ or "ROCR_VISIBLE_DEVICES" in os.environ
)
if IS_ROCM or _inherits_rocm_visibility:
_is_rocm = IS_ROCM or _inherits_rocm_visibility
if not _is_rocm:
# torch.version.hip is a non-empty string on ROCm, None on CUDA.
# AMD SDK / Radeon ROCm wheels can leave torch.version.hip unset but
# still encode "rocm" in torch.__version__, matching detect_hardware().
# Broad except: a probe failure must never crash a training worker.
try:
import torch as _torch
_is_rocm = (
getattr(_torch.version, "hip", None) is not None
or "rocm" in getattr(_torch, "__version__", "").lower()
)
except Exception as e:
logger.debug(
"apply_gpu_ids: torch ROCm probe skipped (%s: %s)",
type(e).__name__,
e,
)
if _is_rocm:
os.environ["HIP_VISIBLE_DEVICES"] = value
os.environ["ROCR_VISIBLE_DEVICES"] = value
_visible_gpu_count = None
if IS_ROCM or _inherits_rocm_visibility:
if _is_rocm:
logger.info("Applied gpu_ids: CUDA_VISIBLE_DEVICES='%s' (rocm)", value)
else:
logger.info("Applied gpu_ids: CUDA_VISIBLE_DEVICES='%s'", value)
@ -1652,8 +2010,6 @@ def safe_num_proc(desired: Optional[int] = None) -> int:
Returns:
A safe integer ≥ 1.
"""
import sys
# Windows and macOS use 'spawn' for multiprocessing -- the overhead of
# re-importing torch/transformers/unsloth per worker is typically slower
# than single-process.
@ -1704,8 +2060,6 @@ def dataset_map_num_proc(desired: Optional[int] = None) -> Optional[int]:
``datasets`` treats ``num_proc=1`` as multiprocessing (creates ``Pool(1)``).
Only ``num_proc=None`` guarantees in-process execution.
"""
import sys
if sys.platform in ("win32", "darwin"):
return None
return safe_num_proc(desired)

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

@ -1156,13 +1156,18 @@ def detect_gguf_model(path: str) -> Optional[str]:
p = Path(path)
# Case 1: direct .gguf file
if p.suffix.lower() == ".gguf" and p.is_file():
if p.suffix.lower() == ".gguf":
if _is_mmproj(p.name):
return None
# Use absolute (not resolve) to preserve symlink names -- e.g.
# Ollama .studio_links/model.gguf -> blobs/sha256-... should
# keep the readable symlink name, not the opaque blob hash.
return str(p.absolute())
# Extension is authoritative: don't gate on is_file()/exists(), which
# can fail in the Windows lock window after llama-server is killed.
try:
is_dir = p.is_dir()
except OSError:
is_dir = False # stat() unavailable in the lock window
if not is_dir:
return str(p.absolute()) # absolute() keeps symlink names readable
# Directory named "*.gguf": fall through to the dir scan below.
# Case 2: directory containing .gguf files (skip mmproj)
if p.is_dir():

2
studio/backend/utils/wheel_utils.py
View file

@ -15,6 +15,7 @@ import urllib.request
from typing import Callable
from utils.native_path_leases import child_env_without_native_path_secret
from utils.subprocess_compat import windows_hidden_subprocess_kwargs
_logger = logging.getLogger(__name__)
@ -106,6 +107,7 @@ def probe_torch_wheel_env(*, timeout: int | None = None) -> dict[str, str] | Non
text = True,
timeout = timeout,
env = child_env_without_native_path_secret(),
**windows_hidden_subprocess_kwargs(),
)
except subprocess.TimeoutExpired:
return None

1
studio/frontend/package.json
View file

@ -12,6 +12,7 @@
"lint": "eslint .",
"preview": "vite preview",
"typecheck": "tsc -b --pretty false",
"i18n:check": "node --experimental-strip-types --no-warnings src/i18n/check-parity.ts",
"biome:check": "biome check",
"biome:fix": "biome check --write"
},

9
studio/frontend/src/app/router.tsx
View file

@ -3,6 +3,7 @@
import { Link, createRouter, useRouterState } from "@tanstack/react-router";
import { Button } from "@/components/ui/button";
import { useT } from "@/i18n";
import { Route as rootRoute } from "./routes/__root";
import { Route as dataRecipesRoute } from "./routes/data-recipes";
import { Route as dataRecipeRoute } from "./routes/data-recipes.$recipeId";
@ -31,7 +32,9 @@ const routeTree = rootRoute.addChildren([
]);
function DefaultNotFound() {
const t = useT();
const pathname = useRouterState({ select: (s) => s.location.pathname });
return (
<div className="flex flex-1 flex-col items-center justify-center gap-4 p-8 text-center">
<img
@ -41,14 +44,14 @@ function DefaultNotFound() {
/>
<div className="flex flex-col items-center gap-1">
<h1 className="font-heading font-semibold text-2xl tracking-tight">
Page not found
{t("shell.notFound.title")}
</h1>
<p className="text-muted-foreground text-sm break-all">
{pathname} does not exist.
{t("shell.notFound.description", { path: pathname })}
</p>
</div>
<Button asChild>
<Link to="/chat">Back to chat</Link>
<Link to="/chat">{t("shell.notFound.backToChat")}</Link>
</Button>
</div>
);

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

@ -6,8 +6,9 @@ import { Navbar } from "@/components/navbar";
import { fetchDeviceType, usePlatformStore } from "@/config/env";
import { SidebarInset, SidebarProvider } from "@/components/ui/sidebar";
import { SettingsDialog, useSettingsDialogStore } from "@/features/settings";
import { useTrainingUnloadGuard } from "@/features/training/hooks/use-training-unload-guard";
import { useTrainingUnloadGuard } from "@/features/training";
import { useSidebarPin } from "@/hooks/use-sidebar-pin";
import { useT, type TranslationKey } from "@/i18n";
import {
Outlet,
createRootRoute,
@ -16,24 +17,25 @@ import {
useRouterState,
} from "@tanstack/react-router";
import { AnimatePresence, motion } from "motion/react";
import { Suspense, useEffect, useLayoutEffect, type ReactNode } from "react";
import { Suspense, useEffect, useLayoutEffect } from "react";
import { AppProvider } from "../provider";
// Type `staticData.title` on every route so the matched-title selector
// below stays type-safe without an inline cast.
declare module "@tanstack/react-router" {
interface StaticDataRouteOption {
title?: string;
titleKey?: TranslationKey;
}
}
// Fallback while a lazy route bundle (Train/Recipes/Export) loads.
// /chat is synchronous and never hits this.
const RouteFallback: ReactNode = (
<div className="flex h-full min-h-0 flex-1 items-center justify-center text-muted-foreground text-sm">
Loading...
</div>
);
function RouteFallback() {
const t = useT();
return (
<div className="flex h-full min-h-0 flex-1 items-center justify-center text-muted-foreground text-sm">
{t("common.loading")}
</div>
);
}
const CHAT_ONLY_ALLOWED = new Set([
"/",
@ -68,6 +70,7 @@ const HIDDEN_NAVBAR_ROUTES = ["/onboarding", "/login", "/change-password"];
const DEFAULT_DOCUMENT_TITLE = "Unsloth Studio";
function RootLayout() {
const t = useT();
const pathname = useRouterState({ select: (s) => s.location.pathname });
const hideNavbar = HIDDEN_NAVBAR_ROUTES.includes(pathname);
const isChatRoute = pathname.startsWith("/chat");
@ -75,24 +78,20 @@ function RootLayout() {
useTrainingUnloadGuard();
// Walk matches deepest-first; each route declares its own title.
const matchedTitle = useMatches({
select: (matches) => {
for (let i = matches.length - 1; i >= 0; i--) {
const title = matches[i].staticData.title;
const { title, titleKey } = matches[i].staticData;
if (titleKey) return t(titleKey);
if (title) return title;
}
return null;
},
});
// `/settings` redirects in `beforeLoad`, so its route never stays
// matched; surface the modal's title via the store instead.
const settingsDialogOpen = useSettingsDialogStore((s) => s.open);
const documentTitle = settingsDialogOpen ? "Settings" : matchedTitle;
const documentTitle = settingsDialogOpen ? t("settings.title") : matchedTitle;
// useLayoutEffect updates the tab title before paint, avoiding a
// one-frame flash of the previous route's title on navigation.
useLayoutEffect(() => {
document.title = documentTitle
? `${documentTitle} - ${DEFAULT_DOCUMENT_TITLE}`
@ -116,7 +115,7 @@ function RootLayout() {
<SettingsDialog />
{hideNavbar ? (
<main className="flex-1">
<Suspense fallback={RouteFallback}>
<Suspense fallback={<RouteFallback />}>
<Outlet />
</Suspense>
</main>
@ -142,7 +141,7 @@ function RootLayout() {
transition={{ duration: 0.15 }}
className={`flex min-h-0 min-w-0 flex-1 basis-0 flex-col ${isChatRoute ? "overflow-hidden" : "overflow-visible"}`}
>
<Suspense fallback={RouteFallback}>
<Suspense fallback={<RouteFallback />}>
<Outlet />
</Suspense>
</motion.div>

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

@ -15,7 +15,7 @@ const StudioPage = lazy(() =>
export const Route = createRoute({
getParentRoute: () => rootRoute,
path: "/studio",
staticData: { title: "Train" },
staticData: { titleKey: "studio.routeTitle" },
beforeLoad: () => requireAuth(),
component: StudioPage,
});

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

@ -49,7 +49,7 @@ import {
Edit03Icon,
Globe02Icon,
HelpCircleIcon,
Logout01Icon,
Logout05Icon,
Search01Icon,
PowerIcon,
PencilEdit02Icon,
@ -90,9 +90,33 @@ import {
useTrainingRuntimeStore,
} from "@/features/training";
import type { TrainingRunSummary } from "@/features/training";
import { useEffect, useRef, useState } from "react";
import { useEffect, useRef, useState, type ReactNode } from "react";
import { toast } from "@/lib/toast";
import { ShutdownDialog } from "@/components/shutdown-dialog";
import { translate, useT, type TranslationKey } from "@/i18n";
const EMPHASIS_MARKER = "__UNSLOTH_I18N_EMPHASIS_MARKER__";
type AppT = ReturnType<typeof useT>;
function renderEmphasizedTranslation(
t: AppT,
key: TranslationKey,
emphasizedValue: string,
): ReactNode {
const translated = t(key, { name: EMPHASIS_MARKER });
const parts = translated.split(EMPHASIS_MARKER);
if (parts.length === 1) return translated;
const nodes: ReactNode[] = [];
parts.forEach((part, index) => {
if (part.length > 0) nodes.push(part);
if (index < parts.length - 1) {
nodes.push(<em key={`emphasis-${index}`}>{emphasizedValue}</em>);
}
});
return nodes;
}
function getTourId(pathname: string): string | null {
if (pathname.startsWith("/studio")) return "studio";
@ -185,6 +209,7 @@ function NavItem({
}
export function AppSidebar() {
const t = useT();
const { isDark, toggleTheme, anchorRef } = useAnimatedThemeToggle();
const { pathname, search } = useRouterState({
select: (s) => ({
@ -204,14 +229,8 @@ export function AppSidebar() {
const chatOnly = usePlatformStore((s) => s.isChatOnly());
const [shutdownOpen, setShutdownOpen] = useState(false);
// Chat collapsible state — open by default, auto-expand on route entry
const isChatRoute = pathname.startsWith("/chat");
const isStudioRoute = pathname === "/studio" || pathname.startsWith("/studio/");
const [chatOpen, setChatOpen] = useState(true);
const [runsOpen, setRunsOpen] = useState(true);
useEffect(() => { if (isChatRoute) setChatOpen(true); }, [isChatRoute]);
useEffect(() => { if (isStudioRoute) setRunsOpen(true); }, [isStudioRoute]);
const scrollRef = useRef<HTMLDivElement | null>(null);
const [scrolled, setScrolled] = useState(false);
@ -290,7 +309,7 @@ export function AppSidebar() {
try {
await renameChatItem(target.item, renameTrimmed);
} catch (err) {
toast.error("Failed to rename chat", {
toast.error(translate("shell.toast.failedToRenameChat"), {
description: err instanceof Error ? err.message : undefined,
});
}
@ -300,7 +319,7 @@ export function AppSidebar() {
const updated = await renameTrainingRun(target.run.id, nextRunDisplayName);
emitTrainingRunUpdated(updated);
} catch (err) {
toast.error("Failed to rename run", {
toast.error(translate("shell.toast.failedToRenameRun"), {
description: err instanceof Error ? err.message : undefined,
});
}
@ -320,14 +339,14 @@ export function AppSidebar() {
try {
await handleDeleteThread(target.item);
} catch (err) {
toast.error("Failed to delete chat", {
toast.error(translate("shell.toast.failedToDeleteChat"), {
description: err instanceof Error ? err.message : undefined,
});
}
return;
}
if (target.run.status === "running") {
toast.error("Cannot delete a running training run");
toast.error(t("shell.toast.cannotDeleteRunningRun"));
return;
}
try {
@ -337,7 +356,7 @@ export function AppSidebar() {
}
emitTrainingRunDeleted(target.run.id);
} catch (err) {
toast.error("Failed to delete run", {
toast.error(translate("shell.toast.failedToDeleteRun"), {
description: err instanceof Error ? err.message : undefined,
});
}
@ -366,7 +385,7 @@ export function AppSidebar() {
});
}}
className="flex items-center gap-[6px] select-none"
aria-label="Unsloth home"
aria-label={t("shell.aria.home")}
>
<img
src="/circle-logo-small.png"
@ -377,7 +396,7 @@ export function AppSidebar() {
unsloth
</span>
<span className="nav-badge ml-0.5 inline-flex items-center justify-center rounded-full border border-nav-beta-border px-[5px] pt-[3px] pb-[2px] text-[8px] font-medium leading-none tracking-[0.04em] text-nav-fg-muted antialiased subpixel-antialiased shadow-[0_1px_2px_rgba(0,0,0,0.06)] dark:shadow-[0_1px_2px_rgba(0,0,0,0.35)]">
BETA
{t("shell.beta")}
</span>
</Link>
{!isMobile && (
@ -387,7 +406,7 @@ export function AppSidebar() {
type="button"
onClick={togglePinned}
className="inline-flex h-[35px] w-[32px] items-center justify-center rounded-[10px] text-nav-icon-idle dark:text-nav-fg-muted transition-colors hover:bg-nav-surface-hover hover:text-black dark:hover:text-white focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring"
aria-label="Close sidebar"
aria-label={t("shell.aria.closeSidebar")}
>
<HugeiconsIcon icon={LayoutAlignLeftIcon} strokeWidth={1.75} className="size-icon" />
</button>
@ -397,7 +416,7 @@ export function AppSidebar() {
sideOffset={6}
className="tooltip-compact"
>
Close sidebar
{t("shell.aria.closeSidebar")}
</TooltipContent>
</Tooltip>
)}
@ -412,7 +431,7 @@ export function AppSidebar() {
type="button"
onClick={togglePinned}
className="inline-flex h-[35px] w-[32px] items-center justify-center rounded-[10px] text-nav-fg transition-colors hover:bg-nav-surface-hover hover:text-black dark:hover:text-white focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring"
aria-label="Open sidebar"
aria-label={t("shell.aria.openSidebar")}
>
<HugeiconsIcon icon={LayoutAlignLeftIcon} strokeWidth={1.75} className="size-icon" />
</button>
@ -422,7 +441,7 @@ export function AppSidebar() {
sideOffset={8}
className="tooltip-compact"
>
Open sidebar
{t("shell.aria.openSidebar")}
</TooltipContent>
</Tooltip>
</div>
@ -434,7 +453,7 @@ export function AppSidebar() {
<SidebarMenu>
<NavItem
icon={PencilEdit02Icon}
label="New Chat"
label={t("shell.navigation.newChat")}
active={false}
disabled={chatDisabled}
onClick={() => {
@ -446,7 +465,7 @@ export function AppSidebar() {
/>
<NavItem
icon={ColumnInsertIcon}
label="Compare"
label={t("shell.navigation.compare")}
active={!!search.compare && !chatItems.some((i) => i.id === search.compare)}
disabled={chatDisabled}
dataTour="chat-compare"
@ -459,7 +478,7 @@ export function AppSidebar() {
/>
<NavItem
icon={Search01Icon}
label="Search"
label={t("shell.navigation.search")}
active={false}
disabled={chatDisabled}
onClick={() => {
@ -477,7 +496,7 @@ export function AppSidebar() {
<SidebarMenu>
<NavItem
icon={TestTubeOutlineIcon}
label="Train"
label={t("shell.navigation.train")}
active={pathname === "/studio" || pathname.startsWith("/studio/")}
disabled={chatOnly}
onClick={() => {
@ -489,7 +508,7 @@ export function AppSidebar() {
<NavItem
icon={ChefHatIcon}
label="Recipes"
label={t("shell.navigation.recipes")}
active={isRecipesRoute}
onClick={() => {
navigate({ to: "/data-recipes" });
@ -499,7 +518,7 @@ export function AppSidebar() {
<NavItem
icon={DownloadSquare01Icon}
label="Export"
label={t("shell.navigation.export")}
active={pathname === "/export" || pathname.startsWith("/export/")}
disabled={chatOnly}
onClick={() => {
@ -513,13 +532,16 @@ export function AppSidebar() {
</SidebarGroup>
<SidebarContent ref={scrollRef} className="gap-0 overflow-y-auto overscroll-contain min-h-0">
{/* Recent Chats — hide on Studio only (Eyera fac13); chatOpen = ec695 clickability */}
{!isStudioRoute && chatItems.length > 0 && (
<Collapsible open={chatOpen} onOpenChange={setChatOpen} asChild>
<Collapsible
key={isChatRoute ? "chat-route" : "non-chat-route"}
defaultOpen
asChild
>
<SidebarGroup className="group-data-[collapsible=icon]:hidden px-0 py-0">
<SidebarGroupLabel className={cn("sidebar-sticky-label", scrolled && "is-scrolled")} asChild>
<CollapsibleTrigger className="cursor-pointer flex w-full items-center justify-between">
Recents
{t("shell.navigation.recents")}
<ChevronDown className="size-3.5 transition-transform duration-200 data-[state=open]:rotate-0 [[data-state=closed]_&]:rotate-[-90deg]" />
</CollapsibleTrigger>
</SidebarGroupLabel>
@ -552,7 +574,7 @@ export function AppSidebar() {
<button
type="button"
onClick={(e) => e.stopPropagation()}
aria-label="Chat options"
aria-label={t("shell.aria.chatOptions")}
className="sidebar-row-action group-hover/recent-item:opacity-100 group-hover/recent-item:pointer-events-auto focus-visible:opacity-100 focus-visible:pointer-events-auto"
>
<span className="sidebar-row-action-glyph">
@ -568,14 +590,14 @@ export function AppSidebar() {
>
<DropdownMenuItem onSelect={() => openRenameChat(item)}>
<HugeiconsIcon icon={Edit03Icon} strokeWidth={1.75} className="size-icon" />
<span>Rename</span>
<span>{t("common.rename")}</span>
</DropdownMenuItem>
<DropdownMenuItem
variant="destructive"
onSelect={() => setConfirmingDelete({ kind: "chat", item })}
>
<HugeiconsIcon icon={Delete02Icon} strokeWidth={1.75} className="size-icon" />
<span>Delete</span>
<span>{t("common.delete")}</span>
</DropdownMenuItem>
</DropdownMenuContent>
</DropdownMenu>
@ -588,13 +610,12 @@ export function AppSidebar() {
</Collapsible>
)}
{/* Recent Runs */}
{isStudioRoute && runItems.length > 0 && !chatOnly && (
<Collapsible open={runsOpen} onOpenChange={setRunsOpen} asChild>
<Collapsible key="studio-runs-route" defaultOpen asChild>
<SidebarGroup className="group-data-[collapsible=icon]:hidden px-0 py-0">
<SidebarGroupLabel className={cn("sidebar-sticky-label", scrolled && "is-scrolled")} asChild>
<CollapsibleTrigger className="cursor-pointer flex w-full items-center justify-between">
Recents
{t("shell.navigation.recents")}
<ChevronDown className="size-3.5 transition-transform duration-200 data-[state=open]:rotate-0 [[data-state=closed]_&]:rotate-[-90deg]" />
</CollapsibleTrigger>
</SidebarGroupLabel>
@ -641,7 +662,7 @@ export function AppSidebar() {
<button
type="button"
onClick={(e) => e.stopPropagation()}
aria-label="Run options"
aria-label={t("shell.aria.runOptions")}
className="sidebar-row-action group-hover/run-item:opacity-100 group-hover/run-item:pointer-events-auto focus-visible:opacity-100 focus-visible:pointer-events-auto"
>
<span className="sidebar-row-action-glyph">
@ -657,7 +678,7 @@ export function AppSidebar() {
>
<DropdownMenuItem onSelect={() => openRenameRun(run)}>
<HugeiconsIcon icon={Edit03Icon} strokeWidth={1.75} className="size-icon" />
<span>Rename</span>
<span>{t("common.rename")}</span>
</DropdownMenuItem>
<DropdownMenuItem
variant="destructive"
@ -667,7 +688,7 @@ export function AppSidebar() {
}
>
<HugeiconsIcon icon={Delete02Icon} strokeWidth={1.75} className="size-icon" />
<span>Delete</span>
<span>{t("common.delete")}</span>
</DropdownMenuItem>
</DropdownMenuContent>
</DropdownMenu>
@ -689,7 +710,7 @@ export function AppSidebar() {
<DropdownMenuTrigger asChild>
<SidebarMenuButton
size="lg"
aria-label={`${displayTitle} account menu`}
aria-label={t("shell.accountMenu", { name: displayTitle })}
className="sidebar-nav-btn !h-[50px] gap-[8px] px-2 py-[9px] rounded-[10px]"
>
<div className="shrink-0">
@ -717,16 +738,16 @@ export function AppSidebar() {
onSelect={() => useSettingsDialogStore.getState().openDialog()}
>
<HugeiconsIcon icon={Settings02Icon} strokeWidth={1.75} className="size-icon" />
<span>Settings</span>
<span>{t("shell.navigation.settings")}</span>
<DropdownMenuShortcut>⌘,</DropdownMenuShortcut>
</DropdownMenuItem>
<DropdownMenuItem
onSelect={() => useSettingsDialogStore.getState().openDialog("api-keys")}
>
<HugeiconsIcon icon={Globe02Icon} strokeWidth={1.75} className="size-[18px]" />
<span>API</span>
<span>{t("shell.navigation.api")}</span>
<span className="ml-auto rounded-[6px] border border-emerald-500/25 bg-emerald-500/10 px-1.5 py-0.5 text-[10px] leading-none font-semibold text-emerald-700 dark:text-emerald-300">
New
{t("common.new")}
</span>
</DropdownMenuItem>
<DropdownMenuItem
@ -734,7 +755,11 @@ export function AppSidebar() {
onSelect={(e) => { e.preventDefault(); toggleTheme(); }}
>
{isDark ? <Sun strokeWidth={1.75} className="size-icon" /> : <Moon strokeWidth={1.75} className="size-icon" />}
<span>{isDark ? "Light Mode" : "Dark Mode"}</span>
<span>
{isDark
? t("shell.navigation.lightMode")
: t("shell.navigation.darkMode")}
</span>
</DropdownMenuItem>
<DropdownMenuItem
disabled={!getTourId(pathname)}
@ -749,7 +774,7 @@ export function AppSidebar() {
}}
>
<HugeiconsIcon icon={CursorInfo02Icon} strokeWidth={1.75} className="size-icon" />
<span>Guided Tour</span>
<span>{t("shell.navigation.guidedTour")}</span>
</DropdownMenuItem>
</DropdownMenuGroup>
<DropdownMenuSeparator className="mx-2.5! my-2.5! h-0! border-t border-border/70 bg-transparent!" />
@ -757,7 +782,7 @@ export function AppSidebar() {
onSelect={() => useSettingsDialogStore.getState().openDialog("about")}
>
<HugeiconsIcon icon={HelpCircleIcon} strokeWidth={1.75} className="size-icon" />
<span>Help</span>
<span>{t("common.help")}</span>
</DropdownMenuItem>
<DropdownMenuItem
onSelect={async () => {
@ -771,12 +796,12 @@ export function AppSidebar() {
void navigate({ to: "/login" });
}}
>
<HugeiconsIcon icon={Logout01Icon} strokeWidth={1.75} className="size-icon" />
<span>Log out</span>
<HugeiconsIcon icon={Logout05Icon} strokeWidth={1.75} className="size-icon" />
<span>{t("shell.navigation.logOut")}</span>
</DropdownMenuItem>
<DropdownMenuItem onSelect={() => setShutdownOpen(true)}>
<HugeiconsIcon icon={PowerIcon} strokeWidth={1.75} className="size-icon" />
<span>Shutdown</span>
<span>{t("common.shutdown")}</span>
</DropdownMenuItem>
</DropdownMenuContent>
</DropdownMenu>
@ -800,20 +825,23 @@ export function AppSidebar() {
<DialogHeader>
<DialogTitle>
{confirmingDelete?.kind === "run"
? "Delete training run"
: "Delete chat"}
? t("shell.dialog.deleteRun.title")
: t("shell.dialog.deleteChat.title")}
</DialogTitle>
<DialogDescription>
{confirmingDelete?.kind === "run" ? (
<>
Are you sure you want to delete this run{" "}
<em>{confirmingDelete.run.display_name ?? confirmingDelete.run.model_name}</em>?
</>
renderEmphasizedTranslation(
t,
"shell.dialog.deleteRun.description",
confirmingDelete.run.display_name ??
confirmingDelete.run.model_name,
)
) : confirmingDelete?.kind === "chat" ? (
<>
Are you sure you want to delete this chat{" "}
<em>{confirmingDelete.item.title}</em>?
</>
renderEmphasizedTranslation(
t,
"shell.dialog.deleteChat.description",
confirmingDelete.item.title,
)
) : null}
</DialogDescription>
</DialogHeader>
@ -823,14 +851,14 @@ export function AppSidebar() {
variant="ghost"
onClick={() => setConfirmingDelete(null)}
>
Cancel
{t("common.cancel")}
</Button>
<Button
type="button"
variant="destructive"
onClick={() => void commitDelete()}
>
Delete
{t("common.delete")}
</Button>
</DialogFooter>
</DialogContent>
@ -844,7 +872,9 @@ export function AppSidebar() {
<DialogContent className="corner-squircle border border-border/60 bg-background/98 shadow-none sm:max-w-md">
<DialogHeader>
<DialogTitle>
{renamingTarget?.kind === "run" ? "Rename run" : "Rename chat"}
{renamingTarget?.kind === "run"
? t("shell.dialog.renameRun.title")
: t("shell.dialog.renameChat.title")}
</DialogTitle>
</DialogHeader>
<Input
@ -858,8 +888,16 @@ export function AppSidebar() {
}}
autoFocus
maxLength={120}
placeholder={renamingTarget?.kind === "run" ? "Run name" : "Chat title"}
aria-label={renamingTarget?.kind === "run" ? "Run name" : "Chat title"}
placeholder={
renamingTarget?.kind === "run"
? t("shell.dialog.renameRun.placeholder")
: t("shell.dialog.renameChat.placeholder")
}
aria-label={
renamingTarget?.kind === "run"
? t("shell.dialog.renameRun.placeholder")
: t("shell.dialog.renameChat.placeholder")
}
className="focus-visible:border-input focus-visible:ring-0"
/>
<DialogFooter className="flex-wrap gap-2 sm:justify-end">
@ -868,14 +906,14 @@ export function AppSidebar() {
variant="ghost"
onClick={() => setRenamingTarget(null)}
>
Cancel
{t("common.cancel")}
</Button>
<Button
type="button"
onClick={() => void commitRename()}
disabled={!renameDirty}
>
Save
{t("common.save")}
</Button>
</DialogFooter>
</DialogContent>

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

@ -238,7 +238,7 @@ export const ComposerAddAttachment: FC = () => {
side="bottom"
variant="ghost"
size="icon"
className="aui-composer-add-attachment size-8.5 rounded-full p-1 font-semibold text-xs hover:bg-muted-foreground/15 dark:border-muted-foreground/15 dark:hover:bg-muted-foreground/30"
className="aui-composer-add-attachment size-8.5 rounded-full p-1 font-semibold text-xs hover:bg-muted-foreground/15 dark:hover:bg-muted-foreground/30"
aria-label="Add Attachment"
>
<PlusIcon className="aui-attachment-add-icon size-5 stroke-[1.5px]" />

36
studio/frontend/src/components/assistant-ui/sources.tsx
View file

@ -263,20 +263,34 @@ const SourcesGroup: FC = () => {
return (
<div className="relative mt-2 mb-3">
{/* Hidden measurement container — renders all badges to measure row positions */}
{/* Hidden measurement container. Renders all badges off-screen so we
can read each child's offsetTop and decide how many fit in two
rows. Wrapped in an absolute, h-0, overflow-hidden box so the
measurement pills do NOT contribute to the viewport's scrollable
overflow region. Without this clip, every hidden source row
adds ~30px to scrollHeight, producing a phantom empty scroll
area below the message: visible to users as unbounded blank
space below the assistant action bar. The inner div still
flex-wraps its children for measurement; offsetTop reads
correctly because the wrapper is positioned (absolute) and the
children's offsetTop is measured relative to it. */}
<div
ref={containerRef}
aria-hidden
className="flex w-full flex-wrap gap-1 invisible absolute pointer-events-none"
className="absolute pointer-events-none overflow-hidden h-0 w-full left-0 top-0"
>
{sources.map((source) => (
<span key={source.id} className="inline-block">
<Source href={source.url}>
<SourceIcon url={source.url} />
<SourceTitle>{source.title || extractDomain(source.url)}</SourceTitle>
</Source>
</span>
))}
<div
ref={containerRef}
className="flex w-full flex-wrap gap-1 invisible"
>
{sources.map((source) => (
<span key={source.id} className="inline-block">
<Source href={source.url}>
<SourceIcon url={source.url} />
<SourceTitle>{source.title || extractDomain(source.url)}</SourceTitle>
</Source>
</span>
))}
</div>
</div>
{/* Visible container */}

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

@ -802,6 +802,7 @@ const ReasoningToggle: FC = () => {
{
isReasoningProvider:
selectedExternalProvider?.isReasoningModel === true,
baseUrl: selectedExternalProvider?.baseUrl ?? null,
},
)
: null;

15
studio/frontend/src/components/assistant-ui/tool-fallback.tsx
View file

@ -101,6 +101,16 @@ const statusIconMap: Record<ToolStatus, ElementType> = {
"requires-action": AlertCircleIcon,
};
const MCP_TOOL_PREFIX = "mcp__";
function formatToolNameForDisplay(toolName: string): string {
if (!toolName.startsWith(MCP_TOOL_PREFIX)) return toolName;
const rest = toolName.slice(MCP_TOOL_PREFIX.length);
const sep = rest.indexOf("__");
if (sep <= 0) return toolName;
return `${rest.slice(0, sep)} · ${rest.slice(sep + 2)}`;
}
function ToolFallbackTrigger({
toolName,
status,
@ -119,6 +129,7 @@ function ToolFallbackTrigger({
const StatusIcon = statusIconMap[statusType];
const label = isCancelled ? "Cancelled tool" : "Used tool";
const displayName = formatToolNameForDisplay(toolName);
return (
<CollapsibleTrigger
@ -160,7 +171,7 @@ function ToolFallbackTrigger({
>
<span className="block truncate">
{label}:{" "}
<span className="font-medium text-foreground/85">{toolName}</span>
<span className="font-medium text-foreground/85">{displayName}</span>
</span>
{isRunning && (
<span
@ -169,7 +180,7 @@ function ToolFallbackTrigger({
className="aui-tool-fallback-trigger-shimmer shimmer pointer-events-none absolute inset-0 block truncate motion-reduce:animate-none"
>
{label}:{" "}
<span className="font-medium text-foreground/85">{toolName}</span>
<span className="font-medium text-foreground/85">{displayName}</span>
</span>
)}
</span>

1
studio/frontend/src/config/training.ts
View file

@ -108,6 +108,7 @@ export const LR_DEFAULT_CPT = 5e-5;
export const DEFAULT_HYPERPARAMS = {
epochs: 3,
contextLength: 2048,
visionImageSize: null as number | null,
learningRate: LR_DEFAULT_LORA,
// null = let backend auto-compute (lr/10 per Unsloth CPT recipe). Only used by CPT.
embeddingLearningRate: null as number | null,

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

File diff suppressed because it is too large Load diff

116
studio/frontend/src/features/chat/api/mcp-servers-api.ts Normal file
View file

@ -0,0 +1,116 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import { authFetch } from "@/features/auth";
import { formatFastApiDetail } from "@/lib/format-fastapi-error";
export interface McpServerConfig {
id: string;
display_name: string;
url: string;
headers: Record<string, string>;
is_enabled: boolean;
use_oauth: boolean;
created_at: string;
updated_at: string;
}
export interface McpServerProbeResult {
ok: boolean;
tool_count: number;
error: string | null;
}
function parseErrorText(status: number, body: unknown): string {
if (body && typeof body === "object") {
const { detail, message } = body as { detail?: unknown; message?: unknown };
const formatted = formatFastApiDetail(detail);
if (formatted) return formatted;
if (typeof message === "string" && message) return message;
}
return `Request failed (${status})`;
}
async function mcpRequest<T>(
path: string,
init?: { method?: string; body?: object },
): Promise<T> {
const response = await authFetch(`/api/mcp/servers${path}`, {
method: init?.method,
headers: init?.body ? { "Content-Type": "application/json" } : undefined,
body: init?.body ? JSON.stringify(init.body) : undefined,
});
// 204 No Content (DELETE) has no body — calling .json() would throw.
if (response.status === 204) return undefined as T;
const json = await response.json().catch(() => null);
if (!response.ok) throw new Error(parseErrorText(response.status, json));
return json as T;
}
export function listMcpServers(): Promise<McpServerConfig[]> {
return mcpRequest("/");
}
export function createMcpServer(payload: {
displayName: string;
url: string;
headers?: Record<string, string>;
isEnabled?: boolean;
useOauth?: boolean;
}): Promise<McpServerConfig> {
return mcpRequest("/", {
method: "POST",
body: {
display_name: payload.displayName,
url: payload.url,
headers: payload.headers ?? null,
is_enabled: payload.isEnabled ?? true,
use_oauth: payload.useOauth ?? false,
},
});
}
export function updateMcpServer(
serverId: string,
payload: {
displayName?: string;
url?: string;
/** null = drop stored headers; omit to leave as-is */
headers?: Record<string, string> | null;
isEnabled?: boolean;
useOauth?: boolean;
},
): Promise<McpServerConfig> {
const body: Record<string, unknown> = {};
if (payload.displayName !== undefined) body.display_name = payload.displayName;
if (payload.url !== undefined) body.url = payload.url;
if (payload.headers !== undefined) body.headers = payload.headers;
if (payload.isEnabled !== undefined) body.is_enabled = payload.isEnabled;
if (payload.useOauth !== undefined) body.use_oauth = payload.useOauth;
return mcpRequest(`/${serverId}`, { method: "PUT", body });
}
export function deleteMcpServer(serverId: string): Promise<void> {
return mcpRequest(`/${serverId}`, { method: "DELETE" });
}
export function refreshMcpServerTools(
serverId: string,
): Promise<McpServerProbeResult> {
return mcpRequest(`/${serverId}/refresh`, { method: "POST" });
}
export function testMcpServer(payload: {
url: string;
headers?: Record<string, string>;
useOauth?: boolean;
}): Promise<McpServerProbeResult> {
return mcpRequest("/test", {
method: "POST",
body: {
url: payload.url,
headers: payload.headers ?? null,
use_oauth: payload.useOauth ?? false,
},
});
}

546
studio/frontend/src/features/chat/chat-mcp-servers-dialog.tsx Normal file
View file

@ -0,0 +1,546 @@
// 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 { useCallback, useEffect, useState } from "react";
import { toast } from "sonner";
import { Delete02Icon, Edit03Icon, PlusSignIcon } from "@hugeicons/core-free-icons";
import { HugeiconsIcon } from "@hugeicons/react";
import { RefreshCwIcon } from "lucide-react";
import { Button } from "@/components/ui/button";
import {
Dialog,
DialogContent,
DialogDescription,
DialogHeader,
DialogTitle,
} from "@/components/ui/dialog";
import { Input } from "@/components/ui/input";
import { Label } from "@/components/ui/label";
import { Spinner } from "@/components/ui/spinner";
import { Switch } from "@/components/ui/switch";
import {
type McpServerConfig,
createMcpServer,
deleteMcpServer,
listMcpServers,
refreshMcpServerTools,
testMcpServer,
updateMcpServer,
} from "./api/mcp-servers-api";
type HeaderRow = { id: string; key: string; value: string };
type FormState = {
displayName: string;
url: string;
headers: HeaderRow[];
useOauth: boolean;
};
const EMPTY_FORM: FormState = {
displayName: "",
url: "",
headers: [],
useOauth: false,
};
function newRowId(): string {
return `r_${Math.random().toString(36).slice(2, 10)}`;
}
function headersFromObject(headers: Record<string, string>): HeaderRow[] {
return Object.entries(headers).map(([k, v]) => ({
id: newRowId(),
key: k,
value: v,
}));
}
function headersToObject(rows: HeaderRow[]): Record<string, string> | undefined {
const out: Record<string, string> = {};
for (const row of rows) {
const key = row.key.trim();
if (!key) continue;
out[key] = row.value;
}
return Object.keys(out).length > 0 ? out : undefined;
}
// A non-HTTP address is a local stdio command. Case-insensitive to match the
// backend's is_stdio(), so all layers split http-vs-command identically.
function isHttpAddress(value: string): boolean {
const trimmed = value.trim().toLowerCase();
return trimmed.startsWith("http://") || trimmed.startsWith("https://");
}
function isValidAddress(value: string): boolean {
const trimmed = value.trim();
if (!trimmed) return false;
if (isHttpAddress(trimmed)) {
try {
const parsed = new URL(trimmed);
return parsed.protocol === "http:" || parsed.protocol === "https:";
} catch {
return false;
}
}
// Anything else is treated as a local command (stdio); the backend gates
// whether stdio servers are allowed on this host. Reject other URL schemes
// only when the command itself is a URL; "://" is fine inside an argument
// (e.g. a database connection string passed to the server).
return !trimmed.split(/\s+/)[0].includes("://");
}
function HeadersEditor({
rows,
onChange,
stdio,
}: {
rows: HeaderRow[];
onChange: (rows: HeaderRow[]) => void;
// stdio servers reuse this editor for environment variables instead of headers.
stdio: boolean;
}) {
const update = (id: string, patch: Partial<HeaderRow>) =>
onChange(rows.map((row) => (row.id === id ? { ...row, ...patch } : row)));
const add = () =>
onChange([...rows, { id: newRowId(), key: "", value: "" }]);
const remove = (id: string) =>
onChange(rows.filter((row) => row.id !== id));
const copy = stdio
? {
label: "Environment variables",
add: "Add variable",
keyPlaceholder: "Variable name",
valuePlaceholder: "Variable value",
remove: "Remove variable",
}
: {
label: "Custom headers",
add: "Add header",
keyPlaceholder: "Header name",
valuePlaceholder: "Header value",
remove: "Remove header",
};
return (
<>
<div className="flex items-center justify-between">
<Label className="text-sm">{copy.label}</Label>
<Button type="button" variant="ghost" size="sm" onClick={add}>
<HugeiconsIcon icon={PlusSignIcon} size={14} />
{copy.add}
</Button>
</div>
{rows.length === 0 ? (
<div className="text-xs text-muted-foreground">
{stdio ? (
"Optional. Environment variables passed to the server process."
) : (
<>
Optional. Add an <code>Authorization</code> header here for servers
that require auth.
</>
)}
</div>
) : (
<div className="flex flex-col gap-2">
{rows.map((row) => (
<div key={row.id} className="flex items-center gap-2">
<Input
value={row.key}
placeholder={copy.keyPlaceholder}
onChange={(e) => update(row.id, { key: e.target.value })}
/>
<Input
value={row.value}
placeholder={copy.valuePlaceholder}
onChange={(e) => update(row.id, { value: e.target.value })}
/>
<Button
type="button"
variant="ghost"
size="icon"
onClick={() => remove(row.id)}
aria-label={copy.remove}
>
<HugeiconsIcon icon={Delete02Icon} size={14} />
</Button>
</div>
))}
</div>
)}
</>
);
}
export interface ChatMcpServersDialogProps {
open: boolean;
onOpenChange: (open: boolean) => void;
}
type View =
| { kind: "list" }
| { kind: "create" }
| { kind: "edit"; id: string };
export function ChatMcpServersDialog({
open,
onOpenChange,
}: ChatMcpServersDialogProps) {
const [servers, setServers] = useState<McpServerConfig[]>([]);
const [loading, setLoading] = useState(false);
const [view, setView] = useState<View>({ kind: "list" });
const [form, setForm] = useState<FormState>(EMPTY_FORM);
const [saving, setSaving] = useState(false);
const [testing, setTesting] = useState(false);
const [refreshingId, setRefreshingId] = useState<string | null>(null);
const refresh = useCallback(async () => {
setLoading(true);
try {
const rows = await listMcpServers();
setServers(rows);
} catch (err) {
toast.error("Failed to load MCP servers", {
description: err instanceof Error ? err.message : String(err),
});
} finally {
setLoading(false);
}
}, []);
useEffect(() => {
if (!open) return;
refresh();
}, [open, refresh]);
function startCreate() {
setView({ kind: "create" });
setForm(EMPTY_FORM);
}
function startEdit(server: McpServerConfig) {
setView({ kind: "edit", id: server.id });
setForm({
displayName: server.display_name,
url: server.url,
headers: headersFromObject(server.headers ?? {}),
useOauth: server.use_oauth ?? false,
});
}
function cancelForm() {
setView({ kind: "list" });
setForm(EMPTY_FORM);
}
async function testConnection() {
const trimmedUrl = form.url.trim();
if (!isValidAddress(trimmedUrl)) {
toast.error("Enter an http(s):// URL or a local command first");
return;
}
setTesting(true);
try {
const result = await testMcpServer({
url: trimmedUrl,
headers: headersToObject(form.headers),
useOauth: form.useOauth,
});
if (result.ok) {
toast.success(
`Connected (${result.tool_count} tool${result.tool_count === 1 ? "" : "s"})`,
);
} else {
toast.error("Connection failed", {
description: result.error ?? "Unknown error",
});
}
} catch (err) {
toast.error("Connection test failed", {
description: err instanceof Error ? err.message : String(err),
});
} finally {
setTesting(false);
}
}
async function submitForm() {
const trimmedName = form.displayName.trim();
const trimmedUrl = form.url.trim();
if (!trimmedName) {
toast.error("Display name is required");
return;
}
if (!trimmedUrl) {
toast.error("URL or command is required");
return;
}
if (!isValidAddress(trimmedUrl)) {
toast.error("Enter an http(s):// URL or a local command");
return;
}
setSaving(true);
try {
const headers = headersToObject(form.headers);
if (view.kind === "edit") {
await updateMcpServer(view.id, {
displayName: trimmedName,
url: trimmedUrl,
headers: headers ?? null,
useOauth: form.useOauth,
});
toast.success("MCP server updated");
} else {
await createMcpServer({
displayName: trimmedName,
url: trimmedUrl,
headers: headers,
useOauth: form.useOauth,
});
toast.success("MCP server added");
}
cancelForm();
await refresh();
} catch (err) {
toast.error("Save failed", {
description: err instanceof Error ? err.message : String(err),
});
} finally {
setSaving(false);
}
}
async function removeServer(server: McpServerConfig) {
const ok = window.confirm(`Delete MCP server "${server.display_name}"?`);
if (!ok) return;
try {
await deleteMcpServer(server.id);
await refresh();
} catch (err) {
toast.error("Delete failed", {
description: err instanceof Error ? err.message : String(err),
});
}
}
async function toggleEnabled(server: McpServerConfig, next: boolean) {
// Optimistic update so the switch doesn't snap back during the round-trip.
setServers((rows) =>
rows.map((row) =>
row.id === server.id ? { ...row, is_enabled: next } : row,
),
);
try {
await updateMcpServer(server.id, { isEnabled: next });
} catch (err) {
setServers((rows) =>
rows.map((row) =>
row.id === server.id ? { ...row, is_enabled: !next } : row,
),
);
toast.error("Update failed", {
description: err instanceof Error ? err.message : String(err),
});
}
}
async function refreshTools(server: McpServerConfig) {
setRefreshingId(server.id);
try {
const result = await refreshMcpServerTools(server.id);
if (result.ok) {
toast.success(
`Refreshed "${server.display_name}" (${result.tool_count} tool${result.tool_count === 1 ? "" : "s"})`,
);
} else {
toast.error(`Refresh failed for "${server.display_name}"`, {
description: result.error ?? "Unknown error",
});
}
} catch (err) {
toast.error("Refresh failed", {
description: err instanceof Error ? err.message : String(err),
});
} finally {
setRefreshingId(null);
}
}
const showForm = view.kind !== "list";
// A local stdio command uses env vars, not headers or OAuth.
const addressIsCommand =
form.url.trim() !== "" && !isHttpAddress(form.url);
return (
<Dialog open={open} onOpenChange={onOpenChange}>
<DialogContent className="max-w-2xl">
<DialogHeader>
<DialogTitle>MCP Servers</DialogTitle>
<DialogDescription>
Register remote (HTTP) or local (stdio command) MCP servers.
</DialogDescription>
</DialogHeader>
{showForm ? (
<div className="flex flex-col gap-4">
<div className="grid gap-2">
<Label htmlFor="mcp-display-name">Display name</Label>
<Input
id="mcp-display-name"
value={form.displayName}
onChange={(e) =>
setForm((prev) => ({ ...prev, displayName: e.target.value }))
}
placeholder="e.g. GitHub MCP"
/>
</div>
<div className="grid gap-2">
<Label htmlFor="mcp-url">URL or command</Label>
<Input
id="mcp-url"
value={form.url}
onChange={(e) =>
setForm((prev) => ({ ...prev, url: e.target.value }))
}
placeholder="https://example.com/mcp or npx -y @modelcontextprotocol/server-filesystem /tmp"
/>
<span className="text-xs text-muted-foreground">
An http(s) URL for a remote server, or a local command to run an
stdio server (desktop app only).
</span>
</div>
{!addressIsCommand && (
<div className="flex items-start justify-between gap-3">
<div className="flex flex-col gap-0.5">
<Label className="text-sm" htmlFor="mcp-oauth">
Use OAuth sign-in
</Label>
<span className="text-xs text-muted-foreground">
For servers that require browser-based authentication
(GitHub, Linear, etc.). A browser window will open on first
connect.
</span>
</div>
<Switch
id="mcp-oauth"
checked={form.useOauth}
onCheckedChange={(useOauth) =>
setForm((prev) => ({ ...prev, useOauth }))
}
/>
</div>
)}
<HeadersEditor
rows={form.headers}
onChange={(headers) => setForm((prev) => ({ ...prev, headers }))}
stdio={addressIsCommand}
/>
<div className="flex items-center justify-between gap-2 pt-2">
<Button
type="button"
variant="outline"
size="sm"
onClick={testConnection}
disabled={testing || saving || !form.url.trim()}
>
{testing ? <Spinner /> : null}
Test connection
</Button>
<div className="flex gap-2">
<Button variant="ghost" onClick={cancelForm} disabled={saving}>
Cancel
</Button>
<Button onClick={submitForm} disabled={saving}>
{saving ? <Spinner /> : null}
{view.kind === "edit" ? "Save changes" : "Add server"}
</Button>
</div>
</div>
</div>
) : (
<div className="flex min-w-0 flex-col gap-3">
<div className="flex justify-end">
<Button size="sm" onClick={startCreate}>
<HugeiconsIcon icon={PlusSignIcon} size={14} />
Add server
</Button>
</div>
{loading ? (
<div className="flex justify-center py-6">
<Spinner />
</div>
) : servers.length === 0 ? (
<div className="rounded-md border border-dashed py-6 text-center text-sm text-muted-foreground">
No MCP servers configured yet.
</div>
) : (
<ul className="flex flex-col divide-y rounded-md border">
{servers.map((server) => (
<li
key={server.id}
className="flex items-center justify-between gap-3 px-3 py-2"
>
<div className="min-w-0 flex-1">
<div className="truncate font-medium">
{server.display_name}
</div>
<div className="truncate text-xs text-muted-foreground">
{server.url}
</div>
</div>
<div className="flex items-center gap-1">
<Switch
checked={server.is_enabled}
onCheckedChange={(next) => toggleEnabled(server, next)}
aria-label="Enable server"
/>
<Button
type="button"
variant="ghost"
size="icon"
onClick={() => refreshTools(server)}
aria-label="Refresh tools"
title="Refresh tools from this server"
disabled={refreshingId === server.id}
>
{refreshingId === server.id ? (
<Spinner />
) : (
<RefreshCwIcon size={14} />
)}
</Button>
<Button
type="button"
variant="ghost"
size="icon"
onClick={() => startEdit(server)}
aria-label="Edit server"
>
<HugeiconsIcon icon={Edit03Icon} size={14} />
</Button>
<Button
type="button"
variant="ghost"
size="icon"
onClick={() => removeServer(server)}
aria-label="Delete server"
>
<HugeiconsIcon icon={Delete02Icon} size={14} />
</Button>
</div>
</li>
))}
</ul>
)}
</div>
)}
</DialogContent>
</Dialog>
);
}

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

@ -728,7 +728,10 @@ export function ChatPage(): ReactElement {
const reasoningCaps = getExternalReasoningCapabilities(
provider?.providerType,
selection.modelId,
{ isReasoningProvider: provider?.isReasoningModel === true },
{
isReasoningProvider: provider?.isReasoningModel === true,
baseUrl: provider?.baseUrl ?? null,
},
);
const state = useChatRuntimeStore.getState();
const preferredEffort = state.reasoningEffort;
@ -769,6 +772,8 @@ export function ChatPage(): ReactElement {
: state.reasoningEffort;
const supportsBuiltinWebSearch = providerSupportsBuiltinWebSearch(
provider?.providerType,
selection.modelId,
provider?.baseUrl,
);
const supportsBuiltinCodeExecution = providerSupportsBuiltinCodeExecution(
provider?.providerType,
@ -968,6 +973,7 @@ export function ChatPage(): ReactElement {
{
isReasoningProvider:
selectedProvider?.isReasoningModel === true,
baseUrl: selectedProvider?.baseUrl ?? null,
},
);
const preferredEffort = store.reasoningEffort;
@ -1009,6 +1015,8 @@ export function ChatPage(): ReactElement {
store.setCheckpoint(value, null);
const supportsBuiltinWebSearch = providerSupportsBuiltinWebSearch(
selectedProvider?.providerType,
selectedExternal?.modelId,
selectedProvider?.baseUrl,
);
const supportsBuiltinCodeExecution = providerSupportsBuiltinCodeExecution(
selectedProvider?.providerType,

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

@ -91,6 +91,8 @@ import {
providerSupportsFastMode,
} from "./provider-capabilities";
import { useChatRuntimeStore } from "./stores/chat-runtime-store";
import { ChatMcpServersDialog } from "./chat-mcp-servers-dialog";
import { listMcpServers } from "./api/mcp-servers-api";
import type { InferenceParams } from "./types/runtime";
export { defaultInferenceParams, type Preset } from "./presets/preset-policy";
@ -1341,6 +1343,12 @@ export function ChatSettingsPanel({
</div>
</CollapsibleSection>
) : null}
{!isExternalModel ? (
<CollapsibleSection label="MCP Servers">
<McpServersSection />
</CollapsibleSection>
) : null}
</div>
</div>
<Dialog
@ -1504,6 +1512,74 @@ function AutoHealToolCallsToggle() {
);
}
function McpServersSection() {
const mcpEnabledForChat = useChatRuntimeStore((s) => s.mcpEnabledForChat);
const setMcpEnabledForChat = useChatRuntimeStore(
(s) => s.setMcpEnabledForChat,
);
const [enabledServerCount, setEnabledServerCount] = useState<number | null>(
null,
);
const [dialogOpen, setDialogOpen] = useState(false);
const [refreshTick, setRefreshTick] = useState(0);
useEffect(() => {
let cancelled = false;
listMcpServers()
.then((rows) => {
if (cancelled) return;
setEnabledServerCount(rows.filter((row) => row.is_enabled).length);
})
.catch(() => {
if (!cancelled) setEnabledServerCount(0);
});
return () => {
cancelled = true;
};
}, [refreshTick]);
return (
<div className="flex flex-col gap-3 pt-1">
<div className="flex items-center justify-between gap-3">
<div className="flex min-w-0 items-center gap-1.5">
<span className="min-w-0 text-[13px] font-medium leading-[1.25] tracking-nav text-nav-fg">
Use MCP Servers
</span>
<InfoHint>
When on, every server marked enabled in the manage dialog is
attached to this chat's tool list.
</InfoHint>
</div>
<Switch
className="panel-switch"
checked={mcpEnabledForChat}
onCheckedChange={setMcpEnabledForChat}
disabled={enabledServerCount === 0 && !mcpEnabledForChat}
/>
</div>
<div className="flex items-center justify-between">
<span className="text-[11px] text-muted-foreground">
{enabledServerCount === null
? "Loading…"
: enabledServerCount === 0
? "No servers configured"
: `${enabledServerCount} server${enabledServerCount === 1 ? "" : "s"} enabled`}
</span>
<Button variant="ghost" size="sm" onClick={() => setDialogOpen(true)}>
Manage…
</Button>
</div>
<ChatMcpServersDialog
open={dialogOpen}
onOpenChange={(next) => {
setDialogOpen(next);
if (!next) setRefreshTick((tick) => tick + 1);
}}
/>
</div>
);
}
function ChatTemplateFields() {
const defaultTemplate = useChatRuntimeStore((s) => s.defaultChatTemplate);
const override = useChatRuntimeStore((s) => s.chatTemplateOverride);

7
studio/frontend/src/features/chat/external-providers.ts
View file

@ -37,6 +37,13 @@ export interface ExternalProviderConfig {
updatedAt: number;
}
// Gemini supports prompt caching, but the wire flow requires a
// separate POST to /v1beta/cachedContents to create the cache before
// the generateContent call can reference it; the boolean Studio
// currently emits on enable_prompt_caching is not enough on its own.
// Until that two-step orchestration ships we keep the picker off so
// the toggle does not silently no-op for Gemini users. See
// https://ai.google.dev/gemini-api/docs/caching.
const PROMPT_CACHING_PROVIDER_TYPES = new Set(["openai", "anthropic"]);
export function supportsProviderPromptCaching(

15
studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts
View file

@ -27,6 +27,7 @@ import {
CHAT_REASONING_ENABLED_KEY,
loadOptionalBool,
type ReasoningEffort,
resolveToolsEnabledOnLoad,
useChatRuntimeStore,
} from "../stores/chat-runtime-store";
import {
@ -698,14 +699,12 @@ export function useChatModelRuntime() {
reasoningEffort: clampedReasoningEffort,
supportsPreserveThinking: loadResponse.supports_preserve_thinking ?? false,
supportsTools,
toolsEnabled:
reloadingSameModel && supportsTools
? stateBeforeUnload.toolsEnabled
: supportsTools,
codeToolsEnabled:
reloadingSameModel && supportsTools
? stateBeforeUnload.codeToolsEnabled
: supportsTools,
...(reloadingSameModel && supportsTools
? {
toolsEnabled: stateBeforeUnload.toolsEnabled,
codeToolsEnabled: stateBeforeUnload.codeToolsEnabled,
}
: resolveToolsEnabledOnLoad(supportsTools)),
kvCacheDtype: loadedKv,
loadedKvCacheDtype: loadedKv,
speculativeType: loadedSpec,

245
studio/frontend/src/features/chat/provider-capabilities.ts
View file

@ -189,7 +189,27 @@ function _inferProviderFromOpenrouterId(
*/
export function providerSupportsBuiltinWebSearch(
providerType: string | null | undefined,
modelId?: string | null | undefined,
baseUrl?: string | null | undefined,
): boolean {
// Gemini ships grounded search via `tools: [{googleSearch: {}}]` on
// every chat-capable model. Most image-tier ids (`-image`,
// `nano-banana`) reject text-tool wiring because the
// responseModalities path is mutually exclusive with text tools, but
// Google explicitly documents Search grounding on the Gemini 3 image
// family (gemini-3-pro-image-preview, gemini-3.1-flash-image-preview,
// nano-banana-pro). Allow Search on those; hide on older image ids.
// Custom Gemini OpenAI-compat proxies (non-Google bases) skip the
// native translator on the backend, so native tool envelopes never
// reach them -- hide the pill there.
if (providerType === "gemini") {
if (isGeminiCustomOpenAICompatBase(baseUrl)) return false;
const normalized = modelId?.trim().toLowerCase() ?? "";
if (normalized && isGeminiImageModel(normalized)) {
return geminiImageModelAllowsGoogleSearch(normalized);
}
return true;
}
return (
providerType === "openai" ||
providerType === "anthropic" ||
@ -319,6 +339,20 @@ export function providerSupportsBuiltinCodeExecution(
normalized.startsWith(prefix),
);
}
if (providerType === "gemini") {
// Gemini's `tools: [{codeExecution: {}}]` is supported on every
// chat-capable model. Image-tier ids (`-image`, `nano-banana`)
// reject text-tool wiring because the inline-image path is
// mutually exclusive with codeExecution. Custom Gemini
// OpenAI-compat proxies skip the native translator on the
// backend, so native codeExecution envelopes do not reach them.
// Wire-up lives in `_stream_gemini` on the backend; output comes
// back inline as executableCode/codeExecutionResult parts. See
// https://ai.google.dev/gemini-api/docs/code-execution.
if (isGeminiCustomOpenAICompatBase(baseUrl)) return false;
if (isGeminiImageModel(normalized)) return false;
return normalized.startsWith("gemini-");
}
return false;
}
@ -351,12 +385,75 @@ export function providerSupportsBuiltinImageGeneration(
modelId: string | null | undefined,
baseUrl?: string | null,
): boolean {
if (providerType !== "openai") return false;
if (!isOpenAICloudBaseUrl(baseUrl)) return false;
const normalized = modelId?.trim().toLowerCase() ?? "";
if (!normalized) return false;
return OPENAI_IMAGE_GENERATION_MODEL_PREFIXES.some((prefix) =>
normalized.startsWith(prefix),
if (providerType === "openai") {
if (!isOpenAICloudBaseUrl(baseUrl)) return false;
return OPENAI_IMAGE_GENERATION_MODEL_PREFIXES.some((prefix) =>
normalized.startsWith(prefix),
);
}
if (providerType === "gemini") {
// Gemini's Nano Banana image-output ids carry either `-image` (e.g.
// `gemini-2.5-flash-image`, `gemini-3.1-flash-image-preview`) or the
// `nano-banana` alias (`nano-banana-pro-preview`). The backend flips
// generationConfig.responseModalities to ["TEXT", "IMAGE"] when one
// is picked, and translates inlineData parts into the same image_b64
// tool_end envelope the OpenAI path emits so the chat UI renders the
// picture inline. Custom Gemini OpenAI-compat proxies skip the
// native translator on the backend, so hide the image pill there.
// See https://ai.google.dev/gemini-api/docs/image-generation.
if (isGeminiCustomOpenAICompatBase(baseUrl)) return false;
return normalized.includes("-image") || normalized.includes("nano-banana");
}
return false;
}
/**
* Whether `modelId` is a Gemini image-output id (Nano Banana family).
* Mirrors the backend's `is_image_picker_model` guard so the frontend
* hides text-only tool pills (web_search, code_execution) for these.
*/
function isGeminiImageModel(modelId: string): boolean {
const m = modelId.toLowerCase();
return m.includes("-image") || m.includes("nano-banana");
}
/**
* Whether the saved Gemini connection points at a custom
* OpenAI-compatible gateway (any non-Google host). The backend
* `_is_openai_compatible` mirrors this to route those connections
* through `/chat/completions` instead of the native translator, so
* native Gemini tool envelopes (googleSearch, codeExecution,
* responseModalities) never reach them. Hide the corresponding
* Studio pills here so the request, builder, and UI agree.
*/
export function isGeminiCustomOpenAICompatBase(
baseUrl: string | null | undefined,
): boolean {
if (!baseUrl) return false;
try {
const host = new URL(baseUrl).hostname.toLowerCase();
return host.length > 0 && host !== "generativelanguage.googleapis.com";
} catch {
return false;
}
}
/**
* Whether the given Gemini image model supports `tools: [{googleSearch: {}}]`.
* Google documents Search grounding on the Gemini 3 image family
* (gemini-3-pro-image-preview, gemini-3.1-flash-image-preview,
* "Nano Banana Pro"); older image ids (gemini-2.5-flash-image) reject
* it with "Search as tool is not enabled for this model".
*/
function geminiImageModelAllowsGoogleSearch(modelId: string): boolean {
const m = modelId.toLowerCase();
return (
m.startsWith("gemini-3-pro-image") ||
m.startsWith("gemini-3.1-flash-image") ||
m.startsWith("nano-banana-pro") ||
m.startsWith("nano-banana-2")
);
}
@ -431,7 +528,20 @@ const PROVIDER_CAPABILITIES: Record<string, ProviderCapabilities> = {
presencePenalty: false,
},
mistral: OPENAI_COMPAT_BASE,
gemini: OPENAI_COMPAT_BASE,
// Gemini's native generationConfig accepts temperature, topP, topK and
// presencePenalty (plus a separate frequencyPenalty we do not surface
// today). minP and repetitionPenalty are not part of the contract --
// see https://ai.google.dev/api/rest/v1beta/GenerationConfig. Backend
// request shaping lives in _stream_gemini in
// studio/backend/core/inference/external_provider.py.
gemini: {
temperature: true,
topP: true,
topK: true,
minP: false,
repetitionPenalty: false,
presencePenalty: true,
},
// Kimi k2.5/k2.6 are reasoning-class — the API locks temperature and
// top_p to fixed defaults and 400s on any other value:
// "invalid temperature: only 1 is allowed for this model".
@ -643,6 +753,119 @@ function resolveKimiReasoningCapabilities(modelId: string): ExternalReasoningCap
return withEnableThinkingStyle();
}
// Gemini's thinking ladder.
// - Gemini 3.x (3 / 3.1 / 3.5, Pro + Flash + Flash-Lite) and the
// gemini-pro-latest / gemini-flash-latest aliases use the new
// `thinkingConfig.thinkingLevel` string field (LOW/MEDIUM/HIGH/
// MINIMAL). Pro tier rejects MINIMAL.
// - Gemini 2.5 Flash + 2.5 Pro stay on the integer
// `thinkingConfig.thinkingBudget` (0=off on Flash, -1=dynamic,
// N>0=cap; Pro rejects 0).
// - 2.5 Flash-Lite: no native thinking surfaced; leave it off.
// - Image-tier ids (`*-image*`, `nano-banana-pro-preview`): image
// generation path -- no reasoning controls.
const GEMINI3_PRO_PREFIXES = [
"gemini-3.5-pro",
"gemini-3.1-pro",
"gemini-3-pro-preview",
"gemini-pro-latest",
];
const GEMINI3_FLASH_PREFIXES = [
"gemini-3.5-flash",
"gemini-3.1-flash",
"gemini-3-flash",
"gemini-flash-latest",
"gemini-flash-lite-latest",
];
const GEMINI25_PRO_PREFIXES = [
"gemini-2.5-pro",
];
const GEMINI25_FLASH_PREFIXES = [
"gemini-2.5-flash",
];
const GEMINI_IMAGE_HINTS = [
"-image",
"nano-banana",
];
function resolveGeminiReasoningCapabilities(
modelId: string,
): ExternalReasoningCapabilities {
const m = modelId.toLowerCase();
if (GEMINI_IMAGE_HINTS.some((h) => m.includes(h))) {
// Image generation; no thinking knob.
return withEnableThinkingStyle();
}
// Gemini 2.5 Flash-Lite supports `thinkingBudget` with `0` = off and
// a positive range starting at 512 (the backend maps "minimal" to
// that floor at external_provider._stream_gemini). Check this branch
// BEFORE the broader `gemini-2.5-flash` prefix.
// https://ai.google.dev/gemini-api/docs/thinking
if (m.startsWith("gemini-2.5-flash-lite")) {
return withReasoningEffortStyle({
supportsReasoning: true,
supportsReasoningOff: true,
reasoningEffortLevels: [
"none",
"minimal",
"low",
"medium",
"high",
"max",
] as const,
});
}
if (GEMINI3_PRO_PREFIXES.some((p) => m.startsWith(p))) {
// Gemini 3.x Pro: thinkingLevel supports low/medium/high per
// https://ai.google.dev/gemini-api/docs/thinking and
// https://docs.cloud.google.com/vertex-ai/generative-ai/docs/models/gemini/3-1-pro.
// Cannot fully disable thinking; "minimal" is rejected on Pro.
return withReasoningEffortStyle({
supportsReasoning: true,
supportsReasoningOff: false,
reasoningEffortLevels: ["low", "medium", "high"] as const,
});
}
if (GEMINI3_FLASH_PREFIXES.some((p) => m.startsWith(p))) {
// Gemini 3 Flash: thinkingLevel minimal/low/medium/high. Minimal
// is the closest to "off" Google offers on Gemini 3.
return withReasoningEffortStyle({
supportsReasoning: true,
supportsReasoningOff: false,
reasoningEffortLevels: [
"minimal",
"low",
"medium",
"high",
] as const,
});
}
if (GEMINI25_PRO_PREFIXES.some((p) => m.startsWith(p))) {
// Gemini 2.5 Pro: thinkingBudget cannot be 0 (API rejects with
// "only works in thinking mode"); backend coerces to a small
// positive budget. The picker still hides the off switch.
return withReasoningEffortStyle({
supportsReasoning: true,
supportsReasoningOff: false,
reasoningEffortLevels: ["low", "medium", "high", "max"] as const,
});
}
if (GEMINI25_FLASH_PREFIXES.some((p) => m.startsWith(p))) {
// Gemini 2.5 Flash: thinkingBudget supports 0 = off cleanly.
return withReasoningEffortStyle({
supportsReasoning: true,
supportsReasoningOff: true,
reasoningEffortLevels: [
"none",
"low",
"medium",
"high",
"max",
] as const,
});
}
return withEnableThinkingStyle();
}
function resolveMistralReasoningCapabilities(modelId: string): ExternalReasoningCapabilities {
if (modelId === "magistral-medium-latest") {
return withReasoningEffortStyle({
@ -665,6 +888,8 @@ function resolveMistralReasoningCapabilities(modelId: string): ExternalReasoning
export interface ExternalReasoningResolveOptions {
/** vLLM connection flagged as a reasoning model in provider config. */
isReasoningProvider?: boolean;
/** Provider base URL; used to detect custom Gemini OAI-compat gateways. */
baseUrl?: string | null;
}
// vLLM has no per-model reasoning signal on OpenAI-compat — pin via user toggle.
@ -740,6 +965,16 @@ export function getExternalReasoningCapabilities(
}
if (isKimiProvider) return resolveKimiReasoningCapabilities(modelForMatching);
if (isMistralProvider) return resolveMistralReasoningCapabilities(modelForMatching);
if (normalizedProvider === "gemini") {
// Custom Gemini OAI-compat gateways (LiteLLM, proxies) route
// through /chat/completions which drops the Gemini-native
// thinkingConfig payload. Hide the native thinking ladder so the
// UI does not advertise a control the backend cannot honor.
if (isGeminiCustomOpenAICompatBase(options?.baseUrl)) {
return withEnableThinkingStyle();
}
return resolveGeminiReasoningCapabilities(modelForMatching);
}
if (!isOpenAIProvider && !isAnthropicProvider) {
return withEnableThinkingStyle();
}

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

@ -396,6 +396,7 @@ export function SharedComposer({
{
isReasoningProvider:
selectedExternalProvider?.isReasoningModel === true,
baseUrl: selectedExternalProvider?.baseUrl ?? null,
},
)
: null;
@ -449,16 +450,36 @@ export function SharedComposer({
const supportsBuiltinWebFetch = providerSupportsBuiltinWebFetch(
selectedExternalProvider?.providerType,
);
const searchDisabled =
!modelLoaded || !(supportsTools || supportsBuiltinWebSearch);
const codeDisabled =
!modelLoaded || !(supportsTools || supportsBuiltinCodeExecution);
// Images pill is only ever lit on OpenAI cloud's Responses-API models.
// No local tool runtime fallback because the only image-generation
// server tool we wire today is OpenAI's; local models cannot dispatch
// it. Hidden entirely when the active model does not advertise it so
// the pill row stays compact for providers without the capability.
// Gemini rejects codeExecution alongside image modalities. Search is
// blocked on older Gemini image ids but allowed on Gemini 3 image
// models -- supportsBuiltinWebSearch already encodes the per-model
// allowance, so we only disable Code unconditionally in Gemini
// image mode.
const isExternalGemini = selectedExternalProvider?.providerType === "gemini";
const imageDisabled = !modelLoaded || !supportsBuiltinImageGeneration;
const imageModeDisablesCode =
isExternalGemini && imageToolsEnabled && !imageDisabled;
// Image-tier Gemini models always reject codeExecution and reject
// web_search on older ids (Gemini 3.x Pro/Flash allow it -- encoded
// in supportsBuiltinWebSearch). Don't let the local `supportsTools`
// runtime flag re-enable a pill the Gemini backend will silently
// drop. Detect "external provider is Gemini AND model is image-tier"
// and gate strictly on the provider builtin support.
const isGeminiImageTier =
isExternalGemini && supportsBuiltinImageGeneration;
const searchDisabled =
!modelLoaded ||
(isGeminiImageTier
? !supportsBuiltinWebSearch
: !(supportsTools || supportsBuiltinWebSearch));
const codeDisabled =
!modelLoaded ||
(isGeminiImageTier
? true
: !(supportsTools || supportsBuiltinCodeExecution)) ||
imageModeDisablesCode;
// Images pill is only ever lit on OpenAI cloud's Responses-API models
// and Gemini Nano Banana family. No local tool runtime fallback.
const showImagePill = supportsBuiltinImageGeneration;
// Fetch pill: Anthropic-only (web_fetch_20250910 / web_fetch_20260209).
const webFetchDisabled = !modelLoaded || !supportsBuiltinWebFetch;
@ -893,7 +914,7 @@ export function SharedComposer({
side="bottom"
variant="ghost"
size="icon"
className="size-8.5 rounded-full p-1 font-semibold text-xs hover:bg-muted-foreground/15 dark:border-muted-foreground/15 dark:hover:bg-muted-foreground/30"
className="size-8.5 rounded-full p-1 font-semibold text-xs hover:bg-muted-foreground/15 dark:hover:bg-muted-foreground/30"
onClick={() => {
// The picker accepts both image and audio. Don't gate the
// button on image-availability — addFiles still filters

26
studio/frontend/src/features/chat/stores/chat-runtime-store.ts
View file

@ -27,6 +27,7 @@ export const CHAT_REASONING_ENABLED_KEY = "unsloth_chat_reasoning_enabled";
export const CHAT_TOOLS_ENABLED_KEY = "unsloth_chat_tools_enabled";
export const CHAT_CODE_TOOLS_ENABLED_KEY = "unsloth_chat_code_tools_enabled";
export const CHAT_IMAGE_TOOLS_ENABLED_KEY = "unsloth_chat_image_tools_enabled";
export const CHAT_MCP_ENABLED_KEY = "unsloth_chat_mcp_enabled";
export const CHAT_WEB_FETCH_TOOLS_ENABLED_KEY =
"unsloth_chat_web_fetch_tools_enabled";
@ -183,6 +184,23 @@ export function loadOptionalBool(key: string): boolean | null {
}
}
/**
* Resolve the web-search / code-execution pill state to apply when a model
* loads. Honors the user's persisted preference so loading a tool-capable
* model never silently re-enables a pill the user turned off; falls back to
* the model's capability only when no preference has been expressed.
*/
export function resolveToolsEnabledOnLoad(supportsTools: boolean): {
toolsEnabled: boolean;
codeToolsEnabled: boolean;
} {
if (!supportsTools) return { toolsEnabled: false, codeToolsEnabled: false };
return {
toolsEnabled: loadOptionalBool(CHAT_TOOLS_ENABLED_KEY) ?? true,
codeToolsEnabled: loadOptionalBool(CHAT_CODE_TOOLS_ENABLED_KEY) ?? true,
};
}
function saveBool(key: string, value: boolean): void {
if (!canUseStorage()) return;
try {
@ -282,6 +300,7 @@ type ChatRuntimeStore = {
toolsEnabled: boolean;
codeToolsEnabled: boolean;
imageToolsEnabled: boolean;
mcpEnabledForChat: boolean;
/**
* Fetch pill state, independent of `toolsEnabled` (Search). Only
* consulted when `providerSupportsBuiltinWebFetch` is true.
@ -349,6 +368,7 @@ type ChatRuntimeStore = {
setToolsEnabled: (enabled: boolean, options?: { persist?: boolean }) => void;
setCodeToolsEnabled: (enabled: boolean) => void;
setImageToolsEnabled: (enabled: boolean) => void;
setMcpEnabledForChat: (enabled: boolean) => void;
setWebFetchToolsEnabled: (enabled: boolean) => void;
setToolStatus: (status: string | null) => void;
setGeneratingStatus: (status: string | null) => void;
@ -599,6 +619,7 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set, get) => ({
toolsEnabled: loadBool(CHAT_TOOLS_ENABLED_KEY, false),
codeToolsEnabled: loadBool(CHAT_CODE_TOOLS_ENABLED_KEY, false),
imageToolsEnabled: loadBool(CHAT_IMAGE_TOOLS_ENABLED_KEY, false),
mcpEnabledForChat: loadBool(CHAT_MCP_ENABLED_KEY, false),
webFetchToolsEnabled: loadBool(CHAT_WEB_FETCH_TOOLS_ENABLED_KEY, false),
toolStatus: null,
generatingStatus: null,
@ -875,6 +896,11 @@ export const useChatRuntimeStore = create<ChatRuntimeStore>((set, get) => ({
saveBool(CHAT_IMAGE_TOOLS_ENABLED_KEY, imageToolsEnabled);
return { imageToolsEnabled };
}),
setMcpEnabledForChat: (mcpEnabledForChat) =>
set(() => {
saveBool(CHAT_MCP_ENABLED_KEY, mcpEnabledForChat);
return { mcpEnabledForChat };
}),
setWebFetchToolsEnabled: (webFetchToolsEnabled) =>
set(() => {
saveBool(CHAT_WEB_FETCH_TOOLS_ENABLED_KEY, webFetchToolsEnabled);

53
studio/frontend/src/features/chat/types/api.ts
View file

@ -219,9 +219,34 @@ export type OpenAIMessageContentPart =
export type OpenAIMessageContent = string | OpenAIMessageContentPart[];
/**
* OpenAI Chat Completions tool_call shape. Assistant turns echo back
* function/tool calls as `tool_calls`; the matching tool result rides
* on a separate `role="tool"` message keyed by `tool_call_id`.
* `extra_content.google.thought_signature` is the Gemini-specific
* round-trip field the backend translator both emits (on `delta.
* tool_calls`) and consumes (when rebuilding the native functionCall
* part on the next turn).
*/
export interface OpenAIToolCallPart {
id?: string;
type?: "function";
function?: {
name?: string;
arguments?: string;
};
extra_content?: unknown;
}
export interface OpenAIChatMessage {
role: "system" | "user" | "assistant";
content: OpenAIMessageContent;
role: "system" | "user" | "assistant" | "tool";
content: OpenAIMessageContent | null;
/** Assistant tool-call deltas, when the turn invoked a function tool. */
tool_calls?: OpenAIToolCallPart[];
/** `role="tool"` only: id matching `assistant.tool_calls[].id`. */
tool_call_id?: string;
/** `role="tool"` only: name of the function that produced the result. */
name?: string;
}
export interface OpenAIChatCompletionsRequest {
@ -262,7 +287,14 @@ export interface OpenAIChatCompletionsRequest {
external_model?: string;
encrypted_api_key?: string;
provider_base_url?: string | null;
enable_prompt_caching?: boolean | null;
/**
* Boolean toggle for OpenAI/Anthropic ephemeral cache_control. For
* Gemini the backend also accepts the cached-content resource name
* (`cachedContents/...`) as a string, which is forwarded as
* `generationConfig.cachedContent` on the native streamGenerateContent
* request.
*/
enable_prompt_caching?: boolean | string | null;
/**
* OpenAI shell-tool container id captured from the prior response in
* this chat thread. When set and the Code pill is on, the backend
@ -292,7 +324,20 @@ export interface OpenAIChatCompletionsRequest {
export interface OpenAIChatDelta {
role?: string;
content?: string;
content?: string | null;
/**
* Streamed assistant tool calls. The Gemini and OpenAI Responses
* translators emit incremental `tool_calls` deltas (function name +
* arguments fragments) so the chat-adapter can render tool cards as
* they arrive.
*/
tool_calls?: OpenAIToolCallPart[];
/**
* Provider-specific passthrough. Gemini ships `thoughtSignature`,
* citations, `native_part`, etc., here so the round-trip can replay
* them on follow-up turns without bleeding into other providers.
*/
extra_content?: Record<string, unknown>;
}
export interface OpenAIChatChunkChoice {

25
studio/frontend/src/features/profile/components/profile-personalization-panel.tsx
View file

@ -5,6 +5,7 @@ import { Button } from "@/components/ui/button";
import { Input } from "@/components/ui/input";
import { Label } from "@/components/ui/label";
import { getAuthToken } from "@/features/auth";
import { useT } from "@/i18n";
import { toastError, toastSuccess } from "@/shared/toast";
import { Camera } from "lucide-react";
import { useMemo, useRef, useState } from "react";
@ -37,6 +38,7 @@ function readPersistedProfile(): { displayName: string; avatarDataUrl: string |
}
export function ProfilePersonalizationPanel() {
const t = useT();
const displayName = useUserProfileStore((s) => s.displayName);
const avatarDataUrl = useUserProfileStore((s) => s.avatarDataUrl);
const setDisplayName = useUserProfileStore((s) => s.setDisplayName);
@ -60,11 +62,11 @@ export function ProfilePersonalizationPanel() {
setDisplayName(trimmed);
const persisted = readPersistedProfile();
if (persisted && persisted.displayName === trimmed) {
toastSuccess("Profile name saved");
toastSuccess(t("settings.profile.nameSaved"));
} else {
toastError(
"Could not persist profile name",
"Name updated for this session, but may not persist after reload.",
t("settings.profile.namePersistErrorTitle"),
t("settings.profile.namePersistErrorDescription"),
);
}
}
@ -78,17 +80,18 @@ export function ProfilePersonalizationPanel() {
setAvatarDataUrl(dataUrl);
const persisted = readPersistedProfile();
if (persisted && persisted.avatarDataUrl === dataUrl) {
toastSuccess("Profile photo updated");
toastSuccess(t("settings.profile.photoUpdated"));
} else {
toastError(
"Could not persist profile photo",
"Photo updated for this session, but may not persist after reload.",
t("settings.profile.photoPersistErrorTitle"),
t("settings.profile.photoPersistErrorDescription"),
);
}
} catch (e) {
const message = e instanceof Error ? e.message : "Could not use this image.";
const message =
e instanceof Error ? e.message : t("settings.profile.imageUseError");
setImageError(message);
toastError("Could not update profile photo", message);
toastError(t("settings.profile.photoUpdateErrorTitle"), message);
}
};
@ -115,7 +118,7 @@ export function ProfilePersonalizationPanel() {
type="button"
onClick={() => fileInputRef.current?.click()}
className="absolute right-0 bottom-0 -translate-x-[15.625%] -translate-y-[15.625%] flex size-8 items-center justify-center rounded-full border border-border bg-background text-foreground shadow-sm transition-colors hover:bg-muted focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring focus-visible:ring-offset-2 focus-visible:ring-offset-background"
aria-label="Change profile picture"
aria-label={t("settings.profile.changePicture")}
>
<Camera className="size-3.5" strokeWidth={2} />
</button>
@ -123,7 +126,7 @@ export function ProfilePersonalizationPanel() {
<div className="flex w-full max-w-[560px] flex-col gap-2">
<Label htmlFor="profile-display-name" className="text-xs font-medium text-muted-foreground">
Display name
{t("settings.profile.displayName")}
</Label>
<div className="flex items-center gap-2">
<Input
@ -142,7 +145,7 @@ export function ProfilePersonalizationPanel() {
className="h-10 min-w-0 flex-1 rounded-full text-sm"
/>
<Button type="button" size="sm" className="h-10 px-5" onClick={saveName} disabled={!hasNameChanges}>
Save
{t("common.save")}
</Button>
</div>
</div>

58
studio/frontend/src/features/settings/components/api-key-row.tsx
View file

@ -14,30 +14,39 @@ import {
MoreHorizontalIcon,
} from "@hugeicons/core-free-icons";
import { HugeiconsIcon } from "@hugeicons/react";
import { useT } from "@/i18n";
import { copyToClipboard } from "@/lib/copy-to-clipboard";
import type { ApiKey } from "../api/api-keys";
function relative(iso: string | null): string {
if (!iso) return "never";
type SettingsT = ReturnType<typeof useT>;
function relative(iso: string | null, t: SettingsT): string {
if (!iso) return t("settings.apiKeys.relativeNever");
const diff = Date.now() - new Date(iso).getTime();
const days = Math.floor(diff / 86400000);
if (days < 1) {
const hours = Math.floor(diff / 3600000);
if (hours < 1) return "just now";
return `${hours}h ago`;
if (hours < 1) return t("settings.apiKeys.relativeJustNow");
return t("settings.apiKeys.relativeHoursAgo", { count: hours });
}
if (days < 30) return `${days}d ago`;
if (days < 365) return `${Math.floor(days / 30)}mo ago`;
return `${Math.floor(days / 365)}y ago`;
if (days < 30) return t("settings.apiKeys.relativeDaysAgo", { count: days });
if (days < 365) {
return t("settings.apiKeys.relativeMonthsAgo", {
count: Math.floor(days / 30),
});
}
return t("settings.apiKeys.relativeYearsAgo", {
count: Math.floor(days / 365),
});
}
function expiresText(iso: string | null): string {
if (!iso) return "never";
function expiresText(iso: string | null, t: SettingsT): string {
if (!iso) return t("settings.apiKeys.relativeNever");
const diff = new Date(iso).getTime() - Date.now();
if (diff < 0) return "expired";
if (diff < 0) return t("settings.apiKeys.expired");
const days = Math.floor(diff / 86400000);
if (days < 1) return "today";
return `in ${days}d`;
if (days < 1) return t("settings.apiKeys.today");
return t("settings.apiKeys.inDays", { count: days });
}
export function ApiKeyRow({
@ -47,6 +56,7 @@ export function ApiKeyRow({
apiKey: ApiKey;
onRevoke: (key: ApiKey) => void;
}) {
const t = useT();
const prefix = `sk-unsloth-${apiKey.key_prefix}…`;
return (
<div className="group flex items-center gap-3 border-b border-border/60 px-1 py-3 last:border-b-0 transition-colors hover:bg-accent/40">
@ -64,11 +74,23 @@ export function ApiKeyRow({
</code>
</div>
<div className="flex flex-wrap gap-x-1.5 text-[11px] text-muted-foreground">
<span>Created {relative(apiKey.created_at)}</span>
<span>
{t("settings.apiKeys.created", {
value: relative(apiKey.created_at, t),
})}
</span>
<span>·</span>
<span>Used {relative(apiKey.last_used_at)}</span>
<span>
{t("settings.apiKeys.used", {
value: relative(apiKey.last_used_at, t),
})}
</span>
<span>·</span>
<span>Expires {expiresText(apiKey.expires_at)}</span>
<span>
{t("settings.apiKeys.expires", {
value: expiresText(apiKey.expires_at, t),
})}
</span>
</div>
</div>
<DropdownMenu>
@ -77,7 +99,7 @@ export function ApiKeyRow({
variant="ghost"
size="sm"
className="size-7 p-0 opacity-0 transition-opacity group-hover:opacity-100 data-[state=open]:opacity-100 max-sm:!opacity-100 max-sm:size-9"
aria-label={`Actions for ${apiKey.name}`}
aria-label={t("settings.apiKeys.actionsFor", { name: apiKey.name })}
>
<HugeiconsIcon icon={MoreHorizontalIcon} className="size-4" />
</Button>
@ -85,14 +107,14 @@ export function ApiKeyRow({
<DropdownMenuContent align="end">
<DropdownMenuItem onClick={async () => { await copyToClipboard(prefix); }}>
<HugeiconsIcon icon={Copy01Icon} className="size-3.5 mr-2" />
Copy prefix
{t("settings.apiKeys.copyPrefix")}
</DropdownMenuItem>
<DropdownMenuItem
onClick={() => onRevoke(apiKey)}
className="text-destructive focus:text-destructive"
>
<HugeiconsIcon icon={Delete02Icon} className="size-3.5 mr-2" />
Revoke token
{t("settings.apiKeys.revokeToken")}
</DropdownMenuItem>
</DropdownMenuContent>
</DropdownMenu>

19
studio/frontend/src/features/settings/components/create-key-form.tsx
View file

@ -3,6 +3,7 @@
import { Button } from "@/components/ui/button";
import { Input } from "@/components/ui/input";
import { useT } from "@/i18n";
import { cn } from "@/lib/utils";
import { useState } from "react";
import { createApiKey } from "../api/api-keys";
@ -21,6 +22,7 @@ export function CreateKeyForm({
onCreated: (rawKey: string) => void;
onError: (message: string) => void;
}) {
const t = useT();
const [name, setName] = useState("");
const [expiry, setExpiry] = useState<number | null>(null);
const [loading, setLoading] = useState(false);
@ -33,8 +35,11 @@ export function CreateKeyForm({
const result = await createApiKey(name.trim(), expiry);
onCreated(result.key);
setName("");
} catch (err) {
onError(err instanceof Error ? err.message : "Couldn't create access token.");
} catch {
// API helpers in ../api/api-keys.ts throw generic English Error
// messages; always use the translated message so zh-CN users do not
// see English text bleed through from internal exceptions.
onError(t("settings.apiKeys.createError"));
} finally {
setLoading(false);
}
@ -49,9 +54,9 @@ export function CreateKeyForm({
<Input
value={name}
onChange={(e) => setName(e.target.value)}
placeholder="Token name (e.g. production)"
placeholder={t("settings.apiKeys.tokenNamePlaceholder")}
className="h-8 min-w-[180px] flex-1 text-sm"
aria-label="New access token name"
aria-label={t("settings.apiKeys.newAccessTokenName")}
/>
<div className="inline-flex items-center rounded-md border border-border bg-background p-0.5">
{EXPIRY_PRESETS.map((p) => {
@ -69,13 +74,15 @@ export function CreateKeyForm({
: "text-muted-foreground hover:text-foreground",
)}
>
{p.label}
{p.value === null ? t("settings.apiKeys.never") : p.label}
</button>
);
})}
</div>
<Button type="submit" size="sm" disabled={loading || !name.trim()}>
{loading ? "Creating…" : "Create token"}
{loading
? t("settings.apiKeys.creating")
: t("settings.apiKeys.createToken")}
</Button>
</div>
</form>

14
studio/frontend/src/features/settings/components/key-reveal-card.tsx
View file

@ -2,6 +2,7 @@
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import { Button } from "@/components/ui/button";
import { useT } from "@/i18n";
import { copyToClipboard } from "@/lib/copy-to-clipboard";
import { cn } from "@/lib/utils";
import { Copy01Icon, Tick02Icon } from "@hugeicons/core-free-icons";
@ -15,6 +16,7 @@ export function KeyRevealCard({
rawKey: string;
onDone: () => void;
}) {
const t = useT();
const [copied, setCopied] = useState(false);
const handleCopy = async () => {
@ -32,7 +34,7 @@ export function KeyRevealCard({
className="size-3.5 text-emerald-600 dark:text-emerald-500"
/>
<span className="text-xs font-medium text-emerald-700 dark:text-emerald-500">
New access token created
{t("settings.apiKeys.newTokenCreated")}
</span>
</div>
<button
@ -43,7 +45,11 @@ export function KeyRevealCard({
"focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring focus-visible:ring-offset-1 focus-visible:ring-offset-background",
copied && "border-emerald-500/40 bg-emerald-500/10",
)}
aria-label={copied ? "Access token copied" : "Copy access token"}
aria-label={
copied
? t("settings.apiKeys.accessTokenCopied")
: t("settings.apiKeys.copyAccessToken")
}
>
<code className="min-w-0 flex-1 break-all text-left text-foreground">
{rawKey}
@ -55,7 +61,7 @@ export function KeyRevealCard({
</button>
<div className="flex items-center justify-between gap-3 pt-0.5">
<p className="text-[11px] text-muted-foreground">
Copy now — this won't be shown again.
{t("settings.apiKeys.copyNow")}
</p>
<Button
type="button"
@ -63,7 +69,7 @@ export function KeyRevealCard({
onClick={onDone}
className="focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring focus-visible:ring-offset-1 focus-visible:ring-offset-background"
>
Done
{t("common.done")}
</Button>
</div>
</div>

46
studio/frontend/src/features/settings/components/language-select.tsx Normal file
View file

@ -0,0 +1,46 @@
// 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 {
Select,
SelectContent,
SelectItem,
SelectTrigger,
SelectValue,
} from "@/components/ui/select";
import {
LOCALES,
isSupportedLocale,
setLocale,
useT,
useLocale,
} from "@/i18n";
export function LanguageSelect() {
const t = useT();
const locale = useLocale();
return (
<Select
value={locale}
onValueChange={(value) => {
if (isSupportedLocale(value)) setLocale(value);
}}
>
<SelectTrigger
aria-label={t("settings.appearance.language.label")}
className="w-40"
size="sm"
>
<SelectValue />
</SelectTrigger>
<SelectContent>
{Object.entries(LOCALES).map(([value, metadata]) => (
<SelectItem key={value} value={value}>
{metadata.nativeLabel}
</SelectItem>
))}
</SelectContent>
</Select>
);
}

16
studio/frontend/src/features/settings/components/theme-segmented.tsx
View file

@ -2,6 +2,7 @@
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import { cn } from "@/lib/utils";
import { useT, type TranslationKey } from "@/i18n";
import {
LaptopIcon,
Moon02Icon,
@ -11,13 +12,18 @@ import { HugeiconsIcon } from "@hugeicons/react";
import { motion, useReducedMotion } from "motion/react";
import { useTheme, type Theme } from "../stores/theme-store";
const OPTIONS: { value: Theme; label: string; icon: typeof Sun02Icon }[] = [
{ value: "light", label: "Light", icon: Sun02Icon },
{ value: "dark", label: "Dark", icon: Moon02Icon },
{ value: "system", label: "System", icon: LaptopIcon },
const OPTIONS: {
value: Theme;
labelKey: TranslationKey;
icon: typeof Sun02Icon;
}[] = [
{ value: "light", labelKey: "settings.appearance.theme.light", icon: Sun02Icon },
{ value: "dark", labelKey: "settings.appearance.theme.dark", icon: Moon02Icon },
{ value: "system", labelKey: "settings.appearance.theme.system", icon: LaptopIcon },
];
export function ThemeSegmented() {
const t = useT();
const { theme, setTheme } = useTheme();
const reduced = useReducedMotion();
return (
@ -49,7 +55,7 @@ export function ThemeSegmented() {
/>
)}
<HugeiconsIcon icon={opt.icon} className="relative z-10 size-3.5" />
<span className="relative z-10">{opt.label}</span>
<span className="relative z-10">{t(opt.labelKey)}</span>
</button>
);
})}

94
studio/frontend/src/features/settings/components/update-studio-instructions.tsx
View file

@ -2,6 +2,7 @@
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import { copyToClipboard } from "@/lib/copy-to-clipboard";
import { useT } from "@/i18n";
import { cn } from "@/lib/utils";
import { Copy01Icon, Tick02Icon } from "@hugeicons/core-free-icons";
import { HugeiconsIcon } from "@hugeicons/react";
@ -29,10 +30,13 @@ export type UpdateInstallSource =
| "unknown";
type UpdateInstallSourceState = UpdateInstallSource | "loading";
function getStudioUpdateInstructionLine(shell: UpdateShell): string {
function getStudioUpdateInstructionLine(
shell: UpdateShell,
t: ReturnType<typeof useT>,
): string {
return shell === "windows"
? "Open PowerShell and run:"
: "Open Terminal and run:";
? t("settings.about.update.openPowerShell")
: t("settings.about.update.openTerminal");
}
function isLocalInstallSource(
@ -59,6 +63,7 @@ function CopyableCommand({
command: string;
copyLabel: string;
}): ReactElement {
const t = useT();
const [copied, setCopied] = useState(false);
const timerRef = useRef<ReturnType<typeof setTimeout> | null>(null);
@ -89,14 +94,26 @@ function CopyableCommand({
value={command}
className="min-w-0 flex-1 bg-transparent px-2 py-1.5 font-mono text-[11px] text-foreground outline-none"
title={command}
aria-label={`${copyLabel} text`}
aria-label={t("settings.about.update.commandText", {
label: copyLabel,
})}
/>
<button
type="button"
onClick={handleCopy}
className="flex shrink-0 items-center justify-center border-l border-border px-2 text-muted-foreground transition-colors hover:bg-muted hover:text-foreground"
title={copied ? "Copied" : "Copy command"}
aria-label={copied ? `${copyLabel} copied` : `Copy ${copyLabel}`}
title={
copied
? t("settings.about.update.copied")
: t("settings.about.update.copyCommand")
}
aria-label={
copied
? t("settings.about.update.commandCopied", { label: copyLabel })
: t("settings.about.update.copyNamedCommand", {
label: copyLabel,
})
}
>
{copied ? (
<HugeiconsIcon
@ -123,7 +140,9 @@ export function UpdateStudioInstructions({
installSource?: UpdateInstallSourceState | null;
showTitle?: boolean;
}): ReactElement {
const [shell, setShell] = useState<UpdateShell>(defaultShell);
const t = useT();
const [shellOverride, setShellOverride] = useState<UpdateShell | null>(null);
const shell = shellOverride ?? defaultShell;
const prefersReducedMotion = useReducedMotion();
const windows = shell === "windows";
const localInstallSource = isLocalInstallSource(installSource);
@ -144,10 +163,6 @@ export function UpdateStudioInstructions({
? { opacity: 1 }
: { opacity: 0, y: -2 };
useEffect(() => {
setShell(defaultShell);
}, [defaultShell]);
return (
<div className={cn("flex flex-col gap-3", className)}>
<div
@ -158,13 +173,13 @@ export function UpdateStudioInstructions({
>
{showTitle ? (
<p className="shrink-0 whitespace-nowrap text-sm font-semibold font-heading">
Update Unsloth Studio
{t("settings.about.update.title")}
</p>
) : null}
<div className="flex shrink-0 items-center gap-0.5 text-[11px]">
<button
type="button"
onClick={() => setShell("windows")}
onClick={() => setShellOverride("windows")}
className={cn(
"px-0.5 py-0.5 font-medium transition-colors",
windows
@ -178,7 +193,7 @@ export function UpdateStudioInstructions({
<span className="text-border">/</span>
<button
type="button"
onClick={() => setShell("unix")}
onClick={() => setShellOverride("unix")}
className={cn(
"px-0.5 py-0.5 font-medium transition-colors",
windows
@ -193,31 +208,28 @@ export function UpdateStudioInstructions({
</div>
{loadingInstallSource ? (
<p className="text-xs text-muted-foreground leading-relaxed">
Checking how Studio was installed…
{t("settings.about.update.checkingInstall")}
</p>
) : localInstallSource ? (
<>
<p className="text-xs text-muted-foreground leading-relaxed">
Source or local install detected. To avoid replacing it with PyPI,
update from the checkout or source you originally installed from.
{t("settings.about.update.localInstallDetected")}
</p>
{checkoutInstallSource ? (
<>
<p className="text-xs text-muted-foreground leading-relaxed">
Pull latest changes from your Unsloth repo checkout, then update
Studio locally:
{t("settings.about.update.pullThenUpdate")}
</p>
<CopyableCommand
command={STUDIO_LOCAL_PULL_CMD}
copyLabel="git pull command"
copyLabel={t("settings.about.update.gitPullCommand")}
/>
<CopyableCommand
command={STUDIO_LOCAL_UPDATE_CMD}
copyLabel="local update command"
copyLabel={t("settings.about.update.localUpdateCommand")}
/>
<p className="text-xs text-muted-foreground leading-relaxed">
If the Studio update command is unavailable, run the local
installer from that checkout:
{t("settings.about.update.localInstallerFallback")}
</p>
<AnimatePresence mode="wait" initial={false}>
<motion.div
@ -233,7 +245,7 @@ export function UpdateStudioInstructions({
? STUDIO_LOCAL_FALLBACK_WINDOWS_CMD
: STUDIO_LOCAL_FALLBACK_UNIX_CMD
}
copyLabel="local installer command"
copyLabel={t("settings.about.update.localInstallerCommand")}
/>
</motion.div>
</AnimatePresence>
@ -242,12 +254,10 @@ export function UpdateStudioInstructions({
{packagedSourceInstall ? (
<>
<p className="text-xs text-muted-foreground leading-relaxed">
This looks like a source or VCS package install. Reinstall from
the original local path or Git URL you used.
{t("settings.about.update.sourceInstallDetected")}
</p>
<p className="text-xs text-muted-foreground leading-relaxed">
If you still have the Unsloth repo checkout, run the local
installer from that checkout:
{t("settings.about.update.repoCheckoutFallback")}
</p>
<AnimatePresence mode="wait" initial={false}>
<motion.div
@ -263,39 +273,37 @@ export function UpdateStudioInstructions({
? STUDIO_LOCAL_FALLBACK_WINDOWS_CMD
: STUDIO_LOCAL_FALLBACK_UNIX_CMD
}
copyLabel="local installer command"
copyLabel={t("settings.about.update.localInstallerCommand")}
/>
</motion.div>
</AnimatePresence>
</>
) : null}
<p className="text-xs text-muted-foreground leading-relaxed">
Restart Studio after updating for changes to take effect.
{t("settings.about.update.restartAfterUpdate")}
</p>
</>
) : unknownInstallSource ? (
<>
<p className="text-xs text-muted-foreground leading-relaxed">
Studio could not detect how it was installed. Check how you
installed Studio first, then choose the matching update path.
{t("settings.about.update.unknownInstall")}
</p>
<p className="text-xs text-muted-foreground leading-relaxed">
For curl or PyPI installs, run:
{t("settings.about.update.curlOrPypi")}
</p>
<CopyableCommand
command={STUDIO_UPDATE_CMD}
copyLabel="update command"
copyLabel={t("settings.about.update.updateCommand")}
/>
<p className="text-xs text-muted-foreground leading-relaxed">
For local checkout installs, update from that checkout instead and
use the local update command:
{t("settings.about.update.localCheckout")}
</p>
<CopyableCommand
command={STUDIO_LOCAL_UPDATE_CMD}
copyLabel="local update command"
copyLabel={t("settings.about.update.localUpdateCommand")}
/>
<p className="text-xs text-muted-foreground leading-relaxed">
Restart Studio after updating for changes to take effect.
{t("settings.about.update.restartAfterUpdate")}
</p>
</>
) : (
@ -309,15 +317,15 @@ export function UpdateStudioInstructions({
transition={fadeTransition}
className="text-xs text-muted-foreground leading-relaxed"
>
{getStudioUpdateInstructionLine(shell)}
{getStudioUpdateInstructionLine(shell, t)}
</motion.p>
</AnimatePresence>
<CopyableCommand
command={STUDIO_UPDATE_CMD}
copyLabel="update command"
copyLabel={t("settings.about.update.updateCommand")}
/>
<p className="text-xs text-muted-foreground leading-relaxed">
If that fails or unsloth studio update is unavailable, run:
{t("settings.about.update.fallbackInstruction")}
</p>
<AnimatePresence mode="wait" initial={false}>
<motion.div
@ -333,12 +341,12 @@ export function UpdateStudioInstructions({
? STUDIO_UPDATE_FALLBACK_WINDOWS_CMD
: STUDIO_UPDATE_FALLBACK_UNIX_CMD
}
copyLabel="fallback command"
copyLabel={t("settings.about.update.fallbackCommand")}
/>
</motion.div>
</AnimatePresence>
<p className="text-xs text-muted-foreground leading-relaxed">
Restart Studio after updating for changes to take effect.
{t("settings.about.update.restartAfterUpdate")}
</p>
</>
)}

24
studio/frontend/src/features/settings/components/usage-examples.tsx
View file

@ -3,6 +3,7 @@
import { cn } from "@/lib/utils";
import { copyToClipboard } from "@/lib/copy-to-clipboard";
import { useT } from "@/i18n";
import {
ArrowUpRight01Icon,
Copy01Icon,
@ -78,6 +79,7 @@ for chunk in response:
}
export function UsageExamples() {
const t = useT();
const [lang, setLang] = useState<Lang>("curl");
const [copied, setCopied] = useState(false);
const snippets = useMemo(
@ -97,17 +99,19 @@ export function UsageExamples() {
return (
<section className="flex min-w-0 max-w-full flex-col">
<h2 className="mb-2 text-sm font-semibold text-foreground">Usage examples</h2>
<h2 className="mb-2 text-sm font-semibold text-foreground">
{t("settings.apiKeys.usageExamples")}
</h2>
<div className="min-w-0 max-w-full overflow-hidden rounded-lg border border-border bg-muted/20">
<div className="flex min-w-0 items-center justify-between gap-2 border-b border-border px-2 py-1.5">
<div className="flex min-w-0 items-center gap-0.5">
{TABS.map((t) => {
const active = lang === t.id;
{TABS.map((tab) => {
const active = lang === tab.id;
return (
<button
key={t.id}
key={tab.id}
type="button"
onClick={() => setLang(t.id)}
onClick={() => setLang(tab.id)}
aria-pressed={active}
className={cn(
"rounded px-2 py-1 text-[11px] font-medium transition-colors focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring",
@ -116,7 +120,9 @@ export function UsageExamples() {
: "text-muted-foreground hover:text-foreground",
)}
>
{t.label}
{tab.id === "tools"
? t("settings.apiKeys.usageTools")
: tab.label}
</button>
);
})}
@ -125,20 +131,20 @@ export function UsageExamples() {
type="button"
onClick={handleCopy}
className="flex items-center gap-1 rounded px-1.5 py-1 text-[11px] text-muted-foreground transition-colors hover:text-foreground focus-visible:outline-none focus-visible:ring-2 focus-visible:ring-ring"
aria-label="Copy snippet"
aria-label={t("settings.apiKeys.copySnippet")}
>
<HugeiconsIcon
icon={copied ? Tick02Icon : Copy01Icon}
className={cn("size-3.5", copied && "text-emerald-600")}
/>
{copied ? "Copied" : "Copy"}
{copied ? t("settings.apiKeys.copied") : t("settings.apiKeys.copy")}
</button>
</div>
<pre className="max-w-full overflow-x-auto whitespace-pre-wrap break-words p-3 font-mono text-[11px] leading-relaxed text-foreground">
{snippets[lang]}
</pre>
<div className="flex flex-wrap items-center gap-x-2 gap-y-1 border-t border-border px-3 py-2 text-[11px] text-muted-foreground">
<span>Setup docs:</span>
<span>{t("settings.apiKeys.setupDocs")}</span>
{DOC_LINKS.map((link) => (
<a
key={link.href}

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