Merge remote-tracking branch 'upstream/main' into feat/model-picker-per-model-config

# Conflicts:
#	studio/frontend/src/features/chat/hooks/use-chat-model-runtime.ts
This commit is contained in:
sneakr 2026-07-11 12:55:22 +02:00
commit a615a40284
84 changed files with 13015 additions and 2157 deletions

View file

@ -527,6 +527,154 @@ case "$MODE" in
echo "[claude] attribution A/B OK (suppressed HIT, header=1 MISS)"
;;
# ── resume: does a launched agent's session survive exit and resume? ────
# Unlike the other modes, this drives the real LAUNCH path (`unsloth start
# <agent> ...`, the interactive default), not the --no-launch recipe. That
# path relocates each agent's home to a throwaway temp dir wiped on exit, so
# a session cannot be resumed -- unless --persist routes it to the stable
# Unsloth agents dir instead. We run one headless turn per pass and check
# whether the turn left a session in a persistent store (deterministic, no
# reliance on the model recalling anything), for a baseline pass and a
# --persist pass, and assert the expected split for this agent.
resume)
CODEWORD="PLATYPUS7"
T1="Remember this codeword for later: ${CODEWORD}. Reply with just the word OK."
T2="What codeword did I ask you to remember? Reply with just that word."
WORK="$WORKDIR_BASE/${AGENT}-resume"
# STABLE_HOME: the stable dir that --no-launch (and --persist) relocate to.
# Read it from a --no-launch probe (which also writes the agent's config
# there). codex/pi relocate their whole home/HOME here; opencode/claude keep
# their session data in a fixed user dir, so STABLE_HOME stays empty for them.
parse_connect
case "$AGENT" in
codex) STABLE_HOME="$(raw_env CODEX_HOME)" ;;
pi) STABLE_HOME="$(raw_env HOME)" ;;
*) STABLE_HOME="" ;;
esac
# The persistent stores a session would land in if it were NOT wiped. We
# count files here before/after each turn; a positive delta means the
# session persisted (is resumable), zero means it went to a wiped temp dir.
resume_tracked_dirs() {
case "$AGENT" in
codex) printf '%s\n' "$HOME/.codex" ;;
opencode) printf '%s\n' "$HOME/.local/share/opencode" "$HOME/.config/opencode" ;;
claude) printf '%s\n' "$HOME/.claude" ;;
pi) printf '%s\n' "$HOME/.pi" ;;
*) : ;;
esac
[ -n "$STABLE_HOME" ] && printf '%s\n' "$STABLE_HOME"
}
count_session_files() {
local total=0 d n
while IFS= read -r d; do
[ -n "$d" ] && [ -d "$d" ] || continue
n="$(find "$d" -type f 2>/dev/null | wc -l)"; total=$((total + n))
done < <(resume_tracked_dirs)
echo "$total"
}
# The headless first-turn subcommand per agent (mirrors file-edit's map),
# forwarded verbatim through the launch path as passthrough args.
set_t1_cmd() {
case "$AGENT" in
claude) T1_CMD=("${CLAUDE_CONNECT_FLAGS[@]}" -p "$T1") ;;
codex) T1_CMD=(exec "$T1") ;;
opencode) T1_CMD=(run "$T1") ;;
pi) T1_CMD=(-p "$T1") ;;
*) guide_fail "resume mode does not cover agent '$AGENT'" ;;
esac
}
# Run one headless turn through the launch path. $1=outfile, $2="" or
# "--persist", rest = the agent subcommand. --yolo auto-approves so no tool
# prompt can hang; --api-key attaches to the already-served CI model.
launch_turn() {
local out="$1" rflag="$2"; shift 2
local flag=(); [ -n "$rflag" ] && flag=("$rflag")
run_timed "$out" unsloth start "$AGENT" "${flag[@]}" --yolo \
--api-key "$UNSLOTH_API_KEY" "$@"
local rc=$?
redact "$out"
return "$rc"
}
# One pass: fresh work dir, one planting turn, set RESULT to PERSISTED/WIPED
# from the session-store delta. Runs in the main shell (not a command
# substitution) so a hang's guide_fail actually fails the job and the
# progress lines reach the CI log. $1 = "" (baseline) or "--persist".
RESULT=""
run_pass() {
local rflag="$1" label="baseline"
[ -n "$rflag" ] && label="resume"
rm -rf "$WORK"; mkdir -p "$WORK"
set_t1_cmd
local out="$LOGS_DIR/${AGENT}-resume-${label}.txt"
local before after rc
before="$(count_session_files)"
pushd "$WORK" >/dev/null || guide_fail "could not enter work dir $WORK"
launch_turn "$out" "$rflag" "${T1_CMD[@]}"; rc=$?
popd >/dev/null || true
after="$(count_session_files)"
echo "[$AGENT] ${label}: session files ${before} -> ${after} (rc=${rc})"
# The turn must succeed for the delta to mean anything: an agent that writes a
# session file then errors would otherwise be misread as PERSISTED. Mirror the
# file-edit mode and fail the pass on a non-zero launch (the flagship codex recall
# below stays WARN-only, driven by its own launch_turn calls).
[ "$rc" -eq 0 ] || { echo "[$AGENT] ${label} transcript (tail):"; tail -30 "$out" 2>/dev/null || true; \
guide_fail "resume ${label} turn for ${AGENT} exited non-zero (rc=${rc})"; }
if [ "$after" -gt "$before" ]; then RESULT="PERSISTED"; else RESULT="WIPED"; fi
}
run_pass ""; BASELINE="$RESULT"
# Only the temp-dir agents (codex/pi) need the --persist pass to prove the fix.
# opencode/claude persist either way, so the baseline already proves it and a
# second full CPU turn only risks a timeout; skip it for them.
case "$AGENT" in
codex|pi) run_pass "--persist"; RESUME="$RESULT" ;;
*) RESUME="n/a (persists either way)" ;;
esac
# Expected: codex/pi relocate their whole home to the temp dir, so a plain
# launch is WIPED and only --persist PERSISTS. opencode/claude keep their
# session data in a fixed user dir, so the baseline already PERSISTS.
case "$AGENT" in
codex|pi) EXPECT_BASELINE="WIPED" ;;
opencode|claude) EXPECT_BASELINE="PERSISTED" ;;
esac
echo "──────────────────────────────────────────────"
echo "[$AGENT] RESUME EXPERIMENT"
echo " baseline (unsloth start ${AGENT}): ${BASELINE} (expected ${EXPECT_BASELINE})"
echo " with --persist (unsloth start ${AGENT} --persist): ${RESUME}"
echo "──────────────────────────────────────────────"
[ "$BASELINE" = "$EXPECT_BASELINE" ] \
|| guide_fail "baseline resume behavior for ${AGENT} was ${BASELINE}, expected ${EXPECT_BASELINE}"
case "$AGENT" in
codex|pi)
[ "$RESUME" = "PERSISTED" ] \
|| guide_fail "--persist did not persist ${AGENT}'s session (got ${RESUME}); the session dir is still not stable" ;;
esac
# Flagship behavioral proof (codex only, WARN-only): after a --persist plant,
# resume the session and check the model actually recalls the codeword. A
# miss is not a failure (the CI model is small); the mechanism gate above is
# the real assertion.
if [ "$AGENT" = "codex" ]; then
rm -rf "$WORK"; mkdir -p "$WORK"
( cd "$WORK" && launch_turn "$LOGS_DIR/codex-resume-plant.txt" "--persist" exec "$T1" ) || true
( cd "$WORK" && launch_turn "$LOGS_DIR/codex-resume-recall.txt" "--persist" exec resume --last "$T2" ) || true
if grep -q "$CODEWORD" "$LOGS_DIR/codex-resume-recall.txt" 2>/dev/null; then
echo "[codex] behavioral recall HIT: resumed session remembered ${CODEWORD}"
else
echo "::warning::[codex] behavioral recall MISS (small CI model); mechanism gate still passed"
fi
fi
echo "[$AGENT] resume OK"
;;
*)
echo "agent-guides-drive.sh: unknown mode '$MODE'" >&2
exit 2

View file

@ -364,7 +364,9 @@ jobs:
tests/utils/test_attention_masks.py \
tests/utils/test_trunc_normal_patch.py \
tests/python/test_fast_language_model_text_only.py \
tests/test_bad_mappings_redirect.py \
tests/test_prefetch_snapshot_scope.py \
tests/test_gemma_2b_mapper_key.py \
--deselect 'tests/utils/test_attention_masks.py::test_run_attention_flash_varlen_receives_window_and_softcap'
# The deselected test monkeypatches flash_attn_varlen_func, which is
# only bound on the module when `flash_attn` is importable. flash_attn

View file

@ -471,6 +471,176 @@ jobs:
redacted-configs/
retention-days: 7
# ═════════════════════════════════════════════════════════════════════
# Job: resume
# Does a conversation started with `unsloth start <agent>` survive exit
# and resume? This drives the REAL launch path (not the --no-launch
# recipe the other jobs use). A plain launch relocates the agent home to
# a temp dir wiped on exit, so codex/pi cannot resume; --persist routes the
# session to the stable Unsloth agents dir so it persists. opencode/claude
# keep their session data in a fixed user dir, so they persist either way.
# Dispatch-only: it is an end-to-end experiment, not a PR gate.
# ═════════════════════════════════════════════════════════════════════
resume:
name: resume (${{ matrix.agent }})
if: github.event_name == 'workflow_dispatch'
runs-on: ubuntu-latest
timeout-minutes: 60
strategy:
fail-fast: false
matrix:
# codex/pi relocate their whole home (resume broken without --persist);
# opencode/claude keep session data in a fixed dir (resume already works).
# One agent from each class proves the split end to end; openclaw/hermes
# share codex's relocation mechanism and are covered by the unit tests.
agent: [codex, opencode, claude, pi]
env:
GGUF_REPO: unsloth/gemma-4-E4B-it-GGUF
GGUF_FILE: gemma-4-E4B-it-UD-Q4_K_XL.gguf
STUDIO_PORT: '18904'
steps:
- uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
with:
persist-credentials: false
- name: Linux deps for llama.cpp prebuilt
run: |
sudo apt-get update
sudo apt-get install -y --no-install-recommends \
libcurl4-openssl-dev libssl-dev jq
- 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'
cache: 'pip'
- name: Restore GGUF model file
id: cache-gguf
uses: actions/cache/restore@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5
continue-on-error: true
with:
path: gguf-cache
key: ${{ runner.os }}-gguf-${{ env.GGUF_REPO }}-${{ env.GGUF_FILE }}-v1
- name: Download GGUF if cache miss
id: download-gguf
if: steps.cache-gguf.outputs.cache-hit != 'true' || steps.cache-gguf.outcome != 'success'
env:
HF_TOKEN: ${{ secrets.HF_TOKEN }}
run: |
python -m pip install --upgrade huggingface_hub
mkdir -p gguf-cache
bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE" gguf-cache
- name: Save GGUF model file
if: always() && steps.download-gguf.outcome == 'success'
uses: actions/cache/save@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5
with:
path: gguf-cache
key: ${{ runner.os }}-gguf-${{ env.GGUF_REPO }}-${{ env.GGUF_FILE }}-v1
- name: Install Studio (--local, --no-torch)
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
HF_TOKEN: ${{ secrets.HF_TOKEN }}
run: |
mkdir -p logs
set -o pipefail
bash install.sh --local --no-torch 2>&1 | tee logs/install.log
- name: Serve unsloth run --disable-tools (gemma-4-E4B)
run: |
unsloth studio reset-password
bash .github/scripts/serve-unsloth-run.sh \
--gguf-file "$GITHUB_WORKSPACE/gguf-cache/${GGUF_FILE}" \
--port "$STUDIO_PORT" --log-dir logs \
--extra "--seed $UNSLOTH_SEED --temp 0" \
--health-timeout 900
- name: Preflight the agent's API dialect (class-a isolation)
env:
AGENT: ${{ matrix.agent }}
run: |
set -uo pipefail
B="$UNSLOTH_BASE_URL"; K="$UNSLOTH_API_KEY"
preflight_fail() {
echo "::error::[server/API regression] agent=$AGENT: $* (preflight failed BEFORE install/connect). Endpoint contract lives in studio/backend/routes/**.";
exit 1
}
code=$(curl -s -o /tmp/pf.json -w '%{http_code}' "$B/v1/models" \
-H "Authorization: Bearer $K") || true
[ "$code" = "200" ] || preflight_fail "/v1/models returned HTTP $code"
case "$AGENT" in
claude)
code=$(curl -s -o /tmp/pf.json -w '%{http_code}' "$B/v1/messages" \
-H "Authorization: Bearer $K" -H 'content-type: application/json' \
--max-time 120 \
-d "{\"model\":\"$UNSLOTH_MODEL_ID\",\"max_tokens\":16,\"messages\":[{\"role\":\"user\",\"content\":\"Hi\"}]}") || true
[ "$code" = "200" ] || preflight_fail "/v1/messages returned HTTP $code"
;;
codex)
code=$(curl -s -o /tmp/pf.json -w '%{http_code}' "$B/v1/responses" \
-H "Authorization: Bearer $K" -H 'content-type: application/json' \
--max-time 120 \
-d "{\"model\":\"$UNSLOTH_MODEL_ID\",\"input\":\"Hi\",\"max_output_tokens\":16,\"stream\":true}") || true
[ "$code" = "200" ] || preflight_fail "/v1/responses returned HTTP $code"
;;
*)
code=$(curl -s -o /tmp/pf.json -w '%{http_code}' "$B/v1/chat/completions" \
-H "Authorization: Bearer $K" -H 'content-type: application/json' \
--max-time 120 \
-d "{\"model\":\"$UNSLOTH_MODEL_ID\",\"max_tokens\":16,\"messages\":[{\"role\":\"user\",\"content\":\"Hi\"}]}") || true
[ "$code" = "200" ] || preflight_fail "/v1/chat/completions returned HTTP $code"
;;
esac
echo "preflight OK for $AGENT"
- name: Install agent CLI (class-b isolation)
env:
AGENT: ${{ matrix.agent }}
run: bash .github/scripts/agent-guides-install.sh "$AGENT"
- name: Resume experiment (launch path)
env:
AGENT: ${{ matrix.agent }}
run: bash .github/scripts/agent-guides-drive.sh resume "$AGENT"
- name: Collect server logs (debug)
if: always()
run: |
mkdir -p logs/studio-logs
cp -r "$HOME/.unsloth/studio/logs/." logs/studio-logs/ 2>/dev/null || true
if [ -n "${UNSLOTH_API_KEY:-}" ]; then
grep -rlF "$UNSLOTH_API_KEY" logs redacted-configs agent-workdir 2>/dev/null | while IFS= read -r f; do
sed -i "s#${UNSLOTH_API_KEY}#<REDACTED>#g" "$f" 2>/dev/null || true
done
fi
- name: Stop Studio
if: always()
run: |
if [ -n "${UNSLOTH_SERVER_PID:-}" ] && [ "${UNSLOTH_SERVER_PID}" != "0" ]; then
kill "${UNSLOTH_SERVER_PID}" 2>/dev/null || true
fi
sleep 2
ss -tln 2>/dev/null | grep ":${STUDIO_PORT}" || true
- name: Upload logs
if: always()
continue-on-error: true
uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
with:
name: resume-${{ matrix.agent }}-log
path: |
logs/
agent-workdir/
redacted-configs/
retention-days: 7
# ═════════════════════════════════════════════════════════════════════
# Job 3: prompt-cache
# (a) curl 2-turn /v1/chat/completions: assert turn-2 cached_tokens > 0

View file

@ -2,8 +2,8 @@
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
# Multi-language supply-chain audit. Triggers:
# - PRs touching any dependency manifest (Python / npm / Cargo) or
# this workflow file,
# - PRs touching any dependency manifest (Python / npm / Cargo), a
# scanner or its allowlist baseline, or this workflow file,
# - push to main / pip,
# - nightly @ 04:13 UTC so newly-published advisories surface even
# when no PR opens,
@ -57,7 +57,9 @@ on:
- 'studio/src-tauri/Cargo.lock'
- 'pyproject.toml'
- 'scripts/scan_packages.py'
- 'scripts/scan_packages_baseline.json'
- 'scripts/scan_npm_packages.py'
- 'scripts/scan_npm_packages_baseline.json'
- '.github/workflows/security-audit.yml'
push:
branches: [main, pip]

View file

@ -444,6 +444,8 @@ jobs:
python - <<'PY'
import json
import os
import time
import urllib.error
import urllib.request
BASE = os.environ["BASE_URL"]
@ -464,8 +466,24 @@ jobs:
"Content-Type": "application/json",
},
)
with urllib.request.urlopen(req, timeout = timeout) as resp:
return resp.status, json.loads(resp.read().decode())
# Shared CI runners stall sporadically, so retry transport-level
# failures only; HTTP status errors surface immediately. Bounded
# to fit the job's timeout-minutes: short probes get 3 full
# attempts, long probes one retry capped at 300s (a healthy
# server answers a retry quickly; a stalled one never does).
attempts = 3 if timeout <= 300 else 2
for attempt in range(attempts):
try:
t = timeout if attempt == 0 else min(timeout, 300)
with urllib.request.urlopen(req, timeout = t) as resp:
return resp.status, json.loads(resp.read().decode())
except urllib.error.HTTPError:
raise
except (TimeoutError, ConnectionError, urllib.error.URLError) as exc:
if attempt == attempts - 1:
raise
print(f"[retry] {path}: {exc!r}", flush = True)
time.sleep(15)
def post_sse(path, body, *, timeout = 600):
"""POST a streaming request and accumulate the assistant
@ -938,6 +956,8 @@ jobs:
import base64
import json
import os
import time
import urllib.error
import urllib.request
from openai import OpenAI
from anthropic import Anthropic
@ -956,8 +976,24 @@ jobs:
"Content-Type": "application/json",
},
)
with urllib.request.urlopen(req, timeout = timeout) as resp:
return resp.status, json.loads(resp.read().decode())
# Shared CI runners stall sporadically, so retry transport-level
# failures only; HTTP status errors surface immediately. Bounded
# to fit the job's timeout-minutes: short probes get 3 full
# attempts, long probes one retry capped at 300s (a healthy
# server answers a retry quickly; a stalled one never does).
attempts = 3 if timeout <= 300 else 2
for attempt in range(attempts):
try:
t = timeout if attempt == 0 else min(timeout, 300)
with urllib.request.urlopen(req, timeout = t) as resp:
return resp.status, json.loads(resp.read().decode())
except urllib.error.HTTPError:
raise
except (TimeoutError, ConnectionError, urllib.error.URLError) as exc:
if attempt == attempts - 1:
raise
print(f"[retry] {path}: {exc!r}", flush = True)
time.sleep(15)
# ── 1. response_format = json_object (JSON mode) ─────────────
# llama.cpp's HTTP server supports OpenAI-compatible JSON

View file

@ -430,6 +430,8 @@ jobs:
python - <<'PY'
import json
import os
import time
import urllib.error
import urllib.request
BASE = os.environ["BASE_URL"]
@ -450,8 +452,24 @@ jobs:
"Content-Type": "application/json",
},
)
with urllib.request.urlopen(req, timeout = timeout) as resp:
return resp.status, json.loads(resp.read().decode())
# Shared CI runners stall sporadically, so retry transport-level
# failures only; HTTP status errors surface immediately. Bounded
# to fit the job's timeout-minutes: short probes get 3 full
# attempts, long probes one retry capped at 300s (a healthy
# server answers a retry quickly; a stalled one never does).
attempts = 3 if timeout <= 300 else 2
for attempt in range(attempts):
try:
t = timeout if attempt == 0 else min(timeout, 300)
with urllib.request.urlopen(req, timeout = t) as resp:
return resp.status, json.loads(resp.read().decode())
except urllib.error.HTTPError:
raise
except (TimeoutError, ConnectionError, urllib.error.URLError) as exc:
if attempt == attempts - 1:
raise
print(f"[retry] {path}: {exc!r}", flush = True)
time.sleep(15)
def post_sse(path, body, *, timeout = 600):
"""POST a streaming request and accumulate the assistant
@ -825,6 +843,8 @@ jobs:
import base64
import json
import os
import time
import urllib.error
import urllib.request
from openai import OpenAI
from anthropic import Anthropic
@ -848,8 +868,24 @@ jobs:
"Content-Type": "application/json",
},
)
with urllib.request.urlopen(req, timeout = timeout) as resp:
return resp.status, json.loads(resp.read().decode())
# Shared CI runners stall sporadically, so retry transport-level
# failures only; HTTP status errors surface immediately. Bounded
# to fit the job's timeout-minutes: short probes get 3 full
# attempts, long probes one retry capped at 300s (a healthy
# server answers a retry quickly; a stalled one never does).
attempts = 3 if timeout <= 300 else 2
for attempt in range(attempts):
try:
t = timeout if attempt == 0 else min(timeout, 300)
with urllib.request.urlopen(req, timeout = t) as resp:
return resp.status, json.loads(resp.read().decode())
except urllib.error.HTTPError:
raise
except (TimeoutError, ConnectionError, urllib.error.URLError) as exc:
if attempt == attempts - 1:
raise
print(f"[retry] {path}: {exc!r}", flush = True)
time.sleep(15)
# ── 1. response_format = json_object (JSON mode) ─────────────
# llama.cpp's HTTP server supports OpenAI-compatible JSON

View file

@ -634,6 +634,8 @@ jobs:
python - <<'PY'
import json
import os
import time
import urllib.error
import urllib.request
BASE = os.environ["BASE_URL"]
@ -656,8 +658,24 @@ jobs:
"Content-Type": "application/json",
},
)
with urllib.request.urlopen(req, timeout = timeout) as resp:
return resp.status, json.loads(resp.read().decode())
# Shared CI runners stall sporadically, so retry transport-level
# failures only; HTTP status errors surface immediately. Bounded
# to fit the job's timeout-minutes: short probes get 3 full
# attempts, long probes one retry capped at 300s (a healthy
# server answers a retry quickly; a stalled one never does).
attempts = 3 if timeout <= 300 else 2
for attempt in range(attempts):
try:
t = timeout if attempt == 0 else min(timeout, 300)
with urllib.request.urlopen(req, timeout = t) as resp:
return resp.status, json.loads(resp.read().decode())
except urllib.error.HTTPError:
raise
except (TimeoutError, ConnectionError, urllib.error.URLError) as exc:
if attempt == attempts - 1:
raise
print(f"[retry] {path}: {exc!r}", flush = True)
time.sleep(15)
def post_sse(path, body, *, timeout = 600):
body = {**body, "stream": True}
@ -1063,6 +1081,8 @@ jobs:
import base64
import json
import os
import time
import urllib.error
import urllib.request
from openai import OpenAI
from anthropic import Anthropic
@ -1082,8 +1102,24 @@ jobs:
"Content-Type": "application/json",
},
)
with urllib.request.urlopen(req, timeout = timeout) as resp:
return resp.status, json.loads(resp.read().decode())
# Shared CI runners stall sporadically, so retry transport-level
# failures only; HTTP status errors surface immediately. Bounded
# to fit the job's timeout-minutes: short probes get 3 full
# attempts, long probes one retry capped at 300s (a healthy
# server answers a retry quickly; a stalled one never does).
attempts = 3 if timeout <= 300 else 2
for attempt in range(attempts):
try:
t = timeout if attempt == 0 else min(timeout, 300)
with urllib.request.urlopen(req, timeout = t) as resp:
return resp.status, json.loads(resp.read().decode())
except urllib.error.HTTPError:
raise
except (TimeoutError, ConnectionError, urllib.error.URLError) as exc:
if attempt == attempts - 1:
raise
print(f"[retry] {path}: {exc!r}", flush = True)
time.sleep(15)
# ── 1. response_format = json_object (JSON mode) ─────────────
status, data = post("/v1/chat/completions", {
@ -1334,42 +1370,75 @@ jobs:
try { Add-MpPreference -ExclusionPath $p -ErrorAction Stop } catch { }
}
- name: Hide Visual Studio + CMake (simulate a host with no build tools)
- name: Prepare no-build-tools simulation
shell: pwsh
run: |
$ErrorActionPreference = 'Stop'
# A Program Files dir can hold a transient handle (Defender / MSBuild node)
# so Rename-Item intermittently fails with "Access is denied"; retry to ride it out.
function Rename-WithRetry($Path, $NewName) {
for ($i = 1; $i -le 6; $i++) {
try { Rename-Item -LiteralPath $Path -NewName $NewName -ErrorAction Stop; return }
catch { if ($i -eq 6) { throw }; Start-Sleep -Seconds 3 }
$root = Join-Path $env:GITHUB_WORKSPACE 'no-build-tools'
$pf = Join-Path $root 'ProgramFiles'
$pfx86 = Join-Path $root 'ProgramFilesx86'
New-Item -ItemType Directory -Force -Path $pf, $pfx86 | Out-Null
$blocked = [System.Collections.Generic.HashSet[string]]::new([StringComparer]::OrdinalIgnoreCase)
foreach ($tool in @('cmake', 'cl.exe')) {
foreach ($cmd in (Get-Command $tool -All -ErrorAction SilentlyContinue)) {
if ($cmd.Source) {
$dir = Split-Path -Parent $cmd.Source
if ($dir) {
[void] $blocked.Add(
[Environment]::ExpandEnvironmentVariables($dir).Trim().Trim('"').TrimEnd('\'))
}
}
}
}
# Rename the Visual Studio install roots (incl. the Installer that holds
# vswhere.exe) so Find-VsBuildTools' vswhere + filesystem scan both miss.
foreach ($d in @("$env:ProgramFiles\Microsoft Visual Studio", "${env:ProgramFiles(x86)}\Microsoft Visual Studio")) {
if (Test-Path -LiteralPath $d) {
Rename-WithRetry $d ((Split-Path $d -Leaf) + '.vsoff')
Write-Host "Hid VS: $d"
}
# Normalized comparison so registry spellings (trailing slash,
# unexpanded %VAR%) still match.
function Test-Blocked([string]$p) {
$n = [Environment]::ExpandEnvironmentVariables($p).Trim().Trim('"').TrimEnd('\')
return $blocked.Contains($n)
}
# Surgically rename each cmake executable on PATH (not its parent dir --
# cmake can share a dir with other shims) so Get-Command cmake fails.
$hidden = @()
foreach ($c in (Get-Command cmake -All -ErrorAction SilentlyContinue)) {
if ($c.Source -and (Test-Path -LiteralPath $c.Source)) {
Rename-WithRetry $c.Source ((Split-Path $c.Source -Leaf) + '.off')
$hidden += $c.Source
Write-Host "Hid cmake: $($c.Source)"
}
$pathParts = $env:Path -split [IO.Path]::PathSeparator |
Where-Object { $_ -and -not (Test-Blocked $_) }
$noBuildToolsPath = $pathParts -join [IO.Path]::PathSeparator
# install.ps1's Refresh-SessionPath and setup.ps1's Refresh-Environment
# rebuild the session Path from these scopes mid-install, so filter
# them too. Originals are saved for the cleanup step.
foreach ($scope in @('Machine', 'User')) {
$orig = [Environment]::GetEnvironmentVariable('Path', $scope)
if (-not $orig) { continue }
Set-Content -LiteralPath (Join-Path $root "orig-path-$scope.txt") -Value $orig -NoNewline
$kept = ($orig -split ';' | Where-Object { $_ -and -not (Test-Blocked $_) }) -join ';'
[Environment]::SetEnvironmentVariable('Path', $kept, $scope)
Write-Host "Filtered $scope Path scope."
}
"NO_BUILD_TOOLS_PROGRAMFILES=$pf" | Out-File -FilePath $env:GITHUB_ENV -Append -Encoding utf8
"NO_BUILD_TOOLS_PROGRAMFILES_X86=$pfx86" | Out-File -FilePath $env:GITHUB_ENV -Append -Encoding utf8
"NO_BUILD_TOOLS_PATH<<NO_BUILD_TOOLS_PATH_EOF" | Out-File -FilePath $env:GITHUB_ENV -Append -Encoding utf8
$noBuildToolsPath | Out-File -FilePath $env:GITHUB_ENV -Append -Encoding utf8
"NO_BUILD_TOOLS_PATH_EOF" | Out-File -FilePath $env:GITHUB_ENV -Append -Encoding utf8
Write-Host "ProgramFiles simulation root: $pf"
Write-Host "ProgramFiles(x86) simulation root: $pfx86"
if ($blocked.Count -gt 0) {
Write-Host "Removed build-tool PATH dirs:"
$blocked | Sort-Object | ForEach-Object { Write-Host " $_" }
} else {
Write-Host "No cmake or cl.exe PATH dirs found to remove."
}
("HIDDEN_CMAKE=" + ($hidden -join '|')) | Out-File -FilePath $env:GITHUB_ENV -Append -Encoding utf8
- name: Assert Visual Studio + CMake are genuinely undetectable
shell: pwsh
run: |
$ErrorActionPreference = 'Stop'
# Set in-script: the runner does not apply step-level env keys with
# parentheses (`ProgramFiles(x86)`), so vswhere still found VS.
if (-not $env:NO_BUILD_TOOLS_PROGRAMFILES) { Write-Error "NO_BUILD_TOOLS_* env missing (Prepare step did not run?)"; exit 1 }
$env:ProgramFiles = $env:NO_BUILD_TOOLS_PROGRAMFILES
${env:ProgramFiles(x86)} = $env:NO_BUILD_TOOLS_PROGRAMFILES_X86
$env:Path = $env:NO_BUILD_TOOLS_PATH
. (Join-Path $env:GITHUB_WORKSPACE 'tests/studio_setup_ps1/Get-FunctionSource.ps1')
$setup = Join-Path $env:GITHUB_WORKSPACE 'studio/setup.ps1'
foreach ($fn in @('Resolve-VsGeneratorFromLabel', 'Find-VsBuildTools')) {
@ -1394,6 +1463,10 @@ jobs:
# Withheld on PR: this step runs checked-out PR code; public GGUF still downloads.
HF_TOKEN: ${{ github.event_name != 'pull_request' && secrets.HF_TOKEN || '' }}
run: |
# Set in-script (see the assert step); child processes inherit these.
$env:ProgramFiles = $env:NO_BUILD_TOOLS_PROGRAMFILES
${env:ProgramFiles(x86)} = $env:NO_BUILD_TOOLS_PROGRAMFILES_X86
$env:Path = $env:NO_BUILD_TOOLS_PATH
New-Item -ItemType Directory -Force -Path logs | Out-Null
$ProgressPreference = 'SilentlyContinue'
& ./install.ps1 --local --no-torch *>&1 | Tee-Object -FilePath logs/install.log
@ -1480,19 +1553,19 @@ jobs:
[ -n "$CONTENT" ] && [ "$CONTENT" != "null" ] || { echo "::error::empty completion"; exit 1; }
echo "Inference OK without Visual Studio: $CONTENT"
- name: Restore Visual Studio + CMake
- name: Clean no-build-tools simulation
if: always()
shell: pwsh
run: |
foreach ($d in @("$env:ProgramFiles\Microsoft Visual Studio", "${env:ProgramFiles(x86)}\Microsoft Visual Studio")) {
$off = "$d.vsoff"
if (Test-Path -LiteralPath $off) { Rename-Item -LiteralPath $off -NewName (Split-Path $d -Leaf); Write-Host "Restored $d" }
}
if ($env:HIDDEN_CMAKE) {
foreach ($src in ($env:HIDDEN_CMAKE -split '\|')) {
if ($src -and (Test-Path -LiteralPath "$src.off")) { Rename-Item -LiteralPath "$src.off" -NewName (Split-Path $src -Leaf) }
$root = Join-Path $env:GITHUB_WORKSPACE 'no-build-tools'
foreach ($scope in @('Machine', 'User')) {
$saved = Join-Path $root "orig-path-$scope.txt"
if (Test-Path -LiteralPath $saved) {
[Environment]::SetEnvironmentVariable('Path', (Get-Content -LiteralPath $saved -Raw), $scope)
Write-Host "Restored $scope Path scope."
}
}
Remove-Item -LiteralPath $root -Recurse -Force -ErrorAction SilentlyContinue
- name: Stop Studio
if: always()
@ -1540,21 +1613,34 @@ jobs:
with:
python-version: '3.12'
- name: Hide Visual Studio
- name: Prepare no-build-tools simulation
shell: pwsh
run: |
$ErrorActionPreference = 'Stop'
# Retry the rename: a Program Files dir can hold a transient handle that
# makes Rename-Item intermittently fail with "Access is denied".
function Rename-WithRetry($Path, $NewName) {
for ($i = 1; $i -le 6; $i++) {
try { Rename-Item -LiteralPath $Path -NewName $NewName -ErrorAction Stop; return }
catch { if ($i -eq 6) { throw }; Start-Sleep -Seconds 3 }
$root = Join-Path $env:GITHUB_WORKSPACE 'no-build-tools'
$pf = Join-Path $root 'ProgramFiles'
$pfx86 = Join-Path $root 'ProgramFilesx86'
New-Item -ItemType Directory -Force -Path $pf, $pfx86 | Out-Null
$blocked = [System.Collections.Generic.HashSet[string]]::new([StringComparer]::OrdinalIgnoreCase)
foreach ($tool in @('cmake', 'cl.exe')) {
foreach ($cmd in (Get-Command $tool -All -ErrorAction SilentlyContinue)) {
if ($cmd.Source) {
$dir = Split-Path -Parent $cmd.Source
if ($dir) { [void] $blocked.Add($dir) }
}
}
}
foreach ($d in @("$env:ProgramFiles\Microsoft Visual Studio", "${env:ProgramFiles(x86)}\Microsoft Visual Studio")) {
if (Test-Path -LiteralPath $d) { Rename-WithRetry $d ((Split-Path $d -Leaf) + '.vsoff'); Write-Host "Hid VS: $d" }
}
$pathParts = $env:Path -split [IO.Path]::PathSeparator |
Where-Object { $_ -and -not $blocked.Contains($_) }
$noBuildToolsPath = $pathParts -join [IO.Path]::PathSeparator
"NO_BUILD_TOOLS_PROGRAMFILES=$pf" | Out-File -FilePath $env:GITHUB_ENV -Append -Encoding utf8
"NO_BUILD_TOOLS_PROGRAMFILES_X86=$pfx86" | Out-File -FilePath $env:GITHUB_ENV -Append -Encoding utf8
"NO_BUILD_TOOLS_PATH<<NO_BUILD_TOOLS_PATH_EOF" | Out-File -FilePath $env:GITHUB_ENV -Append -Encoding utf8
$noBuildToolsPath | Out-File -FilePath $env:GITHUB_ENV -Append -Encoding utf8
"NO_BUILD_TOOLS_PATH_EOF" | Out-File -FilePath $env:GITHUB_ENV -Append -Encoding utf8
- name: Windows CUDA and ROCm prebuilts exist in unslothai/llama.cpp (what GPU users download, no VS)
env:
@ -1577,25 +1663,34 @@ jobs:
echo "Windows CUDA and ROCm prebuilts are available -- GPU users get them without compiling."
- name: The prebuilt resolver runs without Visual Studio
shell: pwsh
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
run: |
$ErrorActionPreference = 'Stop'
# pwsh: bash cannot export `ProgramFiles(x86)`; set in-script so the
# python child inherits the overrides.
$env:ProgramFiles = $env:NO_BUILD_TOOLS_PROGRAMFILES
${env:ProgramFiles(x86)} = $env:NO_BUILD_TOOLS_PROGRAMFILES_X86
$env:Path = $env:NO_BUILD_TOOLS_PATH
# Resolver-only (no GPU on hosted runners, so the host resolves to the
# CPU bundle). The point is that resolution needs no compiler/VS.
python -m pip install --upgrade huggingface_hub
python studio/install_llama_prebuilt.py --resolve-prebuilt latest --output-format json > /tmp/resolve.json || {
echo "::error::resolver exited non-zero"; cat /tmp/resolve.json || true; exit 1; }
cat /tmp/resolve.json
echo "Prebuilt resolver ran with no Visual Studio present."
if ($LASTEXITCODE -ne 0) { Write-Host "::error::pip install huggingface_hub failed"; exit 1 }
python studio/install_llama_prebuilt.py --resolve-prebuilt latest --output-format json > resolve.json
if ($LASTEXITCODE -ne 0) {
Write-Host "::error::resolver exited non-zero"
if (Test-Path resolve.json) { Get-Content resolve.json }
exit 1
}
Get-Content resolve.json
Write-Host "Prebuilt resolver ran with no Visual Studio present."
- name: Restore Visual Studio
- name: Clean no-build-tools simulation
if: always()
shell: pwsh
run: |
foreach ($d in @("$env:ProgramFiles\Microsoft Visual Studio", "${env:ProgramFiles(x86)}\Microsoft Visual Studio")) {
$off = "$d.vsoff"
if (Test-Path -LiteralPath $off) { Rename-Item -LiteralPath $off -NewName (Split-Path $d -Leaf); Write-Host "Restored $d" }
}
Remove-Item -LiteralPath (Join-Path $env:GITHUB_WORKSPACE 'no-build-tools') -Recurse -Force -ErrorAction SilentlyContinue
# ── folded from studio-setup-ps1-vs2026.yml: setup.ps1 unit tests + real-VS detection + vcredist ──
pester:

View file

@ -212,7 +212,7 @@ By default `unsloth studio` binds to `127.0.0.1` (this machine only). To reach i
```bash
unsloth studio --secure -p 8888
```
- `-H 0.0.0.0`: bind the raw port on all network interfaces, reachable from anywhere on the network. Only use this on a trusted network.
- `-H 0.0.0.0`: bind the raw port on all network interfaces, reachable from anywhere on the network. This also starts a public Cloudflare quick tunnel by default, which publishes an internet-reachable `https://*.trycloudflare.com` URL even behind a firewall. Both the raw port and the tunnel expose Studio beyond this machine, so only use this on a network you trust; pass `--no-cloudflare` to drop the public link while keeping the network bind.
```bash
unsloth studio -H 0.0.0.0 -p 8888
```

View file

@ -469,6 +469,17 @@ function Install-UnslothStudio {
param(
[Parameter(Mandatory = $true)][ScriptBlock]$Command
)
# Installer-pinned index installs (torch) must beat an inherited uv mirror
# (#6898): when the command pins an index, clear every uv index env var so
# it wins, then restore in finally. Other installs keep the user's mirror.
$savedUvIndex = $null
if ($Command.ToString() -match '--default-index') {
$savedUvIndex = @{}
foreach ($n in 'UV_DEFAULT_INDEX', 'UV_INDEX_URL', 'UV_INDEX', 'UV_EXTRA_INDEX_URL') {
$savedUvIndex[$n] = [Environment]::GetEnvironmentVariable($n)
Remove-Item "Env:$n" -ErrorAction SilentlyContinue
}
}
$prevEap = $ErrorActionPreference
$ErrorActionPreference = "Continue"
try {
@ -488,6 +499,7 @@ function Install-UnslothStudio {
return [int]$LASTEXITCODE
} finally {
$ErrorActionPreference = $prevEap
if ($savedUvIndex) { foreach ($n in $savedUvIndex.Keys) { if ($null -ne $savedUvIndex[$n]) { Set-Item "Env:$n" $savedUvIndex[$n] } } }
}
}
@ -2200,7 +2212,7 @@ exit 0
# ABI-incompatible torchvision/torchaudio on AMD's per-arch index.
$visionSpec = if ($ROCmGfxArch -and $torchvisionFloorMap.ContainsKey($ROCmGfxArch)) { $torchvisionFloorMap[$ROCmGfxArch] } else { "torchvision" }
$audioSpec = if ($ROCmGfxArch -and $torchaudioFloorMap.ContainsKey($ROCmGfxArch)) { $torchaudioFloorMap[$ROCmGfxArch] } else { "torchaudio" }
$torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch (AMD ROCm)" { uv pip install --python $VenvPython --force-reinstall --index-url $ROCmIndexUrl $torchSpec $visionSpec $audioSpec }
$torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch (AMD ROCm)" { uv pip install --python $VenvPython --force-reinstall --default-index $ROCmIndexUrl $torchSpec $visionSpec $audioSpec }
if ($torchInstallExit -ne 0) {
# Transient AMD-index failure: fall back to a CPU base so the install
# still completes; Studio setup retries ROCm afterwards.
@ -2209,7 +2221,7 @@ exit 0
# torch (e.g. 2.10.0+rocm on gfx110X/gfx90a) that still satisfies the CPU
# torch>= range, so without it uv would keep the ROCm build and only swap
# the companions -- a mismatched venv the flavor-repair block won't fix.
$torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch (CPU fallback)" { uv pip install --python $VenvPython --force-reinstall "torch>=2.4,<2.11.0" torchvision torchaudio --index-url $TorchIndexUrl }
$torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch (CPU fallback)" { uv pip install --python $VenvPython --force-reinstall "torch>=2.4,<2.11.0" torchvision torchaudio --default-index $TorchIndexUrl }
if ($torchInstallExit -ne 0) {
Write-Host "[ERROR] Failed to install PyTorch (ROCm and CPU base both failed, exit code $torchInstallExit)" -ForegroundColor Red
return (Exit-InstallFailure "Failed to install PyTorch (exit code $torchInstallExit)" $torchInstallExit)
@ -2223,7 +2235,7 @@ exit 0
} else {
Write-TauriLog "STEP" "Installing PyTorch"
substep "installing PyTorch ($TorchIndexUrl)..."
$torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch" { uv pip install --python $VenvPython "torch>=2.4,<2.11.0" torchvision torchaudio --index-url $TorchIndexUrl }
$torchInstallExit = Invoke-InstallCommandRetry -Label "install PyTorch" { uv pip install --python $VenvPython "torch>=2.4,<2.11.0" torchvision torchaudio --default-index $TorchIndexUrl }
if ($torchInstallExit -ne 0) {
Write-Host "[ERROR] Failed to install PyTorch (exit code $torchInstallExit)" -ForegroundColor Red
return (Exit-InstallFailure "Failed to install PyTorch (exit code $torchInstallExit)" $torchInstallExit)
@ -2306,7 +2318,7 @@ exit 0
# keeps a stale torch==X+cpu against a CUDA index and setup.ps1 then loops on
# "torch cpu != required cuXXX". Reinstall the right triplet when a GPU build is
# expected: CUDA from $TorchIndexUrl, ROCm from $ROCmIndexUrl (repo.amd.com gfx*
# is a PEP 503 index uv resolves via --index-url, same URL the fresh ROCm install
# is a PEP 503 index uv resolves via --default-index, same URL the fresh ROCm install
# above uses). --no-torch / CPU-only hosts (expected cpu) are no-ops.
if (-not $SkipTorch) {
$expectedTorchTag = Get-ExpectedTorchFlavorTag -TorchIndexUrl $TorchIndexUrl -ROCmIndexUrl $ROCmIndexUrl
@ -2322,7 +2334,7 @@ exit 0
$visionSpec = if ($ROCmGfxArch -and $torchvisionFloorMap.ContainsKey($ROCmGfxArch)) { $torchvisionFloorMap[$ROCmGfxArch] } else { "torchvision" }
$audioSpec = if ($ROCmGfxArch -and $torchaudioFloorMap.ContainsKey($ROCmGfxArch)) { $torchaudioFloorMap[$ROCmGfxArch] } else { "torchaudio" }
substep "PyTorch flavor mismatch (installed $installedTorchTag, need ROCm) -- reinstalling correct build..." "Yellow"
$torchFixExit = Invoke-InstallCommand { uv pip install --python $VenvPython --force-reinstall --index-url $ROCmIndexUrl $rocmSpec $visionSpec $audioSpec }
$torchFixExit = Invoke-InstallCommand { uv pip install --python $VenvPython --force-reinstall --default-index $ROCmIndexUrl $rocmSpec $visionSpec $audioSpec }
if ($torchFixExit -ne 0) {
Write-Host "[ERROR] Failed to reinstall PyTorch with the correct ROCm build (exit code $torchFixExit)" -ForegroundColor Red
return (Exit-InstallFailure "Failed to reinstall PyTorch (ROCm) (exit code $torchFixExit)" $torchFixExit)
@ -2331,7 +2343,7 @@ exit 0
} elseif ($expectedTorchTag -ne 'rocm') {
# CUDA: stale +cpu (or wrong cuXXX) against a CUDA index -> reinstall triplet.
substep "PyTorch flavor mismatch (installed $installedTorchTag, need $expectedTorchTag) -- reinstalling correct build..." "Yellow"
$torchFixExit = Invoke-InstallCommand { uv pip install --python $VenvPython "torch>=2.4,<2.11.0" torchvision torchaudio --index-url $TorchIndexUrl --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio }
$torchFixExit = Invoke-InstallCommand { uv pip install --python $VenvPython "torch>=2.4,<2.11.0" torchvision torchaudio --default-index $TorchIndexUrl --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio }
if ($torchFixExit -ne 0) {
Write-Host "[ERROR] Failed to reinstall PyTorch with the correct CUDA build (exit code $torchFixExit)" -ForegroundColor Red
return (Exit-InstallFailure "Failed to reinstall PyTorch ($expectedTorchTag) (exit code $torchFixExit)" $torchFixExit)

View file

@ -159,6 +159,12 @@ run_maybe_quiet() {
run_install_cmd() {
_label="$1"
shift
# Installer-pinned index installs (torch) must beat an inherited uv mirror
# (#6898): when we pass --default-index, neutralize every uv index env var so
# the pinned index wins. Other installs keep the user's mirror.
case " $* " in
*" --default-index "*) set -- env -u UV_DEFAULT_INDEX -u UV_INDEX_URL -u UV_INDEX -u UV_EXTRA_INDEX_URL "$@" ;;
esac
if _is_verbose; then
"$@" && return 0
_rc=$?
@ -2190,9 +2196,9 @@ _expected_torch_flavor_tag() {
esac
}
# Whether index ($1) supports a plain --index-url reinstall. pytorch.org cuXXX /
# Whether index ($1) supports a plain --default-index reinstall. pytorch.org cuXXX /
# rocmX.Y AND the repo.amd.com gfx* indexes are all PEP 503 simple indexes that uv
# resolves (torch + every transitive dep) via --index-url -- the same URLs the
# resolves (torch + every transitive dep) via --default-index -- the same URLs the
# fresh-install paths above already use -- so a stale wheel is auto-repairable.
# Unknown/odd-mirror leaves -> no, so we warn rather than risk a wrong reinstall.
_torch_index_repairable() {
@ -2744,7 +2750,7 @@ if [ "$_MIGRATED" = true ]; then
substep "repairing ROCm torch (overwritten by dependency resolution)..."
run_install_cmd_retry "repair ROCm torch" uv pip install --python "$_VENV_PY" \
"$TORCH_CONSTRAINT" torchvision torchaudio \
--index-url "$TORCH_INDEX_URL" \
--default-index "$TORCH_INDEX_URL" \
--force-reinstall
fi
;;
@ -2870,7 +2876,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
substep "[WARN] Radeon repo lacks a compatible wheel set for this Python; falling back to ROCm index ($TORCH_INDEX_URL)" "$C_WARN"
run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" \
"$TORCH_CONSTRAINT" torchvision torchaudio \
--index-url "$TORCH_INDEX_URL"
--default-index "$TORCH_INDEX_URL"
else
substep "installing PyTorch from Radeon repo (${_RADEON_BASE_URL})..."
# Pass explicit wheel URLs so the matched trio is
@ -2893,18 +2899,18 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
substep "[WARN] Radeon repo unavailable; falling back to ROCm index ($TORCH_INDEX_URL)" "$C_WARN"
run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" \
"$TORCH_CONSTRAINT" torchvision torchaudio \
--index-url "$TORCH_INDEX_URL"
--default-index "$TORCH_INDEX_URL"
fi
else
substep "[WARN] Radeon GPU detected but could not detect full ROCm version; falling back to ROCm index" "$C_WARN"
run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" \
"$TORCH_CONSTRAINT" torchvision torchaudio \
--index-url "$TORCH_INDEX_URL"
--default-index "$TORCH_INDEX_URL"
fi
else
substep "installing PyTorch ($TORCH_INDEX_URL)..."
run_install_cmd_retry "install PyTorch" uv pip install --python "$_VENV_PY" "$TORCH_CONSTRAINT" torchvision torchaudio \
--index-url "$TORCH_INDEX_URL"
--default-index "$TORCH_INDEX_URL"
fi
# AMD ROCm: install bitsandbytes (once, after torch, for all ROCm paths).
# Gate on SKIP_TORCH=false so a user running with --no-torch on a ROCm
@ -2964,7 +2970,7 @@ elif [ -n "$TORCH_INDEX_URL" ]; then
substep "repairing ROCm torch (overwritten by dependency resolution)..."
run_install_cmd_retry "repair ROCm torch" uv pip install --python "$_VENV_PY" \
"$TORCH_CONSTRAINT" torchvision torchaudio \
--index-url "$TORCH_INDEX_URL" \
--default-index "$TORCH_INDEX_URL" \
--force-reinstall
fi
;;
@ -2999,14 +3005,14 @@ if [ "$SKIP_TORCH" = false ] && [ -n "${TORCH_INDEX_URL:-}" ]; then
_installed_torch_ver=$("$_VENV_PY" -c "import torch; print(torch.__version__)" 2>/dev/null || true)
_installed_torch_tag=""
[ -n "$_installed_torch_ver" ] && _installed_torch_tag=$(_torch_flavor_tag "$_installed_torch_ver")
# Repair when flavor is wrong AND the index is plain --index-url reinstallable
# Repair when flavor is wrong AND the index is plain --default-index reinstallable
# (cuXXX / rocmX.Y / repo.amd.com gfx*); an unknown mirror leaf -> warn only.
if [ -n "$_installed_torch_tag" ] && [ "$_installed_torch_tag" != "$_expected_torch_tag" ] \
&& [ "$(_torch_index_repairable "$TORCH_INDEX_URL")" = "yes" ]; then
substep "PyTorch flavor mismatch (installed $_installed_torch_tag, need $_expected_torch_tag) -- reinstalling correct build..."
run_install_cmd "reinstall PyTorch ($_expected_torch_tag)" uv pip install --python "$_VENV_PY" \
"$TORCH_CONSTRAINT" torchvision torchaudio \
--index-url "$TORCH_INDEX_URL" \
--default-index "$TORCH_INDEX_URL" \
--reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio
_installed_torch_ver=$("$_VENV_PY" -c "import torch; print(torch.__version__)" 2>/dev/null || true)
_installed_torch_tag=""
@ -3017,7 +3023,7 @@ if [ "$SKIP_TORCH" = false ] && [ -n "${TORCH_INDEX_URL:-}" ]; then
substep "[WARN] PyTorch is CPU-only but a $_expected_torch_tag GPU build was expected for this machine." "$C_WARN"
substep "[WARN] Training and GPU inference will run on CPU until this is fixed." "$C_WARN"
substep "[WARN] Re-run this installer, or reinstall the GPU build manually:" "$C_WARN"
substep "[WARN] uv pip install --python \"$_VENV_PY\" \"$TORCH_CONSTRAINT\" torchvision torchaudio --index-url $TORCH_INDEX_URL --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio" "$C_WARN"
substep "[WARN] uv pip install --python \"$_VENV_PY\" \"$TORCH_CONSTRAINT\" torchvision torchaudio --default-index $TORCH_INDEX_URL --reinstall-package torch --reinstall-package torchvision --reinstall-package torchaudio" "$C_WARN"
fi
fi
fi

File diff suppressed because one or more lines are too long

View file

@ -0,0 +1,110 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Standalone free-VRAM probe for the bundled ggml Vulkan backend.
Run in a short-lived subprocess (``python _vulkan_probe.py <bindir>``) so the
Vulkan instance never lives in the long-running backend process. Loads the
bundled ggml Vulkan backend from ``<bindir>`` and prints one
``<idx>\\t<free_bytes>\\t<is_igpu>\\t<total_bytes>`` line per device to stdout.
Indices are ggml's own Vulkan device ordinals, which need not match nvidia-smi
order. ``is_igpu`` (from ggml's device type) is ``1`` for an integrated GPU
sharing system RAM. ``total_bytes`` is the device-local heap; the reader uses
it to reserve absolute headroom on a discrete card (parity with the CUDA/ROCm
fit) and ignores it for an iGPU, whose "VRAM" is shared system RAM.
Uses only the standard library so it stays runnable as a bare script.
"""
import ctypes
import os
import sys
# ggml_backend_dev_type enum (ggml-backend.h): CPU=0, GPU=1, IGPU=2, ...
_GGML_BACKEND_DEVICE_TYPE_IGPU = 2
def _igpu_flags(base, lib, count: int) -> list[bool]:
"""Per-device integrated-GPU flags via ggml's backend registry.
The Vulkan reg enumerates devices in the same order as
``ggml_backend_vk_get_device_memory`` (each context uses ``ctx->device =
i``), so reg index == device ordinal. Returns all-False on any failure so
the reader never over-caps a discrete card.
"""
flags = [False] * count
try:
lib.ggml_backend_vk_reg.restype = ctypes.c_void_p
lib.ggml_backend_vk_reg.argtypes = []
base.ggml_backend_reg_dev_count.restype = ctypes.c_size_t
base.ggml_backend_reg_dev_count.argtypes = [ctypes.c_void_p]
base.ggml_backend_reg_dev_get.restype = ctypes.c_void_p
base.ggml_backend_reg_dev_get.argtypes = [ctypes.c_void_p, ctypes.c_size_t]
base.ggml_backend_dev_type.restype = ctypes.c_int
base.ggml_backend_dev_type.argtypes = [ctypes.c_void_p]
reg = lib.ggml_backend_vk_reg()
if not reg:
return flags
dev_count = base.ggml_backend_reg_dev_count(reg)
for i in range(min(count, dev_count)):
dev = base.ggml_backend_reg_dev_get(reg, i)
if dev:
flags[i] = base.ggml_backend_dev_type(dev) == _GGML_BACKEND_DEVICE_TYPE_IGPU
except Exception:
# Best-effort: any failure degrades to "discrete" so the memory
# readings still get through instead of crashing the probe.
pass
return flags
def main() -> int:
if len(sys.argv) < 2:
return 0
bindir = sys.argv[1]
# Hold add_dll_directory's handle for the rest of main() (the documented
# idiom) so bindir stays on the search path while the sibling ggml DLLs
# resolve below.
_dll_dir = None
if sys.platform == "win32":
base_name, vk_name = "ggml-base.dll", "ggml-vulkan.dll"
try:
_dll_dir = os.add_dll_directory(bindir)
except Exception:
pass
else:
base_name, vk_name = "libggml-base.so", "libggml-vulkan.so"
# RTLD_GLOBAL exposes ggml-base's symbols to ggml-vulkan on POSIX. getattr
# falls back to 0 where the flag doesn't exist (Windows CDLL ignores mode).
_rtld_global = getattr(ctypes, "RTLD_GLOBAL", 0)
try:
base = ctypes.CDLL(os.path.join(bindir, base_name), mode = _rtld_global)
lib = ctypes.CDLL(os.path.join(bindir, vk_name), mode = _rtld_global)
except OSError as e:
print(f"ggml-vulkan load failed: {e}", file = sys.stderr)
return 1
lib.ggml_backend_vk_get_device_count.restype = ctypes.c_int
lib.ggml_backend_vk_get_device_count.argtypes = []
lib.ggml_backend_vk_get_device_memory.restype = None
lib.ggml_backend_vk_get_device_memory.argtypes = [
ctypes.c_int,
ctypes.POINTER(ctypes.c_size_t),
ctypes.POINTER(ctypes.c_size_t),
]
count = lib.ggml_backend_vk_get_device_count()
igpu = _igpu_flags(base, lib, count)
rows = []
for i in range(count):
free, total = ctypes.c_size_t(0), ctypes.c_size_t(0)
lib.ggml_backend_vk_get_device_memory(i, ctypes.byref(free), ctypes.byref(total))
rows.append("%d\t%d\t%d\t%d" % (i, free.value, int(igpu[i]), total.value))
sys.stdout.write("\n".join(rows))
return 0
if __name__ == "__main__":
raise SystemExit(main())

View file

@ -0,0 +1,368 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
"""Admission control for local llama-server generation requests.
The helpers in this module deliberately know nothing about FastAPI, SSE, or the
OpenAI-compatible route shape. They only coordinate how many upstream generation
requests may be active for one llama-server backend and provide a cancellable
FIFO queue for excess requests.
"""
from __future__ import annotations
import asyncio
import os
import threading
from collections import deque
from dataclasses import dataclass
from typing import Deque, Optional
ADMISSION_CONTROL_ENV = "UNSLOTH_OPENAI_COMPAT_ADMISSION_CONTROL"
ADMISSION_QUEUE_TIMEOUT_ENV = "UNSLOTH_OPENAI_COMPAT_ADMISSION_QUEUE_TIMEOUT"
ADMISSION_KEEPALIVE_INTERVAL_ENV = "UNSLOTH_OPENAI_COMPAT_ADMISSION_KEEPALIVE_INTERVAL"
ADMISSION_MAX_QUEUE_ENV = "UNSLOTH_OPENAI_COMPAT_ADMISSION_MAX_QUEUE"
DEFAULT_ADMISSION_ENABLED = True
DEFAULT_ADMISSION_QUEUE_TIMEOUT_S = None
DEFAULT_ADMISSION_KEEPALIVE_INTERVAL_S = 5.0
DEFAULT_ADMISSION_MAX_QUEUE = 64
@dataclass(frozen = True)
class LlamaAdmissionConfig:
enabled: bool = DEFAULT_ADMISSION_ENABLED
queue_timeout_s: Optional[float] = DEFAULT_ADMISSION_QUEUE_TIMEOUT_S
keepalive_interval_s: float = DEFAULT_ADMISSION_KEEPALIVE_INTERVAL_S
max_queue: Optional[int] = DEFAULT_ADMISSION_MAX_QUEUE
@dataclass(frozen = True)
class LlamaAdmissionSnapshot:
key: str
capacity: int
active: int
queued: int
class LlamaAdmissionError(Exception):
def __init__(
self,
message: str,
*,
snapshot: Optional[LlamaAdmissionSnapshot] = None,
):
super().__init__(message)
self.snapshot = snapshot
class LlamaAdmissionQueueFull(LlamaAdmissionError):
pass
class LlamaAdmissionTimeout(LlamaAdmissionError):
pass
class LlamaAdmissionCancelled(LlamaAdmissionError):
pass
def _bool_env(name: str, default: bool) -> bool:
value = os.environ.get(name)
if value is None or not value.strip():
return default
value = value.strip().lower()
if value in {"1", "true", "yes", "on"}:
return True
if value in {"0", "false", "no", "off"}:
return False
return default
def _optional_positive_float_env(name: str, default: Optional[float]) -> Optional[float]:
value = os.environ.get(name)
if value is None or not value.strip():
return default
try:
parsed = float(value.strip())
except ValueError:
return default
return parsed if parsed > 0 else None
def _positive_float_env(name: str, default: float) -> float:
value = os.environ.get(name)
if value is None or not value.strip():
return default
try:
parsed = float(value.strip())
except ValueError:
return default
return parsed if parsed > 0 else default
def _optional_positive_int_env(name: str, default: Optional[int]) -> Optional[int]:
value = os.environ.get(name)
if value is None or not value.strip():
return default
try:
parsed = int(value.strip())
except ValueError:
return default
return parsed if parsed > 0 else None
def llama_admission_config_from_env() -> LlamaAdmissionConfig:
return LlamaAdmissionConfig(
enabled = _bool_env(ADMISSION_CONTROL_ENV, DEFAULT_ADMISSION_ENABLED),
queue_timeout_s = _optional_positive_float_env(
ADMISSION_QUEUE_TIMEOUT_ENV,
DEFAULT_ADMISSION_QUEUE_TIMEOUT_S,
),
keepalive_interval_s = _positive_float_env(
ADMISSION_KEEPALIVE_INTERVAL_ENV,
DEFAULT_ADMISSION_KEEPALIVE_INTERVAL_S,
),
max_queue = _optional_positive_int_env(
ADMISSION_MAX_QUEUE_ENV,
DEFAULT_ADMISSION_MAX_QUEUE,
),
)
@dataclass
class _Waiter:
loop: asyncio.AbstractEventLoop
future: asyncio.Future
cancelled: bool = False
granted_lease: Optional["LlamaAdmissionLease"] = None
class LlamaAdmissionLease:
def __init__(self, queue: Optional["LlamaAdmissionQueue"]):
self._queue = queue
self._released = False
self._release_lock = threading.Lock()
def release(self) -> None:
queue = None
with self._release_lock:
if self._released:
return
self._released = True
queue = self._queue
if queue is not None:
queue.release()
async def __aenter__(self) -> "LlamaAdmissionLease":
return self
async def __aexit__(self, *_args) -> None:
self.release()
class LlamaAdmissionReservation:
def __init__(
self,
*,
queue: Optional["LlamaAdmissionQueue"],
lease: Optional[LlamaAdmissionLease] = None,
waiter: Optional[_Waiter] = None,
snapshot: Optional[LlamaAdmissionSnapshot] = None,
):
self._queue = queue
self._lease = lease
self._waiter = waiter
self.snapshot = snapshot
@property
def is_cancelled(self) -> bool:
return self._lease is None and self._waiter is None
def lease_nowait(self) -> Optional[LlamaAdmissionLease]:
if self._lease is not None:
return self._lease
if self._waiter is None or not self._waiter.future.done():
return None
if self._waiter.future.cancelled():
self._waiter.cancelled = True
self._waiter = None
return None
self._lease = self._waiter.future.result()
self._waiter = None
return self._lease
async def wait(self, timeout_s: float) -> Optional[LlamaAdmissionLease]:
lease = self.lease_nowait()
if lease is not None:
return lease
if self._waiter is None:
return None
waiter = self._waiter
try:
await asyncio.wait_for(asyncio.shield(waiter.future), timeout = timeout_s)
except asyncio.CancelledError:
if waiter.future.cancelled():
waiter.cancelled = True
if self._waiter is waiter:
self._waiter = None
return None
raise
return self.lease_nowait()
def cancel(self) -> None:
lease = self.lease_nowait()
if lease is not None:
lease.release()
self._lease = None
return
if self._queue is not None and self._waiter is not None:
self._queue.cancel(self._waiter)
self._waiter = None
def snapshot_now(self) -> Optional[LlamaAdmissionSnapshot]:
if self._queue is None:
return self.snapshot
return self._queue.snapshot()
class LlamaAdmissionQueue:
def __init__(self, key: str):
self.key = key
self._lock = threading.Lock()
self._active = 0
self._capacity = 1
self._waiters: Deque[_Waiter] = deque()
def reserve(self, *, capacity: int, config: LlamaAdmissionConfig) -> LlamaAdmissionReservation:
capacity = max(1, int(capacity or 1))
if not config.enabled:
return LlamaAdmissionReservation(
queue = None,
lease = LlamaAdmissionLease(None),
snapshot = LlamaAdmissionSnapshot(self.key, capacity, 0, 0),
)
loop = asyncio.get_running_loop()
with self._lock:
self._capacity = capacity
self._prune_waiters_locked()
self._grant_waiters_locked()
if self._active < self._capacity and not self._waiters:
self._active += 1
return LlamaAdmissionReservation(
queue = self,
lease = LlamaAdmissionLease(self),
snapshot = self._snapshot_locked(),
)
if config.max_queue is not None and len(self._waiters) >= config.max_queue:
raise LlamaAdmissionQueueFull(
"llama-server generation queue is full",
snapshot = self._snapshot_locked(),
)
waiter = _Waiter(
loop = loop,
future = loop.create_future(),
)
self._waiters.append(waiter)
return LlamaAdmissionReservation(
queue = self,
waiter = waiter,
snapshot = self._snapshot_locked(),
)
def release(self) -> None:
with self._lock:
if self._active > 0:
self._active -= 1
self._grant_waiters_locked()
def cancel(self, waiter: _Waiter) -> None:
lease_to_release = None
with self._lock:
waiter.cancelled = True
try:
self._waiters.remove(waiter)
except ValueError:
pass
if waiter.granted_lease is not None:
lease_to_release = waiter.granted_lease
waiter.granted_lease = None
if not waiter.future.done():
waiter.loop.call_soon_threadsafe(waiter.future.cancel)
if lease_to_release is not None:
lease_to_release.release()
def snapshot(self) -> LlamaAdmissionSnapshot:
with self._lock:
self._prune_waiters_locked()
return self._snapshot_locked()
def is_idle(self) -> bool:
with self._lock:
self._prune_waiters_locked()
return self._active == 0 and not self._waiters
def _grant_waiters_locked(self) -> None:
self._prune_waiters_locked()
while self._waiters and self._active < self._capacity:
waiter = self._waiters.popleft()
if waiter.cancelled or waiter.future.done():
continue
self._active += 1
lease = LlamaAdmissionLease(self)
waiter.granted_lease = lease
waiter.loop.call_soon_threadsafe(self._deliver_lease, waiter, lease)
def _deliver_lease(self, waiter: _Waiter, lease: LlamaAdmissionLease) -> None:
if waiter.cancelled or waiter.future.done():
waiter.granted_lease = None
if not waiter.future.done():
waiter.future.cancel()
lease.release()
return
try:
waiter.future.set_result(lease)
waiter.granted_lease = None
except asyncio.InvalidStateError:
waiter.granted_lease = None
lease.release()
def _prune_waiters_locked(self) -> None:
self._waiters = deque(
waiter for waiter in self._waiters if not waiter.cancelled and not waiter.future.done()
)
def _snapshot_locked(self) -> LlamaAdmissionSnapshot:
return LlamaAdmissionSnapshot(
key = self.key,
capacity = self._capacity,
active = self._active,
queued = len(self._waiters),
)
_QUEUES_LOCK = threading.Lock()
_QUEUES: dict[str, LlamaAdmissionQueue] = {}
def get_llama_admission_queue(key: str) -> LlamaAdmissionQueue:
with _QUEUES_LOCK:
queue = _QUEUES.get(key)
if queue is None:
queue = LlamaAdmissionQueue(key)
_QUEUES[key] = queue
# base_url carries a fresh ephemeral port on every model load, so
# each load registers a new key. Drop the now-idle queues from prior
# loads so the registry can't grow without bound on a long-running
# server. Queues with in-flight requests are kept until they drain.
for stale_key in [k for k in _QUEUES if k != key and _QUEUES[k].is_idle()]:
del _QUEUES[stale_key]
return queue
def reset_llama_admission_queues() -> None:
with _QUEUES_LOCK:
_QUEUES.clear()

View file

@ -100,6 +100,10 @@ class LlamaServerNotFoundError(RuntimeError):
Subclasses RuntimeError so existing handlers still catch it."""
class _LlamaStreamCancelled(Exception):
"""Internal signal for an expected client/request cancellation."""
# Shared so the from_identifier preflight and the load-time raise stay in sync.
LLAMA_SERVER_NOT_FOUND_DETAIL = (
"This is a GGUF model, but the llama.cpp runtime (llama-server) is not "
@ -1436,6 +1440,50 @@ def _backfill_usage_from_timings(usage, timings):
return out
def _vulkan_lib_filename() -> str:
return "ggml-vulkan.dll" if sys.platform == "win32" else "libggml-vulkan.so"
# Host RAM to leave free on an integrated GPU, matching llama.cpp's own --fit
# margin (default 1024 MiB per device). ggml reports an iGPU's "VRAM" as shared
# system RAM, so hold back the same margin rather than inventing a larger one.
_IGPU_HOST_RESERVE_MIB = 1024
def _apply_igpu_host_reserve_mib(free_mib: int, is_igpu: bool) -> int:
"""Reserve host headroom on an integrated (shared-memory) Vulkan GPU.
An iGPU's reported free "VRAM" is really free system RAM, so sizing
context/offload against all of it would push the host into swap or the OOM
killer. Leave the same margin llama.cpp's --fit uses. ``is_igpu`` comes from
ggml's device type, so a discrete card is never touched; only ever reduces.
"""
if not is_igpu:
return free_mib
return max(0, free_mib - _IGPU_HOST_RESERVE_MIB)
def _llama_lib_dir(binary: str) -> Path:
# The installer exposes llama-server as a top-level entrypoint into build/bin/,
# where the ggml backend libs live, so callers looking for sibling libs (Vulkan
# detection, LD_LIBRARY_PATH, probe bindir) need the real dir. It is normally a
# symlink (resolve() reaches build/bin), but create_exec_entrypoint falls back to
# a shell wrapper (exec "$(dirname "$0")/build/bin/llama-server" "$@") when it
# cannot symlink, and resolve() stops at the wrapper file. Follow the wrapper's
# exec target too, so a wrapper-based install still finds build/bin.
resolved = Path(binary).resolve()
try:
with open(resolved, "rb") as _f:
_head = _f.read(256)
if _head.startswith(b"#!"):
_m = re.search(r'exec "\$\(dirname "\$0"\)/([^"]+)"', _head.decode("utf-8", "ignore"))
if _m:
return (resolved.parent / _m.group(1)).resolve().parent
except OSError:
pass
return resolved.parent
def _is_external_link(path: Path) -> bool:
"""True when ``path`` is a --with-llama-cpp-dir local link: a POSIX symlink
or a Windows directory junction / reparse point. Such a link resolves into
@ -1493,6 +1541,7 @@ class LlamaCppBackend:
self._context_length: Optional[int] = None
self._effective_context_length: Optional[int] = None
self._max_context_length: Optional[int] = None
self._effective_parallel_slots: int = 1
self._chat_template: Optional[str] = None
self._chat_template_override: Optional[str] = None
self._supports_reasoning: bool = False
@ -1678,6 +1727,15 @@ class LlamaCppBackend:
"""Return the effective context length the server is running at."""
return self._effective_context_length or self._context_length
@property
def effective_parallel_slots(self) -> int:
"""Return the serving-slot count the active llama-server actually uses."""
try:
slots = int(getattr(self, "_effective_parallel_slots", 1))
except (TypeError, ValueError):
slots = 1
return max(1, slots)
@property
def max_context_length(self) -> Optional[int]:
"""Return the largest context that fits on this hardware at load time.
@ -1694,6 +1752,16 @@ class LlamaCppBackend:
"""Return the model's native context length from GGUF metadata."""
return self._context_length
def _commit_effective_parallel_slots(self, n_parallel: int) -> None:
try:
slots = int(n_parallel)
except (TypeError, ValueError):
slots = 1
self._effective_parallel_slots = max(1, slots)
def _reset_effective_parallel_slots(self) -> None:
self._effective_parallel_slots = 1
@staticmethod
def _read_rss_bytes(pid: int) -> Optional[int]:
"""Resident set size of ``pid`` in bytes, from /proc/<pid>/status (Linux).
@ -2278,6 +2346,30 @@ class LlamaCppBackend:
return total
@staticmethod
def _is_vulkan_backend(binary: Optional[str] = None) -> bool:
"""True if the installed llama.cpp build is Vulkan-only.
The official prebuilts are single-backend, so the Vulkan ggml lib next
to llama-server identifies a Vulkan build. Keeps the free-memory probe
and GPU pin in ggml's Vulkan device-index space. For a custom
multi-backend build with a CUDA or HIP ggml lib alongside Vulkan, defer
to that backend (torch-usable, better-understood probe/pin).
"""
binary = binary or LlamaCppBackend._find_llama_server_binary()
if not binary:
return False
lib_dir = _llama_lib_dir(binary)
if not (lib_dir / _vulkan_lib_filename()).is_file():
return False
for _backend in ("cuda", "hip"):
sibling = (
f"ggml-{_backend}.dll" if sys.platform == "win32" else f"libggml-{_backend}.so"
)
if (lib_dir / sibling).is_file():
return False
return True
@staticmethod
def _resolve_visible_physical_ids() -> Optional[list[int]]:
"""Physical GPU ids behind the active visibility mask (HIP/ROCR/CUDA on
@ -2440,11 +2532,42 @@ class LlamaCppBackend:
return True
@staticmethod
def _get_gpu_free_memory() -> list[tuple[int, int]]:
def _visible_devices_mask(env_name: str) -> Optional[set[int]]:
"""Physical indices a ``*_VISIBLE_DEVICES`` mask permits, or None if unset.
``if x.strip()`` filters trailing-comma masks ("0,1,"); an empty mask
("") yields an empty set (all devices hidden), distinct from an unset
var (None, no mask). Used by the nvidia-smi probe.
"""
raw = os.environ.get(env_name)
if raw is None:
return None
try:
return set(int(x.strip()) for x in raw.split(",") if x.strip())
except ValueError:
return None
@staticmethod
def _vulkan_pin_args(gpu_indices: Optional[Iterable[int]]) -> list[str]:
"""``--device Vulkan<i>,...`` to pin a Vulkan launch to selected GPUs.
The indices are ggml's compact Vulkan ordinals (as _get_gpu_free_memory
reports and the registry names ``Vulkan<i>``). Pin by that name, NOT via
GGML_VK_VISIBLE_DEVICES: ggml parses that env var in the raw
vkEnumeratePhysicalDevices space (before dropping CPU/llvmpipe devices
and deduplicating ICDs), so a compact ordinal there could select a
different physical device or the CPU rasterizer.
"""
if not gpu_indices:
return []
return ["--device", ",".join(f"Vulkan{i}" for i in gpu_indices)]
@staticmethod
def _get_gpu_free_memory(binary: Optional[str] = None) -> list[tuple[int, int]]:
"""Query free memory per GPU. Returns ``(gpu_index, free_mib)`` sorted by
index; empty if no supported GPU is reachable. Thin wrapper over
``_get_gpu_memory`` for callers that only need free VRAM."""
return [(idx, free) for idx, free, _total in LlamaCppBackend._get_gpu_memory()]
return [(idx, free) for idx, free, _total in LlamaCppBackend._get_gpu_memory(binary)]
@staticmethod
def _apple_metal_memory_budget_bytes() -> int:
@ -2475,7 +2598,7 @@ class LlamaCppBackend:
return int(rec_bytes * _APPLE_UNIFIED_MEMORY_FRACTION)
@staticmethod
def _get_gpu_memory() -> list[tuple[int, int, int]]:
def _get_gpu_memory(binary: Optional[str] = None) -> list[tuple[int, int, int]]:
"""Query free AND total memory per GPU.
Order:
@ -2487,9 +2610,18 @@ class LlamaCppBackend:
probe returned [] on AMD) and NVIDIA hosts missing
``nvidia-smi`` from PATH.
On a Vulkan build the ggml Vulkan probe is authoritative, so the indices
are ggml's compact Vulkan ordinals (the space the pin selects via
``--device Vulkan<i>``). It reports ``total`` for discrete cards and 0
for an iGPU (shared RAM) so the fit falls back to free*frac there.
Otherwise nvidia-smi / torch cover NVIDIA + AMD ROCm.
Returns (gpu_index, free_mib, total_mib) sorted by index; empty if no
supported GPU is reachable. ``total`` lets the fit reserve absolute headroom.
supported GPU is reachable.
"""
binary = binary or LlamaCppBackend._find_llama_server_binary()
if LlamaCppBackend._is_vulkan_backend(binary):
return LlamaCppBackend._get_gpu_free_memory_vulkan(binary)
# ── NVIDIA via nvidia-smi ────────────────────────────────────
try:
result = subprocess.run(
@ -2505,16 +2637,7 @@ class LlamaCppBackend:
**_windows_hidden_subprocess_kwargs(),
)
if result.returncode == 0:
allowed: Optional[set[int]] = None
cvd = os.environ.get("CUDA_VISIBLE_DEVICES")
if cvd is not None:
try:
# `if x.strip()` filters trailing-comma masks ("0,1,").
# Empty mask (CVD="") yields an empty set -> all GPUs
# filtered out, per codebase convention.
allowed = set(int(x.strip()) for x in cvd.split(",") if x.strip())
except ValueError:
pass
allowed = LlamaCppBackend._visible_devices_mask("CUDA_VISIBLE_DEVICES")
gpus: list[tuple[int, int, int]] = []
for line in result.stdout.strip().splitlines():
parts = [p.strip() for p in line.split(",")]
@ -2579,6 +2702,91 @@ class LlamaCppBackend:
logger.debug(f"torch GPU probe failed: {e}")
return []
@staticmethod
def _get_gpu_free_memory_vulkan(binary: Optional[str] = None) -> list[tuple[int, int, int]]:
"""Query free (and total) VRAM per device via the bundled ggml Vulkan backend.
Loads ``libggml-vulkan`` in a short-lived subprocess (no Vulkan instance
in this process) and returns (device_index, free_mib, total_mib) sorted
by index. The index is ggml's compact Vulkan ordinal -- the one the
registry names ``Vulkan<index>`` and load_model pins with ``--device``,
NOT the raw ``GGML_VK_VISIBLE_DEVICES`` space. A user-set
``GGML_VK_VISIBLE_DEVICES`` is honored by ggml (passed through), so the
list already reflects it. iGPUs leave a host-RAM margin (see
``_apply_igpu_host_reserve_mib``) and report total 0; discrete cards pass
their real total through. [] when no Vulkan build or device is reachable.
"""
binary = binary or LlamaCppBackend._find_llama_server_binary()
if not binary:
return []
binary_dir = _llama_lib_dir(binary)
if not (binary_dir / _vulkan_lib_filename()).is_file():
return []
env = child_env_without_native_path_secret()
# Pass any inherited GGML_VK_VISIBLE_DEVICES through to ggml unchanged so
# the probe enumerates the same device list the launch will, named
# Vulkan0..N in the compact order reported here and pinned by that name
# via --device -- probe, mask, and pin stay in one index space. Do NOT
# filter the mask in Python: ggml parses the env var in raw
# vkEnumeratePhysicalDevices space while this probe reports the compact
# post-filter ordinal, so a Python filter would compare mismatched spaces.
if sys.platform != "win32":
# Let the loader resolve sibling ggml libs next to the binary.
existing_ld = env.get("LD_LIBRARY_PATH", "")
env["LD_LIBRARY_PATH"] = (
f"{binary_dir}:{existing_ld}" if existing_ld else str(binary_dir)
)
probe_script = Path(__file__).with_name("_vulkan_probe.py")
try:
result = subprocess.run(
[sys.executable, str(probe_script), str(binary_dir)],
capture_output = True,
text = True,
timeout = 15,
env = env,
**_windows_hidden_subprocess_kwargs(),
)
if result.returncode != 0:
logger.debug(
f"vulkan GPU probe exited {result.returncode}: {result.stderr.strip()}"
)
return []
except Exception as e:
logger.debug(f"vulkan GPU probe failed: {e}")
return []
gpus: list[tuple[int, int, int]] = []
for line in result.stdout.strip().splitlines():
parts = line.split("\t")
if len(parts) != 4:
continue
try:
idx = int(parts[0])
free_mib = int(parts[1]) // (1024 * 1024)
is_igpu = parts[2] == "1"
# iGPU "total" is shared RAM, not a VRAM budget -> keep 0 so the
# fit stays on free*frac (the host reserve below is its
# headroom); a discrete card passes its real total through.
total_mib = 0 if is_igpu else int(parts[3]) // (1024 * 1024)
except ValueError:
continue
capped = _apply_igpu_host_reserve_mib(free_mib, is_igpu)
if capped < free_mib:
logger.info(
f"Vulkan device VK{idx} is an integrated GPU sharing system "
f"RAM; reserving {free_mib - capped}MiB host headroom "
f"({free_mib}->{capped}MiB usable)"
)
gpus.append((idx, capped, total_mib))
gpus.sort(key = lambda g: g[0])
if gpus:
logger.info(
"Vulkan GPU memory detected: "
+ ", ".join(f"VK{idx}={free}MiB" for idx, free, _total in gpus)
)
return gpus
@staticmethod
def _available_system_memory_mib() -> Optional[int]:
"""Available system RAM in MiB (psutil, then /proc/meminfo), or None if
@ -2807,7 +3015,8 @@ class LlamaCppBackend:
def _llama_server_env_for_binary(binary: str) -> dict[str, str]:
"""Build a subprocess env that lets llama-server resolve native libs."""
env = child_env_without_native_path_secret()
binary_dir = str(Path(binary).parent)
# _llama_lib_dir resolves the llama-server symlink to the real build/bin.
binary_dir = str(_llama_lib_dir(binary))
if sys.platform == "win32":
# Ordering: see _build_windows_path_dirs. #5106.
@ -4488,6 +4697,29 @@ class LlamaCppBackend:
cancel_event = cancel_event,
)
def _cached_repo_mtp_drafter(self, hf_repo: str) -> Optional[str]:
"""A drafter already in this repo's local HF cache, reused offline when a
fresh copy can't be fetched. Prefers a repo-root ``mtp-*.gguf`` across all
cached snapshots; else an existing ``MTP/`` copy (any precision -- the
target verifies every drafted token). None if none is cached."""
try:
from utils.models.model_config import _iter_hf_cache_snapshots
roots: list[Path] = []
subdirs: list[Path] = []
for snap in _iter_hf_cache_snapshots(hf_repo): # newest first
for f in sorted(_gguf_snapshot_files(snap)):
if _is_companion_gguf_path(f) and "mmproj" not in f.lower():
(roots if "/" not in f else subdirs).append(snap / f)
# Keep snapshot order (newest first), root before any MTP/ copy, so a
# newer main GGUF pairs with the newest cached drafter, not a stale one.
for cand in roots + subdirs:
if cand.is_file():
return str(cand)
except Exception as e:
logger.debug("Cached MTP drafter lookup failed for %s: %s", hf_repo, e)
return None
def _download_mtp(
self,
*,
@ -4504,11 +4736,25 @@ class LlamaCppBackend:
are intentionally skipped. Returns the local path, or None.
"""
# Offline, reuse any drafter already on disk (a fresh copy can't be
# fetched). Online, _download_companion_gguf/hf_hub_download reuse the
# current cached file and refetch a changed one, so skip the probe here
# rather than pair new weights with a stale draft.
if _hf_env_offline():
cached = self._cached_repo_mtp_drafter(hf_repo)
if cached:
logger.info(f"Reusing cached MTP drafter (offline): {cached}")
return cached
def _pick_mtp(candidates: list[str]) -> Optional[str]:
# Root-level only: MTP/ subdir copies now share the mtp- prefix but
# are explicit-selection, not auto-fetch (they'd sort ahead of root).
mtp_files = sorted(
f
for f in candidates
if f.lower().endswith(".gguf") and Path(f).name.lower().startswith("mtp-")
if f.lower().endswith(".gguf")
and "/" not in f
and Path(f).name.lower().startswith("mtp-")
)
return mtp_files[0] if mtp_files else None
@ -5210,6 +5456,7 @@ class LlamaCppBackend:
# Resolve llama-server now but defer a not-found error: a block-diffusion
# GGUF uses the diffusion runner, and its arch is only known after the header.
binary = self._find_llama_server_binary()
is_vulkan_backend = self._is_vulkan_backend(binary)
# ── Phase 2: download (NO lock held, so cancel can proceed) ──
# mtp_draft_path arrives set for local Gemma loads (detected
@ -5449,7 +5696,8 @@ class LlamaCppBackend:
model_size = gguf_size + mmproj_size
# 2-tuple gpus for existing logic + a total map for the absolute
# per-GPU headroom (correct when the GPU is already partly used).
_gpu_mem = self._get_gpu_memory()
# Pass binary so a Vulkan build probes ggml's Vulkan ordinals.
_gpu_mem = self._get_gpu_memory(binary)
gpus = [(idx, free) for idx, free, _t in _gpu_mem]
total_by_idx = {idx: total for idx, _f, total in _gpu_mem}
@ -6222,7 +6470,12 @@ class LlamaCppBackend:
# cap, not the ROCm-reported VRAM, is the real ceiling); refuse an
# oversize load the OS would otherwise kill mid-flight. Base model
# only: an optional MTP drafter is dropped by the MTP-drop fallback.
if model_size is not None and self._amd_apu_wants_unified_memory(gpu_indices):
# CUDA/ROCm ids only; a Vulkan build's gpu_indices are ggml ordinals.
if (
model_size is not None
and not is_vulkan_backend
and self._amd_apu_wants_unified_memory(gpu_indices)
):
_ram_msg = self._apu_ram_shortfall_message(
model_size, self._available_system_memory_mib()
)
@ -6485,6 +6738,12 @@ class LlamaCppBackend:
", ".join(unsupported_cache_flags),
)
# Vulkan pins via --device (a cmd arg, unlike the env-based
# CUDA/ROCm pin below), emitted BEFORE user extras so llama.cpp's
# last-wins parsing lets a user --device override Studio's pick.
if is_vulkan_backend and gpu_indices is not None:
cmd += LlamaCppBackend._vulkan_pin_args(gpu_indices)
# User pass-through args go last so llama.cpp's last-wins parsing
# lets the user override Studio's auto-set flags. Already
# validated by the route via validate_extra_args().
@ -6536,23 +6795,25 @@ class LlamaCppBackend:
env.setdefault("OMP_NUM_THREADS", "2")
# 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(gpu_indices):
# shared system RAM. setdefault so a user value wins. Not on Vulkan
# (nor DC below): gpu_indices are ggml ordinals, not CUDA/ROCm ids.
if not is_vulkan_backend and self._amd_apu_wants_unified_memory(gpu_indices):
env.setdefault("GGML_CUDA_ENABLE_UNIFIED_MEMORY", "1")
logger.info("AMD unified-memory APU: set GGML_CUDA_ENABLE_UNIFIED_MEMORY=1")
# DC NVIDIA GPUs: FP32 accum (+ P2P / launch queues for multi-GPU).
# See _apply_datacenter_env; opt out with UNSLOTH_DISABLE_DC_TUNING=1.
if self._apply_datacenter_env(env, gpu_indices):
if not is_vulkan_backend and self._apply_datacenter_env(env, gpu_indices):
multi_gpu = self._effective_gpu_count(gpu_indices) > 1
logger.info(
f"Data-center GPU detected: applied DC llama.cpp env tuning (multi_gpu={multi_gpu})"
)
# Pin to selected GPU(s). On ROCm, narrowing only
# CUDA_VISIBLE_DEVICES leaves an AMD child seeing the full
# set, so set HIP_VISIBLE_DEVICES too.
if gpu_indices is not None:
# CUDA_VISIBLE_DEVICES leaves an AMD child seeing the full set, so
# set HIP_VISIBLE_DEVICES too. Vulkan is pinned via --device
# (above), not here.
if gpu_indices is not None and not is_vulkan_backend:
pinned = ",".join(str(i) for i in gpu_indices)
env["CUDA_VISIBLE_DEVICES"] = pinned
try:
@ -6946,6 +7207,7 @@ class LlamaCppBackend:
)
self._healthy = True
self._commit_effective_parallel_slots(n_parallel)
# Commit caller intent only after _healthy=True so a failed start
# can't poison the next inheritance check. None keeps prior, []
@ -7483,6 +7745,7 @@ class LlamaCppBackend:
self._context_length = None
self._effective_context_length = None
self._max_context_length = None
self._reset_effective_parallel_slots()
self._chat_template = None
self._chat_template_override = None
self._supports_reasoning = False
@ -7538,6 +7801,7 @@ class LlamaCppBackend:
# Stop the watchdog before a deliberate kill so a planned reload/unload
# isn't seen as a crash; a real crash never routes through here.
self._stop_mtp_crash_watchdog()
self._reset_effective_parallel_slots()
if self._process is None:
return
try:
@ -8379,7 +8643,7 @@ class LlamaCppBackend:
):
"""Open one streaming POST and let cancel interrupt prefill or reads."""
if cancel_event is not None and cancel_event.is_set():
raise GeneratorExit
raise _LlamaStreamCancelled
_cancel_closed = threading.Event()
_response_ref: list = [None]
@ -8424,13 +8688,13 @@ class LlamaCppBackend:
) as response:
_response_ref[0] = response
if cancel_event is not None and cancel_event.is_set():
raise GeneratorExit
raise _LlamaStreamCancelled
yield response
return
except (httpx.RequestError, RuntimeError):
# Response was closed by the cancel watcher
if cancel_event is not None and cancel_event.is_set():
raise GeneratorExit
raise _LlamaStreamCancelled
raise
finally:
_cancel_closed.set()
@ -8633,6 +8897,8 @@ class LlamaCppBackend:
"finish_reason": _metadata_finish_reason,
}
except _LlamaStreamCancelled:
return
except httpx.ConnectError as e:
# Server already down. If this was an MTP+tensor crash, recover by
# reloading without MTP (scheduled in the background) and fail this
@ -9757,6 +10023,8 @@ class LlamaCppBackend:
break
continue
except _LlamaStreamCancelled:
return
except httpx.ConnectError:
# Mark unresolved provisional cards as failed before raising.
for _pid, _pname in provisional_started_tool_calls.items():
@ -9939,6 +10207,8 @@ class LlamaCppBackend:
if _meta is not None:
yield _meta
except _LlamaStreamCancelled:
return
except httpx.ConnectError:
raise RuntimeError("Lost connection to llama-server")
except Exception as e:

View file

@ -102,16 +102,17 @@ def preferred_mmproj_sibling(siblings: Sequence) -> Optional[object]:
def preferred_mtp_sibling(siblings: Sequence) -> Optional[object]:
"""The separate MTP drafter to fetch with every variant: the repo-root
``mtp-*.gguf`` copy unsloth ships for llama.cpp ``-hf`` auto-discovery
(Gemma 4). Same pick as the loader's drafter resolution (``mtp-`` basename
prefix, first in sort order) so download and load resolve the same file;
the higher-precision ``MTP/`` subdir copies are for explicit selection and
are not auto-fetched. None for repos with the head baked into the main
GGUF (Qwen)."""
(Gemma 4). Same pick as the loader's drafter resolution (root-level
``mtp-`` prefix, first in sort order) so download and load resolve the same
file; the higher-precision ``MTP/`` subdir copies are for explicit
selection and are not auto-fetched. None for repos with the head baked into
the main GGUF (Qwen)."""
# Root-level only: the MTP/ subdir copies now share the mtp- prefix too.
candidates = sorted(
(
s
for s in siblings
if (name := _gguf_rfilename(s)) and name.lower().rsplit("/", 1)[-1].startswith("mtp-")
if (name := _gguf_rfilename(s)) and "/" not in name and name.lower().startswith("mtp-")
),
key = lambda s: getattr(s, "rfilename"),
)

View file

@ -1533,12 +1533,41 @@ class AnthropicToolResultBlock(BaseModel):
tool_use_id: str
content: Union[str, list] = ""
@field_validator("content", mode = "before")
@classmethod
def _coerce_null_content(cls, v):
# Some clients send null content for an empty tool result; the str|list
# union would 400 on it, so treat null as "".
return "" if v is None else v
# Block types the converter translates explicitly. Anything else (thinking /
# redacted_thinking, a provider block a resumed session replays, or a future type)
# is accepted as an unknown block and dropped by the converter, rather than 400-ing
# the whole request on strict validation.
_KNOWN_ANTHROPIC_BLOCK_TYPES = frozenset({"text", "image", "tool_use", "tool_result"})
class AnthropicUnknownBlock(BaseModel):
type: str
model_config = {"extra": "allow"}
@field_validator("type")
@classmethod
def _only_unknown_types(cls, v):
# Known types parse as their typed models above (so a malformed known block
# still fails cleanly); this fallback only catches the rest.
if v in _KNOWN_ANTHROPIC_BLOCK_TYPES:
raise ValueError("known block type handled by its typed model")
return v
AnthropicContentBlock = Union[
AnthropicTextBlock,
AnthropicImageBlock,
AnthropicToolUseBlock,
AnthropicToolResultBlock,
AnthropicUnknownBlock,
]
@ -1583,6 +1612,40 @@ class AnthropicMessage(BaseModel):
role: Literal["user", "assistant"]
content: Union[str, list[AnthropicContentBlock]]
@model_validator(mode = "before")
@classmethod
def _normalize_content(cls, data):
# Role-aware leniency that never silently drops real user input:
# - assistant: a resumed tool-only turn's null content -> "" (str|list would
# 400 on null; "" keeps the converter's `for block in content` safe).
# Unknown blocks (thinking / future types) validate via
# AnthropicUnknownBlock and are dropped by the converter.
# - user: keep strict. Null user content stays None so str|list rejects it
# (400) rather than forwarding an empty prompt; and reject block types the
# converter cannot translate, since it silently skips unknown user blocks
# -- a user turn made only of them would validate yet send no content
# (silent data loss).
if not isinstance(data, dict):
return data
content = data.get("content")
if data.get("role") == "assistant":
# Coerce only an explicit null (resumed tool-only turn). A missing
# content key stays malformed so the required-field check still 400s.
if "content" in data and content is None:
return {**data, "content": ""}
return data
if isinstance(content, list):
for block in content:
btype = (
block.get("type") if isinstance(block, dict) else getattr(block, "type", None)
)
# Guard the value: a non-string type is unsupported too, and a
# membership test on an unhashable value would raise TypeError
# (escaping as a 500 instead of a clean 400).
if not isinstance(btype, str) or btype not in _KNOWN_ANTHROPIC_BLOCK_TYPES:
raise ValueError(f"unsupported content block type {btype!r} in a user message")
return data
class AnthropicTool(BaseModel):
# Client tools have input_schema; server tools may only have type/name.

View file

@ -10,6 +10,7 @@ import binascii
import json
import os
import re
import shutil
from itertools import islice
from pathlib import Path
from typing import Any
@ -59,6 +60,9 @@ UNSTRUCTURED_ALLOWED_EXTS = {".pdf", ".docx", ".txt", ".md"}
SEED_UPLOAD_DIR = seed_uploads_root()
UNSTRUCTURED_UPLOAD_ROOT = unstructured_uploads_root()
_SAFE_ID_RE = re.compile(r"^[a-zA-Z0-9_-]+$")
# Frontend-generated upload namespace (UUID4 hex). Legacy node ids (n1, ...)
# never match: those directories can be shared by several recipes.
_UPLOAD_UID_RE = re.compile(r"^[0-9a-f]{32}$")
def _validate_safe_id(value: str, label: str) -> str:
@ -580,6 +584,39 @@ async def remove_unstructured_file(block_id: str, file_id: str):
return {"status": "ok"}
@router.delete("/seed/unstructured-block/{block_id}")
async def remove_unstructured_block(block_id: str):
"""Delete a block's upload directory; files on disk still count toward its quota.
Only uid-namespaced directories may be bulk-deleted: they have exactly one
owning block. Legacy node-id directories (n1, ...) can be shared by other
recipes, so they are managed file-by-file instead.
"""
_validate_safe_id(block_id, "block_id")
if not _UPLOAD_UID_RE.match(block_id):
raise HTTPException(400, "Invalid block_id: only uid-namespaced blocks can be deleted")
block_dir = (UNSTRUCTURED_UPLOAD_ROOT / block_id).resolve()
if not block_dir.is_relative_to(UNSTRUCTURED_UPLOAD_ROOT.resolve()):
raise HTTPException(400, "Invalid block_id: outside upload root")
if not block_dir.exists():
return {"status": "ok", "deleted": False}
try:
shutil.rmtree(block_dir)
except OSError as exc:
raise log_and_http_error(
exc,
500,
"failed to delete uploaded files",
event = "data_recipe.seed.unstructured_block_delete_failed",
log = logger,
) from exc
if block_dir.exists():
raise HTTPException(500, "failed to delete uploaded files")
return {"status": "ok", "deleted": True}
@router.post("/seed/inspect-upload", response_model = SeedInspectResponse)
def inspect_seed_upload(payload: SeedInspectUploadRequest) -> SeedInspectResponse:
if payload.file_ids is not None:

File diff suppressed because it is too large Load diff

View file

@ -1770,3 +1770,187 @@ class TestAnthropicMessagesToolRouting:
_drive(anthropic_messages(payload, request = None, current_subject = "t"))
assert backend.calls[0][0] == "plain"
def test_resumed_session_thinking_and_null_content_do_not_400():
# A resumed session replays assistant turns with `thinking` (and sometimes null)
# content. Those must be accepted (thinking dropped by the converter), not 400ed.
from pydantic import ValidationError
req = AnthropicMessagesRequest(
model = "x",
max_tokens = 16,
messages = [
{"role": "user", "content": "hi"},
{
"role": "assistant",
"content": [
{"type": "thinking", "thinking": "secret reasoning", "signature": "s"},
{"type": "text", "text": "the answer"},
{"type": "tool_use", "id": "t1", "name": "f", "input": {}},
],
},
{"role": "assistant", "content": None}, # tool-only turn serialized as null
],
)
# Known blocks still parse as their typed models; only the unknown one is loose.
assert type(req.messages[1].content[0]).__name__ == "AnthropicUnknownBlock"
assert type(req.messages[1].content[1]).__name__ == "AnthropicTextBlock"
assert req.messages[2].content == "" # null coerced
openai = anthropic_messages_to_openai([m.model_dump() for m in req.messages])
assistant = next(m for m in openai if m["role"] == "assistant" and m.get("content"))
assert assistant["content"] == "the answer"
assert "secret reasoning" not in json.dumps(openai) # thinking never forwarded
# A malformed KNOWN block still fails cleanly instead of being swallowed.
with pytest.raises(ValidationError):
AnthropicMessagesRequest(
model = "x",
max_tokens = 16,
messages = [{"role": "assistant", "content": [{"type": "tool_use", "name": "f"}]}],
)
def test_user_null_content_rejected():
# The null->"" leniency is assistant-only; a null user content must be rejected
# at the boundary, not coerced into an empty prompt and forwarded to the model.
from pydantic import ValidationError
with pytest.raises(ValidationError):
AnthropicMessagesRequest(
model = "x",
max_tokens = 16,
messages = [{"role": "user", "content": None}],
)
def test_user_unknown_block_rejected_not_silently_dropped():
# The converter skips user blocks it cannot translate, so a user turn whose only
# block is unknown would validate yet forward no content. Reject at the boundary
# to avoid that silent data loss (the assistant fallback is unaffected).
from pydantic import ValidationError
with pytest.raises(ValidationError):
AnthropicMessagesRequest(
model = "x",
max_tokens = 16,
messages = [
{"role": "user", "content": [{"type": "document", "source": {}}]},
],
)
def test_user_translatable_blocks_still_accepted():
# text / image / tool_result are translatable, so a real user message built from
# them must still pass; the unknown-block guard only trips on other types.
req = AnthropicMessagesRequest(
model = "x",
max_tokens = 16,
messages = [
{
"role": "user",
"content": [
{"type": "text", "text": "What is this?"},
{
"type": "image",
"source": {"type": "base64", "media_type": "image/png", "data": "AA"},
},
{"type": "tool_result", "tool_use_id": "t1", "content": "ok"},
],
}
],
)
assert [type(b).__name__ for b in req.messages[0].content] == [
"AnthropicTextBlock",
"AnthropicImageBlock",
"AnthropicToolResultBlock",
]
openai = anthropic_messages_to_openai([m.model_dump() for m in req.messages])
assert any(m["role"] == "tool" and m["tool_call_id"] == "t1" for m in openai)
def test_user_malformed_known_block_still_rejected():
# The guard only allow-lists a user block's *type*; the union still validates its
# shape, so a known-but-malformed block (tool_result without tool_use_id) fails.
from pydantic import ValidationError
with pytest.raises(ValidationError):
AnthropicMessagesRequest(
model = "x",
max_tokens = 16,
messages = [
{"role": "user", "content": [{"type": "tool_result", "content": "x"}]},
],
)
def test_user_content_block_non_string_type_rejected_cleanly():
# A user block whose `type` is a non-string (unhashable list / dict, or a stray
# int) must fail as a clean validation error, not raise TypeError from the
# frozenset membership test and escape as a 500.
from pydantic import ValidationError
for bad_type in ([], {}, 5):
with pytest.raises(ValidationError):
AnthropicMessagesRequest(
model = "x",
max_tokens = 16,
messages = [{"role": "user", "content": [{"type": bad_type}]}],
)
def test_assistant_missing_content_key_still_rejected():
# The null -> "" leniency is only for an EXPLICIT null. An assistant message that
# omits content entirely stays malformed and must fail required-field validation.
from pydantic import ValidationError
with pytest.raises(ValidationError):
AnthropicMessagesRequest(
model = "x",
max_tokens = 16,
messages = [{"role": "assistant"}],
)
# An explicit null is still accepted and coerced (regression guard).
req = AnthropicMessagesRequest(
model = "x",
max_tokens = 16,
messages = [
{"role": "user", "content": "hi"},
{"role": "assistant", "content": None},
],
)
assert req.messages[1].content == ""
def test_resumed_null_assistant_between_users_coalesced_on_messages_route(monkeypatch):
# user -> assistant(null) -> user is now accepted: the null assistant turn coerces
# to "" and is dropped. The route must then coalesce the two remaining user turns
# so a strict GGUF chat template does not 400 on non-alternating roles.
backend = _mock_backend(monkeypatch, context_length = 2048)
class _Req:
state = SimpleNamespace()
url = SimpleNamespace(path = "/v1/messages")
method = "POST"
async def is_disconnected(self):
return False
payload = AnthropicMessagesRequest(
model = "x",
max_tokens = 16,
messages = [
{"role": "user", "content": "first question"},
{"role": "assistant", "content": None},
{"role": "user", "content": "please continue"},
],
)
response = _drive(anthropic_messages(payload, request = _Req(), current_subject = "t"))
assert response.status_code == 200
[(_path, kwargs)] = backend.calls
user_turns = [m for m in kwargs["messages"] if m.get("role") == "user"]
assert len(user_turns) == 1 # the two user turns were merged, not left adjacent
merged = user_turns[0]["content"]
if isinstance(merged, list):
merged = " ".join(p.get("text", "") for p in merged if isinstance(p, dict))
assert "first question" in merged and "please continue" in merged

View file

@ -124,3 +124,100 @@ def test_unstructured_upload_import_errors_stay_generic(monkeypatch, tmp_path, e
assert result.status == "error"
assert result.error == "Text extraction failed."
assert _block_files(seed_route) == []
_TEST_UPLOAD_UID = "0f" * 16
def test_remove_unstructured_block_deletes_directory(monkeypatch, tmp_path):
seed_route = _load_seed_route(monkeypatch, tmp_path)
_run_upload(seed_route, "notes.txt", b"hello", block_id = _TEST_UPLOAD_UID)
assert _block_files(seed_route, _TEST_UPLOAD_UID) != []
result = asyncio.run(seed_route.remove_unstructured_block(_TEST_UPLOAD_UID))
assert result == {"status": "ok", "deleted": True}
assert not (seed_route.UNSTRUCTURED_UPLOAD_ROOT / _TEST_UPLOAD_UID).exists()
def test_remove_unstructured_block_missing_directory_is_ok(monkeypatch, tmp_path):
seed_route = _load_seed_route(monkeypatch, tmp_path)
result = asyncio.run(seed_route.remove_unstructured_block(_TEST_UPLOAD_UID))
assert result == {"status": "ok", "deleted": False}
def test_remove_unstructured_block_rejects_unsafe_ids(monkeypatch, tmp_path):
seed_route = _load_seed_route(monkeypatch, tmp_path)
with pytest.raises(seed_route.HTTPException) as exc:
asyncio.run(seed_route.remove_unstructured_block("../escape"))
assert exc.value.status_code == 400
def test_remove_unstructured_block_rejects_legacy_node_ids(monkeypatch, tmp_path):
seed_route = _load_seed_route(monkeypatch, tmp_path)
_run_upload(seed_route, "notes.txt", b"hello", block_id = "n1")
assert _block_files(seed_route, "n1") != []
with pytest.raises(seed_route.HTTPException) as exc:
asyncio.run(seed_route.remove_unstructured_block("n1"))
assert exc.value.status_code == 400
assert _block_files(seed_route, "n1") != []
def test_remove_unstructured_block_rejects_symlink_escape(monkeypatch, tmp_path):
seed_route = _load_seed_route(monkeypatch, tmp_path)
outside = tmp_path / "outside"
outside.mkdir()
(outside / "victim.txt").write_text("keep me")
root = seed_route.UNSTRUCTURED_UPLOAD_ROOT
root.mkdir(parents = True)
(root / _TEST_UPLOAD_UID).symlink_to(outside)
with pytest.raises(seed_route.HTTPException) as exc:
asyncio.run(seed_route.remove_unstructured_block(_TEST_UPLOAD_UID))
assert exc.value.status_code == 400
assert (outside / "victim.txt").exists()
def test_remove_unstructured_block_fails_if_directory_remains(monkeypatch, tmp_path):
seed_route = _load_seed_route(monkeypatch, tmp_path)
root = seed_route.UNSTRUCTURED_UPLOAD_ROOT
block_dir = root / _TEST_UPLOAD_UID
block_dir.mkdir(parents = True)
(block_dir / "victim.txt").write_text("keep me")
calls = []
def noop_rmtree(path, *args, **kwargs):
calls.append((path, args, kwargs))
monkeypatch.setattr(seed_route.shutil, "rmtree", noop_rmtree)
with pytest.raises(seed_route.HTTPException) as exc:
asyncio.run(seed_route.remove_unstructured_block(_TEST_UPLOAD_UID))
assert calls
assert exc.value.status_code == 500
assert block_dir.exists()
def test_total_upload_quota_is_scoped_per_block(monkeypatch, tmp_path):
seed_route = _load_seed_route(monkeypatch, tmp_path)
monkeypatch.setattr(seed_route, "UNSTRUCTURED_RECIPE_UPLOAD_TOTAL_MAX_BYTES", 10)
first = _run_upload(seed_route, "a.txt", b"123456789")
assert first.status == "ok"
with pytest.raises(seed_route.HTTPException) as exc:
_run_upload(seed_route, "b.txt", b"123")
assert exc.value.status_code == 413
# Another block starts with its own untouched budget.
other = _run_upload(seed_route, "c.txt", b"123", block_id = "other")
assert other.status == "ok"

View file

@ -112,7 +112,7 @@ def test_route_llama_streaming_async_clients_disable_proxy_env():
continue
calls.append(node)
assert len(calls) == 4
assert len(calls) == 5
for call in calls:
assert any(
kw.arg == "trust_env" and isinstance(kw.value, ast.Constant) and kw.value.value is False

View file

@ -55,6 +55,27 @@ def _host(**kw):
return ilp.HostInfo(**base)
def test_force_cpu_clears_all_gpu_attributes_including_intel():
# --cpu-fallback is the "select the CPU prebuilt even when a GPU is present"
# escape hatch. It must drop EVERY GPU attribute, including has_intel_gpu, or
# the planner still prepends the Vulkan asset on an Intel-GPU host.
host = _host(
is_linux = True,
is_x86_64 = True,
has_usable_nvidia = True,
has_physical_nvidia = True,
has_rocm = True,
rocm_gfx_target = "gfx1100",
has_intel_gpu = True,
)
forced = ilp._apply_host_overrides(host, force_cpu = True)
assert forced.has_usable_nvidia is False
assert forced.has_physical_nvidia is False
assert forced.has_rocm is False
assert forced.rocm_gfx_target is None
assert forced.has_intel_gpu is False
def test_macos_upstream_pin_only_for_explicit_pre26_upstream():
pre26 = _host(
system = "Darwin",
@ -313,3 +334,386 @@ def test_sm103_host_drops_cuda128_windows_build():
)
kept_b200 = ilp._drop_blackwell_incapable_windows_cuda(b200, [cuda128, cuda129])
assert [a.name for a in kept_b200] == [cuda128.name, cuda129.name]
def _upstream_release(tag, asset_names):
return {
"tag_name": tag,
"assets": [
{"name": n, "browser_download_url": f"https://example/{n}"} for n in asset_names
],
}
def test_direct_upstream_arm64_intel_prefers_vulkan():
# Auto-detected Intel GPU on Linux arm64 -> Vulkan prebuilt first, CPU
# second (mirrors the x86_64 branch; ggml-org ships the arm64 Vulkan asset).
host = _host(is_linux = True, is_arm64 = True, machine = "aarch64", has_intel_gpu = True)
rel = _upstream_release(
"b9925",
["llama-b9925-bin-ubuntu-vulkan-arm64.tar.gz", "llama-b9925-bin-ubuntu-arm64.tar.gz"],
)
plan = ilp.direct_upstream_release_plan(rel, host, UPSTREAM, "latest")
kinds = [a.install_kind for a in plan.attempts]
assert kinds[0] == "linux-vulkan", kinds
assert "linux-arm64" in kinds
assert plan.attempts[0].name == "llama-b9925-bin-ubuntu-vulkan-arm64.tar.gz"
def test_direct_upstream_intel_with_hidden_nvidia_is_cpu_only():
# A host with a physical NVIDIA hidden via CUDA_VISIBLE_DEVICES (physical
# True, usable False) + an Intel iGPU must NOT get the Vulkan archive even
# when planning directly against upstream: Vulkan ignores CUDA_VISIBLE_DEVICES
# and could grab the reserved card. It falls through to the CPU asset.
host = _host(
is_linux = True,
is_x86_64 = True,
has_intel_gpu = True,
has_physical_nvidia = True,
has_usable_nvidia = False,
)
rel = _upstream_release(
"b9925",
["llama-b9925-bin-ubuntu-vulkan-x64.tar.gz", "llama-b9925-bin-ubuntu-x64.tar.gz"],
)
plan = ilp.direct_upstream_release_plan(rel, host, UPSTREAM, "latest")
assert [a.install_kind for a in plan.attempts] == ["linux-cpu"]
def test_direct_upstream_arm64_without_intel_is_cpu_only():
host = _host(is_linux = True, is_arm64 = True, machine = "aarch64")
rel = _upstream_release(
"b9925",
["llama-b9925-bin-ubuntu-vulkan-arm64.tar.gz", "llama-b9925-bin-ubuntu-arm64.tar.gz"],
)
plan = ilp.direct_upstream_release_plan(rel, host, UPSTREAM, "latest")
assert [a.install_kind for a in plan.attempts] == ["linux-arm64"]
def test_direct_upstream_x86_intel_prefers_vulkan():
host = _host(is_linux = True, is_x86_64 = True, has_intel_gpu = True)
rel = _upstream_release(
"b9925",
["llama-b9925-bin-ubuntu-vulkan-x64.tar.gz", "llama-b9925-bin-ubuntu-x64.tar.gz"],
)
plan = ilp.direct_upstream_release_plan(rel, host, UPSTREAM, "latest")
kinds = [a.install_kind for a in plan.attempts]
assert kinds[0] == "linux-vulkan", kinds
assert "linux-cpu" in kinds
def test_linux_vulkan_health_glob_matches_bare_cpu_lib():
# The widened glob must cover both arch-suffixed (x64) and bare (arm64) CPU
# libs so a valid Vulkan install is not re-flagged unhealthy every check.
choice = ilp.AssetChoice(
repo = UPSTREAM,
tag = "b9925",
name = "llama-b9925-bin-ubuntu-vulkan-arm64.tar.gz",
url = "https://example/x",
source_label = "upstream",
install_kind = "linux-vulkan",
)
groups = ilp.runtime_payload_health_groups(choice)
assert ["libggml-cpu*.so*"] in groups
assert ["libggml-cpu-*.so*"] not in groups
def test_route_to_vulkan_prebuilt_auto_intel_goes_upstream_and_drops_fork_pin():
# Routing fork -> upstream also drops the fork release pin, which is in a
# different tag namespace and would make the upstream resolver miss.
host = _host(is_linux = True, is_x86_64 = True, has_intel_gpu = True)
routed, repo, tag = ilp._route_to_vulkan_prebuilt(host, FORK, "b9596-mix-abc", force_cpu = False)
assert repo == UPSTREAM
assert tag == ""
assert routed.has_intel_gpu is True
def test_route_to_vulkan_prebuilt_preserves_explicit_upstream_pin():
# A pin set WITH an explicit upstream repo is already on upstream -> kept.
host = _host(is_linux = True, is_x86_64 = True, has_intel_gpu = True)
_routed, repo, tag = ilp._route_to_vulkan_prebuilt(host, UPSTREAM, "b9596", force_cpu = False)
assert repo == UPSTREAM
assert tag == "b9596"
def test_route_to_vulkan_prebuilt_cpu_fallback_wins():
# --cpu-fallback suppresses Vulkan routing even for an Intel host.
host = _host(is_linux = True, is_x86_64 = True, has_intel_gpu = True)
routed, repo, tag = ilp._route_to_vulkan_prebuilt(host, FORK, "b9596-mix-abc", force_cpu = True)
assert repo == FORK
assert tag == "b9596-mix-abc"
assert routed is host
def test_route_to_vulkan_prebuilt_hidden_nvidia_not_rerouted():
# A mixed NVIDIA+Intel host that hid NVIDIA (CUDA_VISIBLE_DEVICES=""/-1):
# physical NVIDIA present but not usable. Must NOT auto-route to Vulkan, or
# Vulkan (which ignores CUDA_VISIBLE_DEVICES) could grab the reserved GPU.
host = _host(
is_linux = True,
is_x86_64 = True,
has_intel_gpu = True,
has_physical_nvidia = True,
has_usable_nvidia = False,
)
_routed, repo, _tag = ilp._route_to_vulkan_prebuilt(host, FORK, "", force_cpu = False)
assert repo == FORK
def test_route_to_vulkan_prebuilt_rocm_host_not_rerouted():
# An Intel iGPU alongside a usable ROCm GPU stays on its ROCm/fork path.
host = _host(is_linux = True, is_x86_64 = True, has_intel_gpu = True, has_rocm = True)
_routed, repo, _tag = ilp._route_to_vulkan_prebuilt(host, FORK, "", force_cpu = False)
assert repo == FORK
def test_route_to_vulkan_prebuilt_non_intel_unchanged():
host = _host(is_linux = True, is_x86_64 = True)
routed, repo, _tag = ilp._route_to_vulkan_prebuilt(host, FORK, "", force_cpu = False)
assert repo == FORK
assert routed is host
def test_resolve_prebuilt_intel_host_routes_to_upstream(monkeypatch, capsys):
# The --resolve-prebuilt probe must agree with the install path: an
# auto-detected Intel host resolves against upstream (Vulkan), not the fork.
monkeypatch.setattr(
ilp, "detect_host", lambda: _host(is_linux = True, is_x86_64 = True, has_intel_gpu = True)
)
seen, out = _run_resolve_capture_host(monkeypatch, capsys)
assert seen["repo"] == UPSTREAM
assert out["repo"] == UPSTREAM
# ---------------------------------------------------------------------------
# windows_intel_gpu_in_registry: the in-process Windows Intel probe. A fake
# winreg module stands in for the real registry so the walk runs anywhere.
# ---------------------------------------------------------------------------
class _FakeRegKey:
def __init__(
self,
subkeys = None,
values = None,
denied = False,
):
self.subkeys = subkeys or {}
self.values = values or {}
self.denied = denied
def __enter__(self):
return self
def __exit__(self, *exc):
return False
class _FakeWinreg:
HKEY_LOCAL_MACHINE = object()
def __init__(self, root_key):
self._root_key = root_key
def OpenKey(self, parent, name):
if parent is self.HKEY_LOCAL_MACHINE:
# Pin the production constant: a typo'd class GUID must fail here,
# not silently return the fake tree.
if name != ilp._WINDOWS_DISPLAY_CLASS_KEY:
raise FileNotFoundError(name)
if self._root_key is None:
raise FileNotFoundError(name)
return self._root_key
key = parent.subkeys.get(name)
if key is None:
# Real winreg raises OSError, never KeyError, for a missing key.
raise FileNotFoundError(name)
if key.denied:
raise PermissionError(name)
return key
def QueryInfoKey(self, key):
return (len(key.subkeys), len(key.values), 0)
def EnumKey(self, key, index):
return list(key.subkeys)[index]
def QueryValueEx(self, key, value_name):
if value_name not in key.values:
raise FileNotFoundError(value_name)
return (key.values[value_name], 1)
def _probe_with_display_class(monkeypatch, adapters):
# The helper lazily does `import winreg`; plant the fake in sys.modules the
# same way unsloth_cli/tests/test_start.py fakes it for _refresh_windows_path.
monkeypatch.setitem(sys.modules, "winreg", _FakeWinreg(_FakeRegKey(subkeys = adapters)))
return ilp.windows_intel_gpu_in_registry()
def test_windows_intel_registry_matches_vendor_id(monkeypatch):
assert (
_probe_with_display_class(
monkeypatch,
{
"0000": _FakeRegKey(
values = {
"MatchingDeviceId": r"PCI\VEN_8086&DEV_56A0&SUBSYS_12345678",
"DriverDesc": "Intel(R) Arc(TM) A770 Graphics",
}
),
},
)
is True
)
def test_windows_intel_registry_matches_driver_desc_without_device_id(monkeypatch):
assert (
_probe_with_display_class(
monkeypatch,
{
"0000": _FakeRegKey(values = {"DriverDesc": "Intel(R) UHD Graphics 630"}),
},
)
is True
)
def test_windows_intel_registry_ignores_non_intel_adapters(monkeypatch):
assert (
_probe_with_display_class(
monkeypatch,
{
"0000": _FakeRegKey(
values = {
"MatchingDeviceId": r"PCI\VEN_10DE&DEV_2684",
"DriverDesc": "NVIDIA GeForce RTX 4090",
}
),
"0001": _FakeRegKey(
values = {
"MatchingDeviceId": r"PCI\VEN_1002&DEV_744C",
"DriverDesc": "AMD Radeon RX 7900 XTX",
}
),
},
)
is False
)
def test_windows_intel_registry_skips_restricted_properties_subkey(monkeypatch):
# The real class key carries an ACL-restricted "Properties" subkey and can
# deny access to individual adapter keys; neither may abort the walk.
assert (
_probe_with_display_class(
monkeypatch,
{
"Properties": _FakeRegKey(denied = True),
"0000": _FakeRegKey(denied = True),
"0001": _FakeRegKey(
values = {
"MatchingDeviceId": r"PCI\VEN_8086&DEV_56A0",
}
),
},
)
is True
)
def test_windows_intel_registry_missing_class_key_is_false(monkeypatch):
monkeypatch.setitem(sys.modules, "winreg", _FakeWinreg(None))
assert ilp.windows_intel_gpu_in_registry() is False
def _detect_windows_host(
monkeypatch,
winreg_fake,
powershell_stdout = "",
):
"""Drive the real detect_host() as a GPU-less Windows host with a fake
registry, recording every run_capture invocation. Pins the wiring the
unit tests above cannot see: registry-first, CIM only on a registry miss."""
monkeypatch.setitem(sys.modules, "winreg", winreg_fake)
monkeypatch.setattr(ilp.platform, "system", lambda: "Windows")
monkeypatch.setattr(ilp.platform, "machine", lambda: "AMD64")
for _env in (
"CUDA_VISIBLE_DEVICES",
"HIP_VISIBLE_DEVICES",
"ROCR_VISIBLE_DEVICES",
"HIP_PATH",
"ROCM_PATH",
):
monkeypatch.delenv(_env, raising = False)
monkeypatch.setattr(
ilp.shutil,
"which",
lambda name: "powershell" if name in ("powershell", "pwsh") else None,
)
captured = []
def _fake_run_capture(command, **kwargs):
captured.append(command[0])
if command[0] == "powershell":
return SimpleNamespace(returncode = 0, stdout = powershell_stdout, stderr = "")
return SimpleNamespace(returncode = 1, stdout = "", stderr = "")
monkeypatch.setattr(ilp, "run_capture", _fake_run_capture)
return ilp.detect_host(), captured
def test_detect_host_registry_intel_skips_cim_probe(monkeypatch):
winreg = _FakeWinreg(
_FakeRegKey(
subkeys = {
"0000": _FakeRegKey(values = {"MatchingDeviceId": r"PCI\VEN_8086&DEV_56A0"}),
}
)
)
host, captured = _detect_windows_host(monkeypatch, winreg)
assert host.has_intel_gpu is True
assert "powershell" not in captured
def test_detect_host_cim_fallback_fires_on_registry_miss(monkeypatch):
winreg = _FakeWinreg(
_FakeRegKey(
subkeys = {
"0000": _FakeRegKey(values = {"MatchingDeviceId": r"PCI\VEN_10DE&DEV_2684"}),
}
)
)
host, captured = _detect_windows_host(
monkeypatch, winreg, powershell_stdout = "Intel(R) Arc(TM) A770 Graphics"
)
assert host.has_intel_gpu is True
assert "powershell" in captured
def test_windows_intel_registry_unexpected_error_is_false(monkeypatch):
# The probe is advisory: even a non-OSError bug in the walk must return
# False (deferring to the CIM fallback), never crash detect_host.
class _ExplodingWinreg:
HKEY_LOCAL_MACHINE = object()
def OpenKey(self, parent, name):
raise TypeError(name)
monkeypatch.setitem(sys.modules, "winreg", _ExplodingWinreg())
assert ilp.windows_intel_gpu_in_registry() is False
def test_detect_host_cim_rescues_exploding_registry(monkeypatch):
class _ExplodingWinreg:
HKEY_LOCAL_MACHINE = object()
def OpenKey(self, parent, name):
raise TypeError(name)
host, captured = _detect_windows_host(
monkeypatch, _ExplodingWinreg(), powershell_stdout = "Intel(R) Arc(TM) A770 Graphics"
)
assert host.has_intel_gpu is True
assert "powershell" in captured

View file

@ -0,0 +1,320 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
import asyncio
import os
import sys
import threading
import pytest
_backend = os.path.join(os.path.dirname(__file__), "..")
sys.path.insert(0, _backend)
from core.inference import llama_admission
from core.inference.llama_admission import (
ADMISSION_CONTROL_ENV,
ADMISSION_KEEPALIVE_INTERVAL_ENV,
ADMISSION_MAX_QUEUE_ENV,
ADMISSION_QUEUE_TIMEOUT_ENV,
DEFAULT_ADMISSION_KEEPALIVE_INTERVAL_S,
DEFAULT_ADMISSION_MAX_QUEUE,
DEFAULT_ADMISSION_QUEUE_TIMEOUT_S,
LlamaAdmissionConfig,
LlamaAdmissionQueueFull,
get_llama_admission_queue,
llama_admission_config_from_env,
reset_llama_admission_queues,
)
@pytest.fixture(autouse = True)
def _reset_queues():
reset_llama_admission_queues()
yield
reset_llama_admission_queues()
def test_admission_config_defaults(monkeypatch):
for name in (
ADMISSION_CONTROL_ENV,
ADMISSION_QUEUE_TIMEOUT_ENV,
ADMISSION_KEEPALIVE_INTERVAL_ENV,
ADMISSION_MAX_QUEUE_ENV,
):
monkeypatch.delenv(name, raising = False)
config = llama_admission_config_from_env()
assert config.enabled is True
assert config.queue_timeout_s == DEFAULT_ADMISSION_QUEUE_TIMEOUT_S
assert config.keepalive_interval_s == DEFAULT_ADMISSION_KEEPALIVE_INTERVAL_S
assert config.max_queue == DEFAULT_ADMISSION_MAX_QUEUE
def test_admission_config_env_overrides(monkeypatch):
monkeypatch.setenv(ADMISSION_CONTROL_ENV, "off")
monkeypatch.setenv(ADMISSION_QUEUE_TIMEOUT_ENV, "0")
monkeypatch.setenv(ADMISSION_KEEPALIVE_INTERVAL_ENV, "0.25")
monkeypatch.setenv(ADMISSION_MAX_QUEUE_ENV, "0")
config = llama_admission_config_from_env()
assert config.enabled is False
assert config.queue_timeout_s is None
assert config.keepalive_interval_s == 0.25
assert config.max_queue is None
def test_admission_config_positive_queue_timeout_env(monkeypatch):
monkeypatch.setenv(ADMISSION_QUEUE_TIMEOUT_ENV, "600")
config = llama_admission_config_from_env()
assert config.queue_timeout_s == 600.0
def test_fifo_capacity_one_grants_next_waiter_on_release():
async def _run():
queue = get_llama_admission_queue("http://llama.test")
config = LlamaAdmissionConfig()
first = queue.reserve(capacity = 1, config = config)
second = queue.reserve(capacity = 1, config = config)
third = queue.reserve(capacity = 1, config = config)
first_lease = first.lease_nowait()
assert first_lease is not None
assert second.lease_nowait() is None
assert third.lease_nowait() is None
assert queue.snapshot().queued == 2
first_lease.release()
second_lease = await second.wait(0.1)
assert second_lease is not None
assert third.lease_nowait() is None
second_lease.release()
third_lease = await third.wait(0.1)
assert third_lease is not None
third_lease.release()
snapshot = queue.snapshot()
assert snapshot.active == 0
assert snapshot.queued == 0
asyncio.run(_run())
def test_queue_full_rejects_excess_waiter():
async def _run():
queue = get_llama_admission_queue("http://llama.test")
config = LlamaAdmissionConfig(max_queue = 1)
first = queue.reserve(capacity = 1, config = config)
queued = queue.reserve(capacity = 1, config = config)
assert first.lease_nowait() is not None
assert queued.lease_nowait() is None
with pytest.raises(LlamaAdmissionQueueFull):
queue.reserve(capacity = 1, config = config)
asyncio.run(_run())
def test_disabled_admission_bypasses_active_slot_limit():
async def _run():
queue = get_llama_admission_queue("http://llama.test")
config = LlamaAdmissionConfig(enabled = False)
first = queue.reserve(capacity = 1, config = config)
second = queue.reserve(capacity = 1, config = config)
assert first.lease_nowait() is not None
assert second.lease_nowait() is not None
assert queue.snapshot().active == 0
assert queue.snapshot().queued == 0
asyncio.run(_run())
def test_cancelling_promoted_waiter_releases_slot():
async def _run():
queue = get_llama_admission_queue("http://llama.test")
config = LlamaAdmissionConfig()
first = queue.reserve(capacity = 1, config = config)
second = queue.reserve(capacity = 1, config = config)
first_lease = first.lease_nowait()
first_lease.release()
await asyncio.sleep(0)
second.cancel()
snapshot = queue.snapshot()
assert snapshot.active == 0
assert snapshot.queued == 0
asyncio.run(_run())
def test_cancelling_promoted_waiter_before_delivery_releases_slot():
async def _run():
queue = get_llama_admission_queue("http://llama.test")
config = LlamaAdmissionConfig()
first = queue.reserve(capacity = 1, config = config)
second = queue.reserve(capacity = 1, config = config)
first_lease = first.lease_nowait()
first_lease.release()
second.cancel()
snapshot = queue.snapshot()
assert snapshot.active == 0
assert snapshot.queued == 0
asyncio.run(_run())
def test_external_waiter_future_cancel_invalidates_reservation():
async def _run():
queue = get_llama_admission_queue("http://llama.test")
config = LlamaAdmissionConfig()
first = queue.reserve(capacity = 1, config = config)
second = queue.reserve(capacity = 1, config = config)
first_lease = first.lease_nowait()
assert first_lease is not None
assert second._waiter is not None
second._waiter.future.cancel()
assert second.lease_nowait() is None
assert second.is_cancelled is True
assert await second.wait(0.01) is None
first_lease.release()
snapshot = queue.snapshot()
assert snapshot.active == 0
assert snapshot.queued == 0
asyncio.run(_run())
def test_wait_returns_none_when_waiter_future_cancelled_during_wait():
async def _run():
queue = get_llama_admission_queue("http://llama.test")
config = LlamaAdmissionConfig()
first = queue.reserve(capacity = 1, config = config)
second = queue.reserve(capacity = 1, config = config)
first_lease = first.lease_nowait()
assert first_lease is not None
assert second._waiter is not None
wait_task = asyncio.create_task(second.wait(1.0))
await asyncio.sleep(0)
second._waiter.future.cancel()
assert await asyncio.wait_for(wait_task, timeout = 0.1) is None
assert second.is_cancelled is True
first_lease.release()
snapshot = queue.snapshot()
assert snapshot.active == 0
assert snapshot.queued == 0
asyncio.run(_run())
def test_capacity_increase_promotes_existing_waiter_fifo():
async def _run():
queue = get_llama_admission_queue("http://llama.test")
config = LlamaAdmissionConfig()
first = queue.reserve(capacity = 1, config = config)
second = queue.reserve(capacity = 1, config = config)
first_lease = first.lease_nowait()
assert first_lease is not None
assert second.lease_nowait() is None
assert queue.snapshot().active == 1
assert queue.snapshot().queued == 1
third = queue.reserve(capacity = 2, config = config)
second_lease = await second.wait(0.1)
assert second_lease is not None
assert third.lease_nowait() is None
snapshot = queue.snapshot()
assert snapshot.capacity == 2
assert snapshot.active == 2
assert snapshot.queued == 1
first_lease.release()
third_lease = await third.wait(0.1)
assert third_lease is not None
second_lease.release()
third_lease.release()
snapshot = queue.snapshot()
assert snapshot.active == 0
assert snapshot.queued == 0
asyncio.run(_run())
def test_lease_release_is_idempotent_under_concurrent_calls():
async def _run():
queue = get_llama_admission_queue("http://llama.test")
config = LlamaAdmissionConfig()
reservation = queue.reserve(capacity = 1, config = config)
lease = reservation.lease_nowait()
assert lease is not None
threads = [threading.Thread(target = lease.release) for _ in range(16)]
for thread in threads:
thread.start()
for thread in threads:
thread.join()
snapshot = queue.snapshot()
assert snapshot.active == 0
assert snapshot.queued == 0
asyncio.run(_run())
def test_new_key_evicts_idle_prior_load_queues():
# Each model load carries a fresh ephemeral port, so a new base_url key must
# not leave the drained queues from earlier loads accumulating forever.
get_llama_admission_queue("http://127.0.0.1:1001")
get_llama_admission_queue("http://127.0.0.1:1002")
assert set(llama_admission._QUEUES) == {"http://127.0.0.1:1002"}
get_llama_admission_queue("http://127.0.0.1:1003")
assert set(llama_admission._QUEUES) == {"http://127.0.0.1:1003"}
def test_new_key_retains_in_flight_prior_load_queue():
config = LlamaAdmissionConfig()
busy = get_llama_admission_queue("http://127.0.0.1:2001")
async def _run():
reservation = busy.reserve(capacity = 1, config = config)
lease = reservation.lease_nowait()
assert lease is not None
# A new load must not drop a queue that still has an in-flight request.
get_llama_admission_queue("http://127.0.0.1:2002")
assert set(llama_admission._QUEUES) == {"http://127.0.0.1:2001", "http://127.0.0.1:2002"}
# Once it drains, the next load reclaims it.
lease.release()
get_llama_admission_queue("http://127.0.0.1:2003")
assert set(llama_admission._QUEUES) == {"http://127.0.0.1:2003"}
asyncio.run(_run())

View file

@ -0,0 +1,53 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
import os
import sys
import pytest
_backend = os.path.join(os.path.dirname(__file__), "..")
sys.path.insert(0, _backend)
from core.inference import llama_cpp as llama_cpp_module
from core.inference.llama_cpp import LlamaCppBackend
@pytest.fixture
def backend(monkeypatch):
monkeypatch.setattr(LlamaCppBackend, "_kill_orphaned_servers", lambda self: 0)
monkeypatch.setattr(llama_cpp_module.atexit, "register", lambda *_args, **_kwargs: None)
return LlamaCppBackend()
def test_effective_parallel_slots_initial_value_is_one(backend):
assert backend.effective_parallel_slots == 1
def test_effective_parallel_slots_commit_uses_final_positive_parallel(backend):
backend._commit_effective_parallel_slots(3)
assert backend.effective_parallel_slots == 3
@pytest.mark.parametrize("value", [None, 0, -2, "not-an-int"])
def test_effective_parallel_slots_commit_invalid_value_falls_back_to_one(backend, value):
backend._commit_effective_parallel_slots(value)
assert backend.effective_parallel_slots == 1
def test_effective_parallel_slots_reset_returns_to_one(backend):
backend._commit_effective_parallel_slots(4)
backend._reset_effective_parallel_slots()
assert backend.effective_parallel_slots == 1
def test_effective_parallel_slots_unload_resets_to_one(backend):
backend._commit_effective_parallel_slots(4)
backend.unload_model()
assert backend.effective_parallel_slots == 1

View file

@ -0,0 +1,81 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
import contextlib
import os
import sys
import threading
import httpx
import pytest
_backend = os.path.join(os.path.dirname(__file__), "..")
sys.path.insert(0, _backend)
from core.inference.llama_cpp import LlamaCppBackend, _LlamaStreamCancelled
def _backend_stub() -> LlamaCppBackend:
backend = LlamaCppBackend.__new__(LlamaCppBackend)
backend._process = object()
backend._healthy = True
backend._port = 48848
backend._effective_context_length = 4096
backend._supports_reasoning = False
backend._reasoning_always_on = False
backend._reasoning_style = "enable_thinking"
backend._supports_preserve_thinking = False
return backend
def test_stream_cancel_uses_internal_exception_not_generator_exit():
class FakeResponse:
status_code = 200
def close(self):
pass
class FakeStream:
def __enter__(self):
return FakeResponse()
def __exit__(self, *_args):
return False
class FakeClient:
def stream(self, *_args, **_kwargs):
return FakeStream()
cancel_event = threading.Event()
with pytest.raises(Exception) as exc_info:
with LlamaCppBackend._stream_with_retry(
FakeClient(),
"http://llama.test/v1/chat/completions",
{},
cancel_event,
):
cancel_event.set()
raise httpx.ReadError("client closed")
assert exc_info.type is _LlamaStreamCancelled
assert not issubclass(exc_info.type, GeneratorExit)
def test_generate_chat_completion_swallows_internal_stream_cancel(monkeypatch):
backend = _backend_stub()
@contextlib.contextmanager
def fake_open_stream(*_args, **_kwargs):
raise _LlamaStreamCancelled
monkeypatch.setattr(backend, "_open_stream", fake_open_stream)
chunks = list(
backend.generate_chat_completion(
[{"role": "user", "content": "hi"}],
cancel_event = threading.Event(),
)
)
assert chunks == []

View file

@ -448,6 +448,48 @@ def test_start_update_happy_path(monkeypatch, tmp_path):
assert popen_kwargs["env"]["UNSLOTH_PROGRESS_PERCENT_STEP"] == "5"
def test_start_update_preserves_vulkan_via_env(monkeypatch, tmp_path):
# A Vulkan install (marker asset carries 'vulkan') must re-assert
# UNSLOTH_FORCE_VULKAN on update, or detect_host on a GPU box re-routes to
# CUDA/ROCm and silently replaces the Vulkan build.
install_dir = tmp_path / "llama.cpp"
binary = _write_install(
install_dir,
"b9493",
repo = "ggml-org/llama.cpp",
asset = "llama-b9493-bin-ubuntu-vulkan-x64.tar.gz",
)
monkeypatch.setattr(upd, "_find_binary", lambda: binary)
monkeypatch.setattr(upd, "_installer_script", lambda: tmp_path / "install_llama_prebuilt.py")
monkeypatch.setattr(freshness, "_fetch_latest_release_tag", lambda repo, timeout = 5.0: "b9518")
def _on_start(cmd):
_write_install(
install_dir,
"b9518",
repo = "ggml-org/llama.cpp",
asset = "llama-b9518-bin-ubuntu-vulkan-x64.tar.gz",
)
popen_kwargs: dict = {}
_patch_installer_popen(
monkeypatch,
lines = ["installed\n"],
on_start = _on_start,
captured_kwargs = popen_kwargs,
)
assert upd.start_update()["started"] is True
deadline = time.time() + 10
while time.time() < deadline:
job = upd.get_update_status()["job"]
if job["state"] in ("success", "error"):
break
time.sleep(0.05)
assert job["state"] == "success", job
assert popen_kwargs["env"]["UNSLOTH_FORCE_VULKAN"] == "1"
def test_start_update_reports_full_release_tag(monkeypatch, tmp_path):
install_dir = tmp_path / "llama.cpp"
binary = _write_install(install_dir, "b9595")

View file

@ -0,0 +1,193 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Vulkan free-VRAM reader regression tests on a synthetic probe output.
Covers the post-probe handling in
``LlamaCppBackend._get_gpu_free_memory_vulkan``:
* integrated GPUs (probe reports is_igpu=1) leave a flat per-device host
margin matching llama.cpp's --fit-target, so context auto-sizing can't
over-commit shared RAM, and report total 0 (shared RAM is not a budget),
* discrete GPUs (is_igpu=0) keep their free untouched and pass their real
total through so the fit can reserve absolute headroom,
* an inherited ``GGML_VK_VISIBLE_DEVICES`` is passed through to ggml unchanged
(ggml applies it), not stripped or filtered in Python -- the probe reports
ggml's compact ordinal, which load_model pins with ``--device Vulkan<i>``.
The ggml Vulkan library is never loaded: subprocess.run is mocked to emit
the tab-separated lines the real ``_vulkan_probe.py`` would print.
"""
from __future__ import annotations
import subprocess
import sys
import types as _types
from pathlib import Path
from unittest import mock
import pytest
_BACKEND_DIR = str(Path(__file__).resolve().parent.parent)
if _BACKEND_DIR not in sys.path:
sys.path.insert(0, _BACKEND_DIR)
import importlib as _importlib # noqa: E402
def _maybe_stub(name: str, builder):
try:
_importlib.import_module(name)
except ImportError:
sys.modules[name] = builder()
def _build_loggers_stub():
m = _types.ModuleType("loggers")
m.get_logger = lambda name: __import__("logging").getLogger(name)
return m
_maybe_stub("loggers", _build_loggers_stub)
_maybe_stub("structlog", lambda: _types.ModuleType("structlog"))
from core.inference import llama_cpp as _llama_mod # noqa: E402
from core.inference.llama_cpp import ( # noqa: E402
LlamaCppBackend,
_llama_lib_dir,
_vulkan_lib_filename,
)
MIB = 1024 * 1024
GIB = 1024 * MIB
def _make_vulkan_install(tmp_path: Path) -> str:
"""A binary whose sibling dir holds the Vulkan ggml lib, so the
reader's ``is_vulkan_backend`` sibling-file check passes."""
bindir = tmp_path / "build" / "bin"
bindir.mkdir(parents = True)
binary = bindir / ("llama-server.exe" if sys.platform == "win32" else "llama-server")
binary.write_bytes(b"stub")
(bindir / _vulkan_lib_filename()).write_bytes(b"stub")
return str(binary)
def _mock_probe(rows: list[str], captured_env: dict | None = None):
"""Patch subprocess.run so the _vulkan_probe.py call returns ``rows``
(already tab-formatted), recording the env it was launched with."""
real_run = subprocess.run
def fake_run(cmd, *args, **kwargs):
if isinstance(cmd, list) and any("_vulkan_probe" in str(c) for c in cmd):
if captured_env is not None:
captured_env.clear()
captured_env.update(kwargs.get("env") or {})
return subprocess.CompletedProcess(
args = cmd, returncode = 0, stdout = "\n".join(rows), stderr = ""
)
return real_run(cmd, *args, **kwargs)
return mock.patch("subprocess.run", side_effect = fake_run)
def _row(
idx: int,
free_bytes: int,
is_igpu: int,
total_bytes: int = 0,
) -> str:
return f"{idx}\t{free_bytes}\t{is_igpu}\t{total_bytes}"
def test_integrated_gpu_leaves_host_margin(tmp_path):
binary = _make_vulkan_install(tmp_path)
# iGPU with 30 GiB free; reserve a flat 1024 MiB (llama.cpp --fit-target).
# total stays 0: shared system RAM is not a VRAM budget for the fit.
rows = [_row(0, 30 * GIB, is_igpu = 1, total_bytes = 32 * GIB)]
with _mock_probe(rows):
gpus = LlamaCppBackend._get_gpu_free_memory_vulkan(binary)
assert gpus == [(0, 30 * 1024 - 1024, 0)], gpus
def test_discrete_gpu_free_is_untouched_and_total_passed_through(tmp_path):
binary = _make_vulkan_install(tmp_path)
# 6 GiB free on a partially occupied 24 GiB card: free is untouched and the
# real total flows through so the fit reserves absolute headroom (CUDA/ROCm
# parity) instead of the looser free*frac budget.
rows = [_row(0, 6 * GIB, is_igpu = 0, total_bytes = 24 * GIB)]
with _mock_probe(rows):
gpus = LlamaCppBackend._get_gpu_free_memory_vulkan(binary)
assert gpus == [(0, 6 * 1024, 24 * 1024)], gpus
def test_large_discrete_gpu_is_untouched(tmp_path):
binary = _make_vulkan_install(tmp_path)
# A 48 GiB discrete card stays untouched regardless of size; only the
# iGPU flag triggers the host margin, never a VRAM/RAM ratio.
rows = [_row(0, 47 * GIB, is_igpu = 0, total_bytes = 48 * GIB)]
with _mock_probe(rows):
gpus = LlamaCppBackend._get_gpu_free_memory_vulkan(binary)
assert gpus == [(0, 47 * 1024, 48 * 1024)], gpus
def test_inherited_visible_devices_mask_is_passed_through_to_probe(tmp_path, monkeypatch):
# The mask is NOT stripped or filtered in Python: ggml parses it in raw
# physical-device space while this probe reports the compact post-filter
# ordinal, so mixing spaces would be wrong. It is passed through unchanged
# so ggml applies it to the same device list the launch will enumerate.
binary = _make_vulkan_install(tmp_path)
monkeypatch.setenv("GGML_VK_VISIBLE_DEVICES", "1")
captured: dict = {}
rows = [_row(0, 23 * GIB, is_igpu = 0, total_bytes = 24 * GIB)]
with _mock_probe(rows, captured_env = captured):
LlamaCppBackend._get_gpu_free_memory_vulkan(binary)
assert captured.get("GGML_VK_VISIBLE_DEVICES") == "1", captured
def test_vulkan_pin_args_uses_device_names_not_env_mask():
# Pin by compact device name via --device (the space the probe reports and
# the registry names), never by writing a compact ordinal into the raw
# GGML_VK_VISIBLE_DEVICES index space.
assert LlamaCppBackend._vulkan_pin_args([0]) == ["--device", "Vulkan0"]
assert LlamaCppBackend._vulkan_pin_args([1, 2]) == ["--device", "Vulkan1,Vulkan2"]
assert LlamaCppBackend._vulkan_pin_args(None) == []
assert LlamaCppBackend._vulkan_pin_args([]) == []
def test_vulkan_only_build_is_detected(tmp_path):
binary = _make_vulkan_install(tmp_path)
assert LlamaCppBackend._is_vulkan_backend(binary) is True
def test_multi_backend_build_is_not_vulkan_only(tmp_path):
# A custom build that ships CUDA (or HIP) alongside Vulkan must NOT be
# treated as Vulkan-only, or its CUDA GPU would be probed/pinned as a Vulkan
# device; defer to the CUDA/HIP path instead.
binary = _make_vulkan_install(tmp_path)
cuda = "ggml-cuda.dll" if sys.platform == "win32" else "libggml-cuda.so"
(_llama_lib_dir(binary) / cuda).write_bytes(b"stub")
assert LlamaCppBackend._is_vulkan_backend(binary) is False
@pytest.mark.skipif(sys.platform == "win32", reason = "shell wrapper fallback is POSIX")
def test_shell_wrapper_entrypoint_resolves_to_real_lib_dir(tmp_path):
# create_exec_entrypoint falls back to a #!/bin/sh wrapper at the install root
# when it cannot symlink; _find_llama_server_binary returns that root entrypoint,
# so _llama_lib_dir must follow the wrapper's exec target to build/bin -- else
# _is_vulkan_backend misses libggml-vulkan.so and the Vulkan probe/pin silently
# never engage on a valid Vulkan install.
import os
binary = _make_vulkan_install(tmp_path) # tmp_path/build/bin/llama-server + vulkan lib
bindir = Path(binary).parent
wrapper = tmp_path / "llama-server"
wrapper.write_text('#!/bin/sh\nexec "$(dirname "$0")/build/bin/llama-server" "$@"\n')
os.chmod(wrapper, 0o755)
assert _llama_lib_dir(str(wrapper)) == bindir
assert LlamaCppBackend._is_vulkan_backend(str(wrapper)) is True
if __name__ == "__main__":
raise SystemExit(pytest.main([__file__, "-v"]))

View file

@ -203,3 +203,87 @@ def test_preheader_send_cleanup_on_disconnect_and_cancel():
asyncio.run(_run(False))
asyncio.run(_run(True))
def test_stream_stall_timeout_callable_re_resolved_each_read():
# The OpenAI passthrough passes a callable so the stall bound can switch to
# the short post-terminal grace mid-stream; it must be re-resolved per read,
# not captured once at generator start.
async def _run():
response = SimpleNamespace(request = SimpleNamespace(extensions = {"timeout": {}}))
values = iter([100.0, 2.0])
seen = []
class _Request:
async def is_disconnected(self):
return False
class _Items:
def __init__(self):
self.count = 0
async def __anext__(self):
self.count += 1
if self.count > 3:
raise StopAsyncIteration
return "data: {}"
async for _ in inf_mod._aiter_llama_stream_items(
_Items(),
cancel_event = threading.Event(),
request = _Request(),
response = response,
first_token_deadline = time.monotonic() + 1,
post_first_item_read_timeout_s = lambda: next(values, 5.0),
):
seen.append(response.request.extensions["timeout"].get("read"))
assert len(seen) == 3
# The callable is resolved right after the first item (arming the
# post-first window) and again before each later read, consuming
# successive values.
assert seen[0] == 100.0
assert 1.0 <= seen[1] <= 2.0
assert 4.0 <= seen[2] <= 5.0
asyncio.run(_run())
def test_stream_stall_timeout_disabled_clears_read_timeout():
# UNSLOTH_OPENAI_COMPAT_STREAM_STALL_TIMEOUT=0 disables the stall guard, so
# the callable returns None. Once a chunk has arrived the leftover
# first-token read timeout must be cleared, else a long post-first-chunk gap
# trips a stale deadline the operator asked to turn off.
async def _run():
response = SimpleNamespace(request = SimpleNamespace(extensions = {"timeout": {}}))
seen = []
class _Request:
async def is_disconnected(self):
return False
class _Items:
def __init__(self):
self.count = 0
async def __anext__(self):
self.count += 1
if self.count > 2:
raise StopAsyncIteration
return "data: {}"
async for _ in inf_mod._aiter_llama_stream_items(
_Items(),
cancel_event = threading.Event(),
request = _Request(),
response = response,
first_token_deadline = time.monotonic() + 5,
post_first_item_read_timeout_s = lambda: None,
):
seen.append(response.request.extensions["timeout"].get("read"))
# The first-token path armed a finite read timeout; after the first chunk
# with the guard disabled, it is cleared to None on every subsequent read.
assert seen == [None, None], seen
asyncio.run(_run())

View file

@ -23,7 +23,11 @@ if _BACKEND_DIR not in sys.path:
from hub.utils.download_manifest import ExpectedFile
from hub.utils.gguf import is_mtp_drafter_path
from hub.utils.gguf_plan import build_gguf_variant_plans, plan_from_expected_files
from hub.utils.gguf_plan import (
build_gguf_variant_plans,
plan_from_expected_files,
preferred_mtp_sibling,
)
from utils.models.model_config import (
_is_mtp_drafter,
detect_gguf_model,
@ -37,6 +41,8 @@ from utils.models.model_config import (
DRAFTER_CASES = [
("mtp-gemma-4-12b-it.gguf", True),
("MTP/gemma-4-12b-it-Q8_0-MTP.gguf", True),
# New-scheme MTP/ copies carry the mtp- basename prefix too.
("MTP/mtp-gemma-4-E4B-it-BF16.gguf", True),
("foo/MTP/bar.gguf", True),
("gemma-4-12b-it-Q8_0.gguf", False),
# Baked-in Qwen MTP repos: the head is inside the main GGUF, the file
@ -274,3 +280,178 @@ def test_detect_gguf_model_rejects_mtp_subdir_copy(tmp_path):
assert detect_gguf_model(str(copy)) is None
# Selecting the MTP dir itself must not surface the copies as models.
assert detect_gguf_model(str(sub)) is None
# ── Root drafter wins over new-scheme MTP/ copies ────────────────────
# The MTP/ copies were renamed to share the mtp- basename prefix (e.g.
# MTP/mtp-gemma-4-E4B-it-BF16.gguf). Auto-fetch/load must still resolve the
# small repo-root drafter, not a sort-first MTP/ copy (uppercase precedes
# lowercase, so the subdir path would otherwise win).
NEW_SCHEME_SIBLINGS = [
_sib("gemma-4-12b-it-Q4_K_M.gguf", 4_000, "main-q4"),
_sib("gemma-4-12b-it-Q8_0.gguf", 8_000, "main-q8"),
_sib("mtp-gemma-4-12b-it.gguf", 100, "drafter"),
_sib("MTP/mtp-gemma-4-12b-it-Q8_0.gguf", 100, "mtp-sub-q8"),
_sib("MTP/mtp-gemma-4-12b-it-BF16.gguf", 200, "mtp-sub-bf16"),
_sib("mmproj-F16.gguf", 500, "mmproj"),
]
def test_preferred_mtp_sibling_prefers_root_over_new_scheme_copies():
picked = preferred_mtp_sibling(NEW_SCHEME_SIBLINGS)
assert picked is not None and picked.rfilename == "mtp-gemma-4-12b-it.gguf"
def test_variant_plans_new_scheme_uses_root_drafter():
plans = build_gguf_variant_plans(NEW_SCHEME_SIBLINGS)
assert set(plans) == {"q4_k_m", "q8_0"}
for plan in plans.values():
assert "mtp-gemma-4-12b-it.gguf" in plan.target_filenames
assert not any("MTP/" in name for name in plan.target_filenames)
assert "drafter" in plan.companion_hashes
# Download size = main + mmproj + root drafter (not the 200-byte BF16 copy).
assert plans["q4_k_m"].download_size_bytes == 4_600
def test_download_mtp_prefers_root_over_new_scheme_copies(monkeypatch):
# _pick_mtp is nested; capture it via the companion-download seam.
from core.inference.llama_cpp import LlamaCppBackend
monkeypatch.delenv("HF_HUB_OFFLINE", raising = False) # online: skip reuse probe
captured = {}
def _fake_companion(
*,
hf_repo,
hf_token,
pick,
label,
cancel_event = None,
):
captured["pick"] = pick
return None
b = LlamaCppBackend()
b._download_companion_gguf = _fake_companion
b._download_mtp(hf_repo = "unsloth/gemma-4-E4B-it-qat-mobile-GGUF")
repo_files = [
"MTP/mtp-gemma-4-E4B-it-BF16.gguf",
"MTP/mtp-gemma-4-E4B-it-Q4_0.gguf",
"MTP/mtp-gemma-4-E4B-it-Q8_0.gguf",
"gemma-4-E4B-it-qat-UD-Q2_K_XL.gguf",
"mmproj-F16.gguf",
"mtp-gemma-4-E4B-it.gguf",
]
assert captured["pick"](repo_files) == "mtp-gemma-4-E4B-it.gguf"
# ── Reuse an on-disk drafter offline; fetch fresh online ─────────────
def _seed_snapshot(tmp_path, names):
snap = tmp_path / "snap"
for rel in names:
f = snap / rel
f.parent.mkdir(parents = True, exist_ok = True)
f.write_bytes(b"x")
return snap
def test_download_mtp_reuses_cached_root_drafter_offline(tmp_path, monkeypatch):
import utils.models.model_config as mc
from core.inference.llama_cpp import LlamaCppBackend
monkeypatch.setenv("HF_HUB_OFFLINE", "1")
snap = _seed_snapshot(
tmp_path,
[
"gemma-4-E4B-it-qat-UD-Q2_K_XL.gguf",
"mtp-gemma-4-E4B-it.gguf",
"MTP/mtp-gemma-4-E4B-it-BF16.gguf",
"mmproj-F16.gguf",
],
)
monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda repo: [snap])
got = LlamaCppBackend()._download_mtp(hf_repo = "unsloth/gemma-4-E4B-it-qat-mobile-GGUF")
assert got is not None and Path(got).name == "mtp-gemma-4-E4B-it.gguf"
def test_download_mtp_reuses_cached_subdir_copy_when_no_root_offline(tmp_path, monkeypatch):
# Pre-fix build may have fetched only the MTP/ copy; reuse it offline.
import utils.models.model_config as mc
from core.inference.llama_cpp import LlamaCppBackend
monkeypatch.setenv("HF_HUB_OFFLINE", "1")
snap = _seed_snapshot(
tmp_path,
[
"gemma-4-E4B-it-qat-UD-Q2_K_XL.gguf",
"MTP/mtp-gemma-4-E4B-it-BF16.gguf",
],
)
monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda repo: [snap])
got = LlamaCppBackend()._download_mtp(hf_repo = "unsloth/gemma-4-E4B-it-qat-mobile-GGUF")
assert got is not None and Path(got).name == "mtp-gemma-4-E4B-it-BF16.gguf"
def test_download_mtp_prefers_root_across_snapshots_offline(tmp_path, monkeypatch):
# A newer partial snapshot holds only the MTP/ copy; an older one has the
# root. Must still return the small root, not the large subdir copy.
import utils.models.model_config as mc
from core.inference.llama_cpp import LlamaCppBackend
monkeypatch.setenv("HF_HUB_OFFLINE", "1")
snap_partial = _seed_snapshot(tmp_path / "new", ["MTP/mtp-gemma-4-E4B-it-BF16.gguf"])
snap_full = _seed_snapshot(tmp_path / "old", ["mtp-gemma-4-E4B-it.gguf"])
monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda repo: [snap_partial, snap_full])
got = LlamaCppBackend()._download_mtp(hf_repo = "unsloth/gemma-4-E4B-it-qat-mobile-GGUF")
assert got is not None and Path(got).name == "mtp-gemma-4-E4B-it.gguf"
def test_download_mtp_reuse_follows_snapshot_order_offline(tmp_path, monkeypatch):
# Two snapshots both hold a root drafter; newest-first order must win so a
# fresh main GGUF is not paired with a stale drafter revision.
import utils.models.model_config as mc
from core.inference.llama_cpp import LlamaCppBackend
monkeypatch.setenv("HF_HUB_OFFLINE", "1")
newest = _seed_snapshot(tmp_path / "newest", ["mtp-gemma-4-E4B-it.gguf"])
oldest = _seed_snapshot(tmp_path / "oldest", ["mtp-gemma-4-E4B-it.gguf"])
monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda repo: [newest, oldest])
got = LlamaCppBackend()._download_mtp(hf_repo = "unsloth/gemma-4-E4B-it-qat-mobile-GGUF")
assert got is not None and Path(got).parent.parent.name == "newest"
def test_download_mtp_online_skips_cache_reuse(tmp_path, monkeypatch):
# Online, do not reuse a cached copy: go to the download path so a changed
# drafter is refetched (hf_hub_download checks the current revision).
import utils.models.model_config as mc
from core.inference.llama_cpp import LlamaCppBackend
monkeypatch.delenv("HF_HUB_OFFLINE", raising = False)
snap = _seed_snapshot(tmp_path, ["mtp-gemma-4-E4B-it.gguf"])
monkeypatch.setattr(mc, "_iter_hf_cache_snapshots", lambda repo: [snap])
reached = {}
def _fake_companion(
*,
hf_repo,
hf_token,
pick,
label,
cancel_event = None,
):
reached["hit"] = True
return None
b = LlamaCppBackend()
b._download_companion_gguf = _fake_companion
assert b._download_mtp(hf_repo = "unsloth/gemma-4-E4B-it-qat-mobile-GGUF") is None
assert reached.get("hit") is True

View file

@ -3037,3 +3037,60 @@ def test_acquire_swap_gate_is_cancellation_safe():
inference_route._auto_switch_process_lock.release()
asyncio.run(asyncio.wait_for(main(), timeout = 5))
def test_no_model_loaded_detail_appends_hint_only_when_off(monkeypatch):
# The "no model loaded" errors point at the opt-in auto-switch toggle so a
# request naming a listed-but-unloaded model is self-explanatory -- but only
# when it's off. With it on the name simply didn't resolve, so no hint.
base = "No GGUF model loaded. Load a GGUF model first."
monkeypatch.setattr(settings, "get_openai_auto_switch_enabled", lambda: False)
off = inference_route._no_model_loaded_detail(base)
assert off.startswith(base)
assert "Model auto-switch" in off and "Settings > API" in off
monkeypatch.setattr(settings, "get_openai_auto_switch_enabled", lambda: True)
assert inference_route._no_model_loaded_detail(base) == base
def _run_responses_stream_no_model(monkeypatch, *, enabled, active_model_name):
# Drive _responses_stream's GGUF-not-loaded guard: llama backend unloaded,
# inference backend maybe holding a non-GGUF model. Returns the 400 detail.
from fastapi import HTTPException
from models.inference import ResponsesRequest, ChatMessage
monkeypatch.setattr(settings, "get_openai_auto_switch_enabled", lambda: enabled)
monkeypatch.setattr(
inference_route, "get_llama_cpp_backend", lambda: _FakeBackend(loaded_id = None)
)
monkeypatch.setattr(
inference_route,
"get_inference_backend",
lambda: type("_B", (), {"active_model_name": active_model_name})(),
)
payload = ResponsesRequest(model = "unsloth/Qwen3.5-4B-GGUF", stream = True)
messages = [ChatMessage(role = "user", content = "hi")]
with pytest.raises(HTTPException) as exc:
asyncio.run(inference_route._responses_stream(payload, messages, None))
assert exc.value.status_code == 400
return exc.value.detail
def test_responses_stream_hint_matches_toggle_regardless_of_active_model(monkeypatch):
# Streaming /v1/responses shares the GGUF-only 400 with the other "no model
# loaded" sites, so the auto-switch hint attaches whenever the toggle is
# off -- including while a non-GGUF model is active, since auto-switch
# evicts it to load a resolved GGUF (_maybe_auto_switch_model's resolver
# branch has no active-model guard, unlike its reload-stash branch). Only
# the toggle being on suppresses it.
hinted = _run_responses_stream_no_model(monkeypatch, enabled = False, active_model_name = None)
assert "Model auto-switch" in hinted
on = _run_responses_stream_no_model(monkeypatch, enabled = True, active_model_name = None)
assert "Model auto-switch" not in on
non_gguf_loaded = _run_responses_stream_no_model(
monkeypatch, enabled = False, active_model_name = "unsloth/Llama-3.2-1B-Instruct"
)
assert "Model auto-switch" in non_gguf_loaded

File diff suppressed because it is too large Load diff

View file

@ -515,6 +515,7 @@ class ScriptedClient:
_url,
json = None,
timeout = None,
headers = None,
):
self.posts.append(json)
return httpx.Response(200, json = self.bodies[min(len(self.posts) - 1, len(self.bodies) - 1)])

View file

@ -514,6 +514,12 @@ def _run_update(install_dir: Path, repo: str, asset: Optional[str], script: Path
logger.info("llama update: installing", cmd = " ".join(cmd))
# Stream progress lines into job["progress"].
env = dict(os.environ, UNSLOTH_PROGRESS_PERCENT_STEP = "5")
# Preserve a Vulkan install across updates: detect_host on a CUDA/ROCm
# box would otherwise re-route and silently replace the Vulkan build.
# Re-assert it via the same env flag setup uses (mirrors
# _rocm_install_args).
if asset and "vulkan" in asset.lower():
env["UNSLOTH_FORCE_VULKAN"] = "1"
proc = subprocess.Popen(
cmd,
stdout = subprocess.PIPE,

View file

@ -81,7 +81,6 @@ import {
TestTube01Icon,
ZapIcon,
} from "@hugeicons/core-free-icons";
import { listStoredChatThreads } from "@/features/chat/utils/chat-history-storage";
import {
Tooltip,
TooltipContent,
@ -97,6 +96,7 @@ import {
createChatProject,
deleteChatProject,
deleteChatItem,
listStoredChatThreads,
moveChatItemToProject,
renameChatItem,
renameChatProject,
@ -582,7 +582,14 @@ export function AppSidebar() {
useEffect(() => {
if (!pendingRename) return;
const match = allChatItems.find((i) => i.id === pendingRename.id);
if (match && match.title === pendingRename.title) setPendingRename(null);
if (!match || match.title !== pendingRename.title) return;
queueMicrotask(() => {
setPendingRename((current) =>
current?.id === pendingRename.id && current.title === pendingRename.title
? null
: current,
);
});
}, [allChatItems, pendingRename]);
const [creatingProject, setCreatingProject] = useState(false);
const [projectNameDraft, setProjectNameDraft] = useState("");
@ -680,12 +687,6 @@ export function AppSidebar() {
useState<DeleteTarget | null>(null);
const [deleteProjectFiles, setDeleteProjectFiles] = useState(false);
useEffect(() => {
if (confirmingDelete?.kind !== "project") {
setDeleteProjectFiles(false);
}
}, [confirmingDelete]);
async function commitDelete() {
const target = confirmingDelete;
if (!target) return;

View file

@ -3,27 +3,35 @@
import { Button } from "@/components/ui/button";
import { Progress } from "@/components/ui/progress";
import { useMonitorOverlayStore } from "@/features/settings/stores/monitor-overlay-store";
import { useMonitorOverlayStore } from "@/features/settings";
import { useSystemInfo } from "@/hooks/use-system";
import { useT } from "@/i18n";
import { cn } from "@/lib/utils";
import { CpuIcon, GripVerticalIcon, XIcon } from "lucide-react";
import { motion } from "motion/react";
import { useRef } from "react";
import { AnimatePresence, motion, useDragControls } from "motion/react";
import { type PointerEvent, useMemo, useState } from "react";
function clampPercent(value: number): number {
return Math.max(0, Math.min(100, value));
}
function usageIndicatorClass(percent: number): string {
if (percent >= 90) return "bg-destructive";
if (percent >= 70) return "bg-amber-500";
if (percent >= 90) {
return "bg-destructive";
}
if (percent >= 70) {
return "bg-amber-500";
}
return "bg-primary";
}
function usageTextClass(percent: number): string {
if (percent >= 90) return "text-destructive";
if (percent >= 70) return "text-amber-600 dark:text-amber-400";
if (percent >= 90) {
return "text-destructive";
}
if (percent >= 70) {
return "text-amber-600 dark:text-amber-400";
}
return "text-primary";
}
@ -39,9 +47,18 @@ export function FloatingMonitor() {
const { isOpen, setIsOpen } = useMonitorOverlayStore();
const systemInfo = useSystemInfo({ enabled: isOpen, pollMs: 5000 });
const constraintsRef = useRef<HTMLDivElement>(null);
const [constraintsElement, setConstraintsElement] =
useState<HTMLDivElement | null>(null);
const constraintsRef = useMemo(
() => ({ current: constraintsElement }),
[constraintsElement],
);
const dragControls = useDragControls();
if (!isOpen) return null;
function startDrag(event: PointerEvent<HTMLDivElement>) {
event.preventDefault();
dragControls.start(event);
}
const ramTotal = systemInfo.memory?.total_gb ?? 0;
const ramAvailable = systemInfo.memory?.available_gb ?? 0;
@ -64,99 +81,109 @@ export function FloatingMonitor() {
const hasGpu = (systemInfo.gpu?.available ?? false) && devices.length > 0;
return (
<div
ref={constraintsRef}
className="fixed inset-0 z-50 pointer-events-none"
>
<motion.div
layout={true}
drag={true}
dragConstraints={constraintsRef}
dragElastic={0.1}
dragMomentum={false}
initial={{ opacity: 0, scale: 0.9 }}
animate={{ opacity: 1, scale: 1 }}
exit={{ opacity: 0, scale: 0.9 }}
className="settings-surface fixed bottom-4 right-4 w-64 max-w-[calc(100vw-2rem)] resize overflow-hidden rounded-xl border border-border/70 p-3 shadow-border ring-0 backdrop-blur-sm pointer-events-auto cursor-default select-none"
>
<div className="mb-2 flex items-center justify-between gap-2 border-b border-border/60 pb-2">
<div className="flex min-w-0 flex-1 items-center gap-1.5 truncate text-xs font-semibold text-foreground">
<CpuIcon className="size-3.5 shrink-0 text-primary" />
<span className="truncate">
{t("settings.resources.liveMonitor.title")}
</span>
</div>
<div className="flex items-center gap-1 shrink-0">
<div className="cursor-grab rounded-md px-1 text-muted-foreground/60 transition-colors hover:bg-muted/60 hover:text-muted-foreground active:cursor-grabbing">
<GripVerticalIcon className="size-3.5" />
</div>
<Button
size="icon-xs"
variant="ghost"
className="text-muted-foreground hover:text-foreground"
onClick={() => setIsOpen(false)}
title={t("common.close")}
aria-label={t("common.close")}
>
<XIcon className="size-3" />
</Button>
</div>
</div>
<motion.div
initial={{ opacity: 0, height: 0 }}
animate={{ opacity: 1, height: "auto" }}
exit={{ opacity: 0, height: 0 }}
className="space-y-3 overflow-hidden"
<AnimatePresence>
{isOpen && (
<div
ref={setConstraintsElement}
className="fixed inset-0 z-50 pointer-events-none"
>
<div className="space-y-1">
<div className="flex justify-between text-[11px] font-medium font-mono">
<span>{t("settings.resources.liveMonitor.ram")}</span>
<span className={cn("tabular-nums", usageTextClass(ramPercent))}>
{Math.round(ramPercent)}%
</span>
</div>
<div className="text-xs text-muted-foreground font-mono tabular-nums">
{formatGiB(ramUsed)} / {formatGiB(ramTotal)}
</div>
<Progress
value={ramPercent}
className="mt-1 h-1.5 rounded-full bg-muted"
indicatorClassName={usageIndicatorClass(ramPercent)}
/>
</div>
{hasGpu && (
<div className="space-y-1">
<div className="flex justify-between text-[11px] font-medium font-mono">
<span className="truncate flex-1 pr-2">
{t("settings.resources.liveMonitor.vram")}{" "}
{devices.length > 1
? `(${devices.length} GPUs)`
: `(${devices[0].name ?? "GPU"})`}
<motion.div
drag={true}
dragControls={dragControls}
dragListener={false}
dragConstraints={constraintsRef}
dragElastic={0}
dragMomentum={false}
initial={{ opacity: 0, scale: 0.9 }}
animate={{ opacity: 1, scale: 1 }}
exit={{ opacity: 0, scale: 0.9 }}
className="settings-surface fixed bottom-4 right-4 w-64 max-w-[calc(100vw-2rem)] resize overflow-hidden rounded-xl border border-border/70 p-3 shadow-border ring-0 backdrop-blur-sm pointer-events-auto cursor-default select-none"
>
<div className="mb-2 flex items-center justify-between gap-2 border-b border-border/60 pb-2">
<div className="flex min-w-0 flex-1 items-center gap-1.5 truncate text-xs font-semibold text-foreground">
<CpuIcon className="size-3.5 shrink-0 text-primary" />
<span className="truncate">
{t("settings.resources.liveMonitor.title")}
</span>
<span
className={cn(
"shrink-0 tabular-nums",
usageTextClass(vramPercent),
)}
</div>
<div className="flex items-center gap-1 shrink-0">
<div
onPointerDown={startDrag}
className="touch-none cursor-grab rounded-md px-1 text-muted-foreground/60 transition-colors hover:bg-muted/60 hover:text-muted-foreground active:cursor-grabbing"
>
{Math.round(vramPercent)}%
</span>
<GripVerticalIcon className="size-3.5" />
</div>
<Button
size="icon-xs"
variant="ghost"
className="text-muted-foreground hover:text-foreground"
onClick={() => setIsOpen(false)}
title={t("common.close")}
aria-label={t("common.close")}
>
<XIcon className="size-3" />
</Button>
</div>
<div className="text-xs text-muted-foreground font-mono tabular-nums">
{formatGiB(vramUsed)} / {formatGiB(vramTotal)}
</div>
<Progress
value={vramPercent}
className="mt-1 h-1.5 rounded-full bg-muted"
indicatorClassName={usageIndicatorClass(vramPercent)}
/>
</div>
)}
</motion.div>
</motion.div>
</div>
<motion.div
initial={{ opacity: 0, height: 0 }}
animate={{ opacity: 1, height: "auto" }}
exit={{ opacity: 0, height: 0 }}
className="space-y-3 overflow-hidden"
>
<div className="space-y-1">
<div className="flex justify-between text-[11px] font-medium font-mono">
<span>{t("settings.resources.liveMonitor.ram")}</span>
<span
className={cn("tabular-nums", usageTextClass(ramPercent))}
>
{Math.round(ramPercent)}%
</span>
</div>
<div className="text-xs text-muted-foreground font-mono tabular-nums">
{formatGiB(ramUsed)} / {formatGiB(ramTotal)}
</div>
<Progress
value={ramPercent}
className="mt-1 h-1.5 rounded-full bg-muted"
indicatorClassName={usageIndicatorClass(ramPercent)}
/>
</div>
{hasGpu && (
<div className="space-y-1">
<div className="flex justify-between text-[11px] font-medium font-mono">
<span className="truncate flex-1 pr-2">
{t("settings.resources.liveMonitor.vram")}{" "}
{devices.length > 1
? `(${devices.length} GPUs)`
: `(${devices[0].name ?? "GPU"})`}
</span>
<span
className={cn(
"shrink-0 tabular-nums",
usageTextClass(vramPercent),
)}
>
{Math.round(vramPercent)}%
</span>
</div>
<div className="text-xs text-muted-foreground font-mono tabular-nums">
{formatGiB(vramUsed)} / {formatGiB(vramTotal)}
</div>
<Progress
value={vramPercent}
className="mt-1 h-1.5 rounded-full bg-muted"
indicatorClassName={usageIndicatorClass(vramPercent)}
/>
</div>
)}
</motion.div>
</motion.div>
</div>
)}
</AnimatePresence>
);
}

View file

@ -2,6 +2,7 @@
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import { getAuthToken } from "@/features/auth";
import { resolveInitialConfig } from "@/features/model-picker";
import { projectHasSources } from "@/features/rag/api/rag-api";
import { apiUrl } from "@/lib/api-base";
import { parseParamCountB } from "@/lib/model-size";
@ -63,11 +64,17 @@ import {
listStoredChatThreads,
updateStoredChatThread,
} from "../utils/chat-history-storage";
import {
readLastLocalModelLoad,
recordLastLocalModelLoad,
type LastLocalModelKind,
} from "../utils/last-local-model-load";
import { getImageInputUnavailableReason } from "../utils/image-input-support";
import {
hasClosedThinkTag,
parseAssistantContent,
} from "../utils/parse-assistant-content";
import { resolveLoadMaxSeqLength } from "../presets/preset-policy";
import {
generateAudio,
listCachedGguf,
@ -1309,6 +1316,30 @@ const BIG_ENDIAN_GGUF_FILENAME_RE = /(^|[-_])be(?:[._-]|$)/gi;
const GGUF_KNOWN_QUANT_RE =
/(UD-)?(MXFP[0-9]+(?:_[A-Z0-9]+)*|IQ[0-9]+_[A-Z]+(?:_[A-Z0-9]+)?|TQ[0-9]+_[0-9]+|Q[0-9]+_K_[A-Z]+|Q[0-9]+_[0-9]+|Q[0-9]+_K|BF16|F16|F32)/i;
type AutoLoadCandidate = {
id: string;
kind: LastLocalModelKind;
ggufVariant: string | null;
maxSeqLength: number;
successLabel: string;
};
function autoLoadCandidateKey(
kind: LastLocalModelKind,
id: string,
ggufVariant?: string | null,
): string {
return `${kind}:${id.toLowerCase()}:${(ggufVariant ?? "").toLowerCase()}`;
}
function findCachedRepo<T extends { repo_id: string }>(
repos: T[],
id: string,
): T | undefined {
const normalized = id.toLowerCase();
return repos.find((repo) => repo.repo_id.toLowerCase() === normalized);
}
function hasBigEndianGgufMarker(filename: string, quant?: string | null): boolean {
const normalized = filename.replace(/\\/g, "/").toLowerCase();
const separatorIndex = normalized.lastIndexOf("/");
@ -1357,14 +1388,18 @@ async function autoLoadSmallestModel(): Promise<{
const hfToken = store.hfToken || null;
const trustRemoteCode = store.params.trustRemoteCode ?? false;
const specSettings = resolveSpeculativeSettingsForLoad();
const lastLoaded = readLastLocalModelLoad();
const toastId = toast("Loading a model…", {
description: "Auto-selecting the smallest downloaded model.",
description: lastLoaded
? "Loading last used model."
: "Auto-selecting the smallest downloaded model.",
duration: 5000,
closeButton: true,
});
let blockedByTrustRemoteCode = false;
let hadNonTrustFailure = false;
let loadAttempts = 0;
const skippedAutoLoadCandidates = new Set<string>();
async function canAutoLoad(payload: {
model_path: string;
@ -1389,12 +1424,225 @@ async function autoLoadSmallestModel(): Promise<{
}
return true;
}
async function loadAutoLoadCandidate(
candidate: AutoLoadCandidate,
): Promise<boolean> {
if (loadAttempts >= MAX_AUTO_LOAD_ATTEMPTS) {
return false;
}
const currentStore = useChatRuntimeStore.getState();
const { config } = resolveInitialConfig(candidate.id, candidate.ggufVariant);
const effectiveMaxSeqLength = resolveLoadMaxSeqLength({
modelId: candidate.id,
ggufVariant: candidate.ggufVariant,
isGguf: candidate.kind === "gguf",
customContextLength: config.customContextLength,
ggufContextLength: null,
currentCheckpoint: currentStore.params.checkpoint,
activeGgufVariant: currentStore.activeGgufVariant,
maxSeqLength: candidate.maxSeqLength,
presetSource: currentStore.activePresetSource,
});
const effectiveSpeculativeType =
config.speculativeType ?? specSettings.speculativeType;
const effectiveSpecDraftNMax =
config.specDraftNMax ?? specSettings.specDraftNMax;
const effectiveChatTemplateOverride = config.chatTemplateOverride?.trim()
? config.chatTemplateOverride
: null;
if (
!(await canAutoLoad({
model_path: candidate.id,
max_seq_length: effectiveMaxSeqLength,
is_lora: false,
gguf_variant: candidate.ggufVariant,
}))
) {
skippedAutoLoadCandidates.add(
autoLoadCandidateKey(candidate.kind, candidate.id, candidate.ggufVariant),
);
return false;
}
loadAttempts += 1;
const loadResp = await loadModel({
model_path: candidate.id,
hf_token: hfToken,
max_seq_length: effectiveMaxSeqLength,
load_in_4bit: true,
is_lora: false,
gguf_variant: candidate.ggufVariant,
trust_remote_code: trustRemoteCode,
chat_template_override: effectiveChatTemplateOverride,
cache_type_kv: config.kvCacheDtype,
speculative_type: effectiveSpeculativeType,
spec_draft_n_max: effectiveSpecDraftNMax,
tensor_parallel: config.tensorParallel,
});
saveSpeculativeType(effectiveSpeculativeType);
useChatRuntimeStore
.getState()
.setCheckpoint(candidate.id, candidate.ggufVariant ?? undefined);
const store = useChatRuntimeStore.getState();
store.setModelRequiresTrustRemoteCode(
loadResp.requires_trust_remote_code ?? false,
);
store.setParams({
...store.params,
maxTokens:
candidate.kind === "gguf"
? loadResp.context_length ?? 131072
: effectiveMaxSeqLength,
});
const autoModel: ChatModelSummary = {
id: candidate.id,
name: loadResp.display_name ?? candidate.id,
isVision: loadResp.is_vision ?? false,
isLora: loadResp.is_lora ?? false,
isGguf: loadResp.is_gguf ?? candidate.kind === "gguf",
isAudio: loadResp.is_audio ?? false,
audioType: loadResp.audio_type ?? null,
hasAudioInput: loadResp.has_audio_input ?? false,
};
if (!store.models.some((m) => m.id === candidate.id)) {
store.setModels([...store.models, autoModel]);
}
if (candidate.kind === "gguf") {
useChatRuntimeStore.setState({
ggufContextLength: loadResp.context_length ?? 131072,
ggufMaxContextLength:
loadResp.max_context_length ?? loadResp.context_length ?? 131072,
ggufNativeContextLength: loadResp.native_context_length ?? null,
supportsReasoning: loadResp.supports_reasoning ?? false,
reasoningAlwaysOn: loadResp.reasoning_always_on ?? false,
reasoningEnabled: loadResp.supports_reasoning ?? false,
...reasoningCapsFromLoad(loadResp),
supportsPreserveThinking: loadResp.supports_preserve_thinking ?? false,
supportsTools: loadResp.supports_tools ?? false,
...resolveToolsEnabledOnLoad(loadResp.supports_tools ?? false),
kvCacheDtype: loadResp.cache_type_kv ?? null,
loadedKvCacheDtype: loadResp.cache_type_kv ?? null,
tensorParallel: loadResp.tensor_parallel ?? false,
loadedTensorParallel: loadResp.tensor_parallel ?? false,
defaultChatTemplate: loadResp.chat_template ?? null,
chatTemplateOverride: effectiveChatTemplateOverride,
loadedChatTemplateOverride: effectiveChatTemplateOverride,
customContextLength: null,
loadedIsMultimodal: isMultimodalResponse(loadResp),
loadedIsDiffusion: loadResp.is_diffusion ?? false,
...resolveLoadedSpeculativeSettings(loadResp),
});
} else {
useChatRuntimeStore.setState({
supportsReasoning: loadResp.supports_reasoning ?? false,
reasoningAlwaysOn: loadResp.reasoning_always_on ?? false,
reasoningEnabled: loadResp.supports_reasoning ?? false,
...reasoningCapsFromLoad(loadResp),
supportsPreserveThinking: loadResp.supports_preserve_thinking ?? false,
supportsTools: loadResp.supports_tools ?? false,
...resolveToolsEnabledOnLoad(loadResp.supports_tools ?? false),
kvCacheDtype: loadResp.cache_type_kv ?? null,
loadedKvCacheDtype: loadResp.cache_type_kv ?? null,
tensorParallel: loadResp.tensor_parallel ?? false,
loadedTensorParallel: loadResp.tensor_parallel ?? false,
defaultChatTemplate: loadResp.chat_template ?? null,
chatTemplateOverride: effectiveChatTemplateOverride,
loadedChatTemplateOverride: effectiveChatTemplateOverride,
customContextLength: null,
...resolveLoadedSpeculativeSettings(loadResp),
loadedIsMultimodal: isMultimodalResponse(loadResp),
loadedIsDiffusion: loadResp.is_diffusion ?? false,
});
}
if (!(loadResp.is_lora ?? false)) {
recordLastLocalModelLoad({
id: candidate.id,
kind: candidate.kind,
ggufVariant: candidate.ggufVariant,
});
}
toast.success(candidate.successLabel, { id: toastId });
return true;
}
try {
const [ggufRepos, modelRepos] = await Promise.all([
listCachedGguf().catch(() => []),
listCachedModels().catch(() => []),
]);
if (lastLoaded) {
if (lastLoaded.kind === "gguf") {
const repo = findCachedRepo(ggufRepos, lastLoaded.id);
if (repo && lastLoaded.ggufVariant) {
try {
const variants = await listGgufVariants(repo.repo_id);
const variant = variants.variants.find(
(entry) =>
entry.downloaded &&
entry.quant?.toLowerCase() ===
lastLoaded.ggufVariant?.toLowerCase() &&
isAutoLoadableGgufVariant(entry),
);
if (variant) {
toast("Loading last used model…", {
id: toastId,
description: `${repo.repo_id} (${variant.quant})`,
duration: 5000,
});
if (
await loadAutoLoadCandidate({
id: repo.repo_id,
kind: "gguf",
ggufVariant: variant.quant,
maxSeqLength: 0,
successLabel: `Loaded ${repo.repo_id} (${variant.quant})`,
})
) {
return { loaded: true, blockedByTrustRemoteCode: false };
}
}
} catch {
hadNonTrustFailure = true;
skippedAutoLoadCandidates.add(
autoLoadCandidateKey("gguf", repo.repo_id, lastLoaded.ggufVariant),
);
}
}
} else {
const repo = findCachedRepo(modelRepos, lastLoaded.id);
if (repo) {
try {
toast("Loading last used model…", {
id: toastId,
description: repo.repo_id,
duration: 5000,
});
if (
await loadAutoLoadCandidate({
id: repo.repo_id,
kind: "model",
ggufVariant: null,
maxSeqLength: store.params.maxSeqLength,
successLabel: `Loaded ${repo.repo_id}`,
})
) {
return { loaded: true, blockedByTrustRemoteCode: false };
}
} catch {
hadNonTrustFailure = true;
skippedAutoLoadCandidates.add(
autoLoadCandidateKey("model", repo.repo_id),
);
}
}
}
toast("Loading a model…", {
id: toastId,
description: "Auto-selecting the smallest downloaded model.",
duration: 5000,
});
}
// GGUF first: smallest-total-size repo, then its smallest variant.
if (ggufRepos.length > 0) {
const sorted = [...ggufRepos].sort((a, b) => a.size_bytes - b.size_bytes);
@ -1408,82 +1656,23 @@ async function autoLoadSmallestModel(): Promise<{
if (downloaded.length > 0) {
const variant = downloaded[0];
if (
!(await canAutoLoad({
model_path: repo.repo_id,
max_seq_length: 0,
is_lora: false,
gguf_variant: variant.quant,
}))
skippedAutoLoadCandidates.has(
autoLoadCandidateKey("gguf", repo.repo_id, variant.quant),
)
) {
continue;
}
loadAttempts += 1;
const loadResp = await loadModel({
model_path: repo.repo_id,
hf_token: hfToken,
max_seq_length: 0,
load_in_4bit: true,
is_lora: false,
gguf_variant: variant.quant,
trust_remote_code: trustRemoteCode,
speculative_type: specSettings.speculativeType,
spec_draft_n_max: specSettings.specDraftNMax,
});
saveSpeculativeType(specSettings.speculativeType);
useChatRuntimeStore
.getState()
.setCheckpoint(repo.repo_id, variant.quant);
const store = useChatRuntimeStore.getState();
store.setModelRequiresTrustRemoteCode(
loadResp.requires_trust_remote_code ?? false,
);
store.setParams({
...store.params,
maxTokens: loadResp.context_length ?? 131072,
});
// Add to store so the selector shows the name.
const autoModel: ChatModelSummary = {
id: repo.repo_id,
name: loadResp.display_name ?? repo.repo_id,
isVision: loadResp.is_vision ?? false,
isLora: loadResp.is_lora ?? false,
isGguf: loadResp.is_gguf ?? false,
isAudio: loadResp.is_audio ?? false,
audioType: loadResp.audio_type ?? null,
hasAudioInput: loadResp.has_audio_input ?? false,
};
const existingModels = store.models;
if (!existingModels.some((m) => m.id === repo.repo_id)) {
store.setModels([...existingModels, autoModel]);
if (
await loadAutoLoadCandidate({
id: repo.repo_id,
kind: "gguf",
ggufVariant: variant.quant,
maxSeqLength: 0,
successLabel: `Loaded ${repo.repo_id} (${variant.quant})`,
})
) {
return { loaded: true, blockedByTrustRemoteCode: false };
}
useChatRuntimeStore.setState({
ggufContextLength: loadResp.context_length ?? 131072,
ggufMaxContextLength:
loadResp.max_context_length ??
loadResp.context_length ??
131072,
supportsReasoning: loadResp.supports_reasoning ?? false,
reasoningAlwaysOn: loadResp.reasoning_always_on ?? false,
reasoningEnabled: loadResp.supports_reasoning ?? false,
...reasoningCapsFromLoad(loadResp),
supportsPreserveThinking:
loadResp.supports_preserve_thinking ?? false,
supportsTools: loadResp.supports_tools ?? false,
...resolveToolsEnabledOnLoad(loadResp.supports_tools ?? false),
kvCacheDtype: loadResp.cache_type_kv ?? null,
loadedKvCacheDtype: loadResp.cache_type_kv ?? null,
tensorParallel: loadResp.tensor_parallel ?? false,
loadedTensorParallel: loadResp.tensor_parallel ?? false,
defaultChatTemplate: loadResp.chat_template ?? null,
chatTemplateOverride: null,
loadedChatTemplateOverride: null,
loadedIsMultimodal: isMultimodalResponse(loadResp),
...resolveLoadedSpeculativeSettings(loadResp),
});
toast.success(`Loaded ${repo.repo_id} (${variant.quant})`, {
id: toastId,
});
return { loaded: true, blockedByTrustRemoteCode: false };
}
} catch {
hadNonTrustFailure = true;
@ -1501,64 +1690,23 @@ async function autoLoadSmallestModel(): Promise<{
if (loadAttempts >= MAX_AUTO_LOAD_ATTEMPTS) break;
try {
if (
!(await canAutoLoad({
model_path: repo.repo_id,
max_seq_length: 4096,
is_lora: false,
gguf_variant: null,
}))
skippedAutoLoadCandidates.has(
autoLoadCandidateKey("model", repo.repo_id),
)
) {
continue;
}
loadAttempts += 1;
const sfLoadResp = await loadModel({
model_path: repo.repo_id,
hf_token: hfToken,
max_seq_length: 4096,
load_in_4bit: true,
is_lora: false,
gguf_variant: null,
trust_remote_code: trustRemoteCode,
speculative_type: specSettings.speculativeType,
spec_draft_n_max: specSettings.specDraftNMax,
});
saveSpeculativeType(specSettings.speculativeType);
useChatRuntimeStore.getState().setCheckpoint(repo.repo_id);
const store = useChatRuntimeStore.getState();
store.setModelRequiresTrustRemoteCode(
sfLoadResp.requires_trust_remote_code ?? false,
);
store.setParams({ ...store.params, maxTokens: 4096 });
useChatRuntimeStore.setState({
supportsReasoning: sfLoadResp.supports_reasoning ?? false,
reasoningAlwaysOn: sfLoadResp.reasoning_always_on ?? false,
reasoningEnabled: sfLoadResp.supports_reasoning ?? false,
...reasoningCapsFromLoad(sfLoadResp),
supportsPreserveThinking:
sfLoadResp.supports_preserve_thinking ?? false,
supportsTools: sfLoadResp.supports_tools ?? false,
// Parity with the GGUF branch above.
...resolveToolsEnabledOnLoad(sfLoadResp.supports_tools ?? false),
defaultChatTemplate: sfLoadResp.chat_template ?? null,
chatTemplateOverride: null,
loadedChatTemplateOverride: null,
...resolveLoadedSpeculativeSettings(sfLoadResp),
});
const sfModel: ChatModelSummary = {
id: repo.repo_id,
name: sfLoadResp.display_name ?? repo.repo_id,
isVision: sfLoadResp.is_vision ?? false,
isLora: sfLoadResp.is_lora ?? false,
isGguf: sfLoadResp.is_gguf ?? false,
};
if (!store.models.some((m) => m.id === repo.repo_id)) {
store.setModels([...store.models, sfModel]);
if (
await loadAutoLoadCandidate({
id: repo.repo_id,
kind: "model",
ggufVariant: null,
maxSeqLength: 4096,
successLabel: `Loaded ${repo.repo_id}`,
})
) {
return { loaded: true, blockedByTrustRemoteCode: false };
}
useChatRuntimeStore.setState({
loadedIsMultimodal: isMultimodalResponse(sfLoadResp),
});
toast.success(`Loaded ${repo.repo_id}`, { id: toastId });
return { loaded: true, blockedByTrustRemoteCode: false };
} catch {
hadNonTrustFailure = true;
continue;
@ -1650,6 +1798,11 @@ async function autoLoadSmallestModel(): Promise<{
loadedIsMultimodal: isMultimodalResponse(loadResp),
...resolveLoadedSpeculativeSettings(loadResp),
});
recordLastLocalModelLoad({
id: "unsloth/Qwen3.5-4B-MTP-GGUF",
kind: "gguf",
ggufVariant: "UD-Q4_K_XL",
});
toast.success("Loaded Qwen3.5-4B-MTP (UD-Q4_K_XL)", { id: toastId });
return { loaded: true, blockedByTrustRemoteCode: false };
} catch {

View file

@ -43,6 +43,7 @@ import {
mergeBackendRecommendedInference,
resolveLoadMaxSeqLength,
} from "../presets/preset-policy";
import { recordLastLocalModelLoad } from "../utils/last-local-model-load";
import {
isMultimodalResponse,
} from "../types/api";
@ -804,6 +805,23 @@ export function useChatModelRuntime() {
}
}
await refresh({ signal: abortCtrl.signal });
if (
!isLora &&
!(loadResponse.is_lora ?? false) &&
!nativePathToken &&
!isLocalModelPath(modelId) &&
!isExternalModelId(modelId)
) {
if (loadResponse.is_gguf || isGguf || ggufVariant) {
recordLastLocalModelLoad({
id: modelId,
kind: "gguf",
ggufVariant: ggufVariant ?? null,
});
} else {
recordLastLocalModelLoad({ id: modelId, kind: "model" });
}
}
} catch (error) {
// Skip rollback if user cancelled -- model is already being unloaded.
if (abortCtrl.signal.aborted) throw error;

View file

@ -54,6 +54,7 @@ export { ChatSearchDialog } from "./components/chat-search-dialog";
export { setTrainingCompareHandoff } from "./lib/training-compare-handoff";
export type { ProjectRecord } from "./types";
export { clearAllChats, countAllChats } from "./utils/clear-all-chats";
export { listStoredChatThreads } from "./utils/chat-history-storage";
export { ArtifactCard } from "./artifacts/artifact-card";
export {
useChatArtifactsStore,

View file

@ -0,0 +1,86 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
export type LastLocalModelKind = "gguf" | "model";
export type LastLocalModelLoad = {
id: string;
kind: LastLocalModelKind;
ggufVariant: string | null;
loadedAt: number;
};
const STORAGE_KEY = "unsloth.last-local-model-load.v1";
function storage(): Storage | null {
try {
return typeof localStorage === "undefined" ? null : localStorage;
} catch {
return null;
}
}
function isLastLocalModelKind(value: unknown): value is LastLocalModelKind {
return value === "gguf" || value === "model";
}
export function readLastLocalModelLoad(): LastLocalModelLoad | null {
try {
const raw = storage()?.getItem(STORAGE_KEY);
if (!raw) {
return null;
}
const parsed = JSON.parse(raw) as Partial<LastLocalModelLoad>;
if (
typeof parsed.id !== "string" ||
!parsed.id.trim() ||
!isLastLocalModelKind(parsed.kind) ||
typeof parsed.loadedAt !== "number"
) {
return null;
}
if (
parsed.kind === "gguf" &&
(typeof parsed.ggufVariant !== "string" || !parsed.ggufVariant.trim())
) {
return null;
}
return {
id: parsed.id,
kind: parsed.kind,
ggufVariant:
typeof parsed.ggufVariant === "string" ? parsed.ggufVariant : null,
loadedAt: parsed.loadedAt,
};
} catch {
return null;
}
}
export function recordLastLocalModelLoad(input: {
id: string;
kind: LastLocalModelKind;
ggufVariant?: string | null;
}): void {
const id = input.id.trim();
if (!id) {
return;
}
const ggufVariant = input.ggufVariant?.trim() || null;
if (input.kind === "gguf" && !ggufVariant) {
return;
}
try {
storage()?.setItem(
STORAGE_KEY,
JSON.stringify({
id,
kind: input.kind,
ggufVariant: input.kind === "gguf" ? ggufVariant : null,
loadedAt: Date.now(),
} satisfies LastLocalModelLoad),
);
} catch {
// Ignore disabled storage / quota errors; auto-load falls back to size order.
}
}

View file

@ -494,3 +494,13 @@ export async function removeUnstructuredFile(
throw new Error("Failed to remove file");
}
}
export async function removeUnstructuredBlock(blockId: string): Promise<void> {
const res = await authFetch(
`${DATA_DESIGNER_API_BASE}/seed/unstructured-block/${encodeURIComponent(blockId)}`,
{ method: "DELETE" },
);
if (!res.ok && res.status !== 404) {
throw new Error("Failed to remove uploaded files");
}
}

View file

@ -53,6 +53,11 @@ import {
inspectSeedDataset,
inspectSeedUpload,
} from "../../api";
import { useRecipeStudioStore } from "../../stores/recipe-studio";
import {
makeUnstructuredUploadUid,
resolveUnstructuredUploadBlockId,
} from "../../utils/config-factories";
import { resolveImagePreview } from "../../utils/image-preview";
import type {
GithubItemType,
@ -597,6 +602,41 @@ export function SeedDialog({
const mode = config.seed_source_type ?? "hf";
const previewEmpty = getPreviewEmptyStateCopy(mode);
const queueUploadCleanup = useRecipeStudioStore(
(state) => state.queueUploadCleanup,
);
// config.id collides across recipes (ids reset to n1 on import); use a
// stable per-block uid instead. Generate one synchronously so the first
// rendered drop zone cannot upload under a legacy node id.
const uploadUid = config.unstructured_upload_uid?.trim() ?? "";
const unstructuredFileCount = config.unstructured_file_ids?.length ?? 0;
const generatedUploadUidRef = useRef<string | null>(null);
if (
mode === "unstructured" &&
!uploadUid &&
unstructuredFileCount === 0 &&
generatedUploadUidRef.current === null
) {
generatedUploadUidRef.current = makeUnstructuredUploadUid();
}
const uploadBlockId = resolveUnstructuredUploadBlockId({
configId: config.id,
uploadUid,
generatedUploadUid: generatedUploadUidRef.current,
unstructuredFileCount,
});
useEffect(() => {
if (mode !== "unstructured") return;
if (uploadUid) return;
if (unstructuredFileCount > 0) return;
const nextUid =
generatedUploadUidRef.current ?? makeUnstructuredUploadUid();
generatedUploadUidRef.current = nextUid;
onUpdate({ unstructured_upload_uid: nextUid });
}, [mode, uploadUid, unstructuredFileCount, onUpdate]);
const prevModeRef = useRef(mode);
useEffect(() => {
const prevMode = prevModeRef.current;
@ -720,6 +760,11 @@ export function SeedDialog({
subset: config.hf_subset?.trim() || undefined,
preview_size: 10,
});
// Queue the block's upload directory for deletion after the next
// save; only uid-namespaced directories qualify (single owner).
if (uploadUid && unstructuredFileCount > 0) {
queueUploadCleanup(uploadUid);
}
onUpdate({
hf_path: response.resolved_path,
seed_columns: response.columns,
@ -730,6 +775,7 @@ export function SeedDialog({
hf_split: response.split ?? "",
hf_subset: response.subset ?? "",
local_file_name: "",
unstructured_upload_uid: "",
unstructured_file_ids: [],
unstructured_file_names: [],
unstructured_file_sizes: [],
@ -754,6 +800,11 @@ export function SeedDialog({
content_base64: payload,
preview_size: 10,
});
// Queue the block's upload directory for deletion after the next
// save; only uid-namespaced directories qualify (single owner).
if (uploadUid && unstructuredFileCount > 0) {
queueUploadCleanup(uploadUid);
}
onUpdate({
hf_path: response.resolved_path,
seed_columns: response.columns,
@ -765,6 +816,7 @@ export function SeedDialog({
hf_subset: "",
hf_split: "",
local_file_name: localFile.name,
unstructured_upload_uid: "",
unstructured_file_ids: [],
unstructured_file_names: [],
unstructured_file_sizes: [],
@ -789,7 +841,7 @@ export function SeedDialog({
const { chunkSize, chunkOverlap } = resolveChunking(config);
const response = await inspectSeedUpload({
block_id: config.id,
block_id: uploadBlockId,
file_ids: fileIds,
file_names: fileNames,
preview_size: 10,
@ -827,7 +879,18 @@ export function SeedDialog({
setIsInspecting(false);
}
},
[config, getCurrentLoadKey, localFile, mode, onUpdate, unstructuredFiles],
[
config,
getCurrentLoadKey,
localFile,
mode,
onUpdate,
queueUploadCleanup,
unstructuredFiles,
unstructuredFileCount,
uploadBlockId,
uploadUid,
],
);
useEffect(() => {
@ -997,7 +1060,7 @@ export function SeedDialog({
{mode === "unstructured" && (
<UnstructuredDropZone
blockId={config.id}
blockId={uploadBlockId}
files={unstructuredFiles}
onFilesChange={handleUnstructuredFilesChange}
disabled={isInspecting}

View file

@ -54,11 +54,17 @@ export function UnstructuredDropZone({
}: UnstructuredDropZoneProps) {
const inputRef = useRef<HTMLInputElement>(null);
const filesRef = useRef(files);
const blockIdRef = useRef(blockId);
const mountedRef = useRef(true);
const [isDragOver, setIsDragOver] = useState(false);
useEffect(() => {
filesRef.current = files;
}, [files]);
blockIdRef.current = blockId;
}, [files, blockId]);
useEffect(() => () => {
mountedRef.current = false;
}, []);
const totalSize = files.reduce((sum, f) => sum + f.size, 0);
@ -134,15 +140,32 @@ export function UnstructuredDropZone({
if (entry.status === "uploading" && entry.abortController) {
entry.abortController.abort();
}
if (
const needsServerRemove =
entry.id &&
entry.status === "ok" &&
!deletedIdsRef.current.has(entry.id)
) {
deletedIdsRef.current.add(entry.id);
void removeUnstructuredFile(blockId, entry.id).catch(() => {});
}
!deletedIdsRef.current.has(entry.id);
onFilesChange((prev) => prev.filter((_, i) => i !== index));
if (!needsServerRemove) return;
deletedIdsRef.current.add(entry.id);
removeUnstructuredFile(blockId, entry.id).catch(() => {
// Skip if the drop zone unmounted or its block changed: the id no
// longer belongs here and restoring would leak it into another block.
if (!mountedRef.current || blockIdRef.current !== blockId) return;
// Still exists server-side (counts toward quota); restore it at its
// original position.
deletedIdsRef.current.delete(entry.id);
onFilesChange((prev) => {
const next = [...prev];
next.splice(Math.min(index, next.length), 0, {
id: entry.id,
name: entry.name,
size: entry.size,
status: "ok",
error: "Remove failed — try again",
});
return next;
});
});
},
[blockId, onFilesChange],
);

View file

@ -4,11 +4,13 @@
import { useCallback, useEffect, useMemo, useState } from "react";
import { toastError, toastSuccess } from "@/shared/toast";
import { normalizeNonEmptyName } from "@/utils";
import { removeUnstructuredBlock } from "../api";
import {
buildSignature,
copyTextToClipboard,
formatSavedLabel,
} from "../executions/execution-helpers";
import { useRecipeStudioStore } from "../stores/recipe-studio";
import { importRecipePayload, type RecipeSnapshot } from "../utils/import";
import type { RecipePayloadResult } from "../utils/payload/types";
@ -72,7 +74,10 @@ function stripApiKeys(value: unknown): unknown {
!Array.isArray(output.env)
) {
output.env = Object.fromEntries(
Object.keys(output.env as Record<string, unknown>).map((envKey) => [envKey, ""]),
Object.keys(output.env as Record<string, unknown>).map((envKey) => [
envKey,
"",
]),
);
}
return output;
@ -82,10 +87,7 @@ function inferHfRepoIdFromPath(pathValue: unknown): string {
if (typeof pathValue !== "string") {
return "";
}
const parts = pathValue
.trim()
.split("/")
.filter(Boolean);
const parts = pathValue.trim().split("/").filter(Boolean);
if (parts.length >= 3 && parts[0] === "datasets") {
return `${parts[1]}/${parts[2]}`;
}
@ -126,8 +128,7 @@ function sanitizeSeedForShare(payload: unknown): unknown {
typeof ui?.seed_source_type === "string" ? ui.seed_source_type : null;
const sourceType =
typeof source?.seed_type === "string" ? source.seed_type : null;
const shouldResetHfState =
sourceType === "hf" || uiSourceType === "hf";
const shouldResetHfState = sourceType === "hf" || uiSourceType === "hf";
const shouldResetLocalState =
sourceType === "local" ||
sourceType === "unstructured" ||
@ -144,6 +145,7 @@ function sanitizeSeedForShare(payload: unknown): unknown {
ui.seed_drop_columns = [];
ui.seed_preview_rows = [];
ui.local_file_name = "";
ui.unstructured_upload_uid = "";
ui.unstructured_file_ids = [];
ui.unstructured_file_names = [];
ui.unstructured_file_sizes = [];
@ -165,6 +167,7 @@ function sanitizeSeedForShare(payload: unknown): unknown {
ui.seed_drop_columns = [];
ui.seed_preview_rows = [];
ui.local_file_name = "";
ui.unstructured_upload_uid = "";
ui.unstructured_file_ids = [];
ui.unstructured_file_names = [];
ui.unstructured_file_sizes = [];
@ -174,6 +177,43 @@ function sanitizeSeedForShare(payload: unknown): unknown {
return root;
}
// Delete queued upload directories once a save stops referencing them, so a
// reload before autosave can never leave the saved recipe pointing at
// already-deleted files. Skips any uid the just-saved payload still uses.
function drainQueuedUploadCleanups(
savedPayload: RecipePayloadResult["payload"],
): void {
const pending = useRecipeStudioStore.getState().pendingUploadCleanups;
if (pending.length === 0) {
return;
}
const ui =
savedPayload && typeof savedPayload === "object"
? (savedPayload as { ui?: Record<string, unknown> }).ui
: undefined;
const savedUid =
ui && typeof ui.unstructured_upload_uid === "string"
? ui.unstructured_upload_uid
: "";
const ready = pending.filter((uid) => uid !== savedUid);
if (ready.length === 0) {
return;
}
for (const uid of ready) {
void removeUnstructuredBlock(uid)
.then(() => {
useRecipeStudioStore.setState((state) => ({
pendingUploadCleanups: state.pendingUploadCleanups.filter(
(pendingUid) => pendingUid !== uid,
),
}));
})
.catch((error) => {
console.warn("Failed to clean up uploaded documents:", error);
});
}
}
export function useRecipePersistence({
recipeId,
initialRecipeName,
@ -202,8 +242,10 @@ export function useRecipePersistence({
() => buildSignature(normalizedWorkflowName, currentPayload),
[currentPayload, normalizedWorkflowName],
);
const isDirty = savedSignature.length > 0 && currentSignature !== savedSignature;
const saveTone: SaveTone = !isDirty && Boolean(lastSavedAt) ? "success" : "error";
const isDirty =
savedSignature.length > 0 && currentSignature !== savedSignature;
const saveTone: SaveTone =
!isDirty && Boolean(lastSavedAt) ? "success" : "error";
const savedAtLabel = formatSavedLabel(lastSavedAt);
useEffect(() => {
@ -214,7 +256,9 @@ export function useRecipePersistence({
setLastSavedAt(initialSavedAt);
setCopied(false);
const parsed = importRecipePayload(JSON.stringify(initialPayload));
const parsed = importRecipePayload(JSON.stringify(initialPayload), {
preserveUnstructuredUploads: true,
});
if (parsed.snapshot) {
loadRecipe(parsed.snapshot);
} else {
@ -252,6 +296,7 @@ export function useRecipePersistence({
});
setLastSavedAt(result.updatedAt);
setSavedSignature(buildSignature(nextName, currentPayload));
drainQueuedUploadCleanups(currentPayload);
} catch (error) {
console.error("Save recipe failed:", error);
toastError("Save failed", "Could not save recipe.");
@ -270,11 +315,28 @@ export function useRecipePersistence({
return () => window.clearTimeout(timeoutId);
}, [isDirty, persistRecipe, saveLoading]);
// Drain queued cleanups even when autosave is skipped: a net-zero edit (add
// then remove an unstructured seed before the 800ms debounce) keeps isDirty
// false, so the autosave effect never drains and the queued uid leaks its
// upload dir. Not-dirty means currentPayload equals the saved recipe, and
// drain skips the uid it still references, so only dirs no saved recipe
// points at are deleted (keeps the save-first invariant).
useEffect(() => {
if (!initialRecipeReady || isDirty || saveLoading) {
return;
}
drainQueuedUploadCleanups(currentPayload);
}, [currentPayload, initialRecipeReady, isDirty, saveLoading]);
const copyRecipe = useCallback(async (): Promise<void> => {
setCopied(false);
try {
const safePayload = sanitizeSeedForShare(stripApiKeys(payloadResult.payload));
const ok = await copyTextToClipboard(JSON.stringify(safePayload, null, 2));
const safePayload = sanitizeSeedForShare(
stripApiKeys(payloadResult.payload),
);
const ok = await copyTextToClipboard(
JSON.stringify(safePayload, null, 2),
);
if (!ok) {
throw new Error("Clipboard not available.");
}

View file

@ -36,6 +36,7 @@ import {
} from "../utils/handles";
import type { RecipeSnapshot } from "../utils/import";
import { getLayoutedElements } from "../utils/layout";
import { makeUnstructuredUploadUid } from "../utils/config-factories";
import {
centerModelInfraNodes,
optimizeModelInfraEdgeHandles,
@ -76,6 +77,12 @@ type RecipeStudioState = {
nextId: number;
nextY: number;
fitViewTick: number;
// Upload-uid directories whose owning block dropped them; server-side
// deletion is deferred until a save no longer references them, so a
// reload before autosave cannot leave a saved recipe pointing at
// deleted files.
pendingUploadCleanups: string[];
queueUploadCleanup: (uid: string) => void;
setSheetOpen: (open: boolean) => void;
setSheetView: (view: SheetView) => void;
setProcessors: (processors: RecipeProcessorConfig[]) => void;
@ -137,6 +144,7 @@ const INITIAL_STATE = {
nextId: 3,
nextY: 280,
fitViewTick: 0,
pendingUploadCleanups: [],
} satisfies Pick<
RecipeStudioState,
| "nodes"
@ -154,6 +162,7 @@ const INITIAL_STATE = {
| "nextId"
| "nextY"
| "fitViewTick"
| "pendingUploadCleanups"
>;
function buildAddedNodeState(
@ -269,6 +278,20 @@ function isModelSemanticEdge(
);
}
// Upload uid of a seed block whose server-side directory becomes orphaned
// when the block drops it. Only uid directories qualify (single owner);
// legacy node-id directories can be shared by other recipes.
function seedUploadCleanupUid(config: NodeConfig | undefined): string | null {
if (!config || config.kind !== "seed") {
return null;
}
const uid = config.unstructured_upload_uid?.trim();
if (!uid || !config.unstructured_file_ids?.length) {
return null;
}
return uid;
}
export const useRecipeStudioStore = create<RecipeStudioState>((set, get) => ({
...INITIAL_STATE,
setSheetOpen: (open) => set({ sheetOpen: open }),
@ -278,6 +301,12 @@ export const useRecipeStudioStore = create<RecipeStudioState>((set, get) => ({
setDialogOpen: (open) => set({ dialogOpen: open }),
setExecutionLocked: (locked) => set({ executionLocked: locked }),
resetRecipe: () => set(INITIAL_STATE),
queueUploadCleanup: (uid) =>
set((state) =>
state.pendingUploadCleanups.includes(uid)
? state
: { pendingUploadCleanups: [...state.pendingUploadCleanups, uid] },
),
selectConfig: (id) => set({ activeConfigId: id, dialogOpen: false }),
openConfig: (id) => set({ activeConfigId: id, dialogOpen: true }),
setLayoutDirection: (direction) =>
@ -383,7 +412,18 @@ export const useRecipeStudioStore = create<RecipeStudioState>((set, get) => ({
}
return buildAddedNodeState(state, "sampler", type, position, openDialog);
}),
addSeedNode: (type, position, openDialog = true) =>
addSeedNode: (type, position, openDialog = true) => {
const current = get();
if (!current.executionLocked) {
// The reset below clears the block's upload uid and file list; queue
// its server-side directory for deletion after the next save.
const uid = seedUploadCleanupUid(
Object.values(current.configs).find((config) => config.kind === "seed"),
);
if (uid) {
current.queueUploadCleanup(uid);
}
}
set((state) => {
if (state.executionLocked) {
return state;
@ -413,6 +453,8 @@ export const useRecipeStudioStore = create<RecipeStudioState>((set, get) => ({
hf_token: "",
hf_endpoint: "https://huggingface.co",
local_file_name: "",
unstructured_upload_uid:
nextSourceType === "unstructured" ? makeUnstructuredUploadUid() : "",
unstructured_file_ids: [],
unstructured_file_names: [],
unstructured_file_sizes: [],
@ -446,7 +488,8 @@ export const useRecipeStudioStore = create<RecipeStudioState>((set, get) => ({
activeConfigId: existing.id,
dialogOpen: openDialog,
};
}),
});
},
addLlmNode: (type, position, openDialog = true) =>
set((state) => {
if (state.executionLocked) {
@ -699,6 +742,9 @@ export const useRecipeStudioStore = create<RecipeStudioState>((set, get) => ({
dialogOpen: false,
sheetView: "root",
fitViewTick: state.fitViewTick + 1,
// Queued cleanups belong to the previous recipe; draining them after
// a save of this one could delete files its saved payload still uses.
pendingUploadCleanups: [],
})),
setAuxNodePosition: (id, position) =>
set((state) => {
@ -786,6 +832,17 @@ export const useRecipeStudioStore = create<RecipeStudioState>((set, get) => ({
set(applyUpdate);
},
onNodesChange: (changes) => {
const current = get();
if (!current.executionLocked) {
for (const change of changes) {
if (change.type === "remove") {
const uid = seedUploadCleanupUid(current.configs[change.id]);
if (uid) {
current.queueUploadCleanup(uid);
}
}
}
}
const applyNodesChange = (state: RecipeStudioState) => {
if (state.executionLocked) {
return state;

View file

@ -340,6 +340,8 @@ export type SeedConfig = {
hf_token?: string;
hf_endpoint?: string;
local_file_name?: string;
// ui-only: stable per-block id for uploads, since node ids collide across imports
unstructured_upload_uid?: string;
unstructured_file_ids?: string[];
unstructured_file_names?: string[];
unstructured_file_sizes?: number[];

View file

@ -20,6 +20,46 @@ import type {
} from "../types";
import { nextName } from "./naming";
export function makeUnstructuredUploadUid(): string {
if (typeof globalThis.crypto?.randomUUID === "function") {
return globalThis.crypto.randomUUID().replace(/-/g, "").toLowerCase();
}
if (typeof globalThis.crypto?.getRandomValues === "function") {
const bytes = new Uint8Array(16);
globalThis.crypto.getRandomValues(bytes);
return Array.from(bytes, (byte) => byte.toString(16).padStart(2, "0")).join(
"",
);
}
let uid = "";
while (uid.length < 32) {
uid += Math.floor(Math.random() * 0x100000000)
.toString(16)
.padStart(8, "0");
}
return uid.slice(0, 32);
}
export function resolveUnstructuredUploadBlockId({
configId,
uploadUid,
generatedUploadUid,
unstructuredFileCount,
}: {
configId: string;
uploadUid: string;
generatedUploadUid: string | null;
unstructuredFileCount: number;
}): string {
if (uploadUid) {
return uploadUid;
}
if (generatedUploadUid) {
return generatedUploadUid;
}
return unstructuredFileCount > 0 ? configId : "";
}
export function makeSamplerConfig(
id: string,
samplerType: SamplerType,
@ -368,6 +408,9 @@ export function makeSeedConfig(
hf_token: "",
hf_endpoint: "https://huggingface.co",
local_file_name: "",
...(seedSourceType === "unstructured"
? { unstructured_upload_uid: makeUnstructuredUploadUid() }
: {}),
unstructured_file_ids: [],
unstructured_file_names: [],
unstructured_file_sizes: [],

View file

@ -16,11 +16,7 @@ import type {
} from "../../types";
import { buildEdges } from "./edges";
import { isRecord, parseJson, readString } from "./helpers";
import {
parseColumn,
parseModelConfig,
parseModelProvider,
} from "./parsers";
import { parseColumn, parseModelConfig, parseModelProvider } from "./parsers";
import { parseSeedConfig } from "./parsers/seed-config-parser";
import { buildNodes, parseUi } from "./ui";
import type { ImportResult } from "./types";
@ -43,6 +39,7 @@ type UiInput = {
seed_drop_columns?: unknown;
seed_preview_rows?: unknown;
local_file_name?: unknown;
unstructured_upload_uid?: unknown;
unstructured_file_ids?: unknown;
unstructured_file_names?: unknown;
unstructured_file_sizes?: unknown;
@ -51,6 +48,10 @@ type UiInput = {
advanced_open_by_node?: unknown;
};
type ImportRecipePayloadOptions = {
preserveUnstructuredUploads?: boolean;
};
type UiMarkdownNoteNode = {
name: string;
markdown: string;
@ -90,7 +91,7 @@ function parseProcessors(input: unknown): RecipeProcessorConfig[] {
? templateRaw
: isRecord(templateRaw)
? JSON.stringify(templateRaw, null, 2)
: "{\n \"text\": \"{{ column_name }}\"\n}";
: '{\n "text": "{{ column_name }}"\n}';
processors.push({
id: `p${index + 1}`,
// biome-ignore lint/style/useNamingConvention: api schema
@ -135,9 +136,7 @@ function parseSeedDropColumns(input: unknown): string[] {
return Array.from(values);
}
function parseMcpProviders(
input: unknown,
): Map<string, LlmMcpProviderConfig> {
function parseMcpProviders(input: unknown): Map<string, LlmMcpProviderConfig> {
const providers = new Map<string, LlmMcpProviderConfig>();
if (!Array.isArray(input)) {
return providers;
@ -156,13 +155,12 @@ function parseMcpProviders(
const args = Array.isArray(item.args)
? item.args.map((value) => String(value))
: [];
const envPairs =
isRecord(item.env)
? Object.entries(item.env).map(([key, value]) => ({
key: String(key),
value: String(value),
}))
: [];
const envPairs = isRecord(item.env)
? Object.entries(item.env).map(([key, value]) => ({
key: String(key),
value: String(value),
}))
: [];
providers.set(name, {
id: `mcp-${index + 1}`,
name,
@ -209,7 +207,8 @@ function parseToolConfigs(input: unknown): Map<string, LlmToolConfig> {
allow_tools: allowTools,
// biome-ignore lint/style/useNamingConvention: api schema
max_tool_call_turns:
item.max_tool_call_turns === null || item.max_tool_call_turns === undefined
item.max_tool_call_turns === null ||
item.max_tool_call_turns === undefined
? "5"
: String(item.max_tool_call_turns),
// biome-ignore lint/style/useNamingConvention: api schema
@ -257,7 +256,9 @@ function parseUiMarkdownNoteNodes(input: unknown): UiMarkdownNoteNode[] {
return noteNodes;
}
function parseUiToolProfileNodes(input: unknown): Map<string, Record<string, string[]>> {
function parseUiToolProfileNodes(
input: unknown,
): Map<string, Record<string, string[]>> {
const toolProfiles = new Map<string, Record<string, string[]>>();
if (!Array.isArray(input)) {
return toolProfiles;
@ -312,9 +313,15 @@ function parseAdvancedOpenByNode(input: unknown): Record<string, boolean> {
return out;
}
type AdvancedOpenConfig = LlmConfig | SamplerConfig | SeedConfig | ValidatorConfig;
type AdvancedOpenConfig =
| LlmConfig
| SamplerConfig
| SeedConfig
| ValidatorConfig;
function isAdvancedOpenConfig(config: NodeConfig): config is AdvancedOpenConfig {
function isAdvancedOpenConfig(
config: NodeConfig,
): config is AdvancedOpenConfig {
return (
config.kind === "llm" ||
config.kind === "sampler" ||
@ -350,7 +357,8 @@ function buildToolProfileConfig(
.map((providerName) => mcpProvidersByName.get(providerName))
.flatMap((provider) => (provider ? [cloneMcpProvider(provider)] : [])),
// biome-ignore lint/style/useNamingConvention: ui schema
fetched_tools_by_provider: fetchedToolsByProfileName.get(canonical.tool_alias) ?? {},
fetched_tools_by_provider:
fetchedToolsByProfileName.get(canonical.tool_alias) ?? {},
// biome-ignore lint/style/useNamingConvention: api schema
allow_tools: [...(canonical.allow_tools ?? [])],
// biome-ignore lint/style/useNamingConvention: api schema
@ -360,7 +368,10 @@ function buildToolProfileConfig(
};
}
export function importRecipePayload(input: string): ImportResult {
export function importRecipePayload(
input: string,
options: ImportRecipePayloadOptions = {},
): ImportResult {
const parsed = parseJson(input);
if (!parsed.data || !isRecord(parsed.data)) {
return {
@ -369,9 +380,9 @@ export function importRecipePayload(input: string): ImportResult {
};
}
const recipe = (isRecord(parsed.data.recipe)
? parsed.data.recipe
: parsed.data) as RecipeInput;
const recipe = (
isRecord(parsed.data.recipe) ? parsed.data.recipe : parsed.data
) as RecipeInput;
const ui = isRecord(parsed.data.ui) ? (parsed.data.ui as UiInput) : null;
if (!Array.isArray(recipe.columns)) {
@ -410,21 +421,36 @@ export function importRecipePayload(input: string): ImportResult {
.map((row) => ({ ...row }))
: undefined;
const uiLocalFileName = readString(ui?.local_file_name) ?? undefined;
// Preserve file IDs/names from saved recipes (cleared at share time by sanitizeSeedForShare)
const uiUnstructuredFileIds: string[] = Array.isArray(ui?.unstructured_file_ids)
? (ui.unstructured_file_ids as string[]).filter((v): v is string => typeof v === "string")
: [];
const uiUnstructuredFileNames: string[] = Array.isArray(ui?.unstructured_file_names)
? (ui.unstructured_file_names as string[]).filter((v): v is string => typeof v === "string")
: [];
const uiUnstructuredFileSizes: number[] = Array.isArray(ui?.unstructured_file_sizes)
? (ui.unstructured_file_sizes as number[]).filter((v): v is number => typeof v === "number")
: [];
const preserveUnstructuredUploads =
options.preserveUnstructuredUploads === true;
const uiUnstructuredUploadUid = preserveUnstructuredUploads
? (readString(ui?.unstructured_upload_uid) ?? undefined)
: undefined;
const uiUnstructuredFileIds: string[] =
preserveUnstructuredUploads && Array.isArray(ui?.unstructured_file_ids)
? (ui.unstructured_file_ids as string[]).filter(
(v): v is string => typeof v === "string",
)
: [];
const uiUnstructuredFileNames: string[] =
preserveUnstructuredUploads && Array.isArray(ui?.unstructured_file_names)
? (ui.unstructured_file_names as string[]).filter(
(v): v is string => typeof v === "string",
)
: [];
const uiUnstructuredFileSizes: number[] =
preserveUnstructuredUploads && Array.isArray(ui?.unstructured_file_sizes)
? (ui.unstructured_file_sizes as number[]).filter(
(v): v is number => typeof v === "number",
)
: [];
const uiUnstructuredChunkSize = readStringNumber(ui?.unstructured_chunk_size);
const uiUnstructuredChunkOverlap = readStringNumber(
ui?.unstructured_chunk_overlap,
);
const uiAdvancedOpenByNode = parseAdvancedOpenByNode(ui?.advanced_open_by_node);
const uiAdvancedOpenByNode = parseAdvancedOpenByNode(
ui?.advanced_open_by_node,
);
const uiMarkdownNotes = parseUiMarkdownNoteNodes(ui?.nodes);
const uiToolProfilesByName = parseUiToolProfileNodes(ui?.nodes);
@ -459,11 +485,13 @@ export function importRecipePayload(input: string): ImportResult {
: payloadSeedDropColumns,
seed_preview_rows: uiSeedPreviewRows,
local_file_name: uiLocalFileName,
unstructuredUploadUid: uiUnstructuredUploadUid,
unstructuredFileIds: uiUnstructuredFileIds,
unstructuredFileNames: uiUnstructuredFileNames,
unstructuredFileSizes: uiUnstructuredFileSizes,
unstructured_chunk_size: uiUnstructuredChunkSize,
unstructured_chunk_overlap: uiUnstructuredChunkOverlap,
preserveUnstructuredUploads,
});
if (seedConfig) {
applyAdvancedOpen(seedConfig, uiAdvancedOpenByNode);
@ -567,12 +595,7 @@ export function importRecipePayload(input: string): ImportResult {
const { layouts, auxNodes, edges: uiEdges, layoutDirection } = parseUi(ui);
const resolvedLayoutDirection = layoutDirection ?? "LR";
const nodes = buildNodes(configs, layouts);
const edges = buildEdges(
configs,
nameToId,
uiEdges,
resolvedLayoutDirection,
);
const edges = buildEdges(configs, nameToId, uiEdges, resolvedLayoutDirection);
const auxNodePositions = Object.fromEntries(
auxNodes.flatMap((item) => {
const llmId = nameToId.get(item.llm);
@ -583,10 +606,7 @@ export function importRecipePayload(input: string): ImportResult {
}),
);
const maxY = nodes.reduce(
(acc, node) => Math.max(acc, node.position.y),
0,
);
const maxY = nodes.reduce((acc, node) => Math.max(acc, node.position.y), 0);
return {
errors: [],

View file

@ -197,17 +197,26 @@ export function parseSeedConfig(
seed_drop_columns?: string[];
seed_preview_rows?: Record<string, unknown>[];
local_file_name?: string;
unstructuredUploadUid?: string;
unstructuredFileIds?: string[];
unstructuredFileNames?: string[];
unstructuredFileSizes?: number[];
unstructured_chunk_size?: string;
unstructured_chunk_overlap?: string;
preserveUnstructuredUploads?: boolean;
},
): SeedConfig | null {
if (!seedConfigRaw) {
return null;
}
const parsed = parseSeedSettings(seedConfigRaw);
const parsed = { ...parseSeedSettings(seedConfigRaw) };
if (
parsed.seed_source_type === "unstructured" &&
options?.preserveUnstructuredUploads !== true
) {
parsed.hf_path = "";
parsed.resolved_paths = [];
}
let sourceType: SeedSourceType = "hf";
if (parsed.seed_source_type === "hf") {
sourceType = "hf";
@ -230,6 +239,9 @@ export function parseSeedConfig(
...(options?.local_file_name !== undefined
? { local_file_name: options.local_file_name }
: {}),
...(options?.unstructuredUploadUid
? { unstructured_upload_uid: options.unstructuredUploadUid }
: {}),
...(options?.unstructuredFileIds !== undefined
? { unstructured_file_ids: options.unstructuredFileIds }
: {}),

View file

@ -440,6 +440,9 @@ export function buildRecipePayload(
unstructured_file_names: firstSeed.unstructured_file_names,
unstructured_file_sizes: firstSeed.unstructured_file_sizes,
}),
...(firstSeed?.unstructured_upload_uid?.trim() && {
unstructured_upload_uid: firstSeed.unstructured_upload_uid,
}),
...(firstSeed &&
firstSeed.unstructured_chunk_size !== undefined && {
unstructured_chunk_size: firstSeed.unstructured_chunk_size,

View file

@ -71,6 +71,8 @@ export type RecipePayload = {
seed_preview_rows?: Record<string, unknown>[];
local_file_name?: string;
// biome-ignore lint/style/useNamingConvention: api schema
unstructured_upload_uid?: string;
// biome-ignore lint/style/useNamingConvention: api schema
unstructured_file_ids?: string[];
// biome-ignore lint/style/useNamingConvention: api schema
unstructured_file_names?: string[];

View file

@ -7,6 +7,7 @@ export {
savePersonalization,
} from "./api/personalization";
export { setTheme, useTheme } from "./stores/theme-store";
export { useMonitorOverlayStore } from "./stores/monitor-overlay-store";
export type {
Personalization,
PersonalizationAppearance,

View file

@ -409,7 +409,6 @@ export const en = {
description: "Access Unsloth via the OpenAI-compatible API.",
readDocs: "Read the API docs",
noAccess: "No API access yet.",
newBadge: "New",
accessTokens: "Access tokens",
loadError: "Couldn't load API access.",
createError: "Couldn't create access token.",

View file

@ -296,7 +296,6 @@ export const ja = {
description: "OpenAI互換 API を介して Unsloth にアクセスします。",
readDocs: "API ドキュメントを読む",
noAccess: "まだ API アクセス権がありません。",
newBadge: "新規",
accessTokens: "アクセストークン",
loadError: "API アクセス権を読み込めませんでした。",
createError: "アクセストークンを作成できませんでした。",

View file

@ -363,7 +363,6 @@ export const ptBR = {
"Acesse o Unsloth por meio da API compatível com OpenAI.",
readDocs: "Leia a documentação da API",
noAccess: "Nenhum acesso à API ainda.",
newBadge: "Novo",
accessTokens: "Tokens de acesso",
loadError: "Não foi possível carregar o acesso à API.",
createError: "Não foi possível criar o token de acesso.",

View file

@ -267,7 +267,6 @@ export const zhCN = {
description: "通过兼容 OpenAI 的 API 以编程方式访问 Unsloth。",
readDocs: "阅读 API 文档",
noAccess: "还没有 API 访问权限。",
newBadge: "新",
accessTokens: "访问 token",
loadError: "无法加载 API 访问权限。",
createError: "无法创建访问 token。",

View file

@ -10,6 +10,7 @@ import argparse
import atexit
import errno
import fnmatch
import glob
import hashlib
import json
import os
@ -265,6 +266,7 @@ class HostInfo:
has_physical_nvidia: bool
has_usable_nvidia: bool
has_rocm: bool = False
has_intel_gpu: bool = False
rocm_gfx_target: str | None = None
# (major, minor) from platform.mac_ver(); None off macOS or if unparseable.
# Skips a macos prebuilt whose minimum-OS exceeds this host.
@ -1284,162 +1286,6 @@ def synthetic_checksums_for_release(
)
def parse_direct_linux_release_bundle(
repo: str, release: dict[str, Any]
) -> PublishedReleaseBundle | None:
release_tag = release.get("tag_name")
if not isinstance(release_tag, str) or not release_tag:
return None
assets = release_asset_map(release)
artifacts: list[PublishedLlamaArtifact] = []
inferred_labels: list[str] = []
linux_asset_re = re.compile(
r"^app-(?P<label>.+)-(?P<target>linux-x64(?:-cpu)?|linux-x64-cuda\d+-(?:older|newer|portable))\.tar\.gz$"
)
for asset_name in sorted(assets):
match = linux_asset_re.fullmatch(asset_name)
if not match:
continue
inferred_labels.append(match.group("label"))
target = match.group("target")
if target in {"linux-x64", "linux-x64-cpu"}:
artifacts.append(
PublishedLlamaArtifact(
asset_name = asset_name,
install_kind = "linux-cpu",
runtime_line = None,
coverage_class = None,
supported_sms = [],
min_sm = None,
max_sm = None,
bundle_profile = None,
rank = 1000,
)
)
continue
bundle_profile = target.removeprefix("linux-x64-")
profile = _resolve_linux_bundle_profile(bundle_profile)
if profile is None:
continue
artifacts.append(
PublishedLlamaArtifact(
asset_name = asset_name,
install_kind = "linux-cuda",
runtime_line = str(profile["runtime_line"]),
coverage_class = str(profile["coverage_class"]),
supported_sms = [str(value) for value in profile["supported_sms"]],
min_sm = int(profile["min_sm"]),
max_sm = int(profile["max_sm"]),
bundle_profile = bundle_profile,
rank = int(profile["rank"]),
)
)
if not artifacts:
return None
upstream_tag = (
release_tag
if is_release_tag_like(release_tag)
else inferred_labels[0]
if len(set(inferred_labels)) == 1 and inferred_labels
else release_tag
)
selection_log = [
f"published_release: repo={repo}",
f"published_release: tag={release_tag}",
f"published_release: upstream_tag={upstream_tag}",
"published_release: direct_asset_scan=linux",
]
return PublishedReleaseBundle(
repo = repo,
release_tag = release_tag,
upstream_tag = upstream_tag,
assets = assets,
manifest_asset_name = DEFAULT_PUBLISHED_MANIFEST_ASSET,
artifacts = artifacts,
selection_log = selection_log,
)
def direct_linux_release_plan(
release: dict[str, Any], host: HostInfo, repo: str, requested_tag: str
) -> InstallReleasePlan | None:
bundle = parse_direct_linux_release_bundle(repo, release)
if bundle is None:
return None
if not direct_release_matches_request(
release_tag = bundle.release_tag,
llama_tag = bundle.upstream_tag,
requested_tag = requested_tag,
):
return None
attempts: list[AssetChoice] = []
if host.has_usable_nvidia:
# Prefer the cudart major Studio loads at runtime (torch's bundled
# libcudart), not the newest on disk. Otherwise a stray cuda13
# runtime outranks the torch cuda12 the binary links against.
torch_preference = detect_torch_cuda_runtime_preference(host)
selection = linux_cuda_choice_from_release(
host,
bundle,
preferred_runtime_line = torch_preference.runtime_line,
selection_preamble = torch_preference.selection_log,
)
if selection is not None:
attempts.extend(selection.attempts)
elif not host.has_rocm:
# A ROCm-only host gets no CPU asset: leaving attempts empty lets the
# raise below trigger a HIP source build instead of shipping a CPU
# binary on a GPU host (this ggml-org path has no per-gfx ROCm asset).
cpu_choice = published_asset_choice_for_kind(bundle, "linux-cpu")
if cpu_choice is not None:
attempts.append(cpu_choice)
# NVIDIA hosts whose CUDA selection produced nothing fall through to the
# raise below (mirroring the ROCm policy above): the caller then walks
# back to an older release that still ships a usable CUDA line instead of
# silently installing a CPU binary on a GPU host. Today's walk-back only
# works because partial releases ship no CPU bundle; this keeps it working
# if a future partial release does.
if not attempts:
raise PrebuiltFallback("no compatible Linux prebuilt asset was found")
approved_checksums = synthetic_checksums_for_release(
repo,
bundle.release_tag,
bundle.upstream_tag,
)
resolved_upstream_tag = bundle.upstream_tag
if DEFAULT_PUBLISHED_SHA256_ASSET in bundle.assets and not is_release_tag_like(
bundle.upstream_tag
):
approved_checksums = load_approved_release_checksums(repo, bundle.release_tag)
# Require exact source provenance for branch/pull/commit releases.
# Mirrors validated_checksums_for_bundle so incomplete metadata fails
# closed instead of degrading to the legacy branch-as-tag source
# hydration path this PR eliminates.
if (
not approved_checksums.source_commit
or exact_source_archive_hash(approved_checksums) is None
or source_clone_url_from_checksums(approved_checksums) is None
):
raise PrebuiltFallback(
f"approved checksum asset {DEFAULT_PUBLISHED_SHA256_ASSET} for "
f"{repo}@{bundle.release_tag} did not contain exact source provenance"
)
attempts = apply_approved_hashes(attempts, approved_checksums)
return InstallReleasePlan(
requested_tag = requested_tag,
llama_tag = resolved_upstream_tag,
release_tag = bundle.release_tag,
attempts = attempts,
approved_checksums = approved_checksums,
)
def direct_upstream_release_plan(
release: dict[str, Any], host: HostInfo, repo: str, requested_tag: str
) -> InstallReleasePlan | None:
@ -1482,6 +1328,24 @@ def direct_upstream_release_plan(
install_kind = "windows-hip",
)
)
# Intel (or other non-NVIDIA/non-AMD) GPU: use the Vulkan prebuilt. Gate
# on no PHYSICAL NVIDIA (not just no usable one): a host that hid NVIDIA
# via CUDA_VISIBLE_DEVICES must not reach Vulkan, which ignores that mask
# and could enumerate the reserved card. Falls through to CPU below.
elif host.has_intel_gpu and not host.has_physical_nvidia:
vulkan_asset = f"llama-{release_tag}-bin-win-vulkan-x64.zip"
vulkan_url = assets.get(vulkan_asset)
if vulkan_url:
attempts.append(
AssetChoice(
repo = repo,
tag = release_tag,
name = vulkan_asset,
url = vulkan_url,
source_label = "upstream",
install_kind = "windows-vulkan",
)
)
cpu_asset = f"llama-{release_tag}-bin-win-cpu-x64.zip"
cpu_url = assets.get(cpu_asset)
if cpu_url:
@ -1545,6 +1409,23 @@ def direct_upstream_release_plan(
# ROCm hosts are excluded: this ggml-org path ships no per-gfx ROCm
# asset, so they fall through to the empty-attempts raise (HIP source
# build) rather than silently getting a CPU binary on a GPU host.
# Intel (or other non-NVIDIA/non-AMD) GPU: use the Vulkan prebuilt. The
# elif already excludes usable NVIDIA and ROCm; also require no PHYSICAL
# NVIDIA so a CUDA-hidden card isn't reached through Vulkan (CPU below).
if host.has_intel_gpu and not host.has_physical_nvidia:
vulkan_asset = f"llama-{release_tag}-bin-ubuntu-vulkan-x64.tar.gz"
vulkan_url = assets.get(vulkan_asset)
if vulkan_url:
attempts.append(
AssetChoice(
repo = repo,
tag = release_tag,
name = vulkan_asset,
url = vulkan_url,
source_label = "upstream",
install_kind = "linux-vulkan",
)
)
asset_name = f"llama-{release_tag}-bin-ubuntu-x64.tar.gz"
asset_url = assets.get(asset_name)
if asset_url:
@ -1564,6 +1445,23 @@ def direct_upstream_release_plan(
# selector returned 0 attempts and the installer fell back to a
# source build on every Linux ARM64 host (DGX Spark, Ampere
# Altra, GitHub-hosted ubuntu-24.04-arm runners, etc.).
# Intel (or other non-NVIDIA/non-AMD) GPU: prefer the Vulkan prebuilt,
# mirroring the x86_64 branch. Upstream ships bin-ubuntu-vulkan-arm64.
# No physical NVIDIA: don't reach a CUDA-hidden card through Vulkan.
if host.has_intel_gpu and not host.has_physical_nvidia and not host.has_rocm:
vulkan_asset = f"llama-{release_tag}-bin-ubuntu-vulkan-arm64.tar.gz"
vulkan_url = assets.get(vulkan_asset)
if vulkan_url:
attempts.append(
AssetChoice(
repo = repo,
tag = release_tag,
name = vulkan_asset,
url = vulkan_url,
source_label = "upstream",
install_kind = "linux-vulkan",
)
)
asset_name = f"llama-{release_tag}-bin-ubuntu-arm64.tar.gz"
asset_url = assets.get(asset_name)
if asset_url:
@ -2864,6 +2762,64 @@ def _pick_rocm_gfx_target(out: str) -> str | None:
return _tokens[0]
# Display-adapter device class: one NNNN subkey per installed display driver
# config, each carrying the driver's DriverDesc and PCI MatchingDeviceId.
_WINDOWS_DISPLAY_CLASS_KEY = (
r"SYSTEM\CurrentControlSet\Control\Class\{4d36e968-e325-11ce-bfc1-08002be10318}"
)
def windows_intel_gpu_in_registry() -> bool:
"""Whether the Windows registry lists an Intel display adapter.
In-process Windows counterpart of the Linux DRM vendor-id check (0x8086),
with weaker semantics: the class key lists installed display-driver
configs, which can outlive removed hardware, where sysfs lists present
devices. A stale Intel entry at worst routes to the upstream Vulkan
prebuilt instead of the fork CPU bundle: inference still works (the
Vulkan build runs on CPU when no Vulkan device exists), at the cost of
fork-only extras such as the DiffusionGemma visual server. detect_host's
PowerShell + WMI probe can silently miss a real Intel GPU: a cold
powershell.exe start plus the first CIM query routinely exceeds the 15s
budget on hosts with slow AV scanning or a degraded WMI repository, and
the probe swallows the timeout (#4452, Arc A770 routed to the CPU
prebuilt). Reading the display-adapter class key needs no subprocess and
answers in microseconds. Matches the PCI vendor id in MatchingDeviceId
(ven_8086) or an Intel DriverDesc.
"""
try:
import winreg
except ImportError:
return False
try:
with winreg.OpenKey(winreg.HKEY_LOCAL_MACHINE, _WINDOWS_DISPLAY_CLASS_KEY) as class_key:
for index in range(winreg.QueryInfoKey(class_key)[0]):
try:
name = winreg.EnumKey(class_key, index)
if not name.isdigit():
# "Properties" is ACL-restricted and not an adapter.
continue
with winreg.OpenKey(class_key, name) as adapter_key:
for value_name, needle in (
("MatchingDeviceId", "ven_8086"),
("DriverDesc", "intel"),
):
try:
value, _ = winreg.QueryValueEx(adapter_key, value_name)
except OSError:
continue
if needle in str(value).lower():
return True
except OSError:
continue
except Exception:
# Advisory probe: any unexpected failure must degrade to the CIM
# fallback, never crash the installer (mirrors detect_host's own
# swallow around the CIM probe).
return False
return False
def detect_host() -> HostInfo:
system = platform.system()
machine = platform.machine().lower()
@ -3075,6 +3031,46 @@ def detect_host() -> HostInfo:
# Note: amdhip64.dll presence alone is NOT treated as GPU evidence
# since the HIP SDK can be installed without an AMD GPU.
# Detect an Intel GPU; gates the Vulkan prebuilt. Linux reads the DRM sysfs
# vendor id (0x8086); Windows reads the display-adapter registry class,
# then falls back to the WMI video controller list. Only probed with no
# usable NVIDIA and no ROCm (matching the Vulkan branches), keeping the
# probe (notably the Windows powershell call) off that path.
has_intel_gpu = False
if not has_usable_nvidia and not has_rocm:
if is_linux:
for _vendor_file in glob.glob("/sys/class/drm/card*/device/vendor"):
try:
with open(_vendor_file) as _vf:
if _vf.read().strip().lower() == "0x8086":
has_intel_gpu = True
break
except OSError:
continue
elif is_windows:
# Registry first (in-process; see windows_intel_gpu_in_registry).
# The CIM query stays as the fallback when the registry shows no
# Intel adapter.
has_intel_gpu = windows_intel_gpu_in_registry()
if not has_intel_gpu:
_ps = shutil.which("powershell") or shutil.which("pwsh")
if _ps:
try:
_result = run_capture(
[
_ps,
"-NoProfile",
"-Command",
"Get-CimInstance Win32_VideoController | "
"Select-Object -ExpandProperty Name",
],
timeout = 15,
)
if _result.returncode == 0 and "intel" in _result.stdout.lower():
has_intel_gpu = True
except Exception:
pass
return HostInfo(
system = system,
machine = machine,
@ -3090,6 +3086,7 @@ def detect_host() -> HostInfo:
has_physical_nvidia = has_physical_nvidia,
has_usable_nvidia = has_usable_nvidia,
has_rocm = has_rocm,
has_intel_gpu = has_intel_gpu,
rocm_gfx_target = rocm_gfx_target,
macos_version = macos_version,
)
@ -3126,6 +3123,7 @@ def _apply_host_overrides(
has_physical_nvidia = False,
has_rocm = False,
rocm_gfx_target = None,
has_intel_gpu = False,
)
gfx = _normalize_forwarded_gfx(override_rocm_gfx)
if gfx:
@ -3866,6 +3864,23 @@ def resolve_upstream_asset_choice(host: HostInfo, llama_tag: str) -> AssetChoice
"falling back to source build with HIP support"
)
# Intel (or other non-NVIDIA/non-AMD) GPU: use the Vulkan prebuilt. No
# physical NVIDIA (not just no usable one): a CUDA-hidden card must not
# be reached through Vulkan, which ignores CUDA_VISIBLE_DEVICES.
if host.has_intel_gpu and not host.has_physical_nvidia and not host.has_rocm:
vulkan_name = f"llama-{llama_tag}-bin-ubuntu-vulkan-x64.tar.gz"
if vulkan_name in upstream_assets:
log(f"Intel GPU detected -- using upstream Vulkan prebuilt {vulkan_name}")
return AssetChoice(
repo = UPSTREAM_REPO,
tag = llama_tag,
name = vulkan_name,
url = upstream_assets[vulkan_name],
source_label = "upstream",
install_kind = "linux-vulkan",
)
log("Intel GPU detected but no Vulkan prebuilt found -- falling back to CPU")
upstream_name = f"llama-{llama_tag}-bin-ubuntu-x64.tar.gz"
if upstream_name not in upstream_assets:
raise PrebuiltFallback("upstream Linux CPU asset was not found")
@ -3908,6 +3923,24 @@ def resolve_upstream_asset_choice(host: HostInfo, llama_tag: str) -> AssetChoice
)
log("AMD ROCm detected on Windows but no HIP prebuilt found -- falling back to CPU")
# Intel (or other non-NVIDIA/non-AMD) GPU on Windows: use Vulkan. No
# physical NVIDIA so a CUDA-hidden card isn't reached through Vulkan.
if host.has_intel_gpu and not host.has_physical_nvidia and not host.has_rocm:
vulkan_name = f"llama-{llama_tag}-bin-win-vulkan-x64.zip"
if vulkan_name in upstream_assets:
log(
f"Intel GPU detected on Windows -- using upstream Vulkan prebuilt {vulkan_name}"
)
return AssetChoice(
repo = UPSTREAM_REPO,
tag = llama_tag,
name = vulkan_name,
url = upstream_assets[vulkan_name],
source_label = "upstream",
install_kind = "windows-vulkan",
)
log("Intel GPU detected on Windows but no Vulkan prebuilt found -- falling back to CPU")
upstream_name = f"llama-{llama_tag}-bin-win-cpu-x64.zip"
if upstream_name not in upstream_assets:
raise PrebuiltFallback("upstream Windows CPU asset was not found")
@ -4503,6 +4536,7 @@ def runtime_patterns_for_choice(choice: AssetChoice) -> list[str]:
"linux-arm64-cuda",
"linux-rocm",
"linux-arm64",
"linux-vulkan",
}:
return ["llama-server", "llama-quantize", "llama-diffusion-gemma-visual-server", "lib*.so*"]
if choice.install_kind in {"macos-arm64", "macos-x64"}:
@ -4516,6 +4550,7 @@ def runtime_patterns_for_choice(choice: AssetChoice) -> list[str]:
"windows-cpu",
"windows-cuda",
"windows-hip",
"windows-vulkan",
"windows-rocm",
"windows-arm64",
}:
@ -5731,8 +5766,10 @@ def validate_server(
"linux-cuda",
"linux-arm64-cuda",
"linux-rocm",
"linux-vulkan",
"windows-cuda",
"windows-hip",
"windows-vulkan",
"windows-rocm",
"macos-arm64",
}
@ -6354,6 +6391,20 @@ def runtime_payload_health_groups(choice: AssetChoice) -> list[list[str]]:
["libmtmd.so*"],
["libggml-hip.so*"],
]
if choice.install_kind == "linux-vulkan":
return [
["libllama-common.so*"],
["libllama.so*"],
["libggml.so*"],
["libggml-base.so*"],
# Match the sibling globs (linux-cuda/-rocm): x64 bundles ship
# arch-suffixed libggml-cpu-<variant>.so, arm64 may ship a bare
# libggml-cpu.so; the '-' form missed the latter and re-flagged
# the install unhealthy on every check.
["libggml-cpu*.so*"],
["libmtmd.so*"],
["libggml-vulkan.so*"],
]
if choice.install_kind in {"windows-cpu", "windows-arm64"}:
return [["llama.dll"]]
if choice.install_kind == "windows-cuda":
@ -6373,6 +6424,8 @@ def runtime_payload_health_groups(choice: AssetChoice) -> list[list[str]]:
return groups
if choice.install_kind in {"windows-hip", "windows-rocm"}:
return [["llama.dll"], ["*hip*.dll"]]
if choice.install_kind == "windows-vulkan":
return [["llama.dll"], ["ggml-vulkan.dll"]]
return []
@ -6654,6 +6707,89 @@ def validate_prebuilt_attempts(
raise PrebuiltFallback("no prebuilt bundle passed validation")
def force_vulkan_requested() -> bool:
"""Whether UNSLOTH_FORCE_VULKAN opts this host into the Vulkan llama.cpp
prebuilt instead of its detected CUDA/ROCm backend (e.g. so an AMD user can
run the Vulkan build for inference). Scoped to the llama.cpp backend; the
torch/training stack installs separately and still sees the real GPU.
"""
return os.environ.get("UNSLOTH_FORCE_VULKAN", "").strip().lower() in (
"1",
"true",
"yes",
)
def _vulkan_only_host(host: HostInfo) -> HostInfo:
"""Rewrite ``host`` so the asset selectors take their Vulkan branch.
That branch fires on ``has_intel_gpu and not nvidia and not rocm``, so clear
the CUDA/ROCm flags and raise the integrated-GPU flag. The synthetic flag
never leaves install planning -- it only routes the llama.cpp prebuilt
choice, not the torch/training stack.
"""
return dataclasses_replace(
host,
has_usable_nvidia = False,
has_physical_nvidia = False,
has_rocm = False,
has_intel_gpu = True,
)
def _route_to_vulkan_prebuilt(
host: HostInfo, published_repo: str, published_release_tag: str, *, force_cpu: bool
) -> tuple[HostInfo, str, str]:
"""Point a Vulkan-capable host at the upstream ggml-org Vulkan prebuilt.
The unsloth published repo ships only CUDA/ROCm/CPU assets, so Vulkan comes
from UPSTREAM_REPO. Two triggers route here, both suppressed under
--cpu-fallback (the explicit "give me CPU" last resort wins):
* UNSLOTH_FORCE_VULKAN forces Vulkan over the detected CUDA/ROCm backend;
* an auto-detected Intel GPU with NO physical NVIDIA/ROCm -- the purpose
of the has_intel_gpu probe, since the fork manifest ships no Vulkan asset.
Applied by BOTH the install path and the --resolve-prebuilt probe so the
"is a prebuilt available" answer matches what actually gets installed.
Returns the (possibly rewritten) host, repo, and release tag.
"""
forced = force_vulkan_requested()
# Gate auto-routing on no PHYSICAL NVIDIA, not merely no usable one: a mixed
# NVIDIA+Intel host that hides NVIDIA with CUDA_VISIBLE_DEVICES=""/-1 keeps
# has_physical_nvidia=True while has_usable_nvidia goes False. Vulkan ignores
# CUDA_VISIBLE_DEVICES, so auto-routing such a host would let it grab the
# reserved NVIDIA GPU. An explicit UNSLOTH_FORCE_VULKAN still overrides.
auto_intel = host.has_intel_gpu and not host.has_physical_nvidia and not host.has_rocm
if force_cpu or not (forced or auto_intel):
return host, published_repo, published_release_tag
if host.is_macos:
if forced:
log(
"UNSLOTH_FORCE_VULKAN is set but ignored on macOS "
"(Metal is used; there is no Vulkan prebuilt)"
)
return host, published_repo, published_release_tag
if forced:
log(
"UNSLOTH_FORCE_VULKAN is set; installing the upstream Vulkan "
"llama.cpp prebuilt instead of the detected GPU backend"
)
# Forcing may override a detected NVIDIA/ROCm host, so normalize it to
# Vulkan-only; an auto-detected Intel host already is.
host = _vulkan_only_host(host)
else:
log("Intel GPU detected; installing the upstream Vulkan llama.cpp prebuilt")
# Swapping the fork for upstream invalidates a fork release pin: the two use
# different tag namespaces (fork b9596-mix-<sha> vs upstream b9596), so a
# pinned fork tag would make the upstream resolver query a nonexistent
# release and fall back to source. Drop it and let the upstream resolver
# pick by the requested llama tag. A pin already on an explicit upstream repo
# (repo unchanged here) is preserved.
if published_repo != UPSTREAM_REPO:
published_release_tag = ""
return host, UPSTREAM_REPO, published_release_tag
def diffusion_visual_server_backfill_needed(
install_dir: Path, host: HostInfo, choice: AssetChoice
) -> bool:
@ -6696,6 +6832,9 @@ def install_prebuilt(
override_rocm_gfx = override_rocm_gfx,
force_cpu = force_cpu,
)
host, published_repo, published_release_tag = _route_to_vulkan_prebuilt(
host, published_repo, published_release_tag, force_cpu = force_cpu
)
choice: AssetChoice | None = None
try:
with install_lock(install_lock_path(install_dir)):
@ -6708,7 +6847,9 @@ def install_prebuilt(
f"no existing llama.cpp install detected at {install_dir}; performing fresh prebuilt install"
)
# Single resolver: every fork host selects from the release manifest;
# an explicit ggml-org override selects by asset filename instead.
# an explicit ggml-org override selects by asset filename instead. A
# forced-Vulkan host already has published_repo pointed at
# UPSTREAM_REPO above, so the resolver takes the Vulkan asset branch.
requested_tag, release_plans = resolve_simple_install_release_plans(
llama_tag,
host,
@ -6994,10 +7135,14 @@ def main() -> int:
override_rocm_gfx = args.rocm_gfx,
force_cpu = args.cpu_fallback,
)
repo = args.published_repo
# Same Vulkan routing the install path applies, so the probe's answer
# matches what would install (an Intel/forced-Vulkan host -> upstream).
host, repo, release_tag = _route_to_vulkan_prebuilt(
host, args.published_repo, args.published_release_tag or "", force_cpu = args.cpu_fallback
)
try:
_requested, plans = resolve_simple_install_release_plans(
args.resolve_prebuilt, host, repo, args.published_release_tag or ""
args.resolve_prebuilt, host, repo, release_tag
)
choice = plans[0].attempts[0] if plans and plans[0].attempts else None
if choice is None:

View file

@ -2621,7 +2621,18 @@ function Fast-Install {
param([Parameter(ValueFromRemainingArguments=$true)]$Args_)
if ($UseUv) {
$VenvPy = (Get-Command python).Source
$result = & uv pip install --python $VenvPy @Args_ 2>&1
# An explicit --index-url must win. Inherited uv index env vars otherwise
# override it and pull CPU torch over the CUDA/ROCm build (#6898), so drop
# them only for index-pinned installs; mirrors still apply elsewhere.
$saved = @{}
if (@($Args_) -contains '--index-url') {
foreach ($n in 'UV_DEFAULT_INDEX', 'UV_INDEX_URL', 'UV_INDEX', 'UV_EXTRA_INDEX_URL') {
$saved[$n] = [Environment]::GetEnvironmentVariable($n)
Remove-Item "Env:$n" -ErrorAction SilentlyContinue
}
}
try { $result = & uv pip install --python $VenvPy @Args_ 2>&1 }
finally { foreach ($n in $saved.Keys) { if ($null -ne $saved[$n]) { Set-Item "Env:$n" $saved[$n] } } }
if ($LASTEXITCODE -eq 0) { return }
}
& python -m pip install @Args_ 2>&1

View file

@ -0,0 +1,122 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team.
"""Regression guard for issue #6881: FastSentenceTransformer must preprocess text
like a stock SentenceTransformer for decoder embedding models. ST 5.x infers a
"message" modality for chat-template models (e.g. Qwen/Qwen3-Embedding), so building
via `Transformer(model_name, ...)` chat-wraps inputs and degrades embeddings;
`_create_transformer_module` uses `Transformer.load(...)` instead.
Layers: test_transformer_load_signature_supports_unsloth_kwargs (fast, runs when ST
is importable) and test_fast_sentence_transformer_matches_stock_st (end-to-end parity,
opt-in via UNSLOTH_EMBEDDING_PARITY_MODEL so default CI is unaffected).
"""
from __future__ import annotations
import inspect
import os
import pytest
def test_transformer_load_signature_supports_unsloth_kwargs():
"""Forwards-compat tripwire: a Hub-capable Transformer.load must accept the kwargs
the #6881 fix passes. Legacy ST 3.x/4.x expose load(input_path); the code falls back
to Transformer(...) there, so mirror that gate and skip."""
models = pytest.importorskip("sentence_transformers.models")
load = getattr(models.Transformer, "load", None)
assert callable(load), (
"sentence_transformers Transformer.load is missing; the #6881 fix in "
"unsloth.models.sentence_transformer._create_transformer_module depends on it."
)
params = inspect.signature(load).parameters
accepts_var_kw = any(p.kind is inspect.Parameter.VAR_KEYWORD for p in params.values())
# Mirror _create_transformer_module's hub_capable gate.
hub_capable = accepts_var_kw or any(k in params for k in ("token", "cache_folder", "revision"))
if not hub_capable:
pytest.skip(
"legacy Transformer.load(input_path); production path falls back to Transformer(...)"
)
unsupported = [
k
for k in ("token", "cache_folder", "revision", "trust_remote_code")
if not (accepts_var_kw or k in params)
]
assert not unsupported, (
f"installed sentence_transformers Transformer.load no longer accepts {unsupported} "
f"and has no **kwargs; update _create_transformer_module (#6881) before it silently "
f"falls back to Transformer(...)."
)
def _probe_texts():
return [
"roasted chickpeas in 20 kg bags",
"The capital of France is Paris.",
"A fast brown fox jumps over the lazy dog.",
"recette de tarte aux pommes traditionnelle",
]
def test_fast_sentence_transformer_matches_stock_st():
"""End-to-end: FastSentenceTransformer embeddings and tokenization must match a
stock SentenceTransformer load of the same checkpoint. Opt-in (needs a model) and
GPU-only (FastSentenceTransformer requires CUDA), so it skips on CPU-only runners."""
model_id = os.environ.get("UNSLOTH_EMBEDDING_PARITY_MODEL")
if not model_id:
pytest.skip(
"set UNSLOTH_EMBEDDING_PARITY_MODEL to a chat-template embedding model "
"(HF id or local path) to run the #6881 parity test"
)
torch = pytest.importorskip("torch")
if not torch.cuda.is_available():
pytest.skip("FastSentenceTransformer requires CUDA; skipping on CPU-only runner")
np = pytest.importorskip("numpy")
pytest.importorskip("sentence_transformers")
from sentence_transformers import SentenceTransformer
device = "cuda"
# Prefer bf16 when the GPU supports it: fp16 overflows to NaN on bf16-native
# embedders such as EmbeddingGemma (Gemma3), which would mask real parity.
dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
texts = _probe_texts()
max_seq_length = 256
# Control FIRST, before importing unsloth, so its global import patches never
# touch the stock reference (mirrors the issue's "restart runtime" repro).
ctrl = SentenceTransformer(model_id, device = device, model_kwargs = {"torch_dtype": dtype})
ctrl.max_seq_length = max_seq_length
ctrl_ids = ctrl.tokenize([texts[0]])["input_ids"][0].tolist()
ctrl_emb = np.asarray(
ctrl.encode(texts, normalize_embeddings = True, batch_size = 8), dtype = np.float32
)
import unsloth # noqa: F401
from unsloth import FastSentenceTransformer
fast = FastSentenceTransformer.from_pretrained(
model_id,
max_seq_length = max_seq_length,
dtype = dtype,
load_in_4bit = False,
load_in_16bit = True,
)
fast_ids = fast.tokenize([texts[0]])["input_ids"][0].tolist()
fast_emb = np.asarray(
fast.encode(texts, normalize_embeddings = True, batch_size = 8), dtype = np.float32
)
# Identical tokenization = no chat-template wrapping slipped in (the #6881 defect).
assert fast_ids == ctrl_ids, (
f"tokenization diverged (chat-template wrapping regressed?):\n"
f" stock: {ctrl_ids}\n fast: {fast_ids}"
)
cos = (ctrl_emb * fast_emb).sum(1) / (
np.linalg.norm(ctrl_emb, axis = 1) * np.linalg.norm(fast_emb, axis = 1)
)
assert float(cos.min()) > 0.99, (
f"embedding parity regressed: min cosine {float(cos.min()):.5f} <= 0.99 "
f"(per-text {[round(float(c), 5) for c in cos]})"
)

View file

@ -0,0 +1,46 @@
import ast
from pathlib import Path
def _load_remove_special_tokens():
# Extract remove_special_tokens without importing unsloth (importing unsloth
# needs unsloth_zoo / a GPU). The function is pure Python and uses no imports,
# so it execs cleanly in an empty namespace.
source = Path(__file__).parents[2] / "unsloth" / "chat_templates.py"
tree = ast.parse(source.read_text(encoding = "utf-8"))
funcs = [
node
for node in tree.body
if isinstance(node, ast.FunctionDef) and node.name == "remove_special_tokens"
]
namespace = {}
module = ast.Module(body = funcs, type_ignores = [])
ast.fix_missing_locations(module)
exec(compile(module, str(source), "exec"), namespace)
return namespace["remove_special_tokens"]
class _StubTokenizer:
def __init__(self, bos_token):
self.bos_token = bos_token
def test_no_bos_tokenizer_does_not_crash():
# Tokenizers such as Qwen2 / Qwen2.5, GPT-2, Falcon and GPT-NeoX have no BOS
# token, so tokenizer.bos_token is None. remove_special_tokens must leave the
# prompt untouched instead of raising
# "TypeError: startswith first arg must be str or a tuple of str, not NoneType".
remove_special_tokens = _load_remove_special_tokens()
assert remove_special_tokens(_StubTokenizer(None), "Hello world") == "Hello world"
def test_double_bos_is_stripped():
# A tokenizer with a BOS token still has a single leading BOS removed.
remove_special_tokens = _load_remove_special_tokens()
assert remove_special_tokens(_StubTokenizer("<s>"), "<s>Hello world") == "Hello world"
def test_prompt_without_leading_bos_unchanged():
# A BOS-bearing tokenizer leaves a prompt that does not start with BOS alone.
remove_special_tokens = _load_remove_special_tokens()
assert remove_special_tokens(_StubTokenizer("<s>"), "Hello world") == "Hello world"

View file

@ -0,0 +1,97 @@
import ast
import re
from pathlib import Path
def _load_formatter_builders():
# Extract _parse_combined_prompt and _create_formatter without importing
# unsloth (importing unsloth needs unsloth_zoo / a GPU). Both are pure
# Python and only use the `re` module.
source = Path(__file__).parents[2] / "unsloth" / "chat_templates.py"
tree = ast.parse(source.read_text(encoding = "utf-8"))
wanted = {"_parse_combined_prompt", "_create_formatter"}
funcs = [
node for node in tree.body if isinstance(node, ast.FunctionDef) and node.name in wanted
]
namespace = {"re": re}
module = ast.Module(body = funcs, type_ignores = [])
ast.fix_missing_locations(module)
exec(compile(module, str(source), "exec"), namespace)
return namespace["_parse_combined_prompt"], namespace["_create_formatter"]
class _StubDataset:
def __init__(self, column_names):
self.column_names = column_names
def _render(merged_prompt, columns, batch):
parse, create = _load_formatter_builders()
possible_columns, final_optional_prompts = parse(merged_prompt, _StubDataset(columns))
processor = create(possible_columns, final_optional_prompts, "text")
return processor(batch)["text"]
def test_optional_block_missing_second_column_does_not_render_none():
# A [[...]] block may reference several columns; only the first gates the
# block. A later column that is None must not render as the literal "None".
merged_prompt = "Location: [[{city}, {country}]] end"
out = _render(
merged_prompt,
["city", "country"],
{"city": ["Paris"], "country": [None]},
)
assert out[0] == "Location: Paris, end"
assert "None" not in out[0]
def test_optional_block_all_columns_present_unchanged():
merged_prompt = "Location: [[{city}, {country}]] end"
out = _render(
merged_prompt,
["city", "country"],
{"city": ["Paris"], "country": ["France"]},
)
assert out[0] == "Location: Paris, France end"
def test_optional_block_gating_column_empty_is_dropped():
# When the gating (first) column is empty the whole block is omitted; this
# behaviour is unchanged by the None coercion.
merged_prompt = "Location: [[{city}, {country}]] end"
out = _render(
merged_prompt,
["city", "country"],
{"city": [""], "country": ["France"]},
)
assert out[0] == "Location: end"
def test_single_column_optional_block_gated_out_on_none():
# Single-column blocks were already gated correctly (the sole column is the
# gate); confirm they stay unaffected.
merged_prompt = "Name: [[{name}]]!"
out = _render(merged_prompt, ["name"], {"name": [None, "Bob"]})
assert out == ["Name: !", "Name: Bob!"]
def test_required_column_none_does_not_render_none():
# A required (non-[[...]]) column that is None must not render as the
# literal "None" either; coercion happens at the row source, so both the
# required and optional branches are covered.
merged_prompt = "Location: {city}, {country} end"
out = _render(
merged_prompt,
["city", "country"],
{"city": ["Paris"], "country": [None]},
)
assert out[0] == "Location: Paris, end"
assert "None" not in out[0]
def test_optional_block_falsy_but_present_gating_value_still_renders():
# The gate keeps a block whenever the first column is not "". A falsy but
# real value (0) must not be treated as absent, so the block still renders.
merged_prompt = "Count: [[{n}]]!"
out = _render(merged_prompt, ["n"], {"n": [0]})
assert out[0] == "Count: 0!"

View file

@ -14,6 +14,7 @@ _TESTS_DIR = pathlib.Path(__file__).resolve().parent.parent # tests/
_REPO_ROOT = _TESTS_DIR.parent # unsloth/
_INSTALL_SH = _REPO_ROOT / "install.sh"
_INSTALL_PS1 = _REPO_ROOT / "install.ps1"
_SETUP_PS1 = _REPO_ROOT / "studio" / "setup.ps1"
_NO_TORCH_RT = _REPO_ROOT / "studio" / "backend" / "requirements" / "no-torch-runtime.txt"
@ -109,6 +110,56 @@ class TestStructuralInstallPs1Unchanged:
assert '"torch>=2.4,<2.11.0"' in self._ps1
class TestInstallPs1UvDefaultIndex:
"""Installer-managed torch indexes must override inherited uv defaults."""
_ps1 = _read(_INSTALL_PS1)
def test_torch_installs_use_default_index(self):
assert "--default-index $TorchIndexUrl" in self._ps1
assert "--default-index $ROCmIndexUrl" in self._ps1
def test_torch_installs_do_not_use_deprecated_index_url(self):
assert "--index-url $TorchIndexUrl" not in self._ps1
assert "--index-url $ROCmIndexUrl" not in self._ps1
def test_torch_installs_neutralize_all_uv_index_env_vars(self):
# Extra-index vars outrank --default-index, so pinned installs must clear them.
for var in ("UV_DEFAULT_INDEX", "UV_INDEX_URL", "UV_INDEX", "UV_EXTRA_INDEX_URL"):
assert var in self._ps1
assert 'Remove-Item "Env:$n"' in self._ps1
class TestSetupPs1FastInstallIndex:
"""setup.ps1 Fast-Install must neutralize inherited uv indexes when pinning."""
_ps1 = _read(_SETUP_PS1)
def test_fast_install_clears_all_uv_index_env_vars(self):
for var in ("UV_DEFAULT_INDEX", "UV_INDEX_URL", "UV_INDEX", "UV_EXTRA_INDEX_URL"):
assert var in self._ps1
# Must truly remove the vars (child sees no value), not set them empty.
assert 'Remove-Item "Env:$n"' in self._ps1
class TestInstallShUvDefaultIndex:
"""Linux/Mac installer torch indexes must override inherited uv defaults."""
_sh = _read(_INSTALL_SH)
def test_torch_installs_use_default_index(self):
assert '--default-index "$TORCH_INDEX_URL"' in self._sh
def test_torch_installs_do_not_use_deprecated_index_url(self):
assert '--index-url "$TORCH_INDEX_URL"' not in self._sh
def test_torch_installs_neutralize_all_uv_index_env_vars(self):
# --default-index installs run with all uv index env vars unset via `env -u`.
assert (
"env -u UV_DEFAULT_INDEX -u UV_INDEX_URL -u UV_INDEX -u UV_EXTRA_INDEX_URL" in self._sh
)
# Group 2 -- Shell snippet tests (bash subprocess, mocked python)
class TestTorchConstraintShell:
"""Test the TORCH_CONSTRAINT block via bash with mocked python minor versions."""

View file

@ -312,7 +312,7 @@ if [ "$SKIP_TORCH" = true ]; then
else
echo "==> Installing PyTorch ($TORCH_INDEX_URL)..."
uv pip install --python "$_VENV_PY" "torch>=2.4,<2.11.0" torchvision torchaudio \
--index-url "$TORCH_INDEX_URL"
--default-index "$TORCH_INDEX_URL"
fi
TORCH_EOF

View file

@ -2684,68 +2684,6 @@ class TestBlackwellCuda124Exclusion:
assert kept == [cpu]
# N.1c3. direct_linux_release_plan -- no silent CPU on NVIDIA hosts
class TestDirectLinuxNvidiaCpuGate:
"""A linux-cpu-only release on an NVIDIA host must raise (caller walks back to a usable CUDA line), not silently CPU-install. CPU-only hosts keep the CPU bundle."""
def _bundle_cpu_only(self):
return make_release(
[
make_artifact(
"llama-b8508-bin-ubuntu-x64.tar.gz",
install_kind = "linux-cpu",
runtime_line = None,
coverage_class = None,
supported_sms = [],
min_sm = None,
max_sm = None,
bundle_profile = None,
),
]
)
def _patch(self, monkeypatch):
monkeypatch.setattr(
INSTALL_LLAMA_PREBUILT,
"parse_direct_linux_release_bundle",
lambda repo, release: self._bundle_cpu_only(),
)
monkeypatch.setattr(
INSTALL_LLAMA_PREBUILT,
"detect_torch_cuda_runtime_preference",
lambda host: CudaRuntimePreference(runtime_line = None, selection_log = []),
)
monkeypatch.setattr(
INSTALL_LLAMA_PREBUILT,
"detected_linux_runtime_lines",
lambda: (["cuda13"], {"cuda13": ["/usr/local/cuda/lib64"]}),
)
def test_nvidia_host_without_cuda_line_raises_for_walkback(self, monkeypatch):
self._patch(monkeypatch)
host = make_host(driver_cuda_version = (13, 1), compute_caps = ["100"])
with pytest.raises(PrebuiltFallback, match = "no compatible Linux prebuilt"):
INSTALL_LLAMA_PREBUILT.direct_linux_release_plan(
{"tag_name": "b8508"}, host, "unslothai/llama.cpp", "latest"
)
def test_cpu_host_still_gets_cpu_bundle(self, monkeypatch):
self._patch(monkeypatch)
host = make_host(
nvidia_smi = None,
driver_cuda_version = None,
compute_caps = [],
has_physical_nvidia = False,
has_usable_nvidia = False,
)
plan = INSTALL_LLAMA_PREBUILT.direct_linux_release_plan(
{"tag_name": "b8508"}, host, "unslothai/llama.cpp", "latest"
)
assert [a.install_kind for a in plan.attempts] == ["linux-cpu"]
class TestLinuxPublishedAttemptsNvidiaCpuGate:
"""Live fork-manifest path: an NVIDIA host whose CUDA selection finds nothing gets an empty attempt list (source-builds with CUDA), not the manifest CPU bundle. CPU-only hosts still get the CPU bundle."""

View file

@ -1264,22 +1264,124 @@ with sync_playwright() as p:
# placeholder, and /api/health goes unreachable shortly after.
# ─────────────────────────────────────────────────────
step("Shutdown via account menu")
# Re-login with NEW2 for a valid /api/shutdown token (CLI rotation
# invalidated the old one). The stale token can make the SPA auth guard
# abort this goto with ERR_ABORTED, or redirect to the same /login URL
# ("interrupted by another navigation"); resolve on domcontentloaded and
# tolerate either -- the pw-field wait below confirms we are on /login.
_tolerated_nav = ("ERR_ABORTED", "interrupted by another navigation")
# Start fresh after the CLI rotation invalidates this browser session.
# Stay in the SAME context: macOS Chromium runs --single-process, where
# closing the last context kills the browser and a second context cannot
# be created. Open the new page before closing the old one; the context
# init script covers the new page.
try:
page.goto(f"{BASE}/login", wait_until = "domcontentloaded", timeout = 60_000)
ctx.clear_cookies()
except Exception as exc:
if not any(t in str(exc) for t in _tolerated_nav):
raise
info(f"goto /login interrupted ({exc!r}); password-field wait will confirm /login")
pw_field = page.locator("#password")
pw_field.wait_for(state = "visible", timeout = 60_000)
pw_field.fill(NEW2)
page.locator('button[type="submit"]').click()
info(f"WARN clearing stale session cookies failed: {exc!r}")
# Auth tokens live in localStorage, and /login's guest guard redirects on
# their mere presence, so drop them before navigating.
try:
page.evaluate(
"['unsloth_auth_token', 'unsloth_auth_refresh_token']"
".forEach((key) => localStorage.removeItem(key))"
)
except Exception as exc:
info(f"WARN clearing stale auth tokens failed: {exc!r}")
_fresh_page = ctx.new_page()
_fresh_page.set_default_timeout(60_000)
_fresh_page.on("pageerror", lambda e: page_errors.append(str(e)))
_fresh_page.on("console", _on_console)
try:
page.close()
except Exception:
pass
page = _fresh_page
# Re-login with NEW2 for a valid /api/shutdown token. Route changes can
# still abort or interrupt this navigation, so the field wait below is the
# final confirmation that we reached /login.
_tolerated_nav = ("ERR_ABORTED", "interrupted by another navigation")
# A slow CI runner can make this re-login navigation time out even with the
# server healthy, so retry the whole goto/wait/fill/submit sequence (mirrors
# the change-password retry above). wait_for_health is a diagnostic pre-gate.
wait_for_health(BASE, timeout = 30.0, info = info)
relogin_err: Exception | None = None
for _relogin_attempt in range(3):
try:
try:
page.goto(f"{BASE}/login", wait_until = "domcontentloaded", timeout = 60_000)
except Exception as exc:
if not any(t in str(exc) for t in _tolerated_nav):
raise
info(f"goto /login interrupted ({exc!r}); password-field wait will confirm /login")
pw_field = page.locator("#password")
pw_field.wait_for(state = "visible", timeout = 60_000)
pw_field.fill(NEW2)
# Wait on the login POST so a transient 4xx/5xx is caught and retried
# here, not swallowed until the out-of-loop composer wait.
status, _ = click_and_wait_for_response(
page,
url_substr = "/api/auth/login",
method = "POST",
do_click = lambda: page.locator('button[type="submit"]').click(),
timeout_ms = 30_000,
info = lambda m: print(f"[ui] {m}", flush = True),
)
if status is not None and status >= 400:
raise AssertionError(
f"login POST returned {status}; see console_errors={console_errors[:1]!r}"
)
relogin_err = None
break
except Exception as e:
relogin_err = e
try:
cur_url = page.url
except Exception:
cur_url = "<page closed>"
print(
f"[ui] re-login attempt {_relogin_attempt + 1} failed: "
f"{type(e).__name__}: {str(e)[:200]}; page.url={cur_url}; "
f"page_errors={len(page_errors)} console_errors={len(console_errors)}",
flush = True,
)
if console_errors:
print(
f"[ui] first console.error: {console_errors[0][:200]!r}",
flush = True,
)
if page_errors:
print(f"[ui] first pageerror: {page_errors[0][:200]!r}", flush = True)
try:
shoot(f"18-relogin-attempt-{_relogin_attempt + 1}-fail")
except Exception:
pass
if _relogin_attempt < 2:
# ERR_NO_BUFFER_SPACE needs the OS to recover socket
# buffers; back off 5s then 15s before retrying.
if "ERR_NO_BUFFER_SPACE" in str(e):
backoff_s = 5 if _relogin_attempt == 0 else 15
print(
f"[ui] ENOBUFS detected; sleeping {backoff_s}s "
f"before retry to let OS recover socket buffers...",
flush = True,
)
time.sleep(backoff_s)
# Replace the page if it died; otherwise next iteration's
# page.goto() handles the reload.
old_page = page
page = recover_or_replace_page(
page,
ctx,
default_timeout_ms = 60_000,
info = lambda m: print(f"[ui] recovery: {m}", flush = True),
)
# A freshly created replacement page loses the pageerror/console
# listeners; re-attach so error tracking survives recovery.
if page is not old_page:
page.on("pageerror", lambda e: page_errors.append(str(e)))
page.on("console", _on_console)
if relogin_err is not None:
raise relogin_err
# Composer mount confirms the rotated session is authenticated. Kept OUTSIDE the
# retry: the loop breaks right after submit, so we never re-goto /login once login
# has set tokens -- that would hit the guest guard, redirect to /chat, and make a
# merely-slow composer look like a broken login.
composer = page.locator('textarea[aria-label="Message input"]')
composer.wait_for(state = "visible", timeout = 60_000)
shoot("18-relogin-with-NEW2")

View file

@ -17,7 +17,7 @@ from pathlib import Path
SOURCE_PATH = Path(__file__).resolve().parents[2] / "studio" / "backend" / "routes" / "inference.py"
SRC = SOURCE_PATH.read_text()
SRC = SOURCE_PATH.read_text(encoding = "utf-8")
_TREE = ast.parse(SRC)
@ -166,16 +166,20 @@ def test_chat_completions_streams_avoid_starlette_task_group():
def test_openai_passthrough_stream_avoids_starlette_task_group():
top = _async_function("_openai_passthrough_stream")
functions = [
_async_function("_openai_passthrough_stream"),
_async_function("_openai_passthrough_stream_admitted"),
]
legacy_calls = []
same_task_calls = 0
for sub in ast.walk(top):
if not (isinstance(sub, ast.Call) and isinstance(sub.func, ast.Name)):
continue
if sub.func.id == "StreamingResponse":
legacy_calls.append(sub.lineno)
if sub.func.id == "_SameTaskStreamingResponse":
same_task_calls += 1
for fn in functions:
for sub in ast.walk(fn):
if not (isinstance(sub, ast.Call) and isinstance(sub.func, ast.Name)):
continue
if sub.func.id == "StreamingResponse":
legacy_calls.append(sub.lineno)
if sub.func.id == "_SameTaskStreamingResponse":
same_task_calls += 1
assert not legacy_calls, (
"OpenAI passthrough streams must use _SameTaskStreamingResponse, "
"not Starlette's legacy task-group StreamingResponse. Lines: "
@ -197,7 +201,7 @@ def test_direct_llama_server_streams_install_disconnect_watcher():
"openai_completions",
"_responses_stream",
"_anthropic_passthrough_stream",
"_openai_passthrough_stream",
"_openai_passthrough_stream_admitted",
}
missing = [
name

View file

@ -0,0 +1,47 @@
"""Regression test for BAD_MAPPINGS redirecting oversized dynamic quants.
get_model_name previously applied BAD_MAPPINGS only to the resolver's output,
but several listed names (the `-unsloth-bnb-4bit` dynamic quants, plus any name
the resolver doesn't map) come back as None, so their BAD_MAPPINGS entries were
dead and the oversized model loaded. Asserting over every entry catches all of
them. The mapper table and the resolver have no heavy imports of their own,
so we exec the import-free mapper module and ast-extract the resolver functions
rather than importing unsloth (which needs a GPU).
"""
import ast
import os
_MODELS = os.path.join(os.path.dirname(__file__), os.pardir, "unsloth", "models")
def _load_get_model_name():
mapper_ns = {}
with open(os.path.join(_MODELS, "mapper.py"), encoding = "utf-8") as f:
exec(compile(f.read(), "mapper.py", "exec"), mapper_ns)
with open(os.path.join(_MODELS, "loader_utils.py"), encoding = "utf-8") as f:
tree = ast.parse(f.read())
namespace = dict(mapper_ns)
namespace["SUPPORTS_FOURBIT"] = True
namespace["_env_says_offline"] = lambda: True
namespace["_get_new_mapper"] = lambda: ({}, {}, {})
wanted = {"__get_model_name", "_resolve_with_mappers", "get_model_name"}
for node in tree.body:
if isinstance(node, ast.Assign) and any(
getattr(target, "id", None) == "BAD_MAPPINGS" for target in node.targets
):
exec(compile(ast.Module([node], []), "<bad_mappings>", "exec"), namespace)
elif isinstance(node, ast.FunctionDef) and node.name in wanted:
exec(compile(ast.Module([node], []), node.name, "exec"), namespace)
return namespace["get_model_name"], namespace["BAD_MAPPINGS"]
def test_bad_mappings_redirect_every_listed_name():
get_model_name, bad_mappings = _load_get_model_name()
assert bad_mappings, "BAD_MAPPINGS should not be empty"
for name, expected in bad_mappings.items():
assert get_model_name(name, load_in_4bit = True) == expected, name

View file

@ -0,0 +1,63 @@
# Copyright 2023-present Daniel Han-Chen & the Unsloth team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""`get_lora_parameters` must not treat a `weight_scale` as a quant state for a weight that is
already dequantized to bf16 (e.g. a compressed-tensors layer at forward time). Otherwise the
bnb fast_gemv / fast_dequantize path reads a missing `absmax` and crashes.
"""
from types import SimpleNamespace
import pytest
import torch
# unsloth.kernels.utils imports bitsandbytes unconditionally, so skip the whole module up
# front on runners without it (e.g. CPU-only) before importing unsloth, otherwise collection
# errors instead of producing a skip. Any other import error still surfaces as a failure.
pytest.importorskip("bitsandbytes")
import unsloth # noqa: F401 (sets UNSLOTH_IS_PRESENT before transformers)
from unsloth.kernels.utils import get_lora_parameters_bias, _FP8_WEIGHT_DTYPES
_FP8 = _FP8_WEIGHT_DTYPES[0] if _FP8_WEIGHT_DTYPES else None
def _proj(weight, weight_scale = None):
proj = SimpleNamespace(weight = weight, bias = None, merged = False)
if weight_scale is not None:
proj.weight_scale = weight_scale
return proj
def test_bf16_weight_scale_not_used_as_quant_state():
"""A bf16 weight carrying a weight_scale (compressed-tensors) -> quant state must be None."""
proj = _proj(torch.randn(4, 4, dtype = torch.bfloat16), torch.rand(2, 2))
W, W_quant = get_lora_parameters_bias(proj)[:2]
assert W_quant is None
def test_fp8_weight_keeps_scale():
"""An actual fp8 weight still resolves its weight_scale as the quant state."""
if _FP8 is None:
pytest.skip("no float8 dtype in this torch build")
scale = torch.rand(2, 2)
proj = _proj(torch.randn(4, 4).to(_FP8), scale)
W, W_quant = get_lora_parameters_bias(proj)[:2]
assert W_quant is scale
def test_plain_bf16_has_no_quant_state():
proj = _proj(torch.randn(4, 4, dtype = torch.bfloat16))
W, W_quant = get_lora_parameters_bias(proj)[:2]
assert W_quant is None

View file

@ -0,0 +1,249 @@
from __future__ import annotations
import ast
from contextlib import nullcontext
from pathlib import Path
from types import SimpleNamespace
import pytest
REPO_ROOT = Path(__file__).resolve().parents[1]
FP8_SOURCE = REPO_ROOT / "unsloth" / "kernels" / "fp8.py"
class _FakeDeviceModule:
def __init__(self, device_count: int) -> None:
self._device_count = device_count
self.device_calls = []
def device_count(self) -> int:
return self._device_count
def device(self, device):
self.device_calls.append(device)
return ("device-context", device)
class _FakeTorch:
Tensor = object
def __init__(
self,
cuda_device_count: int,
xpu_device_count: int = 0,
) -> None:
self.cuda = _FakeDeviceModule(cuda_device_count)
self.xpu = _FakeDeviceModule(xpu_device_count)
class _LaunchVisitor(ast.NodeVisitor):
def __init__(self) -> None:
self.guarded_launches: set[str] = set()
self.unguarded_launches: set[str] = set()
self._inside_fp8_device_context = 0
def visit_With(self, node: ast.With) -> None:
enters_context = any(
isinstance(item.context_expr, ast.Call)
and isinstance(item.context_expr.func, ast.Name)
and item.context_expr.func.id == "_fp8_triton_device_context"
for item in node.items
)
if enters_context:
self._inside_fp8_device_context += 1
for statement in node.body:
self.visit(statement)
if enters_context:
self._inside_fp8_device_context -= 1
def visit_Call(self, node: ast.Call) -> None:
launch_name = self._triton_launch_name(node)
if launch_name is not None:
if self._inside_fp8_device_context:
self.guarded_launches.add(launch_name)
else:
self.unguarded_launches.add(launch_name)
self.generic_visit(node)
@staticmethod
def _triton_launch_name(node: ast.Call) -> str | None:
if isinstance(node.func, ast.Name) and node.func.id == "triton_quantize_fp8_block":
return node.func.id
if not isinstance(node.func, ast.Subscript):
return None
if not isinstance(node.func.value, ast.Name):
return None
return node.func.value.id
def _load_device_context_helper(fake_torch: _FakeTorch):
source = FP8_SOURCE.read_text()
tree = ast.parse(source)
for node in tree.body:
if isinstance(node, ast.FunctionDef) and node.name == "_fp8_triton_device_context":
namespace = {"torch": fake_torch, "nullcontext": nullcontext}
exec(ast.get_source_segment(source, node), namespace)
return namespace["_fp8_triton_device_context"]
raise AssertionError("_fp8_triton_device_context was not found")
def test_fp8_device_context_selects_cuda_tensor_device_on_multi_gpu() -> None:
fake_torch = _FakeTorch(cuda_device_count = 2)
helper = _load_device_context_helper(fake_torch)
tensor = SimpleNamespace(device = SimpleNamespace(type = "cuda"))
context = helper(tensor)
assert context == ("device-context", tensor.device)
assert fake_torch.cuda.device_calls == [tensor.device]
def test_fp8_device_context_is_noop_for_single_cuda_device() -> None:
fake_torch = _FakeTorch(cuda_device_count = 1)
helper = _load_device_context_helper(fake_torch)
tensor = SimpleNamespace(device = SimpleNamespace(type = "cuda"))
context = helper(tensor)
assert isinstance(context, nullcontext)
assert fake_torch.cuda.device_calls == []
def test_fp8_device_context_selects_xpu_tensor_device_on_multi_gpu() -> None:
fake_torch = _FakeTorch(cuda_device_count = 0, xpu_device_count = 2)
helper = _load_device_context_helper(fake_torch)
tensor = SimpleNamespace(device = SimpleNamespace(type = "xpu"))
context = helper(tensor)
assert context == ("device-context", tensor.device)
assert fake_torch.xpu.device_calls == [tensor.device]
def test_fp8_device_context_is_noop_for_single_xpu_device() -> None:
fake_torch = _FakeTorch(cuda_device_count = 0, xpu_device_count = 1)
helper = _load_device_context_helper(fake_torch)
tensor = SimpleNamespace(device = SimpleNamespace(type = "xpu"))
context = helper(tensor)
assert isinstance(context, nullcontext)
assert fake_torch.xpu.device_calls == []
def test_fp8_device_context_is_noop_for_non_cuda_tensor() -> None:
fake_torch = _FakeTorch(cuda_device_count = 8)
helper = _load_device_context_helper(fake_torch)
tensor = SimpleNamespace(device = SimpleNamespace(type = "cpu"))
context = helper(tensor)
assert isinstance(context, nullcontext)
assert fake_torch.cuda.device_calls == []
def test_fp8_triton_launches_enter_tensor_device_context() -> None:
tree = ast.parse(FP8_SOURCE.read_text())
function_names = {node.name for node in ast.walk(tree) if isinstance(node, ast.FunctionDef)}
assert "_fp8_triton_device_context" in function_names
visitor = _LaunchVisitor()
visitor.visit(tree)
expected_launches = {
"weight_dequant_kernel",
"act_quant_kernel",
"_w8a8_block_fp8_matmul",
"triton_quantize_fp8_block",
}
assert expected_launches <= visitor.guarded_launches
assert not (expected_launches & visitor.unguarded_launches)
def _require_two_cuda_devices():
torch = pytest.importorskip("torch")
pytest.importorskip("triton")
if not torch.cuda.is_available() or torch.cuda.device_count() < 2:
pytest.skip("requires at least two CUDA devices")
return torch
def test_weight_dequant_block_runs_on_tensor_device_when_current_device_differs() -> None:
torch = _require_two_cuda_devices()
from unsloth.kernels.fp8 import weight_dequant_block
previous_device = torch.cuda.current_device()
try:
torch.cuda.set_device(0)
x = torch.arange(256 * 256, device = "cuda:1", dtype = torch.float32).reshape(256, 256)
scales = torch.tensor([[1.0, 2.0], [3.0, 4.0]], device = "cuda:1", dtype = torch.float32)
actual = weight_dequant_block(x, scales, block_size = 128, dtype = torch.float32)
expanded_scales = scales.repeat_interleave(128, dim = 0).repeat_interleave(128, dim = 1)
expected = x * expanded_scales
assert actual.device == x.device
assert torch.cuda.current_device() == 0
torch.testing.assert_close(actual, expected)
finally:
torch.cuda.set_device(previous_device)
def test_act_quant_runs_on_tensor_device_when_current_device_differs() -> None:
torch = _require_two_cuda_devices()
if not hasattr(torch, "float8_e4m3fn"):
pytest.skip("requires torch.float8_e4m3fn")
if torch.cuda.get_device_capability(1)[0] < 9:
pytest.skip("requires FP8-capable CUDA hardware")
from unsloth.kernels.fp8 import act_quant
previous_device = torch.cuda.current_device()
try:
torch.cuda.set_device(0)
x = torch.arange(256, device = "cuda:1", dtype = torch.float32).reshape(2, 128)
y, scales = act_quant(x, block_size = 128)
assert y.device == x.device
assert scales.device == x.device
assert torch.cuda.current_device() == 0
finally:
torch.cuda.set_device(previous_device)
def test_w8a8_block_fp8_matmul_triton_runs_on_tensor_device_when_current_device_differs() -> None:
torch = _require_two_cuda_devices()
if not hasattr(torch, "float8_e4m3fn"):
pytest.skip("requires torch.float8_e4m3fn")
if torch.cuda.get_device_capability(1)[0] < 9:
pytest.skip("requires FP8-capable CUDA hardware")
from unsloth.kernels.fp8 import w8a8_block_fp8_matmul_triton
previous_device = torch.cuda.current_device()
try:
torch.cuda.set_device(0)
A = torch.ones((128, 128), device = "cuda:1", dtype = torch.float32).to(torch.float8_e4m3fn)
B = torch.ones((128, 128), device = "cuda:1", dtype = torch.float32).to(torch.float8_e4m3fn)
As = torch.ones((128, 1), device = "cuda:1", dtype = torch.float32)
Bs = torch.ones((1, 1), device = "cuda:1", dtype = torch.float32)
actual = w8a8_block_fp8_matmul_triton(
A,
B,
As,
Bs,
block_size = [128, 128],
output_dtype = torch.float32,
)
expected = torch.full((128, 128), 128.0, device = "cuda:1", dtype = torch.float32)
assert actual.device == A.device
assert torch.cuda.current_device() == 0
torch.testing.assert_close(actual, expected)
finally:
torch.cuda.set_device(previous_device)

View file

@ -0,0 +1,358 @@
# Copyright 2023-present Daniel Han-Chen & the Unsloth team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Restoring dropped block-fp8 `weight_scale_inv` tensors on load (#6200).
Some block-scale fp8 checkpoints leave a Linear (e.g. `mlp.gate_proj`) unconverted, so its raw
quantized values land in a plain bf16 weight and its `weight_scale_inv` is dropped, producing a
garbage un-scaled weight. `_restore_dropped_fp8_scales` dequantizes such orphaned weights in place
using the scale from the checkpoint. Runs offline on CPU with synthetic checkpoints.
"""
import json
import os
import tempfile
from types import SimpleNamespace
import torch
from torch import nn
from safetensors.torch import save_file
# Import unsloth first to set UNSLOTH_IS_PRESENT env var.
import unsloth
from unsloth.models.loader_utils import _restore_dropped_fp8_scales, _FP8_DTYPES
_SHARD = "model-00001-of-00001.safetensors"
_FP8 = _FP8_DTYPES[0] if _FP8_DTYPES else None
def _write_checkpoint(
path,
tensors,
filename = _SHARD,
include_index = True,
):
save_file(tensors, os.path.join(path, filename))
if include_index:
weight_map = {name: filename for name in tensors}
with open(os.path.join(path, "model.safetensors.index.json"), "w") as f:
json.dump({"weight_map": weight_map}, f)
def _fp8_config(block = (2, 2)):
return SimpleNamespace(
quantization_config = {
"quant_method": "fp8",
"weight_block_size": list(block),
}
)
def _fp8_anchor():
"""A module carrying a real fp8 weight, so the model looks like a genuine fp8 load."""
m = nn.Linear(2, 2, bias = False)
m.weight = nn.Parameter(torch.randn(2, 2).to(_FP8), requires_grad = False)
return m
def _bf16_linear(out_f, in_f, raw):
m = nn.Linear(in_f, out_f, bias = False).to(torch.bfloat16)
with torch.no_grad():
m.weight.copy_(raw)
return m
def _expand(scale, block, shape):
bs0, bs1 = block
expanded = scale.repeat_interleave(bs0, dim = 0).repeat_interleave(bs1, dim = 1)
return expanded[: shape[0], : shape[1]]
def test_restore_dequantizes_orphaned_scale():
"""A plain bf16 weight whose scale was dropped is dequantized in place."""
if _FP8 is None:
return
torch.manual_seed(0)
raw = torch.randn(4, 4, dtype = torch.bfloat16)
scale = torch.rand(2, 2, dtype = torch.float32) + 0.1
model = nn.Module()
model.config = _fp8_config((2, 2))
model.anchor = _fp8_anchor()
model.layer = _bf16_linear(4, 4, raw)
with tempfile.TemporaryDirectory() as d:
_write_checkpoint(
d,
{
"layer.weight": raw.to(torch.float32),
"layer.weight_scale_inv": scale,
},
)
restored, skipped = _restore_dropped_fp8_scales(model, d, local_files_only = True)
assert restored == 1
expected = (raw.to(torch.float32) * _expand(scale, (2, 2), (4, 4))).to(torch.bfloat16)
assert torch.equal(model.layer.weight.data, expected)
def test_skips_already_fp8_weight():
"""A correctly converted fp8 weight is skipped, never double-scaled."""
if _FP8 is None:
return
weight = torch.randn(4, 4).to(_FP8)
before = weight.clone()
model = nn.Module()
model.config = _fp8_config((2, 2))
model.layer = nn.Linear(4, 4, bias = False)
model.layer.weight = nn.Parameter(weight, requires_grad = False)
with tempfile.TemporaryDirectory() as d:
_write_checkpoint(d, {"layer.weight_scale_inv": torch.rand(2, 2)})
restored, skipped = _restore_dropped_fp8_scales(model, d, local_files_only = True)
assert restored == 0 and skipped == 1
assert torch.equal(model.layer.weight.data.float(), before.float())
def test_skips_offloaded_meta_weight():
"""A disk-offloaded layer (weight on the meta device) is skipped without error or restore."""
if _FP8 is None:
return
raw = torch.randn(4, 4, dtype = torch.bfloat16)
scale = torch.rand(2, 2, dtype = torch.float32) + 0.1
model = nn.Module()
model.config = _fp8_config((2, 2))
model.anchor = _fp8_anchor()
model.layer = nn.Linear(4, 4, bias = False)
# Simulate an offloaded weight living on the meta device.
model.layer.weight = nn.Parameter(
torch.empty(4, 4, dtype = torch.bfloat16, device = "meta"), requires_grad = False
)
with tempfile.TemporaryDirectory() as d:
_write_checkpoint(
d,
{
"layer.weight": raw.to(torch.float32),
"layer.weight_scale_inv": scale,
},
)
restored, skipped = _restore_dropped_fp8_scales(model, d, local_files_only = True)
assert restored == 0
assert model.layer.weight.device.type == "meta"
def test_noop_when_fully_dequantized():
"""If the model has no fp8 weights at all (e.g. load_in_16bit dequantize), do not rescale."""
raw = torch.randn(4, 4, dtype = torch.bfloat16)
scale = torch.rand(2, 2, dtype = torch.float32) + 0.1
model = nn.Module()
model.config = _fp8_config((2, 2))
model.layer = _bf16_linear(4, 4, raw) # no fp8 anchor -> looks dequantized
with tempfile.TemporaryDirectory() as d:
_write_checkpoint(d, {"layer.weight_scale_inv": scale})
restored, skipped = _restore_dropped_fp8_scales(model, d, local_files_only = True)
assert (restored, skipped) == (0, 0)
assert torch.equal(model.layer.weight.data, raw) # untouched
def test_non_block_divisible_shape():
"""Block scale is expanded then sliced to a non-divisible weight shape."""
if _FP8 is None:
return
raw = torch.randn(3, 4, dtype = torch.bfloat16)
scale = torch.rand(2, 2, dtype = torch.float32) + 0.1
model = nn.Module()
model.config = _fp8_config((2, 2))
model.anchor = _fp8_anchor()
model.layer = _bf16_linear(3, 4, raw) # weight shape [3, 4]
with tempfile.TemporaryDirectory() as d:
_write_checkpoint(d, {"layer.weight_scale_inv": scale})
restored, skipped = _restore_dropped_fp8_scales(model, d, local_files_only = True)
assert restored == 1
expected = (raw.to(torch.float32) * _expand(scale, (2, 2), (3, 4))).to(torch.bfloat16)
assert torch.equal(model.layer.weight.data, expected)
def test_transposed_scale_layout():
"""A scale stored in the transposed block grid is transposed before use."""
if _FP8 is None:
return
raw = torch.randn(4, 2, dtype = torch.bfloat16) # weight [4, 2] -> grid (2, 1)
scale_correct = torch.rand(2, 1, dtype = torch.float32) + 0.1
scale_stored = scale_correct.t().contiguous() # stored transposed as (1, 2)
model = nn.Module()
model.config = _fp8_config((2, 2))
model.anchor = _fp8_anchor()
model.layer = _bf16_linear(4, 2, raw)
with tempfile.TemporaryDirectory() as d:
_write_checkpoint(d, {"layer.weight_scale_inv": scale_stored})
restored, _ = _restore_dropped_fp8_scales(model, d, local_files_only = True)
assert restored == 1
expected = (raw.to(torch.float32) * _expand(scale_correct, (2, 2), (4, 2))).to(torch.bfloat16)
assert torch.equal(model.layer.weight.data, expected)
def test_single_file_checkpoint_without_index():
"""Unsharded model.safetensors (no index) is still scanned for dropped scales."""
if _FP8 is None:
return
raw = torch.randn(4, 4, dtype = torch.bfloat16)
scale = torch.rand(2, 2, dtype = torch.float32) + 0.1
model = nn.Module()
model.config = _fp8_config((2, 2))
model.anchor = _fp8_anchor()
model.layer = _bf16_linear(4, 4, raw)
with tempfile.TemporaryDirectory() as d:
_write_checkpoint(
d, {"layer.weight_scale_inv": scale}, filename = "model.safetensors", include_index = False
)
restored, _ = _restore_dropped_fp8_scales(model, d, local_files_only = True)
assert restored == 1
expected = (raw.to(torch.float32) * _expand(scale, (2, 2), (4, 4))).to(torch.bfloat16)
assert torch.equal(model.layer.weight.data, expected)
def test_scalar_block_size_config():
"""A scalar weight_block_size (not a list) is handled without error."""
if _FP8 is None:
return
raw = torch.randn(4, 4, dtype = torch.bfloat16)
scale = torch.rand(2, 2, dtype = torch.float32) + 0.1
model = nn.Module()
model.config = SimpleNamespace(
quantization_config = {"quant_method": "fp8", "weight_block_size": 2}
)
model.anchor = _fp8_anchor()
model.layer = _bf16_linear(4, 4, raw)
with tempfile.TemporaryDirectory() as d:
_write_checkpoint(d, {"layer.weight_scale_inv": scale})
restored, _ = _restore_dropped_fp8_scales(model, d, local_files_only = True)
assert restored == 1
def test_text_only_prefix_mapping():
"""Checkpoint keys with a language_model prefix match the stripped text-only module names."""
if _FP8 is None:
return
raw = torch.randn(2, 2, dtype = torch.bfloat16)
scale = torch.rand(1, 1, dtype = torch.float32) + 0.1
model = nn.Module()
model.config = _fp8_config((2, 2))
model.anchor = _fp8_anchor()
model.model = nn.Module()
model.model.gate_proj = _bf16_linear(2, 2, raw) # module lacks the language_model prefix
with tempfile.TemporaryDirectory() as d:
# checkpoint key carries the language_model wrapper the text-only load stripped
_write_checkpoint(d, {"model.language_model.gate_proj.weight_scale_inv": scale})
restored, _ = _restore_dropped_fp8_scales(model, d, local_files_only = True)
assert restored == 1
expected = (raw.to(torch.float32) * _expand(scale, (2, 2), (2, 2))).to(torch.bfloat16)
assert torch.equal(model.model.gate_proj.weight.data, expected)
def test_skips_variant_load():
"""A variant load (variant="fp8") is skipped to avoid applying default-checkpoint scales."""
if _FP8 is None:
return
raw = torch.randn(4, 4, dtype = torch.bfloat16)
scale = torch.rand(2, 2, dtype = torch.float32) + 0.1
model = nn.Module()
model.config = _fp8_config((2, 2))
model.anchor = _fp8_anchor()
model.layer = _bf16_linear(4, 4, raw)
with tempfile.TemporaryDirectory() as d:
_write_checkpoint(d, {"layer.weight_scale_inv": scale})
result = _restore_dropped_fp8_scales(model, d, local_files_only = True, variant = "fp8")
assert result == (0, 0)
assert torch.equal(model.layer.weight.data, raw) # untouched
def test_vlm_language_model_model_alias():
"""A checkpoint key language_model.model.* matches a model.language_model.* module."""
if _FP8 is None:
return
raw = torch.randn(2, 2, dtype = torch.bfloat16)
scale = torch.rand(1, 1, dtype = torch.float32) + 0.1
model = nn.Module()
model.config = _fp8_config((2, 2))
model.anchor = _fp8_anchor()
model.model = nn.Module()
model.model.language_model = nn.Module()
model.model.language_model.gate_proj = _bf16_linear(
2, 2, raw
) # -> model.language_model.gate_proj
with tempfile.TemporaryDirectory() as d:
_write_checkpoint(d, {"language_model.model.gate_proj.weight_scale_inv": scale})
restored, _ = _restore_dropped_fp8_scales(model, d, local_files_only = True)
assert restored == 1
expected = (raw.to(torch.float32) * _expand(scale, (2, 2), (2, 2))).to(torch.bfloat16)
assert torch.equal(model.model.language_model.gate_proj.weight.data, expected)
def test_noop_without_scale_keys():
if _FP8 is None:
return
model = nn.Module()
model.config = _fp8_config((2, 2))
model.anchor = _fp8_anchor()
model.layer = _bf16_linear(4, 4, torch.randn(4, 4, dtype = torch.bfloat16))
with tempfile.TemporaryDirectory() as d:
_write_checkpoint(d, {"layer.weight": torch.randn(4, 4)})
assert _restore_dropped_fp8_scales(model, d, local_files_only = True) == (0, 0)
def test_noop_without_index_or_single_file():
if _FP8 is None:
return
model = nn.Module()
model.config = _fp8_config((2, 2))
model.anchor = _fp8_anchor()
model.layer = _bf16_linear(4, 4, torch.randn(4, 4, dtype = torch.bfloat16))
with tempfile.TemporaryDirectory() as d:
assert _restore_dropped_fp8_scales(model, d, local_files_only = True) == (0, 0)
def test_noop_when_not_block_fp8():
"""A non-fp8 (or non-block) quantization config is ignored."""
scale = torch.rand(2, 2)
model = nn.Module()
model.config = SimpleNamespace(quantization_config = {"quant_method": "compressed-tensors"})
model.layer = nn.Linear(4, 4, bias = False)
with tempfile.TemporaryDirectory() as d:
_write_checkpoint(d, {"layer.weight_scale_inv": scale})
assert _restore_dropped_fp8_scales(model, d, local_files_only = True) == (0, 0)

View file

@ -0,0 +1,46 @@
"""Regression test for the duplicate ``unsloth/gemma-2b-bnb-4bit`` key in
``unsloth/models/mapper.py``.
The 4bit instruction-tuned Gemma 2B entry was accidentally keyed with the base
model's repo name, so ``__INT_TO_FLOAT_MAPPER`` held two identical
``unsloth/gemma-2b-bnb-4bit`` keys. Python keeps only the last value for a
duplicate literal key, so the base 4bit repo resolved to the *instruct* model,
the base model lost its reverse (4x-faster) mapping, and
``unsloth/gemma-2b-it-bnb-4bit`` was never registered at all.
``mapper.py`` has no imports, so we exec it directly and inspect the built
mappers without importing ``unsloth`` (which requires a GPU).
"""
import os
MAPPER_PATH = os.path.join(os.path.dirname(__file__), os.pardir, "unsloth", "models", "mapper.py")
def _load_mappers():
with open(MAPPER_PATH) as f:
source = f.read()
namespace = {}
exec(compile(source, MAPPER_PATH, "exec"), namespace)
return namespace
def test_gemma_2b_base_and_instruct_4bit_are_distinct():
namespace = _load_mappers()
int_to_float = namespace["INT_TO_FLOAT_MAPPER"]
float_to_int = namespace["FLOAT_TO_INT_MAPPER"]
# The base 4bit repo must resolve to the base model, not the instruct one.
assert int_to_float["unsloth/gemma-2b-bnb-4bit"] == "unsloth/gemma-2b"
# The instruct 4bit repo must be registered and resolve to the instruct model.
assert "unsloth/gemma-2b-it-bnb-4bit" in int_to_float
assert int_to_float["unsloth/gemma-2b-it-bnb-4bit"] == "unsloth/gemma-2b-it"
# The base model must reverse-map back to the base 4bit repo.
assert float_to_int["unsloth/gemma-2b"] == "unsloth/gemma-2b-bnb-4bit"
assert float_to_int["google/gemma-2b"] == "unsloth/gemma-2b-bnb-4bit"
# The instruct model must reverse-map to the instruct 4bit repo.
assert float_to_int["unsloth/gemma-2b-it"] == "unsloth/gemma-2b-it-bnb-4bit"
assert float_to_int["google/gemma-2b-it"] == "unsloth/gemma-2b-it-bnb-4bit"

View file

@ -257,6 +257,39 @@ def test_recompute_helper_scales_on_cpu():
), "_unsloth_recompute_inv_freq must return vanilla inv_freq when unscaled."
def test_extended_rope_scaling_keeps_llama3_and_carries_theta():
# Long-context extension keeps native llama3, but falls back to linear for every other
# type (the patched attention constructor only rebuilds linear/llama3/longrope), and the
# linear dict carries rope_theta so transformers v5 does not fall back to base 10000.
from types import SimpleNamespace
from unsloth.models.llama import _extended_rope_scaling
# llama3 model: keep native scaling, do not synthesize linear.
scaling, native = _extended_rope_scaling(_make_config(LLAMA3_ROPE_SCALING), 2.0)
assert (
scaling is None and native == "llama3"
), "must keep native llama3 scaling instead of overwriting it with linear."
# yarn is not rebuildable by the patcher -> keep the safe linear fallback, not native.
yarn = SimpleNamespace(rope_scaling = {"rope_type": "yarn", "factor": 2.0}, rope_theta = 500000.0)
scaling, _ = _extended_rope_scaling(yarn, 2.0)
assert scaling == {
"type": "linear",
"factor": 2.0,
"rope_theta": 500000.0,
}, f"yarn must fall back to linear (patcher cannot rebuild it), got {scaling}."
# plain RoPE with theta only under v5 rope_parameters: linear must carry rope_theta.
v5 = SimpleNamespace(rope_parameters = {"rope_type": "default", "rope_theta": 1000000.0})
scaling, _ = _extended_rope_scaling(v5, 2.0)
assert scaling == {
"type": "linear",
"factor": 2.0,
"rope_theta": 1000000.0,
}, f"linear override dropped rope_theta on v5 (got {scaling}); base would fall back to 10000."
def test_extended_rotary_reads_config_factor():
# LlamaExtendedRotaryEmbedding must honor the config factor, not hardcode 8
# (Llama-3.2 uses 32); otherwise the subclass path re-drops scaling (#2405).

View file

@ -19,6 +19,8 @@ ST_TAGS = [
"v5.2.3",
"v5.3.0",
"v5.4.1",
"v5.5.1",
"v5.6.0",
"master",
]
@ -120,6 +122,42 @@ def test_st_transformer_base_class_either_path(tag: str):
)
# Transformer.load classmethod: unsloth builds saved-ST modules through it (#6881).
@pytest.mark.parametrize("tag", ST_TAGS)
def test_st_transformer_load_accepts_unsloth_kwargs(tag: str):
"""unsloth builds saved ST models via Transformer.load(...) so the saved
modality_config is honored (#6881). If .load stops accepting the hub kwargs it
passes (and has no **kwargs), update the fix before it silently regresses. Not
locating .load is a SKIP (may be inherited); the live test guards the install."""
candidates = [
"sentence_transformers/models/Transformer.py",
"sentence_transformers/models/transformer.py",
"sentence_transformers/base/modules/transformer.py",
"sentence_transformers/base/modules/module.py",
]
for p in candidates:
src = fetch_text("UKPLab/sentence-transformers", tag, p)
if src is None or not has_def(src, "load", "func"):
continue
m = re.search(r"def\s+load\s*\((.*?)\)\s*(?:->[^:]*)?:", src, re.S)
if m is None:
continue
sig = m.group(1)
accepts_var_kw = "**" in sig
missing = [
kw
for kw in ("token", "cache_folder", "revision", "trust_remote_code")
if not (accepts_var_kw or re.search(rf"\b{re.escape(kw)}\b", sig))
]
assert not missing, (
f"{tag}: Transformer.load in {p} no longer accepts {missing} and has no "
f"**kwargs; update unsloth.models.sentence_transformer._create_transformer_module "
f"(#6881) before it silently falls back to Transformer(...)."
)
return
pytest.skip(f"{tag}: Transformer.load not locatable in {candidates} (may be inherited)")
# sentence_transformers.util: import_from_string + load_dir_path helpers unsloth calls.
@pytest.mark.parametrize("tag", ST_TAGS)
def test_st_util_helpers(tag: str):

View file

@ -26,6 +26,7 @@ already_imported = [mod for mod in critical_modules if mod in sys.modules]
# Fix some issues before importing other packages
from .import_fixes import (
fix_message_factory_issue,
fix_torch_check_is_size,
check_fbgemm_gpu_version,
disable_broken_causal_conv1d,
disable_broken_vllm,
@ -72,6 +73,7 @@ fix_bitsandbytes_rocm_arch_detection()
disable_broken_causal_conv1d()
disable_broken_vllm()
fix_message_factory_issue()
fix_torch_check_is_size()
check_fbgemm_gpu_version()
torchvision_compatibility_check()
fix_diffusers_warnings()
@ -81,6 +83,7 @@ del fix_bitsandbytes_rocm_arch_detection
del disable_broken_causal_conv1d
del disable_broken_vllm
del fix_message_factory_issue
del fix_torch_check_is_size
del check_fbgemm_gpu_version
del torchvision_compatibility_check
del fix_diffusers_warnings
@ -173,6 +176,7 @@ from .import_fixes import (
fix_vllm_guided_decoding_params,
fix_vllm_pdl_blackwell,
fix_triton_compiled_kernel_missing_attrs,
fix_dynamo_config_thread_visibility,
patch_trunc_normal_precision_issue,
ignore_logger_messages,
patch_ipykernel_hf_xet,
@ -203,6 +207,10 @@ fix_vllm_guided_decoding_params()
fix_trl_vllm_ascend()
fix_vllm_pdl_blackwell()
fix_triton_compiled_kernel_missing_attrs()
# Must run before unsloth_zoo's patch_torch_compile and the gpt-oss temporary
# patches raise the dynamo recompile limits, so those settings reach the
# autograd worker threads on torch >= 2.12.
fix_dynamo_config_thread_visibility()
patch_trunc_normal_precision_issue()
ignore_logger_messages()
patch_ipykernel_hf_xet()
@ -233,6 +241,7 @@ del fix_vllm_guided_decoding_params
del fix_trl_vllm_ascend
del fix_vllm_pdl_blackwell
del fix_triton_compiled_kernel_missing_attrs
del fix_dynamo_config_thread_visibility
del patch_trunc_normal_precision_issue
del ignore_logger_messages
del patch_ipykernel_hf_xet

View file

@ -2103,8 +2103,9 @@ def get_chat_template(
def remove_special_tokens(tokenizer, prompt):
# Removes double BOS token
if prompt.startswith(tokenizer.bos_token):
prompt = prompt[len(tokenizer.bos_token):]
bos_token = getattr(tokenizer, "bos_token", None)
if bos_token is not None and prompt.startswith(bos_token):
prompt = prompt[len(bos_token):]
return prompt
@ -2185,7 +2186,16 @@ def _create_formatter(possible_columns, final_optional_prompts, user_column_name
texts = []
for row_idx in range(n_rows):
row_values = {column: examples[column][row_idx] for column in columns}
# Coerce missing (None) columns to "" so they do not render as the
# literal string "None" in the emitted text. In a [[...]] block only
# the first column gates the block, so a later column can still be
# None here; required columns can be None too. Coercing at the source
# covers both; since None is now "", the gate below only needs to
# test for "" (an empty first column still drops the block).
row_values = {
column: ("" if (value := examples[column][row_idx]) is None else value)
for column in columns
}
formatter_values = {}
for formatter_template in formatter_templates:
@ -2196,7 +2206,7 @@ def _create_formatter(possible_columns, final_optional_prompts, user_column_name
continue
_, optional_name, prompt, needed_columns = formatter_template
if row_values[needed_columns[0]] not in (None, ""):
if row_values[needed_columns[0]] != "":
prompt_values = {column: row_values[column] for column in needed_columns}
formatter_values[optional_name] = prompt.format(**prompt_values)
else:

View file

@ -172,6 +172,10 @@ if not UNSLOTH_ENABLE_LOGGING:
# Deprecation warnings from torchao
warnings.filterwarnings("ignore", message = "`int4_weight_only` is deprecated")
warnings.filterwarnings("ignore", message = "`int8_weight_only` is deprecated")
# torch._check_is_size FutureWarning (called by bitsandbytes 4-bit dequant)
warnings.filterwarnings(
"ignore", message = r"_check_is_size will be removed", category = FutureWarning
)
# TorchAO deprecated import paths (https://github.com/pytorch/ao/issues/2752)
warnings.filterwarnings(
@ -253,6 +257,30 @@ if not UNSLOTH_ENABLE_LOGGING:
)
def fix_torch_check_is_size():
"""Shim torch._check_is_size if a future torch removes it (bitsandbytes 4-bit
dequant calls it). The FutureWarning is silenced in suppress_cuda_printf."""
try:
import torch
if hasattr(torch, "_check_is_size"):
return
def _check_is_size(
i,
message = None,
*,
max = None,
):
torch._check(i >= 0, message)
if max is not None:
torch._check(i <= max, message)
torch._check_is_size = _check_is_size
except Exception:
return
# Fix up AttributeError: 'MessageFactory' object has no attribute 'GetPrototype'
# MUST do this at the start primarily due to tensorflow causing issues
def fix_message_factory_issue():
@ -1064,6 +1092,135 @@ def fix_triton_compiled_kernel_missing_attrs():
)
def fix_dynamo_config_thread_visibility():
"""torch 2.12 made torch._dynamo/_inductor config overrides thread-local
(ContextVars), so `config.recompile_limit = 1024` set on the main thread is
invisible to the autograd worker threads that run backward. Gradient
checkpointing recompiles fullgraph gpt-oss kernels there against the default
limit of 8, raising FailOnRecompileLimitHit at step 0. Mirror direct config
assignments into the process-global entry default (torch <= 2.11 semantics).
config.patch(...) and config.load_config(...) also assign via __setattr__ but
are thread-local by design, so skip mirroring while inside one (tracked per
thread). No-op below torch 2.12 and on any torch without this internal layout.
"""
try:
import torch
if Version(torch.__version__) < Version("2.12.0"):
return
import torch._dynamo.config as _dynamo_config
from torch.utils._config_module import ConfigModule
from contextvars import ContextVar
except Exception:
return
try:
probe = getattr(_dynamo_config, "_config", {}).get("recompile_limit", None)
if probe is None or not isinstance(getattr(probe, "user_override", None), ContextVar):
# Overrides are not context-local on this torch; nothing to fix.
return
original_setattr = ConfigModule.__setattr__
if getattr(original_setattr, "__unsloth_patched__", False):
return
except Exception:
return
mirrored_modules = ("torch._dynamo.config", "torch._inductor.config")
# config.patch(...) and config.load_config(...) also assign via __setattr__, but
# their writes are thread-local by design; a per-thread depth counter marks them
# so they are not mirrored into the process-global default.
import threading
_scoped_depth = threading.local()
def _in_scoped_write():
return getattr(_scoped_depth, "n", 0) > 0
def _bump(delta):
_scoped_depth.n = getattr(_scoped_depth, "n", 0) + delta
original_patch = ConfigModule.patch
if not getattr(original_patch, "__unsloth_patched__", False):
@functools.wraps(original_patch)
def _patched_patch(self, *args, **kwargs):
ctx = original_patch(self, *args, **kwargs)
try:
cls = type(ctx) # patch() builds a fresh ConfigPatch class each call
if not getattr(cls, "__unsloth_patch_wrapped__", False):
_enter0, _exit0 = cls.__enter__, cls.__exit__
def _enter(s, _e = _enter0):
_bump(1)
try:
return _e(s)
finally:
_bump(-1)
def _exit(
s,
*a,
_x = _exit0,
):
_bump(1)
try:
return _x(s, *a)
finally:
_bump(-1)
cls.__enter__, cls.__exit__ = _enter, _exit
cls.__unsloth_patch_wrapped__ = True
except Exception:
pass
return ctx
_patched_patch.__unsloth_patched__ = True
ConfigModule.patch = _patched_patch
# load_config restores a saved config by calling setattr per key (thread-local).
original_load_config = getattr(ConfigModule, "load_config", None)
if callable(original_load_config) and not getattr(
original_load_config, "__unsloth_patched__", False
):
@functools.wraps(original_load_config)
def _patched_load_config(self, *args, **kwargs):
_bump(1)
try:
return original_load_config(self, *args, **kwargs)
finally:
_bump(-1)
_patched_load_config.__unsloth_patched__ = True
ConfigModule.load_config = _patched_load_config
@functools.wraps(original_setattr)
def _patched_setattr(self, name, value):
original_setattr(self, name, value)
if _in_scoped_write():
return # transient patch / load_config write: keep it thread-local
# Aliases (cache_size_limit -> recompile_limit) re-enter with the real name.
if self.__dict__.get("__name__", None) in mirrored_modules:
try:
entry = self.__dict__["_config"].get(name, None)
if entry is not None and entry.alias is None:
entry.default = value
except Exception:
pass
_patched_setattr.__unsloth_patched__ = True
ConfigModule.__setattr__ = _patched_setattr
# No replay of existing overrides: unsloth installs this before it sets any
# dynamo/inductor config, so the wrapper mirrors every later assignment. Replaying
# would also bake a still-active config.patch override into the global default.
logger.info(
"Unsloth: Patched torch config modules so dynamo/inductor settings "
"(e.g. recompile_limit) apply across threads on torch >= 2.12."
)
def patch_trunc_normal_precision_issue():
"""
Patch torch.nn.init.trunc_normal_ for low precision tensors to run init in fp32.
@ -1323,8 +1480,7 @@ def fix_vllm_pdl_blackwell():
if patched:
logger.info(
f"Unsloth: Applied PDL fix for SM100 ({sm100_gpu_name}) - "
f"patched: {', '.join(patched)}"
f"Unsloth: Applied PDL fix for SM100 ({sm100_gpu_name}) - patched: {', '.join(patched)}"
)
else:
# Just set the env var - vLLM might be an older version without supports_pdl

View file

@ -12,6 +12,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
import os
from contextlib import nullcontext
import torch
import torch.nn as nn
import triton
@ -24,6 +25,15 @@ from unsloth_zoo.temporary_patches.common import torch_compile
torch_matmul = torch.matmul
def _fp8_triton_device_context(tensor: torch.Tensor):
if tensor.device.type == "cuda" and torch.cuda.device_count() > 1:
return torch.cuda.device(tensor.device)
if tensor.device.type == "xpu" and hasattr(torch, "xpu") and torch.xpu.device_count() > 1:
return torch.xpu.device(tensor.device)
return nullcontext()
try:
from transformers.integrations.finegrained_fp8 import FP8Linear
except:
@ -95,7 +105,8 @@ def weight_dequant_block(
triton.cdiv(M, meta["BLOCK_SIZE"]),
triton.cdiv(N, meta["BLOCK_SIZE"]),
)
weight_dequant_kernel[grid](x, s, y, M, N, BLOCK_SIZE = block_size)
with _fp8_triton_device_context(x):
weight_dequant_kernel[grid](x, s, y, M, N, BLOCK_SIZE = block_size)
return y
@ -149,7 +160,8 @@ def act_quant(x: torch.Tensor, block_size: int = 128) -> tuple[torch.Tensor, tor
def grid(meta):
return (triton.cdiv(x.numel(), meta["BLOCK_SIZE"]),)
act_quant_kernel[grid](x, y, s, BLOCK_SIZE = block_size)
with _fp8_triton_device_context(x):
act_quant_kernel[grid](x, y, s, BLOCK_SIZE = block_size)
return y, s
@ -274,32 +286,33 @@ def w8a8_block_fp8_matmul_triton(
def grid(META):
return (triton.cdiv(M, META["BLOCK_SIZE_M"]) * triton.cdiv(N, META["BLOCK_SIZE_N"]),)
_w8a8_block_fp8_matmul[grid](
A,
B,
C,
As,
Bs,
M,
N,
K,
block_n,
block_k,
A.stride(-2),
A.stride(-1),
B.stride(1),
B.stride(0),
C.stride(-2),
C.stride(-1),
As.stride(-2),
As.stride(-1),
Bs.stride(1),
Bs.stride(0),
BLOCK_SIZE_M = BLOCK_SIZE_M,
BLOCK_SIZE_N = BLOCK_SIZE_N,
BLOCK_SIZE_K = BLOCK_SIZE_K,
GROUP_SIZE_M = 8,
)
with _fp8_triton_device_context(A):
_w8a8_block_fp8_matmul[grid](
A,
B,
C,
As,
Bs,
M,
N,
K,
block_n,
block_k,
A.stride(-2),
A.stride(-1),
B.stride(1),
B.stride(0),
C.stride(-2),
C.stride(-1),
As.stride(-2),
As.stride(-1),
Bs.stride(1),
Bs.stride(0),
BLOCK_SIZE_M = BLOCK_SIZE_M,
BLOCK_SIZE_N = BLOCK_SIZE_N,
BLOCK_SIZE_K = BLOCK_SIZE_K,
GROUP_SIZE_M = 8,
)
return C
@ -311,13 +324,14 @@ def torchao_block_matmul(
block_size: tuple[int, int],
output_dtype: torch.dtype = torch.bfloat16,
):
out = torchao_blockwise_gemm(
act_q.contiguous(),
act_scale.contiguous(),
weight_q.contiguous(),
weight_scale.contiguous(),
block_size = block_size[1],
)
with _fp8_triton_device_context(act_q):
out = torchao_blockwise_gemm(
act_q.contiguous(),
act_scale.contiguous(),
weight_q.contiguous(),
weight_scale.contiguous(),
block_size = block_size[1],
)
return out.to(output_dtype)
@ -540,7 +554,8 @@ class FP8_fbgemm_block_linear(torch.autograd.Function):
f"Weight shape {weight.shape} and scales shape {weight_scale.shape} is not compatible with block size {bs_n, bs_k}"
)
xq, xs = triton_quantize_fp8_block(X, bs_m, bs_n, None)
with _fp8_triton_device_context(X):
xq, xs = triton_quantize_fp8_block(X, bs_m, bs_n, None)
# TODO: WARNING - diverges from baseline for high X values, producing
# gibberish / high starting loss. Do not use until resolved; kept for a
# future headstart.

View file

@ -282,6 +282,21 @@ def QUANT_STATE(W):
return getattr(W, "quant_state", None)
# fp8 weight dtypes. A `weight_scale` / `weight_scale_inv` should only be treated as a
# quant state when the weight itself is still fp8. compressed-tensors layers expose an
# already-dequantized bf16 weight at forward time while keeping a `weight_scale` around;
# reading that as a quant state routes a bf16 weight into the bitsandbytes fast_gemv /
# fast_dequantize path, which then reads a missing `absmax` and crashes.
_FP8_WEIGHT_DTYPES = tuple(
dtype
for dtype in (
getattr(torch, "float8_e4m3fn", None),
getattr(torch, "float8_e5m2", None),
)
if dtype is not None
)
def get_lora_parameters(proj):
"""Return (weight, weight quant_state, lora A, lora B, lora scale).
With QAT enabled, also fake-quantizes the base layer and lora weights.
@ -298,9 +313,11 @@ def get_lora_parameters(proj):
if weight_fake_quantizer is not None:
W = weight_fake_quantizer(W)
# Get quant state for 4bit or FP8
# Get quant state for 4bit or FP8. Only fall back to a weight_scale(_inv) when the
# weight is still fp8; a bf16 weight (e.g. a decompressed compressed-tensors layer)
# must not carry a scale as its quant state or fast_gemv will crash on it.
W_quant = getattr(W, "quant_state", None)
if W_quant is None:
if W_quant is None and W.dtype in _FP8_WEIGHT_DTYPES:
W_quant = getattr(base_layer, "weight_scale_inv", None)
if W_quant is None:
W_quant = getattr(base_layer, "weight_scale", None)
@ -349,9 +366,11 @@ def get_lora_parameters_bias(proj):
) # (proj.base_layer if hasattr(proj, "base_layer") else proj)
W = base_layer.weight
# Get quant state for 4bit or FP8
# Get quant state for 4bit or FP8. Only fall back to a weight_scale(_inv) when the
# weight is still fp8; a bf16 weight (e.g. a decompressed compressed-tensors layer)
# must not carry a scale as its quant state or fast_gemv will crash on it.
W_quant = getattr(W, "quant_state", None)
if W_quant is None:
if W_quant is None and W.dtype in _FP8_WEIGHT_DTYPES:
W_quant = getattr(base_layer, "weight_scale_inv", None)
if W_quant is None:
W_quant = getattr(base_layer, "weight_scale", None)

View file

@ -28,7 +28,11 @@ from ._utils import (
is_bfloat16_supported,
get_quant_type,
)
from .loader_utils import _exclude_rope_inv_freq_from_ddp, _get_fp8_mode_and_check_settings
from .loader_utils import (
_exclude_rope_inv_freq_from_ddp,
_get_fp8_mode_and_check_settings,
_restore_dropped_fp8_scales,
)
from ..utils.packing import (
get_packed_info_from_kwargs,
mask_packed_sequence_boundaries,
@ -1651,6 +1655,26 @@ def _rope_scaling_as_dict(rope_scaling):
return {}
def _extended_rope_scaling(config, factor):
"""RoPE scaling to extend a model past its native window. Keeps native llama3 as-is
(linear extension is far worse for long context); everything else gets linear. Returns
(scaling_or_None, type): None keeps llama3. The linear dict carries rope_theta so
transformers v5 (which stores it under rope_parameters) keeps the real base, not 10000.
Only llama3 is preserved because patch_llama_rope_scaling can only rebuild linear/llama3/
longrope and its longrope branch needs a top-level original_max_position_embeddings."""
existing = _rope_scaling_as_dict(
getattr(config, "rope_scaling", None) or getattr(config, "rope_parameters", None) or {}
)
existing_type = existing.get("rope_type") or existing.get("type")
if existing_type == "llama3":
return None, existing_type
return {
"type": "linear",
"factor": factor,
"rope_theta": _get_rope_theta(config),
}, existing_type
def _llama3_inv_freq_from_config(
config,
rope_scaling,
@ -2518,34 +2542,33 @@ class FastLlamaModel:
max_seq_length = model_max_seq_length
if (rope_scaling is None) and (max_seq_length > model_max_seq_length):
rope_scaling = max_seq_length / model_max_seq_length
factor = max_seq_length / model_max_seq_length
if fast_inference:
raise NotImplementedError(
"Unsloth: Fast inference does not yet work with RoPE Scaling."
)
logger.warning_once(
f"Unsloth: {model_name} can only handle sequence lengths of at most "
f"{model_max_seq_length}.\nBut with kaiokendev's RoPE scaling of "
f"{round(rope_scaling, 3)}, it can be magically be extended to "
f"{max_seq_length}!"
)
# Warn RoPE scaling isn't allowed
if not has_rope_scaling:
raise RuntimeError(
f"However, {model_name} doesn't support RoPE Scaling!\n"
"Please file a feature request at https://github.com/unslothai/unsloth."
linear_scaling, native_type = _extended_rope_scaling(model_config, factor)
if linear_scaling is not None:
logger.warning_once(
f"Unsloth: {model_name} can only handle sequence lengths of at most "
f"{model_max_seq_length}.\nBut with kaiokendev's RoPE scaling of "
f"{round(factor, 3)}, it can be magically be extended to "
f"{max_seq_length}!"
)
if not has_rope_scaling:
raise RuntimeError(
f"However, {model_name} doesn't support RoPE Scaling!\n"
"Please file a feature request at https://github.com/unslothai/unsloth."
)
kwargs["rope_scaling"] = linear_scaling
else:
# Native llama3 scaling already handles long context; just widen the window.
logger.warning_once(
f"Unsloth: extending {model_name} to {max_seq_length} using its native "
f"{native_type} RoPE scaling."
)
rope_scaling = {
"type": "linear",
"factor": rope_scaling,
}
# Add to kwargs
kwargs["rope_scaling"] = rope_scaling
from .loader_utils import (
check_and_disable_bitsandbytes_loading,
@ -2659,6 +2682,18 @@ class FastLlamaModel:
offload_embedding = False,
fast_inference = fast_inference,
)
# Re-apply block-fp8 weight_scale_inv tensors transformers dropped on load (#6200).
_restore_dropped_fp8_scales(
model,
model_name,
local_files_only = kwargs.get("local_files_only", False),
token = token,
# Weights load from the default branch (revision not forwarded), so read scales from there too.
revision = None,
subfolder = kwargs.get("subfolder"),
cache_dir = kwargs.get("cache_dir"),
variant = kwargs.get("variant"),
)
elif not fast_inference:
if user_config is not None:
# Transformers 5.x @strict model init rejects extra kwargs next
@ -2697,6 +2732,18 @@ class FastLlamaModel:
offload_embedding = False,
fast_inference = False,
)
# Re-apply block-fp8 weight_scale_inv tensors transformers dropped on load (#6200).
_restore_dropped_fp8_scales(
model,
model_name,
local_files_only = kwargs.get("local_files_only", False),
token = token,
# Weights load from the default branch (revision not forwarded), so read scales from there too.
revision = None,
subfolder = kwargs.get("subfolder"),
cache_dir = kwargs.get("cache_dir"),
variant = kwargs.get("variant"),
)
model.fast_generate = make_fast_generate_wrapper(model.generate)
model.fast_generate_batches = None
else:

View file

@ -239,6 +239,11 @@ def get_model_name(
and new_model_name.lower() in BAD_MAPPINGS
):
new_model_name = BAD_MAPPINGS[new_model_name.lower()]
elif new_model_name is None and model_name.lower() in BAD_MAPPINGS:
# Some bad names (e.g. the `-unsloth-bnb-4bit` dynamic quants) are keys
# of the mappers, not values, so the resolver returns None for them and
# the remap above is skipped; remap the input name directly instead.
new_model_name = BAD_MAPPINGS[model_name.lower()]
if (
new_model_name is None
@ -362,6 +367,287 @@ def _tag_model_with_fp8_torchao_config(model: torch.nn.Module, fp8_mode: str):
pass
_FP8_DTYPES = tuple(
dtype
for dtype in (getattr(torch, "float8_e4m3fn", None), getattr(torch, "float8_e5m2", None))
if dtype is not None
)
def _fp8_block_size_from_config(model):
"""Return the [block_out, block_in] block size of an fp8 checkpoint, or None if not block-fp8."""
config = getattr(model, "config", None)
quant = getattr(config, "quantization_config", None)
if quant is None:
return None
if hasattr(quant, "to_dict"):
quant = quant.to_dict()
if not isinstance(quant, dict):
return None
if quant.get("quant_method") != "fp8":
return None
block = quant.get("weight_block_size")
if not block:
return None
if isinstance(block, (int, float)):
block = [block, block]
elif isinstance(block, (list, tuple)):
if len(block) == 1:
block = [block[0], block[0]]
elif len(block) < 2:
return None
else:
return None
return [int(block[0]), int(block[1])]
def _load_fp8_weight_map(
model_name,
local_files_only,
token,
revision = None,
subfolder = None,
cache_dir = None,
):
"""Return the checkpoint's tensor->file map, using the same snapshot the load used.
Prefers the sharded `model.safetensors.index.json`; falls back to a single `model.safetensors`
(every tensor maps to that one file) so unsharded checkpoints are covered too.
"""
def _local_path(filename):
return (
os.path.join(model_name, subfolder, filename)
if subfolder
else os.path.join(model_name, filename)
)
def _remote_path(filename):
from huggingface_hub import hf_hub_download
return hf_hub_download(
model_name,
filename,
revision = revision,
subfolder = subfolder,
cache_dir = cache_dir,
local_files_only = local_files_only,
token = token,
)
index_file = "model.safetensors.index.json"
single_file = "model.safetensors"
is_local = os.path.isdir(model_name)
# Sharded checkpoint.
if is_local and os.path.exists(_local_path(index_file)):
index_path = _local_path(index_file)
elif not is_local:
try:
index_path = _remote_path(index_file)
except Exception:
index_path = None
else:
index_path = None
if index_path is not None:
import json
with open(index_path, "r") as f:
return json.load(f).get("weight_map", None)
# Unsharded single file: map every tensor to it.
try:
if is_local and os.path.exists(_local_path(single_file)):
single_path = _local_path(single_file)
elif not is_local:
single_path = _remote_path(single_file)
else:
return None
from safetensors import safe_open
with safe_open(single_path, framework = "pt") as f:
return {key: single_file for key in f.keys()}
except Exception:
return None
def _resolve_fp8_shard(
model_name,
shard,
local_files_only,
token,
revision = None,
subfolder = None,
cache_dir = None,
):
"""Resolve a checkpoint shard filename to a local path (repo id or local dir)."""
if os.path.isdir(model_name):
return (
os.path.join(model_name, subfolder, shard)
if subfolder
else os.path.join(model_name, shard)
)
from huggingface_hub import hf_hub_download
return hf_hub_download(
model_name,
shard,
revision = revision,
subfolder = subfolder,
cache_dir = cache_dir,
local_files_only = local_files_only,
token = token,
)
def _match_fp8_module(module_by_name, base):
"""Resolve a checkpoint module name to a live module, allowing for VLM key remappings.
VLM loads can name the text tower differently from the checkpoint keys: `text_only=True`
strips the `language_model.` wrapper (so `model.language_model.layers.*` -> `model.layers.*`),
and full VLM loads may expose `model.language_model.*` while the checkpoint stores
`language_model.model.*`. Try the raw key first, then a few safe remappings.
"""
if base in module_by_name:
return module_by_name[base]
candidates = []
if "language_model." in base:
candidates.append(base.replace("language_model.", "", 1)) # text-only: drop wrapper
if "language_model.model." in base:
candidates.append(base.replace("language_model.model.", "model.language_model.", 1))
if base.startswith("language_model."):
candidates.append("model." + base) # add model. prefix
for candidate in candidates:
if candidate in module_by_name:
return module_by_name[candidate]
return None
def _restore_dropped_fp8_scales(
model,
model_name,
*,
local_files_only = False,
token = None,
revision = None,
subfolder = None,
cache_dir = None,
variant = None,
):
"""Re-apply block-fp8 `weight_scale_inv` tensors that transformers dropped on load.
On some block-scale fp8 checkpoints (e.g. Qwen3.6-27B-FP8, issue #6200) transformers fails to
convert a Linear (such as `mlp.gate_proj`) to an fp8 module, loading the raw quantized values
into a plain bf16 weight and discarding its `weight_scale_inv` as an unexpected key. The weight
is then used un-scaled, producing a garbage model. For every checkpoint scale whose live weight
is not fp8, dequantize the orphaned weight in place. Modules that were converted correctly keep
an fp8 weight and are skipped, so a healthy checkpoint is a no-op. Returns (restored, skipped).
"""
try:
block = _fp8_block_size_from_config(model)
if block is None or not _FP8_DTYPES:
return (0, 0)
# A variant load reads variant-named files; skip to avoid applying default scales to them.
if variant:
return (0, 0)
# No fp8 params means the checkpoint was dequantized on purpose (e.g. load_in_16bit);
# re-applying a scale would corrupt those already-correct 16bit weights, so do nothing.
if not any(p.dtype in _FP8_DTYPES for p in model.parameters()):
return (0, 0)
weight_map = _load_fp8_weight_map(
model_name, local_files_only, token, revision, subfolder, cache_dir
)
if not weight_map:
return (0, 0)
scale_keys = {k: v for k, v in weight_map.items() if k.endswith(".weight_scale_inv")}
if not scale_keys:
return (0, 0)
module_by_name = dict(model.named_modules())
bs0, bs1 = block
restored = 0
skipped = 0
failed = 0
offloaded = 0
shard_cache = {}
for scale_key, shard in scale_keys.items():
base = scale_key[: -len(".weight_scale_inv")]
module = _match_fp8_module(module_by_name, base)
if module is None:
continue
weight = getattr(module, "weight", None)
if not isinstance(weight, torch.Tensor) or weight.ndim != 2:
continue
if weight.device.type == "meta":
# Disk-offloaded layer: weight lives on meta until forward, so it cannot be
# scaled in place here. Count and warn rather than silently leave it unscaled.
offloaded += 1
continue
if weight.dtype in _FP8_DTYPES:
# Correctly converted fp8 module: the fp8 path already handles the scale.
skipped += 1
continue
# Errors after this point are per-tensor: warn and continue, never abort or hide them.
try:
if shard not in shard_cache:
from safetensors import safe_open
shard_path = _resolve_fp8_shard(
model_name,
shard,
local_files_only,
token,
revision,
subfolder,
cache_dir,
)
shard_cache[shard] = safe_open(shard_path, framework = "pt")
scale = shard_cache[shard].get_tensor(scale_key).to(torch.float32)
out_features, in_features = weight.shape
out_blocks = (out_features + bs0 - 1) // bs0
in_blocks = (in_features + bs1 - 1) // bs1
if tuple(scale.shape) == (out_blocks, in_blocks):
pass
elif tuple(scale.shape) == (in_blocks, out_blocks) and out_blocks != in_blocks:
# Transposed block layout: same handling as the fp8 forward path.
scale = scale.t().contiguous()
else:
# Shape does not match the block grid: skip rather than apply a wrong scale.
continue
scale = scale.to(weight.device)
with torch.no_grad():
if out_features % bs0 == 0 and in_features % bs1 == 0:
# Memory-frugal path: multiply block views in place against the broadcast
# fp32 scale, avoiding a full expanded scale and fp32 copy that could OOM.
# The in-place multiply promotes to fp32, matching the fallback exactly.
module.weight.data.view(out_blocks, bs0, in_blocks, bs1).mul_(
scale[:, None, :, None]
)
else:
scale_expanded = scale.repeat_interleave(bs0, dim = 0).repeat_interleave(
bs1, dim = 1
)[:out_features, :in_features]
module.weight.data = (weight.to(torch.float32) * scale_expanded).to(
weight.dtype
)
restored += 1
except Exception:
failed += 1
continue
if restored > 0:
print(f"Unsloth: Restored {restored} dropped FP8 weight_scale_inv tensor(s) on load")
if failed > 0:
print(f"Unsloth: {failed} dropped FP8 weight_scale_inv tensor(s) could not be restored")
if offloaded > 0:
print(
f"Unsloth: {offloaded} dropped FP8 weight_scale_inv tensor(s) skipped because the "
"layer is disk-offloaded; load without disk offload so the scales can be restored"
)
return (restored, skipped)
except Exception:
return (0, 0)
def check_and_disable_bitsandbytes_loading(
model_config,
load_in_4bit = True,

View file

@ -134,7 +134,7 @@ __INT_TO_FLOAT_MAPPER = \
"unsloth/gemma-7b-it",
"google/gemma-7b-it",
),
"unsloth/gemma-2b-bnb-4bit" : (
"unsloth/gemma-2b-it-bnb-4bit" : (
"unsloth/gemma-2b-it",
"google/gemma-2b-it",
),

View file

@ -990,7 +990,17 @@ class FastSentenceTransformer(FastModel):
return None
@staticmethod
def _create_transformer_module(model_name, model, tokenizer, max_seq_length, trust_remote_code):
def _create_transformer_module(
model_name,
model,
tokenizer,
max_seq_length,
trust_remote_code,
token = None,
cache_dir = None,
revision = None,
module_subfolder = "",
):
"""Helper to create and configure a Transformer module."""
from sentence_transformers.models import Transformer
@ -1077,7 +1087,45 @@ class FastSentenceTransformer(FastModel):
elif "tokenizer_args" in transformer_init_params:
transformer_kwargs["tokenizer_args"] = trust_remote_code_kwargs.copy()
transformer_module = Transformer(model_name, **transformer_kwargs)
# Build via Transformer.load so the saved modality_config is honored: plain
# Transformer(...) makes ST 5.x infer a "message" modality for chat-template
# models (e.g. Qwen3-Embedding), chat-wrapping inputs and degrading embeddings
# (#6881). Only use .load when it resolves a Hub id (accepts the kwargs or
# **kwargs); legacy ST 3.x/4.x load(input_path) is local-only with no modality
# bug, so fall back to the constructor.
transformer_module = None
transformer_load = getattr(Transformer, "load", None)
has_modules_json = (
FastSentenceTransformer._module_path(
model_name, token, cache_dir = cache_dir, revision = revision
)
is not None
)
if callable(transformer_load) and has_modules_json:
load_params = inspect.signature(transformer_load).parameters
accepts_var_kw = any(
p.kind is inspect.Parameter.VAR_KEYWORD for p in load_params.values()
)
hub_capable = accepts_var_kw or any(
key in load_params for key in ("token", "cache_folder", "revision")
)
if hub_capable:
load_kwargs = {
"token": token,
"cache_folder": cache_dir,
"revision": revision,
"trust_remote_code": trust_remote_code,
**transformer_kwargs,
}
# Resolve config/tokenizer from the module's saved subfolder
# (modules.json "path"), like stock ST; "" (root) is a no-op.
if module_subfolder:
load_kwargs["subfolder"] = module_subfolder
if not accepts_var_kw:
load_kwargs = {k: v for k, v in load_kwargs.items() if k in load_params}
transformer_module = Transformer.load(model_name, **load_kwargs)
if transformer_module is None:
transformer_module = Transformer(model_name, **transformer_kwargs)
finally:
# Restore original Auto* loading immediately
AutoModel.from_pretrained = original_model_from_pretrained
@ -1191,6 +1239,10 @@ class FastSentenceTransformer(FastModel):
tokenizer,
max_seq_length,
trust_remote_code,
token,
cache_dir,
revision,
module_subfolder = module_config.get("path") or "",
)
modules[name] = transformer_module
else:
@ -1226,7 +1278,14 @@ class FastSentenceTransformer(FastModel):
)
transformer_module = FastSentenceTransformer._create_transformer_module(
model_name, model, tokenizer, max_seq_length, trust_remote_code
model_name,
model,
tokenizer,
max_seq_length,
trust_remote_code,
token,
cache_dir,
revision,
)
modules["0"] = transformer_module

View file

@ -41,7 +41,11 @@ from ._utils import (
set_task_config_attr,
)
from ._utils import *
from .loader_utils import _exclude_rope_inv_freq_from_ddp, _get_fp8_mode_and_check_settings
from .loader_utils import (
_exclude_rope_inv_freq_from_ddp,
_get_fp8_mode_and_check_settings,
_restore_dropped_fp8_scales,
)
from ..save import patch_saving_functions
from ..models.loader_utils import is_distributed
from unsloth_zoo.gradient_checkpointing import (
@ -1192,6 +1196,17 @@ class FastBaseModel:
offload_embedding = offload_embedding,
fast_inference = fast_inference,
)
# Re-apply block-fp8 weight_scale_inv tensors transformers dropped on load (#6200).
_restore_dropped_fp8_scales(
model,
model_name,
local_files_only = local_files_only,
token = token,
revision = kwargs.get("revision"),
subfolder = kwargs.get("subfolder"),
cache_dir = kwargs.get("cache_dir"),
variant = kwargs.get("variant"),
)
if hasattr(model, "generate"):
model.fast_generate = make_fast_generate_wrapper(model.generate)
model.fast_generate_batches = error_out_no_vllm

View file

@ -133,6 +133,21 @@ _YOLO_OPTION = typer.Option(
"flag/config. Any of the three spellings works for any agent."
),
)
_PERSIST_OPTION = typer.Option(
False,
"--persist/--no-persist",
help = (
"Keep this agent's Unsloth-managed session dir so you can resume it later. "
"codex/openclaw/hermes/pi have their whole home relocated into an Unsloth dir "
"that is a throwaway temp dir (wiped on exit) by default; with --persist it "
"lives under the Unsloth agents dir and survives, so their own resume can reopen "
"it. claude and opencode keep sessions in your own stores (~/.claude, "
"~/.local/share/opencode), so they already resume regardless. To reopen a "
"session, pass the agent's own resume command through, e.g. "
"`unsloth start codex --persist resume` or `claude --resume <id>`; those flow to "
"the agent unchanged."
),
)
# Per-agent CLI flag for "run tools without prompting". opencode and openclaw have no
# such flag (config only) and are handled in their config writers, so they are absent.
@ -989,6 +1004,12 @@ def _refresh_windows_path() -> None:
os.environ["PATH"] = os.pathsep.join(entries)
def _install_source(install_hint: str) -> Optional[str]:
"""The first http(s) URL an install hint fetches, or None (e.g. an npm install)."""
match = re.search(r"https?://[^\s'\")]+", install_hint)
return match.group(0) if match else None
def _install_agent(name: str, install_hint: str) -> Optional[str]:
# Missing agent under --launch: offer to run its documented install command, then
# re-resolve it on PATH. Consent-based (we never auto-run a remote install script
@ -997,7 +1018,18 @@ def _install_agent(name: str, install_hint: str) -> Optional[str]:
if not sys.stdin.isatty():
return None
typer.echo(f"`{name}` is not installed.")
if not typer.confirm(f"Install it now with `{install_hint}`?", default = False):
# Make the supply-chain risk explicit before the prompt: these are the vendors'
# own installers (curl | bash, irm | iex, npm), run with the user's privileges,
# and nothing checks a signature or hash on the fetched content. Naming the source
# turns a blind "yes" into informed consent.
source = _install_source(install_hint)
warning = (
f"This will download and RUN a script from {source} with your privileges"
if source
else f"This will RUN `{install_hint}` with your privileges"
)
typer.secho(f"{warning}; there is no signature or hash check.", fg = "yellow", err = True)
if not typer.confirm(f"Install `{name}` now with `{install_hint}`?", default = False):
return None
# Run each hint through the shell it is written for: PowerShell (irm | iex, or npm)
# on Windows, /bin/sh (curl | bash, or npm) everywhere else.
@ -1116,15 +1148,20 @@ def _agents_config_root() -> Path:
@contextlib.contextmanager
def _session_config(agent: str, launch: bool):
def _session_config(
agent: str,
launch: bool,
persist: bool = False,
):
"""Yield a private directory for an agent's session config (never the user's own).
launch: an ephemeral temp dir removed after the agent process exits, so nothing
persists. no-launch: a stable Unsloth-owned dir (the printed recipe is run later
on this machine), reused across runs. Either way the user's real ~/.<agent>
config is left untouched.
launch (default): an ephemeral temp dir removed after the agent process exits, so
nothing persists. no-launch: a stable Unsloth-owned dir (the printed recipe is run
later on this machine), reused across runs. persist (from --persist): use that same
stable dir even for a launch, so the agent's session survives the exit and can be
resumed next time. Either way the user's real ~/.<agent> config is left untouched.
"""
if launch:
if launch and not persist:
path = Path(tempfile.mkdtemp(prefix = f"unsloth-{agent}-"))
try:
yield path
@ -1436,6 +1473,7 @@ def claude(
tensor_parallel: bool = _TENSOR_PARALLEL_OPTION,
serve: bool = _SERVE_OPTION,
yolo: bool = _YOLO_OPTION,
persist: bool = _PERSIST_OPTION,
):
"""Point Claude Code at the running Studio server and start it."""
base, key, entry = _connect(
@ -1480,6 +1518,9 @@ def claude(
# --yolo (or its aliases) maps to Claude's own --dangerously-skip-permissions.
# IS_SANDBOX is left unset on purpose: Claude refuses bypass mode as root unless a
# sandbox is detected, and we don't want to falsely claim one on the user's host.
# claude keeps its history in ~/.claude/projects, which --settings/env never
# relocate, so a session already survives exit; resume it with `claude --continue`
# or `--resume <id>` passed through.
command = [
"claude",
"--model",
@ -1516,6 +1557,7 @@ def codex(
tensor_parallel: bool = _TENSOR_PARALLEL_OPTION,
serve: bool = _SERVE_OPTION,
yolo: bool = _YOLO_OPTION,
persist: bool = _PERSIST_OPTION,
):
"""Point OpenAI Codex at the running Studio server and start it."""
base, key, entry = _connect(
@ -1541,7 +1583,7 @@ def codex(
*_yolo_command_flags("codex", yolo),
*ctx.args,
]
with _session_config("codex", launch) as home:
with _session_config("codex", launch, persist = persist) as home:
write_codex_config(base, entry, home)
env = {_CODEX_ENV_KEY: key, "CODEX_HOME": str(home)}
_run(base, entry, env, command, launch = launch, install_hint = "npm install -g @openai/codex")
@ -1559,6 +1601,7 @@ def openclaw(
tensor_parallel: bool = _TENSOR_PARALLEL_OPTION,
serve: bool = _SERVE_OPTION,
yolo: bool = _YOLO_OPTION,
persist: bool = _PERSIST_OPTION,
):
"""Point OpenClaw at the running Studio server and start it."""
base, key, entry = _connect(
@ -1584,7 +1627,7 @@ def openclaw(
if os.name == "nt"
else "curl -fsSL https://openclaw.ai/install.sh | bash"
)
with _session_config("openclaw", launch) as cfg:
with _session_config("openclaw", launch, persist = persist) as cfg:
config_path = cfg / "openclaw.json"
# key lives in the config, not the env; --yolo writes the exec policy here too.
write_openclaw_config(base, key, entry, config_path, yolo = yolo)
@ -1605,6 +1648,7 @@ def opencode(
tensor_parallel: bool = _TENSOR_PARALLEL_OPTION,
serve: bool = _SERVE_OPTION,
yolo: bool = _YOLO_OPTION,
persist: bool = _PERSIST_OPTION,
):
"""Point OpenCode at the running Studio server and start it."""
base, key, entry = _connect(
@ -1628,7 +1672,9 @@ def opencode(
command = ["opencode", "--model", opencode_model]
else:
command = ["opencode"]
with _session_config("opencode", launch) as cfg:
# opencode keeps sessions in ~/.local/share/opencode (never relocated), so resume
# already survives exit; reopen the last one by passing `opencode --continue` through.
with _session_config("opencode", launch, persist = persist) as cfg:
config_path = cfg / "opencode.json"
# OPENCODE_CONFIG is an overlay (loaded between the user's global and project
# configs), so this adds the Unsloth provider/model for the session without
@ -1680,6 +1726,7 @@ def hermes(
tensor_parallel: bool = _TENSOR_PARALLEL_OPTION,
serve: bool = _SERVE_OPTION,
yolo: bool = _YOLO_OPTION,
persist: bool = _PERSIST_OPTION,
):
"""Point Hermes (Nous Research) at the running Studio server and start it."""
base, key, entry = _connect(
@ -1691,7 +1738,7 @@ def hermes(
)
command = ["hermes", *_yolo_command_flags("hermes", yolo), *ctx.args]
install_hint = _hermes_install_hint()
with _session_config("hermes", launch) as home:
with _session_config("hermes", launch, persist = persist) as home:
# HERMES_HOME relocates hermes' whole home dir (config.yaml, sessions, state)
# like CODEX_HOME, so the user's ~/.hermes is left untouched for the session.
write_hermes_config(base, entry, home / "config.yaml")
@ -1711,6 +1758,7 @@ def pi(
tensor_parallel: bool = _TENSOR_PARALLEL_OPTION,
serve: bool = _SERVE_OPTION,
yolo: bool = _YOLO_OPTION,
persist: bool = _PERSIST_OPTION,
):
"""Point Pi (coding agent) at the running Studio server and start it."""
base, key, entry = _connect(
@ -1735,7 +1783,7 @@ def pi(
# --ignore-scripts matches Pi's documented install recipe (its README notes Pi needs
# no install scripts), so accepting the prompt skips dependency lifecycle scripts.
install_hint = "npm install -g --ignore-scripts @earendil-works/pi-coding-agent"
with _session_config("pi", launch) as home:
with _session_config("pi", launch, persist = persist) as home:
# Pi resolves its config dir from PI_CODING_AGENT_DIR first (getAgentDir() prefers
# it over $HOME/.pi/agent), so pin it at the session dir: an inherited
# PI_CODING_AGENT_DIR in the user's shell would otherwise send Pi to their real

View file

@ -128,6 +128,32 @@ def test_install_agent_uses_powershell_on_windows(monkeypatch):
assert ran == [["powershell", "-NoProfile", "-Command", install_hint]]
def test_install_agent_warns_and_names_remote_source(monkeypatch, capsys):
# Before the confirm, a remote installer must name the URL it fetches so the
# user consents to a specific source rather than blindly accepting.
monkeypatch.setattr(start.os, "name", "nt")
monkeypatch.setattr(start.sys, "stdin", SimpleNamespace(isatty = lambda: True))
monkeypatch.setattr(start.typer, "confirm", lambda *a, **k: False) # decline: nothing runs
hint = "& ([scriptblock]::Create((irm https://hermes-agent.nousresearch.com/install.ps1))) -SkipSetup"
assert start._install_agent("hermes", hint) is None
err = capsys.readouterr().err
assert "https://hermes-agent.nousresearch.com/install.ps1" in err
assert "download and RUN" in err
assert "signature or hash" in err
def test_install_agent_warns_for_package_installer(monkeypatch, capsys):
# An npm-style installer has no URL to fetch, but still runs with the user's
# privileges, so the warning names the command instead.
monkeypatch.setattr(start.os, "name", "posix")
monkeypatch.setattr(start.sys, "stdin", SimpleNamespace(isatty = lambda: True))
monkeypatch.setattr(start.typer, "confirm", lambda *a, **k: False)
assert start._install_agent("codex", "npm install -g @openai/codex") is None
err = capsys.readouterr().err
assert "npm install -g @openai/codex" in err
assert "with your privileges" in err
def test_hermes_install_hint_is_windows_native_on_windows(monkeypatch):
monkeypatch.setattr(start.os, "name", "nt")
@ -2522,3 +2548,145 @@ def test_session_config_no_launch_preserves_existing_state(fake_studio, tmp_path
with start._session_config("codex", launch = False) as home2:
assert home2 == home
assert (home2 / "sessions" / "live.sqlite").read_text() == "state"
# ── --persist: persist the agent session so it can be resumed ────────────────
def test_session_config_persist_uses_stable_dir_and_survives(monkeypatch, tmp_path):
# --persist routes a launch to the stable Unsloth agents dir (the one --no-launch
# already uses) instead of a throwaway temp dir, and never wipes it on exit.
monkeypatch.setattr(start, "_agents_config_root", lambda: tmp_path / "agents")
with start._session_config("codex", launch = True, persist = True) as home:
assert home == tmp_path / "agents" / "codex"
(home / "marker").write_text("kept")
assert home.exists()
assert (home / "marker").read_text() == "kept"
def test_session_config_default_launch_is_ephemeral():
# Default launch (no --persist) still uses a throwaway temp dir wiped on exit.
with start._session_config("codex", launch = True) as home:
assert home.exists()
assert "unsloth-codex-" in home.name
assert not home.exists()
# The temp-dir agents: --persist points each one's home/state env at the stable dir;
# without it, at an ephemeral temp path. opencode is handled separately (only its
# config overlay is relocated; its session data was never in the temp dir).
_RESUME_ENV_VAR = {
"codex": "CODEX_HOME",
"openclaw": "OPENCLAW_STATE_DIR",
"hermes": "HERMES_HOME",
"pi": "HOME",
}
def _capture_launch(monkeypatch, argv):
captured = {}
def run(
command,
env = None,
**kwargs,
):
captured["command"] = command
captured["env"] = env
return SimpleNamespace(returncode = 0)
monkeypatch.setattr(start.subprocess, "run", run)
result = CliRunner().invoke(start.start_app, argv)
assert result.exit_code == 0, result.output
return captured
@pytest.mark.parametrize("agent", sorted(_RESUME_ENV_VAR))
def test_resume_persists_agent_home_to_stable_dir(agent, fake_studio, tmp_path, monkeypatch):
monkeypatch.setattr(start.shutil, "which", lambda _: f"/usr/local/bin/{agent}")
captured = _capture_launch(monkeypatch, [agent, "--persist"])
stable = tmp_path / "agents" / agent
assert captured["env"][_RESUME_ENV_VAR[agent]] == str(stable)
# The stable dir survives the agent exit, so the session can be resumed.
assert stable.exists()
@pytest.mark.parametrize("agent", sorted(_RESUME_ENV_VAR))
def test_default_launch_home_is_ephemeral(agent, fake_studio, tmp_path, monkeypatch):
monkeypatch.setattr(start.shutil, "which", lambda _: f"/usr/local/bin/{agent}")
captured = _capture_launch(monkeypatch, [agent])
home = captured["env"][_RESUME_ENV_VAR[agent]]
assert f"unsloth-{agent}-" in home
assert str(tmp_path / "agents") not in home
def test_resume_opencode_config_in_stable_dir(fake_studio, tmp_path, monkeypatch):
# opencode's session data lives in ~/.local/share/opencode (never relocated), so
# resume already survives exit; --persist also stabilizes its config overlay dir.
monkeypatch.setattr(start.shutil, "which", lambda _: "/usr/local/bin/opencode")
captured = _capture_launch(monkeypatch, ["opencode", "--persist"])
stable = tmp_path / "agents" / "opencode"
assert captured["env"]["OPENCODE_CONFIG"] == str(stable / "opencode.json")
assert stable.exists()
def test_persist_bare_codex_launch_has_no_resume_token(fake_studio, monkeypatch):
# A bare `--persist` only persists the session dir; it must NOT auto-append a native
# resume token, or the very first launch (no session yet) would send codex down its
# no-session error path. The user resumes explicitly: `unsloth start codex --persist resume`.
monkeypatch.setattr(start.shutil, "which", lambda _: "/usr/local/bin/codex")
captured = _capture_launch(monkeypatch, ["codex", "--persist"])
assert "resume" not in captured["command"]
# command[0] is the resolved executable path; assert the argv after it.
assert captured["command"][1:] == ["--oss", "--profile", start._CODEX_PROFILE]
def test_persist_bare_opencode_launch_has_no_resume_token(fake_studio, monkeypatch):
monkeypatch.setattr(start.shutil, "which", lambda _: "/usr/local/bin/opencode")
captured = _capture_launch(monkeypatch, ["opencode", "--persist"])
assert "--continue" not in captured["command"]
assert captured["command"][1:] == ["--model", f"{start._OPENCODE_PROVIDER}/{MODEL['id']}"]
def test_persist_bare_claude_launch_has_no_resume_token(fake_studio, monkeypatch):
monkeypatch.setattr(start.shutil, "which", lambda _: "/usr/local/bin/claude")
monkeypatch.setattr(start, "_claude_flags", lambda: [])
captured = _capture_launch(monkeypatch, ["claude", "--persist"])
assert "--continue" not in captured["command"]
assert captured["command"][1:] == ["--model", MODEL["id"]]
def test_resume_with_passthrough_does_not_auto_append(fake_studio, monkeypatch):
# When the caller drives their own subcommand, --persist only persists the dir; it
# must not inject a resume token that would collide with the user's command.
monkeypatch.setattr(start.shutil, "which", lambda _: "/usr/local/bin/codex")
captured = _capture_launch(monkeypatch, ["codex", "--persist", "exec", "hello"])
assert "resume" not in captured["command"]
assert captured["command"][-2:] == ["exec", "hello"]
def test_default_launch_has_no_resume_token(fake_studio, monkeypatch):
monkeypatch.setattr(start.shutil, "which", lambda _: "/usr/local/bin/codex")
captured = _capture_launch(monkeypatch, ["codex"])
assert "resume" not in captured["command"]
def test_resume_persist_only_agents_have_no_resume_token(fake_studio, monkeypatch):
# openclaw/hermes persist their session dir but have no non-interactive resume
# selector, so --persist must not append a token; their own picker resumes.
for agent in ("openclaw", "hermes"):
monkeypatch.setattr(start.shutil, "which", lambda _, a = agent: f"/usr/local/bin/{a}")
captured = _capture_launch(monkeypatch, [agent, "--persist"])
assert "resume" not in captured["command"]
assert "--continue" not in captured["command"]
def test_native_resume_flag_passes_through_unchanged(fake_studio, monkeypatch):
# The persistence flag is --persist, NOT --resume, so an agent's own
# `--resume <id>` (e.g. `unsloth start claude --resume <guid>`) still flows
# through to the agent verbatim and is not swallowed as a Studio option.
monkeypatch.setattr(start.shutil, "which", lambda _: "/usr/local/bin/claude")
monkeypatch.setattr(start, "_claude_flags", lambda: [])
captured = _capture_launch(monkeypatch, ["claude", "--resume", "some-session-guid"])
assert captured["command"][-2:] == ["--resume", "some-session-guid"]
# Studio never auto-appends its own resume token when the user drives resume.
assert captured["command"].count("--resume") == 1
assert "--continue" not in captured["command"]