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

Merge branch 'main' into daniel/studio-sliding-window-compaction

Browse source
This commit is contained in:
Lee Jackson 2026-06-16 08:58:51 +01:00 • committed by GitHub
commit 2ec4ff93e2
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1113 changed files with 172789 additions and 37452 deletions
Download patch file Download diff file Expand all files Collapse all files

8
.git-blame-ignore-revs Normal file
View file

@ -0,0 +1,8 @@
# Commits listed here are skipped by `git blame` so that bulk, whitespace-only
# changes don't obscure the real authorship of a line.
#
# GitHub honors this file automatically. To use it locally, run once:
# git config blame.ignoreRevsFile .git-blame-ignore-revs
# chore(studio/frontend): normalize line endings to LF
c50b8ab910f5aa56dd7ae0022d2c7b96bfe3384a

11
.gitattributes vendored
View file

@ -1,2 +1,13 @@
# Normalize Python files to LF line endings
*.py text eol=lf
# Always check out shell scripts with LF endings. Without this rule a Windows
# clone (core.autocrlf=true) rewrites them to CRLF, and the trailing \r breaks
# them when run in WSL/Linux (e.g. `set -e` -> "set: Illegal option -").
*.sh text eol=lf
# Normalize Studio frontend sources to LF. Scoped to the frontend tree (rather
# than repo-wide *.ts/*.tsx/... rules) so the policy can't force LF on files
# elsewhere. text=auto lets Git detect and leave binary assets (logos, fonts)
# untouched while text files (.ts/.tsx/.json/.html/.svg/...) are stored as LF.
studio/frontend/** text=auto eol=lf

20
.github/CODEOWNERS vendored
View file

@ -6,10 +6,10 @@
/unsloth/models/rl_replacements.py @Datta0 @pluesclues @danielhanchen
/unsloth/trainer.py @danielhanchen
/unsloth/models/sentence_transformer.py @Etherll @danielhanchen
/unsloth/save.py @rolandtannous @danielhanchen
/unsloth/save.py @danielhanchen
/unsloth/tokenizer_utils.py @mmathew23 @danielhanchen
/unsloth/chat_templates.py @rolandtannous @danielhanchen
/unsloth/ollama_template_mappers.py @rolandtannous @danielhanchen
/unsloth/chat_templates.py @danielhanchen
/unsloth/ollama_template_mappers.py @danielhanchen
/unsloth/kernels/moe/*.py @Datta0
/unsloth/import_fixes.py @danielhanchen
/unsloth/device_type.py @danielhanchen
@ -45,14 +45,14 @@
/unsloth/utils/hf_hub.py @mmathew23
/unsloth/utils/packing.py @mmathew23
/cli/ @rolandtannous @Manan17
/studio/frontend/ @Shine1i @rolandtannous @Manan17
/cli/ @Manan17
/studio/frontend/ @Shine1i @Manan17
/studio/frontend/public/ @Shine1i
/studio/backend/ @rolandtannous
/studio/backend/core/data_recipe/ @rolandtannous
/studio/backend/tests/ @rolandtannous @danielhanchen
/tests/ @rolandtannous @danielhanchen
/scripts/ @rolandtannous @danielhanchen
/studio/backend/
/studio/backend/core/data_recipe/
/studio/backend/tests/ @danielhanchen
/tests/ @danielhanchen
/scripts/ @danielhanchen
# Snapshot data for the notebook linter / Colab oracle. Drift in these
# files changes the pin floor for every Unsloth notebook, so refreshes

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

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

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

@ -269,7 +269,8 @@ jobs:
tests/saving/test_patch_saving_none_tokenizer.py \
tests/saving/test_fix_sentencepiece_gguf_robustness.py \
tests/utils/test_attention_masks.py \
tests/utils/test_trunc_normal_patch.py
tests/utils/test_trunc_normal_patch.py \
tests/python/test_fast_language_model_text_only.py
python -m pytest --collect-only -q "$RUNNER_TEMP/unsloth-zoo/tests/"
- name: import_fixes drift detectors (18 tests, HARD GATE)
@ -332,12 +333,21 @@ jobs:
run: |
python -m pytest -v --tb=short tests/test_callback_signature_drift.py
- name: generation correctness guards (HARD GATE)
# Deterministic CPU guards, each validated to fail on its pre-fix code:
# leftpad = batched left-padded generation (#1066/#3699, fixed by
# #2216 + #4100; staging proof: unsloth-staging-2 PRs 170/172);
# rope_scaling_drift = config.rope_scaling dropped by replaced rotary
# classes (#2405). AST checks run first so import breakage cannot mask them.
run: |
python -m pytest -v --tb=short \
tests/utils/test_prepare_inputs_leftpad.py \
tests/utils/test_rope_scaling_drift.py
- name: unsloth Bucket-A — CPU tests not in Repo tests (CPU)
# 16 tests across 5 files. They live inside tests/saving/ and
# tests/utils/, both of which Repo tests (CPU) excludes via --ignore
# because their sibling files need real GPUs / real HF weights.
# The five files below are pure-Python + AST/protobuf/regex tests
# that run cleanly on CPU. Env inherited from the job block.
# CPU tests across 6 files under tests/saving/, tests/utils/, tests/python/
# that Repo tests (CPU) --ignores. AST/protobuf/regex plus tiny CPU model
# loads; run cleanly here (transformers/torch installed).
run: |
python -m pytest -q --tb=short \
tests/saving/test_save_shell_injection.py \
@ -345,11 +355,12 @@ jobs:
tests/saving/test_fix_sentencepiece_gguf_robustness.py \
tests/utils/test_attention_masks.py \
tests/utils/test_trunc_normal_patch.py \
tests/python/test_fast_language_model_text_only.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
# requires CUDA + dev toolchain, which the CPU-only ubuntu-latest
# runner does not have. The other 15 Bucket-A tests pass cleanly.
# runner does not have. The other Bucket-A tests pass cleanly.
- name: unsloth_zoo @ ${{ env.UNSLOTH_ZOO_REF }} — full pytest (CPU)
# 106 of 111 test_* in unsloth_zoo are CPU-only. The two CUDA-skip
@ -990,8 +1001,10 @@ jobs:
# First seen on transformers >=5,<6; each represents a slow
# or recursive source-rewriter path the zoo can address.
"beit": "TimeoutError: compile exceeds per-model budget",
"deepseek_ocr2": "TimeoutError: compile exceeds per-model budget",
"sam": "TimeoutError: compile exceeds per-model budget",
"sam_hq": "TimeoutError: compile exceeds per-model budget",
"deepseek_ocr2": "TimeoutError: compile exceeds per-model budget",
}

61
.github/workflows/cross-platform-parity-ci.yml vendored Normal file
View file

@ -0,0 +1,61 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
# Runs tests/python/test_cross_platform_parity.py on Windows and macOS.
#
# Why: that test is the guard that install.sh and install.ps1 stay in
# sync, but today it only runs on ubuntu-latest (auto-discovered by
# studio-backend-ci.yml's "Repo tests (CPU)" job). The test reads both
# installer scripts, and on Windows Path.read_text() defaults to the
# cp1252 locale encoding, so a non-cp1252 byte in install.sh (it already
# contains a U+274C) raises UnicodeDecodeError there even though Linux and
# macOS default to UTF-8. The reads were pinned to encoding="utf-8" in
# #6166; this job keeps that from silently regressing by exercising the
# test on the platforms it claims parity for. Pure pytest, no GPU,
# sub-second, so the matrix is cheap.
name: Cross-platform parity
on:
pull_request:
paths:
- 'install.sh'
- 'install.ps1'
- 'tests/python/test_cross_platform_parity.py'
- '.github/workflows/cross-platform-parity-ci.yml'
push:
branches: [main]
paths:
- 'install.sh'
- 'install.ps1'
- 'tests/python/test_cross_platform_parity.py'
- '.github/workflows/cross-platform-parity-ci.yml'
workflow_dispatch:
concurrency:
group: ${{ github.workflow }}-${{ github.ref }}
cancel-in-progress: true
permissions:
contents: read
jobs:
parity:
name: parity (${{ matrix.os }})
strategy:
fail-fast: false
matrix:
os: [windows-latest, macos-latest]
runs-on: ${{ matrix.os }}
timeout-minutes: 10
steps:
- uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2
with:
persist-credentials: false
- uses: actions/setup-python@a309ff8b426b58ec0e2a45f0f869d46889d02405 # v6.2.0
with:
python-version: '3.12'
cache: 'pip'
- run: python -m pip install -U pip pytest
- name: Cross-platform parity test
run: python -m pytest tests/python/test_cross_platform_parity.py -q

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

@ -79,6 +79,64 @@ jobs:
run: |
ruff check unsloth unsloth_cli studio tests cli.py unsloth-cli.py
- name: Import-hoist verifier self-test
# scripts/verify_import_hoist.py is a scope-aware (LEGB) AST
# resolver that gates import-hoisting / alias-rename refactors
# against two bugs ruff and pyflakes both miss:
# 1. dangling alias -- `from a import b as _b` hoisted to
# `from a import b` but a leftover `_b` reference now
# resolves to nothing (or to some other module-level `_b`).
# 2. rename clash -- `_b -> b` silently re-points at a
# different object already named `b` in that scope.
# This step runs the tool's 8 negative-control cases so a
# regression in the verifier itself fails before we trust it on
# a diff. Hermetic, stdlib-only, sub-second. Hard gate.
run: |
python scripts/verify_import_hoist.py --self-test
- name: Import-hoist / alias-rename safety (changed Python files)
# Runs the verifier in compare mode on every in-place-modified
# .py in the PR: parses each file BEFORE (base branch) and AFTER
# (this diff), resolves every name load, and fails on a BLOCKER
# (dangling alias / rename clash / re-pointed import). INFO
# findings (a helper relocated to another file) do not fail.
#
# --diff-filter=M (in-place edits only) is deliberate: that is
# exactly where a hoist refactor lives, and it skips brand-new
# files whose re-export imports would otherwise look "unused".
#
# Diff against the true merge-base, not the base tip. A two-dot
# diff against the tip re-lints every file the base branch
# changed after the PR branched, comparing newer base code
# (BEFORE) against the PR's older snapshot (AFTER) - a
# time-reversed comparison that flags the base branch's own
# refactors as blockers on PRs that never touched those files.
# The compare API returns the merge-base without needing local
# history, and fetching that single commit by SHA keeps the
# shallow (fetch-depth: 1) clone.
if: github.event_name == 'pull_request'
env:
GH_TOKEN: ${{ github.token }}
run: |
MERGE_BASE=$(gh api \
"repos/${{ github.repository }}/compare/${{ github.event.pull_request.base.sha }}...${{ github.event.pull_request.head.sha }}" \
--jq .merge_base_commit.sha)
git fetch --no-tags --depth=1 origin "$MERGE_BASE"
mapfile -t CHANGED < <(
git diff --name-only --diff-filter=M \
"$MERGE_BASE" HEAD -- '*.py' \
| grep -vE '(^|/)(unsloth_compiled_cache|node_modules|build|dist)/' || true
)
if [ "${#CHANGED[@]}" -eq 0 ]; then
echo "no in-place-modified Python files to check"
exit 0
fi
printf 'merge base: %s\n' "$MERGE_BASE"
printf 'checking %d file(s):\n' "${#CHANGED[@]}"
printf ' %s\n' "${CHANGED[@]}"
python scripts/verify_import_hoist.py \
--before "$MERGE_BASE" --after HEAD "${CHANGED[@]}"
- name: No leftover debugger / pdb / breakpoint calls
# Catches the "I'll just stick a breakpoint() here" mistake
# before it ships. AST-based so commented-out debugger

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

@ -231,33 +231,14 @@ jobs:
tests/studio/test_is_mlx_dispatch_gate.py \
tests/studio/test_mlx_training_worker_behaviors.py
# Studio prebuilt llama.cpp install + GGUF inference. Drives the
# exact path Studio's setup.sh takes on macOS: invokes
# studio/install_llama_prebuilt.py with --published-repo
# ggml-org/llama.cpp and --published-release-tag b9049 (the
# latest llama.cpp release at the time this step was added; bump
# via UNSLOTH_LLAMA_TAG / DEFAULT_LLAMA_TAG when refreshing).
# The installer downloads llama-b9049-bin-macos-arm64.tar.gz,
# which is the universal Apple Silicon (arm64) build -- the
# same artifact works on M1/M2/M3/M4 because llama.cpp compiles
# against the ARMv8.2 baseline.
#
# The b9049 release also publishes:
# - llama-b9049-bin-macos-arm64-kleidiai.tar.gz
# KleidiAI dispatches at runtime; on M1 it falls back where
# ISA features (e.g. I8MM) are missing, so this asset also
# runs on M1 -- Studio just doesn't choose it by default.
# - llama-b9049-bin-macos-x64.tar.gz
# Intel-only; would only run on M1 via Rosetta 2 emulation,
# which we explicitly avoid.
# - iOS XCFramework
# iOS-app build artifact, unrelated to a macOS desktop CI.
#
# After install, downloads a small published GGUF
# (unsloth/gemma-3-270m-it-GGUF, Q4_K_M) from HuggingFace and
# runs the prebuilt llama-cli on it. Asserts the prompt echo
# appears in stdout. If the install fails OR the binary exits
# non-zero, that's an Unsloth/Studio bug.
# Studio prebuilt llama.cpp install + GGUF inference. Mirrors the
# path Studio's setup.sh takes on macOS since #5963: plan against
# the unslothai/llama.cpp fork's latest release, which ships the
# bin-macos-arm64 bundle plus the llama-prebuilt-manifest.json the
# default policy reads. After install, downloads a small published
# GGUF (unsloth/gemma-3-270m-it-GGUF, Q4_K_M) and validates
# llama-server /completion end to end. An install failure or a
# non-zero binary exit is an Unsloth/Studio bug.
- name: Studio prebuilt llama.cpp install + GGUF inference (Mac M1)
env:
HF_TOKEN: ${{ secrets.HF_TOKEN }}
@ -272,20 +253,12 @@ jobs:
set -euo pipefail
INSTALL_DIR="$HOME/.unsloth-studio-prebuilt-test/llama.cpp"
rm -rf "$INSTALL_DIR"
# --simple-policy is required when --published-repo points
# at upstream ggml-org/llama.cpp; that repo doesn't ship the
# llama-prebuilt-manifest.json asset Studio's default policy
# expects, so the simple platform-specific policy maps
# Darwin+arm64 -> bin-macos-arm64 directly. studio/setup.sh
# passes both --published-repo ggml-org/llama.cpp AND
# --simple-policy automatically on macOS, so this CI step
# exercises the same code path users hit when they run
# `curl -fsSL https://unsloth.ai/install.sh | sh`.
# Mirror studio/setup.sh on macOS (the install.sh user path):
# it plans against the unslothai/llama.cpp fork's latest
# release with no policy or tag flags.
python studio/install_llama_prebuilt.py \
--install-dir "$INSTALL_DIR" \
--published-repo ggml-org/llama.cpp \
--published-release-tag b9049 \
--simple-policy
--published-repo unslothai/llama.cpp
# Studio bundles only llama-server + llama-quantize from the
# prebuilt (not llama-cli) -- inference goes through

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

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

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

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

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

@ -77,7 +77,7 @@ jobs:
path: hf-cache
# Same key as studio-ui-smoke.yml so the two jobs share a
# single GGUF download across CI.
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v1
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2
- name: Prime HF_HOME with the GGUF
id: prime-hf
@ -88,17 +88,19 @@ jobs:
python -m pip install --upgrade huggingface_hub
mkdir -p hf-cache
bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE"
bash .github/scripts/hf-download-with-retry.sh ggml-org/models tinyllamas/stories260K.gguf
- name: Save HF_HOME for ${{ env.GGUF_REPO }}
if: always() && steps.prime-hf.outcome == 'success'
uses: actions/cache/save@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5
with:
path: hf-cache
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v1
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2
- name: Install Studio (--local, --no-torch)
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
HF_TOKEN: ${{ secrets.HF_TOKEN }}
run: |
mkdir -p logs
set -o pipefail

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

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

38
.github/workflows/studio-frontend-ci.yml vendored
View file

@ -17,6 +17,8 @@ on:
- 'studio/frontend/**'
- 'scripts/check_frontend_dep_removal.py'
- 'tests/studio/test_frontend_dep_removal.py'
- 'scripts/sync_allow_scripts_pins.py'
- 'tests/studio/test_sync_allow_scripts_pins.py'
- '.github/workflows/studio-frontend-ci.yml'
push:
branches: [main, pip]
@ -60,6 +62,19 @@ jobs:
with:
node-version: '22'
# node 22 bundles npm 10.x, which predates allowScripts. Move to the
# 11.x line and fail loudly if the gate is still missing, so the
# strict flag below can never silently degrade into a warning.
- name: Upgrade npm to 11.x (allowScripts enforcement)
working-directory: ${{ github.workspace }}
run: |
npm install -g npm@^11 --no-fund --no-audit
V=$(npm -v)
case "$V" in
11.1[6-9].*|11.[2-9][0-9].*|1[2-9].*) echo "npm $V has allowScripts" ;;
*) echo "::error::npm $V lacks allowScripts (need >=11.16)"; exit 1 ;;
esac
# Run the structural lockfile scan BEFORE npm ci. A compromised
# tarball runs its `prepare` / `postinstall` during `npm ci`,
# so any catch has to fire upstream of that. The scanner is
@ -68,14 +83,23 @@ jobs:
working-directory: ${{ github.workspace }}
run: python3 scripts/lockfile_supply_chain_audit.py
# Dependency bumps strand the version-pinned allowScripts entries.
# The paired pre-commit hook auto-fixes PRs; this is the backstop.
- name: allowScripts pins must match the lockfile
working-directory: ${{ github.workspace }}
run: |
python3 tests/studio/test_sync_allow_scripts_pins.py
python3 scripts/sync_allow_scripts_pins.py --check
- name: Lockfile must agree with package.json (npm ci is strict)
# Lifecycle scripts (esbuild native-binary postinstall, etc.) are
# required for `vite build`. The pre-install lockfile structural
# audit (lockfile_supply_chain_audit.py) is the practical defence
# against the npm postinstall-dropper class -- it fires BEFORE any
# tarball runs, on the injection pattern itself rather than an
# advisory-DB lookup.
run: npm ci --no-fund --no-audit
# The vite 8 chain (rolldown, lightningcss, tailwind oxide) ships napi
# binaries with no install scripts. The only script-bearing deps are
# covered by `allowScripts` in package.json (npm >=11.16, default in
# npm 12). The pre-install lockfile audit above stays the first line
# of defence -- it fires before any tarball can run code.
# --strict-allow-scripts: any unreviewed install script hard-fails
# the job; the sync hook keeps the pins fresh after bumps.
run: npm ci --strict-allow-scripts --no-fund --no-audit
- name: npm ci must not have modified the working tree
working-directory: ${{ github.workspace }}

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

@ -20,7 +20,7 @@
# enable_tools / enabled_tools, and enable_thinking on/off.
#
# 3. JSON, images
# gemma-4-E2B-it UD-IQ3_XXS (~2.4 GiB) + mmproj-F16 (~986 MiB).
# Qwen3-VL-2B-Instruct UD-Q4_K_XL (~1.1 GiB) + mmproj-F16 (~780 MiB).
# response_format JSON-schema decoding and OpenAI image_url
# (data URI) plus Anthropic source/base64 image inputs.
#
@ -91,7 +91,7 @@ jobs:
continue-on-error: true
with:
path: hf-cache
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v1
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2
- name: Prime HF_HOME with the GGUF
id: prime-hf
@ -102,17 +102,19 @@ jobs:
python -m pip install --upgrade huggingface_hub
mkdir -p hf-cache
bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE"
bash .github/scripts/hf-download-with-retry.sh ggml-org/models tinyllamas/stories260K.gguf
- name: Save HF_HOME for ${{ env.GGUF_REPO }}
if: always() && steps.prime-hf.outcome == 'success'
uses: actions/cache/save@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5
with:
path: hf-cache
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v1
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2
- name: Install Studio (--local, --no-torch)
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
HF_TOKEN: ${{ secrets.HF_TOKEN }}
run: |
mkdir -p logs
set -o pipefail
@ -296,6 +298,8 @@ jobs:
- name: Upload logs
# Always upload so green runs are still reviewable.
if: always()
# Diagnostic only: a transient artifact-service drop must not fail a green job.
continue-on-error: true
uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
with:
name: openai-anthropic-log
@ -373,6 +377,7 @@ jobs:
- name: Install Studio (--local, --no-torch)
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
HF_TOKEN: ${{ secrets.HF_TOKEN }}
run: |
mkdir -p logs
set -o pipefail
@ -771,6 +776,8 @@ jobs:
- name: Upload logs
# Always upload so green runs are still reviewable.
if: always()
# Diagnostic only: a transient artifact-service drop must not fail a green job.
continue-on-error: true
uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
with:
name: tool-calling-log
@ -787,9 +794,15 @@ jobs:
runs-on: ubuntu-latest
timeout-minutes: 30
env:
GGUF_REPO: unsloth/gemma-4-E2B-it-GGUF
GGUF_VARIANT: UD-IQ3_XXS
GGUF_FILE: gemma-4-E2B-it-UD-IQ3_XXS.gguf
GGUF_REPO: unsloth/Qwen3-VL-2B-Instruct-GGUF
# UD-Q4_K_XL, not UD-IQ2_XXS: at 2-bit the temp-0 answer to the JSON
# step's capital-of-France probe flips with the host's SIMD kernels
# (GitHub runners deterministically answered France while other CPUs
# answer Paris; seeds do not rescue it, 1/5 Paris at temp 0.7). The
# Q4 quant answered Paris 13/13 across temps and seeds on the same
# runners, so the hard Paris assertion below stays reliable.
GGUF_VARIANT: UD-Q4_K_XL
GGUF_FILE: Qwen3-VL-2B-Instruct-UD-Q4_K_XL.gguf
MMPROJ_FILE: mmproj-F16.gguf
STUDIO_PORT: '18890'
HF_HOME: ${{ github.workspace }}/hf-cache
@ -819,7 +832,7 @@ jobs:
continue-on-error: true
with:
path: hf-cache
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-${{ env.MMPROJ_FILE }}-v1
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-${{ env.MMPROJ_FILE }}-v2
- name: Prime HF_HOME with the GGUF + mmproj
id: prime-hf
@ -831,17 +844,19 @@ jobs:
mkdir -p hf-cache
bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE"
bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$MMPROJ_FILE"
bash .github/scripts/hf-download-with-retry.sh ggml-org/models tinyllamas/stories260K.gguf
- name: Save HF_HOME for ${{ env.GGUF_REPO }} (model + mmproj)
if: always() && steps.prime-hf.outcome == 'success'
uses: actions/cache/save@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5
with:
path: hf-cache
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-${{ env.MMPROJ_FILE }}-v1
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-${{ env.MMPROJ_FILE }}-v2
- name: Install Studio (--local, --no-torch)
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
HF_TOKEN: ${{ secrets.HF_TOKEN }}
run: |
mkdir -p logs
set -o pipefail
@ -884,13 +899,23 @@ jobs:
-H 'content-type: application/json' \
-d "{\"username\":\"unsloth\",\"password\":\"$NEW\"}" | jq -r .access_token)
echo "API_KEY=$TOKEN" >> "$GITHUB_ENV"
# Load the GGUF (mmproj is auto-detected via the HF repo
# lookup, the cached file is pulled out of HF_HOME).
curl -fs -X POST "http://127.0.0.1:${STUDIO_PORT}/api/inference/load" \
-H "Authorization: Bearer $TOKEN" -H 'content-type: application/json' \
--max-time 900 \
-d "{\"model_path\":\"$GGUF_REPO\",\"gguf_variant\":\"$GGUF_VARIANT\",\"is_lora\":false,\"max_seq_length\":2048}" \
| jq '{status, display_name, is_vision}'
# Retry: llama-server startup can race process teardown after a
# failed attempt. Keep curl out of a pipe so HTTP failures are not
# masked by jq.
LOAD_OK=0
for attempt in 1 2 3; do
HTTP=$(curl -s -o /tmp/load.json -w '%{http_code}' \
-X POST "http://127.0.0.1:${STUDIO_PORT}/api/inference/load" \
-H "Authorization: Bearer $TOKEN" -H 'content-type: application/json' \
--max-time 900 \
-d "{\"model_path\":\"$GGUF_REPO\",\"gguf_variant\":\"$GGUF_VARIANT\",\"is_lora\":false,\"max_seq_length\":2048}")
if [ "$HTTP" = "200" ]; then LOAD_OK=1; break; fi
echo "::warning::/api/inference/load attempt $attempt returned $HTTP; response:"
cat /tmp/load.json || true
sleep 10
done
[ "$LOAD_OK" = "1" ] || { echo "::error::/api/inference/load failed 3 attempts"; exit 22; }
jq '{status, display_name, is_vision}' /tmp/load.json
- name: JSON schema decoding + image input
env:
@ -1043,6 +1068,8 @@ jobs:
- name: Upload logs
# Always upload so green runs are still reviewable.
if: always()
# Diagnostic only: a transient artifact-service drop must not fail a green job.
continue-on-error: true
uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
with:
name: json-images-log

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

@ -62,7 +62,7 @@ jobs:
continue-on-error: true
with:
path: hf-cache
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v1
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2
- name: Prime HF_HOME with the GGUF
id: prime-hf
@ -73,29 +73,26 @@ jobs:
python -m pip install --upgrade huggingface_hub
mkdir -p hf-cache
bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE"
bash .github/scripts/hf-download-with-retry.sh ggml-org/models tinyllamas/stories260K.gguf
- name: Save HF_HOME for ${{ env.GGUF_REPO }}
if: always() && steps.prime-hf.outcome == 'success'
uses: actions/cache/save@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5
with:
path: hf-cache
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v1
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2
- 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: Assert install.sh used the Mac llama.cpp prebuilt
run: |
if grep -q "falling back to source build" logs/install.log; then
echo "::error::install.sh fell back to source-build llama.cpp on Mac. Studio must install the prebuilt llama-bNNNN-bin-macos-arm64 on Apple Silicon."
grep -E "llama-prebuilt|llama.cpp" logs/install.log | tail -60
exit 1
fi
- name: Assert llama.cpp loads on this macOS
run: bash .github/scripts/assert-llama-loads.sh
- name: Install pyjwt for the JWT-expiry forge test
run: pip install 'pyjwt>=2.6'

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

@ -85,7 +85,7 @@ jobs:
continue-on-error: true
with:
path: hf-cache
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v1
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2
- name: Prime HF_HOME with the GGUF
id: prime-hf
@ -96,6 +96,7 @@ jobs:
python -m pip install --upgrade huggingface_hub
mkdir -p hf-cache
bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE"
bash .github/scripts/hf-download-with-retry.sh ggml-org/models tinyllamas/stories260K.gguf
# Save partial caches on cancel/timeout -- hf download resumes by
# content hash. `outcome != skipped` keeps cache-hit a no-op.
@ -104,23 +105,19 @@ jobs:
uses: actions/cache/save@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5
with:
path: hf-cache
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v1
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2
- 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: Assert install.sh used the Mac llama.cpp prebuilt
run: |
if grep -q "falling back to source build" logs/install.log; then
echo "::error::install.sh fell back to source-build llama.cpp on Mac. Studio must install the prebuilt llama-bNNNN-bin-macos-arm64 on Apple Silicon."
grep -E "llama-prebuilt|llama.cpp" logs/install.log | tail -60
exit 1
fi
- name: Assert llama.cpp loads on this macOS
run: bash .github/scripts/assert-llama-loads.sh
- name: Install OpenAI + Anthropic Python SDKs
run: pip install 'openai>=1.50' 'anthropic>=0.40'
@ -294,6 +291,8 @@ jobs:
- name: Upload logs
# Always upload so green runs are still reviewable.
if: always()
# Diagnostic only: a transient artifact-service drop must not fail a green job.
continue-on-error: true
uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
with:
name: openai-anthropic-log
@ -364,18 +363,14 @@ jobs:
- 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: Assert install.sh used the Mac llama.cpp prebuilt
run: |
if grep -q "falling back to source build" logs/install.log; then
echo "::error::install.sh fell back to source-build llama.cpp on Mac. Studio must install the prebuilt llama-bNNNN-bin-macos-arm64 on Apple Silicon."
grep -E "llama-prebuilt|llama.cpp" logs/install.log | tail -60
exit 1
fi
- name: Assert llama.cpp loads on this macOS
run: bash .github/scripts/assert-llama-loads.sh
- name: Reset auth + boot Studio (API-only, default tool policy)
# We deliberately use the API-only mode rather than
@ -659,6 +654,8 @@ jobs:
- name: Upload logs
# Always upload so green runs are still reviewable.
if: always()
# Diagnostic only: a transient artifact-service drop must not fail a green job.
continue-on-error: true
uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
with:
name: tool-calling-log
@ -755,18 +752,14 @@ jobs:
- 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: Assert install.sh used the Mac llama.cpp prebuilt
run: |
if grep -q "falling back to source build" logs/install.log; then
echo "::error::install.sh fell back to source-build llama.cpp on Mac. Studio must install the prebuilt llama-bNNNN-bin-macos-arm64 on Apple Silicon."
grep -E "llama-prebuilt|llama.cpp" logs/install.log | tail -60
exit 1
fi
- name: Assert llama.cpp loads on this macOS
run: bash .github/scripts/assert-llama-loads.sh
- name: Install OpenAI + Anthropic Python SDKs
run: pip install 'openai>=1.50' 'anthropic>=0.40'
@ -1040,6 +1033,8 @@ jobs:
- name: Upload logs
# Always upload so green runs are still reviewable.
if: always()
# Diagnostic only: a transient artifact-service drop must not fail a green job.
continue-on-error: true
uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
with:
name: json-images-log

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

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

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

@ -62,7 +62,7 @@ jobs:
continue-on-error: true
with:
path: hf-cache
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v1
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2
- name: Prime HF_HOME with the GGUF
id: prime-hf
@ -73,29 +73,26 @@ jobs:
python -m pip install --upgrade huggingface_hub
mkdir -p hf-cache
bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE"
bash .github/scripts/hf-download-with-retry.sh ggml-org/models tinyllamas/stories260K.gguf
- name: Save HF_HOME for ${{ env.GGUF_REPO }}
if: always() && steps.prime-hf.outcome == 'success'
uses: actions/cache/save@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5
with:
path: hf-cache
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v1
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2
- 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: Assert install.sh used the Mac llama.cpp prebuilt
run: |
if grep -q "falling back to source build" logs/install.log; then
echo "::error::install.sh fell back to source-build llama.cpp on Mac. Studio must install the prebuilt llama-bNNNN-bin-macos-arm64 on Apple Silicon."
grep -E "llama-prebuilt|llama.cpp" logs/install.log | tail -60
exit 1
fi
- name: Assert llama.cpp loads on this macOS
run: bash .github/scripts/assert-llama-loads.sh
- name: Install Playwright + Chromium
# No --with-deps on Mac: that flag installs Linux apt packages.

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

@ -62,30 +62,19 @@ jobs:
- 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: Assert install.sh used the Mac llama.cpp prebuilt
run: |
# Mac install must take the prebuilt path. Source-build
# fallback here is an Unsloth bug.
if grep -q "falling back to source build" logs/install.log; then
echo "::error::install.sh fell back to source-build llama.cpp on Mac. Studio must install the prebuilt llama-bNNNN-bin-macos-arm64 on Apple Silicon."
grep -E "llama-prebuilt|llama.cpp" logs/install.log | tail -60
exit 1
fi
if ! grep -qE "prebuilt installed and validated|prebuilt up to date and validated|bin-macos-arm64" logs/install.log; then
echo "::error::no Mac prebuilt llama.cpp marker in install.log."
grep -E "llama-prebuilt|llama.cpp" logs/install.log | tail -60
exit 1
fi
echo "install.sh installed the Mac prebuilt llama.cpp"
- name: Assert llama.cpp loads on this macOS
run: bash .github/scripts/assert-llama-loads.sh
- name: First update should be a no-op (prebuilt already validated)
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
HF_TOKEN: ${{ secrets.HF_TOKEN }}
run: |
set -o pipefail
unsloth studio update --local 2>&1 | tee logs/update.log
@ -104,6 +93,7 @@ jobs:
- name: Second update must also be a no-op
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
HF_TOKEN: ${{ secrets.HF_TOKEN }}
run: |
set -o pipefail
unsloth studio update --local 2>&1 | tee logs/update2.log

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

@ -76,7 +76,7 @@ jobs:
continue-on-error: true
with:
path: hf-cache
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v1
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2
- name: Prime HF_HOME with the GGUF
id: prime-hf
@ -87,17 +87,19 @@ jobs:
python -m pip install --upgrade huggingface_hub
mkdir -p hf-cache
bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE"
bash .github/scripts/hf-download-with-retry.sh ggml-org/models tinyllamas/stories260K.gguf
- name: Save HF_HOME for ${{ env.GGUF_REPO }}
if: always() && steps.prime-hf.outcome == 'success'
uses: actions/cache/save@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5
with:
path: hf-cache
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v1
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2
- name: Install Studio (--local, --no-torch)
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
HF_TOKEN: ${{ secrets.HF_TOKEN }}
run: |
mkdir -p logs
set -o pipefail

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

@ -71,6 +71,7 @@ jobs:
# prebuilt path falls back to source build.
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
HF_TOKEN: ${{ secrets.HF_TOKEN }}
run: |
mkdir -p logs
set -o pipefail
@ -85,6 +86,7 @@ jobs:
# idempotency regressed.
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
HF_TOKEN: ${{ secrets.HF_TOKEN }}
run: |
set -o pipefail
unsloth studio update --local 2>&1 | tee logs/update.log
@ -107,6 +109,7 @@ jobs:
# the first one.
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
HF_TOKEN: ${{ secrets.HF_TOKEN }}
run: |
set -o pipefail
unsloth studio update --local 2>&1 | tee logs/update2.log

6
.github/workflows/studio-windows-api-smoke.yml vendored
View file

@ -69,7 +69,7 @@ jobs:
continue-on-error: true
with:
path: hf-cache
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v1
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2
- name: Prime HF_HOME with the GGUF
id: prime-hf
@ -80,13 +80,14 @@ jobs:
python -m pip install --upgrade huggingface_hub
mkdir -p hf-cache
bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE"
bash .github/scripts/hf-download-with-retry.sh ggml-org/models tinyllamas/stories260K.gguf
- name: Save HF_HOME for ${{ env.GGUF_REPO }}
if: always() && steps.prime-hf.outcome == 'success'
uses: actions/cache/save@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5
with:
path: hf-cache
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v1
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2
- name: Pre-install Windows tweaks (npm 11 + Defender exclusions)
shell: pwsh
@ -123,6 +124,7 @@ jobs:
shell: pwsh
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
HF_TOKEN: ${{ secrets.HF_TOKEN }}
run: |
New-Item -ItemType Directory -Force -Path logs | Out-Null
# *>&1 captures Write-Host (Information stream) output;

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

@ -13,7 +13,7 @@
# 2. Tool calling Tests
# Qwen3.5-2B UD-Q4_K_XL (~890 MiB).
# 3. JSON, images
# gemma-4-E2B-it UD-Q4_K_XL + mmproj-F16 (~3.4 GiB total).
# Qwen3-VL-2B-Instruct UD-IQ2_XXS + mmproj-F16 (~1.4 GiB total).
# Within the 14 GB windows-latest SSD budget.
name: Windows Studio GGUF CI
@ -65,6 +65,18 @@ jobs:
with:
persist-credentials: false
# Fast GPU-free gate: parse setup.ps1 and run the Resolve-CudaToolkit unit
# test (deferred Windows CUDA Toolkit check) before the heavy GGUF smoke.
- name: setup.ps1 unit test (Resolve-CudaToolkit)
shell: pwsh
run: |
$errs = $null
[void][System.Management.Automation.Language.Parser]::ParseFile(
(Resolve-Path studio/setup.ps1).Path, [ref]$null, [ref]$errs)
if ($errs) { $errs | ForEach-Object { $_.ToString() }; exit 1 }
Write-Host "setup.ps1 parsed with no errors"
pwsh -NoProfile -File tests/studio/test_resolve_cuda_toolkit.ps1
- uses: actions/setup-node@48b55a011bda9f5d6aeb4c2d9c7362e8dae4041e # v6.4.0
with:
node-version: '22'
@ -89,7 +101,7 @@ jobs:
continue-on-error: true
with:
path: hf-cache
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v1
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2
- name: Prime HF_HOME with the GGUF
id: prime-hf
@ -102,6 +114,7 @@ jobs:
python -m pip install --upgrade huggingface_hub
mkdir -p hf-cache
bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE"
bash .github/scripts/hf-download-with-retry.sh ggml-org/models tinyllamas/stories260K.gguf
- name: Save HF_HOME cache for ${{ env.GGUF_REPO }}
# Only write a fresh cache entry when we actually rebuilt the
@ -111,7 +124,7 @@ jobs:
uses: actions/cache/save@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5
with:
path: hf-cache
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v1
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2
- name: Pre-install Windows tweaks (npm 11 + Defender exclusions)
shell: pwsh
@ -148,6 +161,7 @@ jobs:
shell: pwsh
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
HF_TOKEN: ${{ secrets.HF_TOKEN }}
run: |
New-Item -ItemType Directory -Force -Path logs | Out-Null
# *>&1 captures Write-Host (Information stream) output;
@ -360,6 +374,9 @@ jobs:
- name: Collect llama-server logs
if: always()
# A transient Windows DLL-init crash (0xC0000142) in this diagnostic
# copy must not fail an otherwise-green job.
continue-on-error: true
shell: bash
# Copy llama-server's own stdout/stderr (teed by Studio under
# ~/.unsloth/studio/logs/llama-server/) into the workspace so
@ -373,6 +390,8 @@ jobs:
- name: Upload logs
if: always()
# Diagnostic only: a transient artifact-service drop must not fail a green job.
continue-on-error: true
uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
with:
name: windows-openai-anthropic-log
@ -487,6 +506,7 @@ jobs:
shell: pwsh
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
HF_TOKEN: ${{ secrets.HF_TOKEN }}
run: |
New-Item -ItemType Directory -Force -Path logs | Out-Null
# *>&1 captures Write-Host (Information stream) output;
@ -788,6 +808,9 @@ jobs:
- name: Collect llama-server logs
if: always()
# A transient Windows DLL-init crash (0xC0000142) in this diagnostic
# copy must not fail an otherwise-green job.
continue-on-error: true
shell: bash
# Copy llama-server's own stdout/stderr (teed by Studio under
# ~/.unsloth/studio/logs/llama-server/) into the workspace so
@ -801,6 +824,8 @@ jobs:
- name: Upload logs
if: always()
# Diagnostic only: a transient artifact-service drop must not fail a green job.
continue-on-error: true
uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
with:
name: windows-tool-calling-log
@ -821,9 +846,9 @@ jobs:
run:
shell: bash
env:
GGUF_REPO: unsloth/gemma-4-E2B-it-GGUF
GGUF_VARIANT: UD-Q4_K_XL
GGUF_FILE: gemma-4-E2B-it-UD-Q4_K_XL.gguf
GGUF_REPO: unsloth/Qwen3-VL-2B-Instruct-GGUF
GGUF_VARIANT: UD-IQ2_XXS
GGUF_FILE: Qwen3-VL-2B-Instruct-UD-IQ2_XXS.gguf
MMPROJ_FILE: mmproj-F16.gguf
STUDIO_PORT: '18899'
HF_HOME: ${{ github.workspace }}/hf-cache
@ -857,7 +882,7 @@ jobs:
continue-on-error: true
with:
path: hf-cache
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-${{ env.MMPROJ_FILE }}-v1
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-${{ env.MMPROJ_FILE }}-v2
- name: Prime HF_HOME with the GGUF + mmproj
id: prime-hf
@ -869,13 +894,14 @@ jobs:
mkdir -p hf-cache
bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE"
bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$MMPROJ_FILE"
bash .github/scripts/hf-download-with-retry.sh ggml-org/models tinyllamas/stories260K.gguf
- name: Save HF_HOME cache for ${{ env.GGUF_REPO }} (model + mmproj)
if: always() && steps.prime-hf.outcome == 'success'
uses: actions/cache/save@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5
with:
path: hf-cache
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-${{ env.MMPROJ_FILE }}-v1
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-${{ env.MMPROJ_FILE }}-v2
- name: Pre-install Windows tweaks (npm 11 + Defender exclusions)
shell: pwsh
@ -912,6 +938,7 @@ jobs:
shell: pwsh
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
HF_TOKEN: ${{ secrets.HF_TOKEN }}
run: |
New-Item -ItemType Directory -Force -Path logs | Out-Null
# *>&1 captures Write-Host (Information stream) output;
@ -1101,7 +1128,7 @@ jobs:
)
data_uri = f"data:image/png;base64,{PNG_64X64_RED_B64}"
# On Windows + the gemma-4-E2B mmproj, llama.cpp's vision
# On Windows + the Qwen3-VL mmproj, llama.cpp's vision
# path runs on CPU (no Metal involvement). The wrapper is
# kept for resilience but the vision path is expected to
# work on Windows; an exception here is a real regression.
@ -1186,6 +1213,9 @@ jobs:
- name: Collect llama-server logs
if: always()
# A transient Windows DLL-init crash (0xC0000142) in this diagnostic
# copy must not fail an otherwise-green job.
continue-on-error: true
shell: bash
# Copy llama-server's own stdout/stderr (teed by Studio under
# ~/.unsloth/studio/logs/llama-server/) into the workspace so
@ -1199,6 +1229,8 @@ jobs:
- name: Upload logs
if: always()
# Diagnostic only: a transient artifact-service drop must not fail a green job.
continue-on-error: true
uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7.0.1
with:
name: windows-json-images-log

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

@ -85,7 +85,7 @@ jobs:
continue-on-error: true
with:
path: hf-cache
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v1
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2
- name: Prime HF_HOME with the GGUF
id: prime-hf
@ -96,13 +96,14 @@ jobs:
python -m pip install --upgrade huggingface_hub
mkdir -p hf-cache
bash .github/scripts/hf-download-with-retry.sh "$GGUF_REPO" "$GGUF_FILE"
bash .github/scripts/hf-download-with-retry.sh ggml-org/models tinyllamas/stories260K.gguf
- name: Save HF_HOME for ${{ env.GGUF_REPO }}
if: always() && steps.prime-hf.outcome == 'success'
uses: actions/cache/save@27d5ce7f107fe9357f9df03efb73ab90386fccae # v5.0.5
with:
path: hf-cache
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v1
key: ${{ runner.os }}-hf-${{ env.GGUF_REPO }}-${{ env.GGUF_VARIANT }}-v2
- name: Pre-install Windows tweaks (npm 11 + Defender exclusions)
shell: pwsh
@ -143,6 +144,7 @@ jobs:
shell: pwsh
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
HF_TOKEN: ${{ secrets.HF_TOKEN }}
run: |
New-Item -ItemType Directory -Force -Path logs | Out-Null
# *>&1 redirects ALL PowerShell streams (stdout, stderr,

3
.github/workflows/studio-windows-update-smoke.yml vendored
View file

@ -133,6 +133,7 @@ jobs:
shell: pwsh
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
HF_TOKEN: ${{ secrets.HF_TOKEN }}
run: |
New-Item -ItemType Directory -Force -Path logs | Out-Null
# *>&1 captures Write-Host (Information stream) output;
@ -179,6 +180,7 @@ jobs:
- name: First update should be a no-op (prebuilt already validated)
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
HF_TOKEN: ${{ secrets.HF_TOKEN }}
run: |
set -o pipefail
unsloth studio update --local 2>&1 | tee logs/update.log
@ -197,6 +199,7 @@ jobs:
- name: Second update must also be a no-op
env:
GH_TOKEN: ${{ secrets.GITHUB_TOKEN }}
HF_TOKEN: ${{ secrets.HF_TOKEN }}
run: |
set -o pipefail
unsloth studio update --local 2>&1 | tee logs/update2.log

17
.pre-commit-config.yaml
View file

@ -1,6 +1,6 @@
repos:
- repo: https://github.com/astral-sh/ruff-pre-commit
rev: v0.15.13
rev: v0.15.17
hooks:
- id: ruff
args:
@ -14,5 +14,20 @@ repos:
entry: scripts/run_ruff_format.py
language: python
types: [python]
# Mirror ruff's [tool.ruff] extend-exclude so this hook does not
# half-process files ruff itself skips (which produced churn).
exclude: '(chat_templates|ollama_template_mappers|_auto_install|mapper)\.py$'
additional_dependencies:
- ruff==0.6.9
# Re-pins allowScripts entries after dependency bumps. pre-commit.ci
# pushes the fix to PR branches, Dependabot's included, so stale pins
# heal without a human in the loop.
- id: sync-allow-scripts-pins
name: Sync allowScripts pins with the frontend lockfile
# `python <script>` not a direct exec: autofix commits can drop the
# executable bit, which kills shebang-style entries.
entry: python scripts/sync_allow_scripts_pins.py
args: [--fix]
language: python
files: ^studio/frontend/(package\.json|package-lock\.json)$
pass_filenames: false

6
CONTRIBUTING.md
View file

@ -27,3 +27,9 @@ Your support extends beyond code:
Finally, please be mindful of our [Code of Conduct](https://github.com/unslothai/unsloth/blob/main/CODE_OF_CONDUCT.md) to ensure a welcoming and inclusive environment for everyone.
Thank you so much for reading and we hope you have lots of fun using Unsloth! 🦥
## Pull Request Guidelines
- Keep PRs focused on a single change
- Include a concise description and motivation
- Link related issues when applicable

45
README.md
View file

@ -72,10 +72,13 @@ Unsloth Studio (Beta) works on **Windows, Linux, WSL** and **macOS**.
```bash
curl -fsSL https://unsloth.ai/install.sh | sh
```
Use the same command to update.
#### Windows:
```powershell
irm https://unsloth.ai/install.ps1 | iex
```
Use the same command to update.
#### Launch
```bash
@ -83,9 +86,6 @@ unsloth studio -p 8888
```
For cloud or global access, add `-H 0.0.0.0`. By default, Unsloth is accessible only locally.
#### Update
To update, use the same install commands above or use `unsloth studio update`.
#### Docker
Use our [Docker image](https://hub.docker.com/r/unsloth/unsloth) ```unsloth/unsloth``` container. Run:
```bash
@ -171,7 +171,9 @@ unsloth studio -p 8888
```
Then to update :
```bash
unsloth studio update
cd unsloth && git pull
./install.sh --local
unsloth studio -p 8888
```
#### Developer installs: Windows PowerShell:
@ -184,7 +186,9 @@ unsloth studio -p 8888
```
Then to update :
```bash
unsloth studio update
cd unsloth && git pull
./install.sh --local
unsloth studio -p 8888
```
#### Nightly: MacOS, Linux, WSL:
@ -202,7 +206,7 @@ unsloth studio -p 8888
#### Nightly: Windows:
Run in Windows Powershell:
```bash
```powershell
git clone https://github.com/unslothai/unsloth.git
cd unsloth
git checkout nightly
@ -215,6 +219,35 @@ Then to launch every time:
unsloth studio -p 8888
```
#### Advanced launch options
Installer options can be passed as environment variables. On macOS, Linux and WSL place the variable after the pipe so the shell passes it to `sh`; on Windows set it with `$env:` before piping to `iex`.
Skip PyTorch (GGUF-only mode):
```bash
curl -fsSL https://unsloth.ai/install.sh | UNSLOTH_NO_TORCH=1 sh
```
```powershell
$env:UNSLOTH_NO_TORCH=1; irm https://unsloth.ai/install.ps1 | iex
```
Pin the Python version:
```bash
curl -fsSL https://unsloth.ai/install.sh | UNSLOTH_PYTHON=3.12 sh
```
```powershell
$env:UNSLOTH_PYTHON='3.12'; irm https://unsloth.ai/install.ps1 | iex
```
Install to a custom location with `UNSLOTH_STUDIO_HOME`:
```bash
curl -fsSL https://unsloth.ai/install.sh | UNSLOTH_STUDIO_HOME=/abs/path sh
```
```powershell
$env:UNSLOTH_STUDIO_HOME='C:\path'; irm https://unsloth.ai/install.ps1 | iex
```
Cap Studio's native CPU thread pools on high-core hosts: `UNSLOTH_CPU_THREADS=8 unsloth studio -p 8888`.
#### Uninstall
The recommended way to fully remove Unsloth Studio is the matching uninstall script for your OS. It stops any running servers, removes the install dir, the launcher data dir, the desktop shortcut, and any platform-specific entries (macOS `.app` bundle + Launch Services on Mac; Start Menu, `HKCU\Software\Unsloth` registry key and user `PATH` entries on Windows):

818
install.ps1
View file

File diff suppressed because it is too large Load diff

856
install.sh
View file

File diff suppressed because it is too large Load diff

9
pyproject.toml
View file

@ -26,6 +26,7 @@ classifiers = [
]
dependencies = [
"typer",
"rich",
"pydantic",
"pyyaml",
"nest-asyncio",
@ -54,6 +55,7 @@ studio = [
"frontend/.git*",
"backend/requirements/**/*",
"backend/plugins/**/*",
"backend/assets/**/*.jinja",
"backend/core/data_recipe/oxc-validator/*.json",
"backend/core/data_recipe/oxc-validator/*.mjs",
]
@ -69,7 +71,7 @@ triton = [
]
huggingfacenotorch = [
"unsloth_zoo>=2026.5.4",
"unsloth_zoo>=2026.6.5",
"wheel>=0.42.0",
"packaging",
"numpy",
@ -90,7 +92,7 @@ huggingfacenotorch = [
]
huggingface = [
"unsloth[huggingfacenotorch]",
"unsloth_zoo>=2026.5.4",
"unsloth_zoo>=2026.6.5",
"torchvision",
"unsloth[triton]",
]
@ -580,7 +582,7 @@ colab-ampere-torch220 = [
"flash-attn>=2.6.3 ; ('linux' in sys_platform)",
]
colab-new = [
"unsloth_zoo>=2026.5.4",
"unsloth_zoo>=2026.6.5",
"packaging",
"tyro",
"transformers>=4.51.3,!=4.52.0,!=4.52.1,!=4.52.2,!=4.52.3,!=4.53.0,!=4.54.0,!=4.55.0,!=4.55.1,!=4.57.0,!=4.57.4,!=4.57.5,!=5.0.0,!=5.1.0,<=5.5.0",
@ -1308,6 +1310,7 @@ repository = "https://github.com/unslothai/unsloth"
[tool.ruff]
target-version = "py311"
line-length = 100
force-exclude = true
extend-exclude = [
"*chat_templates.py",

393
scripts/check_frontend_dep_removal.py
View file

@ -43,7 +43,7 @@ DEP_FIELDS = (
"optionalDependencies",
)
# Sources where seeing a package name does NOT count as usage.
# Files where seeing a package name does NOT count as usage.
EXPECTED_NOISE_FILES = {
"studio/frontend/package.json",
"studio/frontend/package-lock.json",
@ -51,19 +51,15 @@ EXPECTED_NOISE_FILES = {
"studio/backend/core/data_recipe/oxc-validator/package-lock.json",
}
# Only quoted-string occurrences in these file types can be module specifiers.
JS_LIKE_EXT = re.compile(
r"\.(ts|tsx|js|jsx|mjs|cjs|html|htm|css|scss|sass|json|jsonc)$"
)
# Files where JS-syntactic import patterns (static/dynamic/require/re-export)
# could be a real module reference. Markdown gets a separate gate (.mdx is
# File types where a quoted string can be a module specifier.
JS_LIKE_EXT = re.compile(r"\.(ts|tsx|js|jsx|mjs|cjs|html|htm|css|scss|sass|json|jsonc)$")
# Files where JS import patterns could be a real module reference (.mdx is
# real ESM; .md code fences are not).
SCRIPT_LIKE_EXT = re.compile(r"\.(ts|tsx|js|jsx|mjs|cjs|mdx)$")
STYLE_EXT = re.compile(r"\.(css|scss|sass)$")
HTML_EXT = re.compile(r"\.(html|htm)$")
TS_LIKE_EXT = re.compile(r"\.(ts|tsx|mts|cts|mdx)$")
# Files where a removed package's CLI binary could be invoked (npx, bunx,
# yarn dlx, pnpm exec, or a bare `pkg --flag` shell call).
# Files where a removed package's CLI binary could be invoked.
COMMAND_LIKE_EXT = re.compile(r"(\.(ya?ml|sh|ps1|bat)$|(^|/)Dockerfile[^/]*$)")
GREP_INCLUDES = [
@ -104,7 +100,7 @@ GREP_EXCLUDES = [
"--exclude-dir=venv",
]
# A pip-installed playwright reference is the PyPI package, not npm.
# A pip-installed playwright ref is the PyPI package, not npm.
PIP_PLAYWRIGHT = re.compile(
r"(pip\s+install\s+['\"]?playwright"
r"|python\s+-m\s+playwright"
@ -155,9 +151,8 @@ def all_decl_names(pkg: dict) -> set[str]:
def _resolve_install_path(parent_path: str, name: str, pkgs: dict) -> str | None:
"""Walk up the nested node_modules chain from `parent_path` to find
where `name` actually resolves. Mirrors Node module resolution.
"""
"""Walk up the nested node_modules chain from `parent_path` to find where
`name` resolves, mirroring Node module resolution."""
parts = parent_path.split("/node_modules/")
for i in range(len(parts), 0, -1):
prefix = "/node_modules/".join(parts[:i])
@ -170,11 +165,8 @@ def _resolve_install_path(parent_path: str, name: str, pkgs: dict) -> str | None
def _deps_of(meta: dict) -> dict:
"""Deps npm actually installs. Optional peers are skipped: npm only
installs them when another package declares the same dep, so for the
purpose of "is this package still reachable" they cannot keep a
removed top-level dep alive on their own.
"""
"""Deps npm actually installs. Optional peers are skipped: they can't keep
a removed top-level dep reachable on their own."""
out = {}
for field in ("dependencies", "optionalDependencies"):
out.update(meta.get(field) or {})
@ -187,10 +179,8 @@ def _deps_of(meta: dict) -> dict:
def reachable_from_head(head_pkg: dict, lock: dict) -> set[str]:
"""BFS the lockfile dep graph starting from `head_pkg`'s top-level
declared deps. Returns the set of lockfile install paths that survive.
Stale lockfile entries (orphaned by the new package.json) are excluded.
"""
"""BFS the lockfile dep graph from `head_pkg`'s top-level deps. Returns the
surviving install paths, excluding stale (orphaned) lockfile entries."""
pkgs = lock.get("packages", {})
if not pkgs:
return set()
@ -217,23 +207,17 @@ def reachable_from_head(head_pkg: dict, lock: dict) -> set[str]:
def classify(pkg: str, file: str, content: str) -> str | None:
"""Return why `content` references `pkg`, or None.
`content` may span multiple lines (for multi-line imports/exports);
each pattern uses re.DOTALL where it matters. The bare-spec
regexes use a word-boundary check on the package name so that
`foobar` does not match `foo`.
File-type gating: JS-syntactic patterns only fire on .ts/.tsx/.js/.jsx/
.mjs/.cjs/.mdx files, so an `import x from "pkg"` snippet inside a
Python test fixture or a Markdown code block is not mistaken for a
real npm usage. CSS patterns only fire on .css/.scss/.sass. HTML
patterns only fire on .html/.htm.
`content` may span multiple lines (multi-line imports/exports use re.DOTALL).
Bare-spec regexes word-boundary the package name so `foobar` doesn't match
`foo`. File-type gating restricts JS patterns to .ts/.tsx/.js/.jsx/.mjs/
.cjs/.mdx, CSS to .css/.scss/.sass, HTML to .html/.htm, so a snippet inside
a Python fixture or Markdown code block isn't mistaken for real npm usage.
"""
if file in EXPECTED_NOISE_FILES:
return None
esc = re.escape(pkg)
# Subpath gate: after the package name, the next char must be either
# the closing quote, `/`, or end-of-string. Prevents foo matching foobar.
# Subpath gate: pkg must be followed by quote, `/`, or end-of-string.
sub = r"(?:/[^'\"`]*)?"
flags_dotall = re.DOTALL | re.MULTILINE
@ -243,68 +227,51 @@ def classify(pkg: str, file: str, content: str) -> str | None:
is_html = bool(HTML_EXT.search(file))
is_ts = bool(TS_LIKE_EXT.search(file))
# If the file is none of script / style / html / json (which is the
# quoted-string fallback surface) and is not an mdx file, no classify
# rule applies. This is what gates out Python fixtures, Markdown code
# blocks, shell snippets, etc.
# Gate out Python fixtures, Markdown code blocks, shell snippets, etc.
is_json = file.endswith(".json") or file.endswith(".jsonc")
if not (is_script or is_style or is_html or is_json):
return None
# CSS @import is checked first so it does not collide with the
# side-effect-import regex below.
# CSS @import first so it doesn't collide with side-effect-import below.
if is_style and re.search(rf"@import\s+['\"]{esc}{sub}['\"]", content):
return "css_import"
# Static imports: handle multi-line `import { ... } from "pkg"` by
# allowing arbitrary content (newlines included) between `import`
# and `from`. The non-greedy match plus the required `from` keeps
# this scoped to a single statement.
# Static imports, including multi-line `import { ... } from "pkg"`.
if is_script and re.search(
rf"(?<!@)\bimport\b[^;'\"]*?\bfrom\s+['\"]{esc}{sub}['\"]",
content,
flags_dotall,
):
return "static_import"
# Side-effect import: `import "pkg"` (no `from`). The negative
# lookbehind rules out CSS `@import` lines.
# Side-effect import `import "pkg"` (no `from`); lookbehind rules out @import.
if is_script and re.search(rf"(?<!@)\bimport\s+['\"]{esc}{sub}['\"]", content):
return "side_effect_import"
# Dynamic import: `import("pkg")` and `await import("pkg")`.
if is_script and re.search(rf"\bimport\(\s*['\"]{esc}{sub}['\"]\s*\)", content):
return "dynamic_import"
# require / require.resolve
if is_script and re.search(
rf"\brequire(?:\.resolve)?\(\s*['\"]{esc}{sub}['\"]\s*\)", content
):
if is_script and re.search(rf"\brequire(?:\.resolve)?\(\s*['\"]{esc}{sub}['\"]\s*\)", content):
return "require"
# Re-exports: `export * from "pkg"`, `export { x } from "pkg"`,
# `export type { Foo } from "pkg"`. Multi-line supported.
# Re-exports: `export * from`, `export { x } from`, `export type { Foo } from`.
if is_script and re.search(
rf"\bexport\b[^;'\"]*?\bfrom\s+['\"]{esc}{sub}['\"]",
content,
flags_dotall,
):
return "re_export"
# HTML script / link. Match the package name as a complete path
# segment bounded by a quote / `#` / `?` or a subpath `/`, so
# `/node_modules/foo-extra/...` is NOT treated as usage of `foo`.
# HTML script / link. Match pkg as a complete path segment so
# `/node_modules/foo-extra/...` is not treated as usage of `foo`.
html_pkg = rf"{esc}(?:/[^'\"#?]*)?(?=['\"#?])"
if is_html and re.search(
rf"<script[^>]*src\s*=\s*['\"][^'\"]*/{html_pkg}", content
):
if is_html and re.search(rf"<script[^>]*src\s*=\s*['\"][^'\"]*/{html_pkg}", content):
return "html_script"
if is_html and re.search(rf"<link[^>]*href\s*=\s*['\"][^'\"]*/{html_pkg}", content):
return "html_link"
# TypeScript triple-slash
if is_ts and re.search(
rf"///\s*<reference\s+types\s*=\s*['\"]{esc}{sub}['\"]", content
):
if is_ts and re.search(rf"///\s*<reference\s+types\s*=\s*['\"]{esc}{sub}['\"]", content):
return "tsc_triple_slash"
# new URL("pkg/...", import.meta.url)
if is_script and re.search(rf"\bnew\s+URL\(\s*['\"]{esc}{sub}['\"]", content):
return "new_url"
# CSS url(...). Accept quoted ("pkg/x") AND unquoted (pkg/x) variants,
# bounded by a path-segment lookahead so `pkg-extra` does not match.
# CSS url(...), quoted and unquoted, bounded so `pkg-extra` doesn't match.
if is_style and re.search(
rf"\burl\(\s*['\"]?(?:[^)'\"\s]+/)?{esc}(?:/[^)'\"`]*)?['\"]?\s*\)",
content,
@ -317,20 +284,18 @@ def classify(pkg: str, file: str, content: str) -> str | None:
if is_script and re.search(rf"@import\(\s*['\"]{esc}{sub}['\"]\s*\)", content):
return "jsdoc_import"
# Bare quoted-string fallback (config plugin lists, vite aliases,
# tsconfig paths, biome config plugin arrays, shadcn registries).
# tsconfig paths, biome plugin arrays, shadcn registries).
if not JS_LIKE_EXT.search(file):
return None
# Boundary: pkg must be followed by `'`, `"`, or `/` to avoid
# matching `foo` inside `foobar`.
# pkg must be followed by `'`, `"`, or `/` so `foo` doesn't match `foobar`.
if re.search(rf"['\"]{esc}(?:['\"]|/)", content):
return "string_literal"
return None
def lockfile_root_sync(head_pkg: dict, head_lock: dict) -> list[str]:
"""Return a list of warnings if package-lock.json's <root> dep map
disagrees with package.json (i.e., npm install was not re-run).
"""
"""Warn if package-lock.json's <root> dep map disagrees with package.json
(i.e. npm install was not re-run)."""
warnings = []
if not head_lock:
return warnings
@ -358,10 +323,8 @@ def lockfile_root_sync(head_pkg: dict, head_lock: dict) -> list[str]:
def types_orphan_warnings(head_pkg: dict) -> list[str]:
"""Flag @types/<X> deps where <X> is no longer declared anywhere
in package.json. Removing X without also dropping @types/X leaves
dangling type packages.
"""
"""Flag @types/<X> deps where <X> is no longer declared in package.json,
which leaves dangling type packages."""
decl = set()
for f in DEP_FIELDS:
decl.update((head_pkg.get(f) or {}).keys())
@ -369,9 +332,7 @@ def types_orphan_warnings(head_pkg: dict) -> list[str]:
for name in decl:
if not name.startswith("@types/"):
continue
# @types/foo provides types for `foo`
# @types/foo-bar provides types for `foo-bar`
# @types/scope__pkg provides types for `@scope/pkg`
# @types/scope__pkg provides types for @scope/pkg.
target = name[len("@types/") :]
if "__" in target:
scope, sub = target.split("__", 1)
@ -394,8 +355,7 @@ _PKG_JSON_SKIP_KEYS = {
"bundledDependencies",
}
# Top-level fields whose contents are never package references. We walk
# everything else recursively.
# Top-level fields whose contents are never package references.
_PKG_JSON_OPAQUE_KEYS = {
"browserslist", # browser queries
"keywords", # free-form strings
@ -437,20 +397,12 @@ _PKG_JSON_OPAQUE_KEYS = {
def package_json_extra_refs(pkg: dict, target: str) -> list[str]:
"""Walk every key/value in package.json EXCEPT the dep declaration
blocks, and return citations for string values or dict keys that
equal `target` (or `target/subpath`).
"""Walk package.json (except dep declaration blocks) and return citations
for string values or dict keys equal to `target` (or `target/subpath`).
Catches the patterns the public dep-checker tools commonly miss:
- `overrides` / `resolutions` / `pnpm.overrides` keys
- `pnpm.patchedDependencies` keys
- `peerDependenciesMeta` keys
- `prettier`: "@my/prettier-config"
- `eslintConfig.extends`: ["..."] / "..."
- `stylelint.extends` / `stylelint.plugins`
- `babel.presets` / `babel.plugins`
- `jest.preset` / `jest.setupFiles` / `jest.transform`
- `commitlint.extends`, `renovate.extends`, `remarkConfig.plugins`
Catches refs that public dep-checkers commonly miss: overrides/resolutions/
pnpm.overrides keys, pnpm.patchedDependencies, peerDependenciesMeta,
prettier, eslintConfig.extends, stylelint, babel, jest, commitlint, etc.
"""
target_sub = target + "/"
cites: list[str] = []
@ -461,14 +413,11 @@ def package_json_extra_refs(pkg: dict, target: str) -> list[str]:
def walk(obj: object, path: str) -> None:
if isinstance(obj, dict):
for k, v in obj.items():
# Skip top-level dep declaration fields entirely.
if path == "" and k in _PKG_JSON_SKIP_KEYS:
continue
# Top-level fields whose contents are never package refs.
if path == "" and k in _PKG_JSON_OPAQUE_KEYS:
continue
# Inside `overrides` / `resolutions` / etc., the KEY itself
# is a package reference.
# Inside overrides/resolutions/etc., the KEY is a package ref.
if matches(k):
cites.append(f"{path}.{k}" if path else k)
walk(v, f"{path}.{k}" if path else k)
@ -484,9 +433,8 @@ def package_json_extra_refs(pkg: dict, target: str) -> list[str]:
def build_bin_to_pkg(head_lock: dict) -> dict[str, str]:
"""Map a binary name (e.g. 'vite', 'tsc', 'eslint') to the package
that provides it. Built from each lockfile entry's `bin` field.
"""
"""Map a binary name (e.g. 'vite', 'eslint') to its providing package,
from each lockfile entry's `bin` field."""
out: dict[str, str] = {}
if not head_lock:
return out
@ -505,70 +453,47 @@ def build_bin_to_pkg(head_lock: dict) -> dict[str, str]:
_SCRIPT_TOKENIZE = re.compile(r"\s*(?:&&|\|\||;|\|(?!\|))\s*")
# Wrappers that delegate to a real CLI in the same shell word list.
# After stripping env prefixes and (optionally) `npx`/`pnpm exec`/`yarn dlx`/
# `bunx`, if the leading token is one of these we advance past the
# wrapper's own flags and any further env-prefix tokens, then re-check.
# `cross-env` is the common one; `dotenv-cli` / `dotenvx` use `--` as a
# separator. Wrappers that operate on named npm-scripts (concurrently,
# npm-run-all, run-s, run-p, wireit, turbo, nx) intentionally aren't
# here -- they reference script names, not bin names, so the real bin
# is in the *target* script's chunk which we already tokenize.
# Wrappers that delegate to a real CLI in the same shell word list; we skip
# past them and their flags to find the wrapped bin. Script-name wrappers
# (concurrently, npm-run-all, turbo, nx) are excluded: they reference script
# names, so the real bin lives in the target script's chunk we already tokenize.
_SCRIPT_WRAPPERS = {"cross-env", "dotenv", "dotenvx", "env-cmd"}
_ENV_PREFIX_RE = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*=")
def _next_real_bin(words: list[str], idx: int) -> str | None:
"""Walk `words` from `idx`, peeling env-prefix tokens, the leading
package-manager runner (`npx`, `pnpm exec`, etc.), and the known
wrapper bins. Return the next token that looks like the real CLI
binary, or None if the chunk has nothing to look up.
Recursion depth is bounded by the chunk's word count, so the loop
cannot run away on a pathological wrapper chain.
"""
"""Walk `words` from `idx`, peeling env-prefix tokens, the package-manager
runner (npx, pnpm exec, etc.), and known wrapper bins. Return the next
real CLI binary, or None. Bounded by the chunk's word count."""
seen_wrappers: set[str] = set()
while idx < len(words):
# 1. env-prefix run: `FOO=bar BAZ="a b" cmd ...`. shlex has
# already collapsed quoted values into one word, so this
# tokenizer is safe for them.
# 1. env-prefix run `FOO=bar BAZ="a b" cmd ...` (shlex pre-collapsed).
while idx < len(words) and _ENV_PREFIX_RE.match(words[idx]):
idx += 1
if idx >= len(words):
return None
first = words[idx]
# 2. Package-manager runner: `npx <pkg> args`, `pnpm exec <pkg>`,
# `yarn dlx <pkg>`, `bunx <pkg>`. Strip and continue (so the
# wrapped command goes through the same unwrap loop).
# 2. Package-manager runner (npx/pnpm exec/yarn dlx/bunx): strip and
# continue so the wrapped command re-enters the unwrap loop.
if first in {"npx", "pnpx", "bunx"} and idx + 1 < len(words):
idx += 1
continue
if (
first in {"pnpm", "yarn"}
and idx + 2 < len(words)
and words[idx + 1] in {"exec", "dlx"}
):
if first in {"pnpm", "yarn"} and idx + 2 < len(words) and words[idx + 1] in {"exec", "dlx"}:
idx += 2
continue
# 3. Wrapper bin (cross-env, dotenv, etc.). Skip the wrapper's
# own flags and any subsequent env-prefix tokens, then re-loop.
bin_token = first.removeprefix("./node_modules/.bin/").removeprefix(
"node_modules/.bin/"
)
# 3. Wrapper bin (cross-env, dotenv): skip its flags and env prefixes.
bin_token = first.removeprefix("./node_modules/.bin/").removeprefix("node_modules/.bin/")
if bin_token in _SCRIPT_WRAPPERS and bin_token not in seen_wrappers:
seen_wrappers.add(bin_token)
idx += 1
# cross-env / env-cmd: no flags; just more env-prefix tokens.
# dotenv / dotenvx: skip `-e <file>` style flags and the
# optional `--` separator before the wrapped command.
# dotenv/dotenvx use `-e <file>` flags and an optional `--`.
while idx < len(words):
tok = words[idx]
if tok.startswith("-") and tok != "--":
idx += 1
# `-e .env` style: also skip the flag's argument
# when it does not look like another flag.
# `-e .env`: also skip the flag's argument.
if (
idx < len(words)
and not words[idx].startswith("-")
@ -585,21 +510,13 @@ def _next_real_bin(words: list[str], idx: int) -> str | None:
return None
def scripts_bin_refs(
head_pkg: dict, bin_to_pkg: dict[str, str]
) -> dict[str, list[str]]:
"""Return `{package_name: ['scripts.X: cmd', ...]}` listing every
package referenced via its bin name in package.json scripts.
def scripts_bin_refs(head_pkg: dict, bin_to_pkg: dict[str, str]) -> dict[str, list[str]]:
"""Return `{package_name: ['scripts.X: cmd', ...]}` for every package
referenced via its bin name in package.json scripts.
Each script value is split on shell separators (`&&`, `||`, `;`,
`|`). Within each chunk, `_next_real_bin()` unwraps env prefixes,
package-manager runners (`npx` / `pnpm exec` / `yarn dlx` / `bunx`),
and wrapper bins like `cross-env` / `dotenv` so that
`cross-env CI=1 biome check` correctly credits `biome` to its
declaring package.
Tokenization uses shlex.split so quoted env values
(`FOO="a b" biome`) survive unbroken.
Each script is split on shell separators; `_next_real_bin()` unwraps env
prefixes, package-manager runners, and wrapper bins so `cross-env CI=1
biome check` credits `biome`. Uses shlex.split so quoted env values survive.
"""
import shlex
@ -615,7 +532,7 @@ def scripts_bin_refs(
try:
words = shlex.split(chunk, posix = True)
except ValueError:
# Unbalanced quotes -- fall back to plain split.
# Unbalanced quotes: fall back to plain split.
words = chunk.split()
if not words:
continue
@ -629,11 +546,8 @@ def scripts_bin_refs(
def tsconfig_compiler_types_refs() -> set[str]:
"""Read studio/frontend/tsconfig*.json and return the set of
package names referenced in compilerOptions.types arrays. These are
implicitly loaded by tsc and count as a real use even though they
have no explicit import.
"""
"""Return package names in tsconfig*.json compilerOptions.types arrays.
These are implicitly loaded by tsc and count as real uses."""
out: set[str] = set()
base = REPO_ROOT / "studio/frontend"
for name in ("tsconfig.json", "tsconfig.app.json", "tsconfig.node.json"):
@ -651,30 +565,16 @@ def tsconfig_compiler_types_refs() -> set[str]:
for t in types:
if not isinstance(t, str):
continue
# `vite/client` resolves to `vite` package.
pkg = (
t.split("/", 1)[0]
if not t.startswith("@")
else "/".join(t.split("/", 2)[:2])
)
# `vite/client` resolves to the `vite` package.
pkg = t.split("/", 1)[0] if not t.startswith("@") else "/".join(t.split("/", 2)[:2])
out.add(pkg)
return out
def enumerate_dep_usage(head_pkg: dict, head_lock: dict) -> dict[str, list]:
"""For every declared dep, classify whether it appears used. Returns
a dict with these categories:
- used: has at least one detected usage in src/,
config files, scripts.bin, package.json
field refs, or tsconfig types
- unused: no detected usage anywhere
- type_pkg_kept: @types/X where X is still declared
- type_pkg_orphan: @types/X where X is no longer declared
(or X is removed) -- candidate for removal
Each entry is the package name. The categorisation is opinionated;
`unused` is a CANDIDATE list, not a guarantee. The caller should
verify before deletion.
"""For every declared dep, classify usage into a dict of package-name lists:
used, unused, type_pkg_kept (@types/X with X declared), type_pkg_orphan
(@types/X with X gone). `unused` is a CANDIDATE list; verify before deletion.
"""
decl = all_decl_names(head_pkg)
bin_to_pkg = build_bin_to_pkg(head_lock) if head_lock else {}
@ -703,11 +603,8 @@ def enumerate_dep_usage(head_pkg: dict, head_lock: dict) -> dict[str, list]:
# Real-source-usage check
hits = find_usage(name)
used = bool(hits)
# CLI usage in shell / workflow / Dockerfile surfaces. Skip for
# `@types/*` packages because they never expose a CLI binary and
# the unscoped-tail bin name candidate would scan workflow files
# for the bare runtime name (a removed `@types/foo` would look
# for invocations of `foo`).
# CLI usage in shell/workflow/Dockerfile. Skipped for @types/* (no CLI
# binary; the bare-name bin candidate would false-match the runtime).
if not used and not name.startswith("@types/") and find_command_usage(name):
used = True
# Bin scripts
@ -727,28 +624,15 @@ def enumerate_dep_usage(head_pkg: dict, head_lock: dict) -> dict[str, list]:
def find_imports_without_decl(head_pkg: dict) -> list[tuple[str, int, str]]:
"""Reverse check: find bare-specifier imports in studio/frontend/src
that don't correspond to any declared package.json dep. Catches the
case where someone adds an import but forgets the dep declaration.
Returns (file, line, spec) tuples.
Match shapes covered:
import "pkg"
import Foo from "pkg"
import { Foo } from "pkg"
import type { Foo } from "pkg"
const x = require("pkg")
const x = await import("pkg")
"""Reverse check: find bare-specifier imports in studio/frontend/src with
no matching package.json dep (import added but dep declaration forgotten).
Covers import/require/dynamic-import shapes. Returns (file, line, spec).
"""
decl = set()
for f in DEP_FIELDS:
decl.update((head_pkg.get(f) or {}).keys())
# Also: anything tsconfig path-aliases (just '@/...' here) is internal.
# The capture group is the specifier; the leading alternation accepts
# any of: `from "..."`, bare side-effect `import "..."`,
# `import("..."), or `require("...")`. We exclude relative paths and
# the `@/` alias prefix by requiring the first char of the specifier
# to be neither `.` nor `/`.
# Exclude relative paths and the `@/` alias by requiring the specifier's
# first char to be neither `.` nor `/`. Capture group is the specifier.
pattern = (
r"(?:\bfrom\s+|"
r"\bimport\s+(?:\(\s*)?|"
@ -774,7 +658,7 @@ def find_imports_without_decl(head_pkg: dict) -> list[tuple[str, int, str]]:
file, ln, content = m.group(1), int(m.group(2)), m.group(3)
for spec_match in re.finditer(pattern, content):
spec = spec_match.group(1)
# Resolve to package name (strip subpath)
# Resolve to package name (strip subpath).
if spec.startswith("@"):
parts = spec.split("/", 2)
pkg_name = "/".join(parts[:2]) if len(parts) >= 2 else spec
@ -782,7 +666,7 @@ def find_imports_without_decl(head_pkg: dict) -> list[tuple[str, int, str]]:
pkg_name = spec.split("/", 1)[0]
if pkg_name in decl:
continue
# Internal aliases like '@/foo' or starts with builtin names
# Internal aliases like '@/foo' or builtin names.
if pkg_name == "@":
continue
if pkg_name in {
@ -820,20 +704,17 @@ _file_lines_cache: dict[str, list[str]] = {}
def _read_file(path: str) -> list[str]:
if path not in _file_lines_cache:
try:
_file_lines_cache[path] = (
Path(path).read_text(errors = "replace").splitlines()
)
_file_lines_cache[path] = Path(path).read_text(errors = "replace").splitlines()
except (OSError, UnicodeDecodeError):
_file_lines_cache[path] = []
return _file_lines_cache[path]
def find_usage(pkg: str) -> list[Hit]:
"""Return real usages of `pkg`. Filters pip-playwright separately.
"""Return real usages of `pkg` (pip-playwright filtered separately).
For each filename returned by grep, also feed a multi-line window
around the matching line into classify() so multi-line imports
(`import {\n a\n} from "pkg"`) get picked up.
For each grep hit, also feed a multi-line window into classify() so
multi-line imports get picked up.
"""
rows = grep_repo(re.escape(pkg))
hits = []
@ -844,10 +725,8 @@ def find_usage(pkg: str) -> list[Hit]:
# Try the single-line classify first.
kind = classify(pkg, file, content)
if not kind:
# Multi-line window: a generous 25 lines above + the line +
# 25 below so Prettier's one-import-per-line formatting for
# 12-20+ named imports still includes the `import` keyword
# in the same window as the `from "pkg"` clause.
# Multi-line window (25 lines each side) so Prettier's
# one-import-per-line formatting still pairs `import` with `from`.
lines = _read_file(file)
lo = max(0, lineno - 26)
hi = min(len(lines), lineno + 25)
@ -863,28 +742,21 @@ def find_usage(pkg: str) -> list[Hit]:
def _candidate_bin_names(pkg: str) -> set[str]:
"""Names a removed package's CLI could be invoked under in shell
scripts and workflow files. Most npm CLIs use the package name
(`vite`, `eslint`, `playwright`); scoped CLI packages commonly
expose an unscoped binary name (`@biomejs/biome` -> `biome`).
"""
"""Bin names a removed package's CLI could be invoked under. Most npm CLIs
use the package name; scoped ones expose an unscoped bin (@biomejs/biome ->
biome)."""
return {pkg, pkg.rsplit("/", 1)[-1]}
def find_command_usage(pkg: str) -> list[Hit]:
"""Find package CLI invocations in shell / workflow / Dockerfile
surfaces: `npx pkg`, `bunx pkg`, `pnpm exec pkg`, `yarn dlx pkg`,
or a bare `pkg --flag`. Returns Hit("command_bin").
Detection is bounded to COMMAND_LIKE_EXT files so a JS string that
happens to contain `npx foo` inside a TS test fixture is not
mistaken for a real invocation.
"""Find package CLI invocations in shell/workflow/Dockerfile surfaces (npx,
bunx, pnpm exec, yarn dlx, or bare `pkg --flag`). Bounded to
COMMAND_LIKE_EXT so `npx foo` in a TS fixture isn't mistaken for real use.
"""
bins = sorted(_candidate_bin_names(pkg), key = len, reverse = True)
esc_bins = "|".join(re.escape(b) for b in bins)
# grep ERE pattern (POSIX classes for whitespace/word boundaries).
# Build without f-strings to avoid f-string-vs-{} confusion with the
# POSIX `[[:space:]]` literals and trailing `})}` boundary class.
# grep ERE pattern. Built without f-strings to avoid clashing with the
# POSIX `[[:space:]]` literals.
grep_pat = (
r"(^|[[:space:]:;&|(\[])"
r"(npx[[:space:]]+|pnpm[[:space:]]+exec[[:space:]]+"
@ -916,10 +788,8 @@ def find_command_usage(pkg: str) -> list[Hit]:
def types_target_name(pkg: str) -> str | None:
"""Strip `@types/` prefix and decode the npm scope-encoding so the
return value matches the runtime package name. `@types/foo` -> `foo`,
`@types/foo__bar` -> `@foo/bar`. Returns None for non-@types packages.
"""
"""Strip `@types/` and decode scope-encoding to the runtime package name
(`@types/foo__bar` -> `@foo/bar`). None for non-@types packages."""
if not pkg.startswith("@types/"):
return None
target = pkg[len("@types/") :]
@ -930,11 +800,8 @@ def types_target_name(pkg: str) -> str | None:
def find_types_runtime_usage(pkg: str, tsc_types: set[str]) -> list[Hit]:
"""For a removed `@types/X`, find usages of `X` itself: explicit
`/// <reference types="X" />`, `tsconfig.compilerOptions.types: ["X"]`,
and runtime `import "X"` shapes. The whole point of `@types/X` is to
type one of those; if any are present, the type package must stay.
"""
"""For a removed `@types/X`, find usages of `X` itself (triple-slash
reference, tsconfig types, runtime import). If any exist, @types/X stays."""
target = types_target_name(pkg)
if target is None:
return []
@ -952,18 +819,14 @@ def find_types_runtime_usage(pkg: str, tsc_types: set[str]) -> list[Hit]:
def main() -> int:
p = argparse.ArgumentParser(
description = __doc__, formatter_class = argparse.RawTextHelpFormatter
)
p = argparse.ArgumentParser(description = __doc__, formatter_class = argparse.RawTextHelpFormatter)
p.add_argument(
"--base",
default = "origin/main",
help = "git ref to diff against (default: origin/main). "
"Examples: HEAD~1, main, a-tag, a-sha.",
)
p.add_argument(
"--base-pkg", help = "optional override: read base package.json from this path"
)
p.add_argument("--base-pkg", help = "optional override: read base package.json from this path")
p.add_argument(
"--base-lock",
help = "optional override: read base package-lock.json from this path. "
@ -1023,10 +886,9 @@ def main() -> int:
return 2
head_lock = read_pkg_file(head_lock_path)
# Base lockfile is best-effort. We use it only to recover the
# bin -> package mapping for packages the PR is removing -- so a
# `scripts.biome:check` cite still fires when `@biomejs/biome` is
# being dropped and the head lockfile no longer has it.
# Base lockfile is best-effort: only used to recover the bin -> package
# mapping for packages the PR removes, so a scripts.biome cite still fires
# when @biomejs/biome is dropped from the head lockfile.
if args.base_lock:
base_lock_path = Path(args.base_lock)
base_lock = read_pkg_file(base_lock_path) if base_lock_path.exists() else {}
@ -1037,9 +899,8 @@ def main() -> int:
head_names = all_decl_names(head_pkg)
removed = sorted(base_names - head_names)
# All hygiene checks compute up front so they can run on both the
# removal-present and removal-empty paths (so `--strict` actually
# fails when only hygiene issues exist).
# Hygiene checks compute up front so they run on both the removal-present
# and removal-empty paths (so --strict fails on hygiene-only issues).
sync_warns = lockfile_root_sync(head_pkg, head_lock)
types_warns = types_orphan_warnings(head_pkg)
missing_imports = find_imports_without_decl(head_pkg)
@ -1057,9 +918,7 @@ def main() -> int:
print(f" - {w}")
print()
if missing_imports:
print(
f"Imports without a matching package.json dep ({len(missing_imports)}):"
)
print(f"Imports without a matching package.json dep ({len(missing_imports)}):")
for file, ln, spec in missing_imports[:20]:
print(f" - {file}:{ln} imports '{spec}'")
print()
@ -1097,19 +956,14 @@ def main() -> int:
return 1
return 0
print(
f"Checking {len(removed)} removed package(s) from studio/frontend/package.json"
)
print(f"Checking {len(removed)} removed package(s) from studio/frontend/package.json")
print(f"Base: {args.base} Head: working tree")
print()
reachable_paths = reachable_from_head(head_pkg, head_lock) if head_lock else set()
# bin -> package map: start from the head lockfile, then layer the
# base lockfile's entries on top for packages this PR is removing.
# A correct removal updates the head lockfile to drop node_modules/foo,
# so build_bin_to_pkg(head_lock) loses the mapping; we recover it
# from the base lockfile so `scripts.biome:check` still flags as a
# usage when `@biomejs/biome` is being dropped.
# bin -> package map from the head lockfile, layering base-lockfile entries
# for removed packages so scripts.biome still flags when @biomejs/biome is
# dropped (head lockfile no longer maps it).
bin_to_pkg = build_bin_to_pkg(head_lock) if head_lock else {}
base_bin_to_pkg = build_bin_to_pkg(base_lock) if base_lock else {}
removed_set = set(removed)
@ -1121,15 +975,12 @@ def main() -> int:
def reachable_install_paths(name: str) -> tuple[str | None, list[str]]:
"""Return (top_level_path, nested_paths). top_level is what bare
`import "name"` from src/ actually resolves to; nested copies are
only visible inside the parent package that nested them.
"""
`import "name"` resolves to; nested copies are only visible inside
their parent package."""
top = f"node_modules/{name}"
top_path = top if top in reachable_paths else None
nested = sorted(
p
for p in reachable_paths
if p != top and p.endswith(f"/node_modules/{name}")
p for p in reachable_paths if p != top and p.endswith(f"/node_modules/{name}")
)
return top_path, nested
@ -1138,8 +989,7 @@ def main() -> int:
hits = find_usage(name)
# CLI invocations in shell scripts / workflows / Dockerfiles.
hits.extend(find_command_usage(name))
# @types/X is "used" if X is referenced as a type or as a
# runtime import elsewhere in the repo.
# @types/X is "used" if X is referenced as a type or runtime import.
hits.extend(find_types_runtime_usage(name, tsc_types))
for cite in script_refs.get(name, []):
hits.append(Hit("studio/frontend/package.json", 0, "script_bin", cite))
@ -1147,9 +997,8 @@ def main() -> int:
hits.append(Hit("studio/frontend/package.json", 0, "pkg_json_field", cite))
top, nested = reachable_install_paths(name)
importable_top_level = top is not None
# Source imports of bare specifier `name` resolve ONLY to top-level
# node_modules/<name>. Nested copies under another package are
# invisible to src/ files.
# Bare specifier `name` resolves ONLY to top-level node_modules/<name>;
# nested copies are invisible to src/ files.
if hits and not importable_top_level:
status = "FAIL"
elif hits and importable_top_level:
@ -1177,9 +1026,7 @@ def main() -> int:
_print_hygiene()
if failures:
print(
f"FAIL: {len(failures)} removed package(s) still referenced and not resolvable"
)
print(f"FAIL: {len(failures)} removed package(s) still referenced and not resolvable")
for name, _ in failures:
print(f" - {name}")
return 1

88
scripts/check_new_install_scripts.py
View file

@ -4,33 +4,18 @@
"""Diff two `package-lock.json` files and flag NEW install-script deps.
A package with `"hasInstallScript": true` runs `preinstall` / `install` /
`postinstall` lifecycle hooks every time `npm ci` lays it down. Every
npm supply-chain compromise of the last 18 months (Shai-Hulud,
TanStack, axios-style, ArmorCode hijacks) leveraged exactly this lever:
the attacker publishes a new malicious version of a dep we already
trust, and the post-install hook runs the next time CI installs.
A `"hasInstallScript": true` package runs preinstall/install/postinstall
hooks on every `npm ci` -- the lever behind recent npm supply-chain
compromises (attacker publishes a malicious version of a trusted dep).
This refuses to land a newly-introduced install-script dep without a
maintainer eyeball; pre-existing ones are not re-flagged.
This scanner refuses to allow a newly-introduced install-script dep to
land without a maintainer eyeball on the lifecycle script body.
Existing install-script deps are NOT re-flagged -- if `node-gyp` has
been in the lockfile since day one, it's not part of this PR's threat
model. Only new entries are surfaced.
Supports lockfileVersion 1 (recursive `dependencies`) and 2/3 (flat
`packages` with `node_modules/.../node_modules/...` nesting). For each
new entry we best-effort fetch the registry metadata to recover the
postinstall command body; the finding is still emitted if unreachable.
Supports lockfileVersion 1 (`dependencies` key, recursive), 2 and 3
(flat `packages` key with `node_modules/<a>/node_modules/<b>` nesting
for transitive entries). For each NEW install-script package we
attempt a stdlib-only fetch of
`https://registry.npmjs.org/<name>/<version>` to recover the actual
postinstall command body. If the network is blocked we still emit the
finding -- the lifecycle command body is informational, not
load-bearing.
Exit codes
==========
0 no newly-added install-script deps
1 one or more newly-added install-script deps; listed on stderr
2 internal error (missing lockfile, malformed JSON, etc.)
Exit codes: 0 = none; 1 = one or more (on stderr); 2 = internal error.
"""
from __future__ import annotations
@ -53,9 +38,7 @@ HIGH = "HIGH"
class Finding:
__slots__ = ("severity", "name", "version", "kind", "detail")
def __init__(
self, severity: str, name: str, version: str, kind: str, detail: str
) -> None:
def __init__(self, severity: str, name: str, version: str, kind: str, detail: str) -> None:
self.severity = severity
self.name = name
self.version = version
@ -70,21 +53,14 @@ class Finding:
)
# ─────────────────────────────────────────────────────────────────────
# Lockfile parsing.
# ─────────────────────────────────────────────────────────────────────
def _strip_nm_prefix(key: str) -> str:
"""Convert a v2/v3 `packages` key into a bare package name.
`node_modules/foo` -> `foo`; `node_modules/foo/node_modules/bar` ->
`bar`. The empty key (`""`) is the project root and returns "".
"""
"""Convert a v2/v3 `packages` key into a bare package name (leaf after last `node_modules/`)."""
if not key:
return ""
# Use the LAST `node_modules/` segment so transitives map to their
# leaf name, matching how npm install resolves a postinstall.
# LAST node_modules/ segment so transitives map to their leaf name.
marker = "node_modules/"
idx = key.rfind(marker)
if idx == -1:
@ -93,14 +69,9 @@ def _strip_nm_prefix(key: str) -> str:
def _collect_install_script_entries(lock: dict) -> dict[str, str]:
"""Walk a parsed lockfile and return {package_name: version} for
every entry with `hasInstallScript: true` (v2/v3) OR a
non-empty `scripts.preinstall|install|postinstall` (v1).
"""Return {name@version: name} for entries with hasInstallScript (v2/v3) or a lifecycle script (v1).
The same package may appear at multiple versions in a single
lockfile (de-duplicated copies under different parents); we key by
`name@version` so we don't lose either copy. Returns a dict keyed
by `name@version` -> the same string for convenience.
Keyed by name@version so dup copies at different versions aren't lost.
"""
seen: dict[str, str] = {}
version = lock.get("lockfileVersion")
@ -120,10 +91,7 @@ def _collect_install_script_entries(lock: dict) -> dict[str, str]:
ver = entry.get("version") or "<unversioned>"
seen[f"{name}@{ver}"] = name
# v1 also embeds a `dependencies` tree; v2/v3 carry both for
# backwards-compat but `packages` is canonical for them. For v1
# there is no `hasInstallScript` flag, so look for a non-empty
# `scripts.preinstall|install|postinstall` directly.
# v1 has no hasInstallScript flag; detect lifecycle scripts directly.
def _walk_v1(deps: dict, depth: int = 0) -> None:
if depth > 64 or not isinstance(deps, dict):
return
@ -135,8 +103,6 @@ def _collect_install_script_entries(lock: dict) -> dict[str, str]:
isinstance(scripts, dict) and scripts.get(hook)
for hook in ("preinstall", "install", "postinstall")
)
# v1 also sets `requires` only on the parent, no flag, so
# the lifecycle-script presence is the only signal.
if lifecycle:
ver = entry.get("version") or "<unversioned>"
seen[f"{name}@{ver}"] = name
@ -157,19 +123,11 @@ def _load_lockfile(path: Path) -> dict:
raise ValueError(f"{path}: not valid JSON: {exc}") from exc
# ─────────────────────────────────────────────────────────────────────
# Registry lookup for the postinstall command body (best-effort).
# ─────────────────────────────────────────────────────────────────────
def _fetch_registry_scripts(name: str, version: str) -> dict[str, str] | None:
"""Return {hook: command} for any of preinstall / install /
postinstall published in the registry metadata for this name@ver.
Returns None on any error (network blocked, 404, malformed JSON).
Never raises; the caller treats absence as "could not enrich, emit
finding anyway".
"""
"""Return {hook: command} for lifecycle hooks in registry metadata; None on any error (never raises)."""
safe_name = urllib.parse.quote(name, safe = "@/")
url = f"{REGISTRY_BASE}{safe_name}/{urllib.parse.quote(version)}"
try:
@ -192,9 +150,7 @@ def _fetch_registry_scripts(name: str, version: str) -> dict[str, str] | None:
return keep or None
# ─────────────────────────────────────────────────────────────────────
# Diff.
# ─────────────────────────────────────────────────────────────────────
def diff_new_install_scripts(base_lock: dict, head_lock: dict) -> list[Finding]:
@ -205,10 +161,7 @@ def diff_new_install_scripts(base_lock: dict, head_lock: dict) -> list[Finding]:
if key in base:
continue # pre-existing install-script dep; not in scope
name = head[key]
# key is "name@version"; rsplit("@", 1) handles scoped names.
version = (
key[len(name) + 1 :] if key.startswith(name + "@") else "<unversioned>"
)
version = key[len(name) + 1 :] if key.startswith(name + "@") else "<unversioned>"
scripts = _fetch_registry_scripts(name, version)
if scripts:
detail = "; ".join(f"{h}={cmd!r}" for h, cmd in scripts.items())
@ -230,16 +183,13 @@ def diff_new_install_scripts(base_lock: dict, head_lock: dict) -> list[Finding]:
return findings
# ─────────────────────────────────────────────────────────────────────
# CLI.
# ─────────────────────────────────────────────────────────────────────
def main(argv: list[str] | None = None) -> int:
parser = argparse.ArgumentParser(
description = (
"Diff two package-lock.json files and refuse any newly-"
"added install-script dep."
"Diff two package-lock.json files and refuse any newly-added install-script dep."
),
)
parser.add_argument(

481
scripts/enforce_kwargs_spacing.py
View file

@ -1,5 +1,10 @@
#!/usr/bin/env python3
"""Ensure keyword arguments use spaces around '=', prune redundant pass statements."""
"""Ensure keyword arguments use spaces around '=', prune redundant pass statements,
drop the blank line after a short indented import block, merge adjacent same-line
string literals, normalize def-signature magic commas (pre-ruff) so a def with
>= 3 params and a default goes one-per-line while everything else stays
collapsible, and collapse a short multi-line assert onto one line (pre-ruff) by
stripping the magic trailing comma that holds it open."""
from __future__ import annotations
@ -15,13 +20,8 @@ from pathlib import Path
def _atomic_write_text(path: Path, data: str, encoding: str) -> None:
"""Write ``data`` to ``path`` atomically.
Stages a tmp file in the same directory (so it's on the same
filesystem as the destination), fsyncs, then `os.replace`s into
place. A crash mid-write therefore leaves either the previous
content or the fully new content -- never a truncated source file.
"""
"""Write ``data`` to ``path`` atomically via same-dir tmp + fsync + os.replace,
so a crash mid-write leaves either the old or full new content, never a truncation."""
dirpath = str(path.parent) or "."
fd, tmp_path = tempfile.mkstemp(prefix=".kwargs_fix.", dir=dirpath)
try:
@ -123,9 +123,7 @@ def remove_redundant_passes(text: str) -> tuple[str, bool]:
lines = text.splitlines(keepends=True)
changed = False
for node in sorted(
redundant, key=lambda item: (item.lineno, item.col_offset), reverse=True
):
for node in sorted(redundant, key=lambda item: (item.lineno, item.col_offset), reverse=True):
start = node.lineno - 1
end = (node.end_lineno or node.lineno) - 1
if start >= len(lines):
@ -139,7 +137,7 @@ def remove_redundant_passes(text: str) -> tuple[str, bool]:
lines[start] = segment if segment.strip() else ""
continue
# Defensive fall-back for unexpected multi-line 'pass'.
# Fall-back for unexpected multi-line 'pass'.
prefix = lines[start][: node.col_offset]
lines[start] = prefix if prefix.strip() else ""
for idx in range(start + 1, end):
@ -160,7 +158,441 @@ def remove_redundant_passes(text: str) -> tuple[str, bool]:
return "".join(result_lines), changed
def process_file(path: Path) -> bool:
def remove_blank_after_short_import(text: str) -> tuple[str, bool]:
"""Drop blank line(s) after an import block in a small nested suite.
In an indented suite of <= 3 statements (never module level), when consecutive
imports are followed across blank lines (nothing else) by another statement,
remove those blanks. A comment in the gap blocks the rule. Removing blank lines
never changes the AST.
"""
try:
tree = ast.parse(text)
except SyntaxError:
return text, False
lines = text.splitlines(keepends=True)
import_types = (ast.Import, ast.ImportFrom)
drop: set[int] = set() # 1-based physical line numbers to delete
def suites_of(node: ast.AST) -> list[list[ast.stmt]]:
if isinstance(node, ast.Module):
return [] # module-level import spacing is left alone
out: list[list[ast.stmt]] = []
for attr in ("body", "orelse", "finalbody"):
val = getattr(node, attr, None)
if isinstance(val, list) and val and all(isinstance(s, ast.stmt) for s in val):
out.append(val)
return out
for node in ast.walk(tree):
for suite in suites_of(node):
if len(suite) > 3: # only small blocks
continue
i = 0
while i < len(suite):
if not isinstance(suite[i], import_types):
i += 1
continue
j = i
while j + 1 < len(suite) and isinstance(suite[j + 1], import_types):
j += 1
if j + 1 < len(suite): # an import block followed by another statement
last_imp, nxt = suite[j], suite[j + 1]
gap = range((last_imp.end_lineno or last_imp.lineno) + 1, nxt.lineno)
nums = [n for n in gap if 1 <= n <= len(lines)]
if nums and all(lines[n - 1].strip() == "" for n in nums):
drop.update(nums)
i = j + 1
if not drop:
return text, False
kept = [ln for idx, ln in enumerate(lines, start=1) if idx not in drop]
return "".join(kept), True
_STRING_TRIVIA = (tokenize.NL, tokenize.NEWLINE, tokenize.COMMENT, tokenize.INDENT, tokenize.DEDENT)
_DEF_MIN_PARAMS_FOR_MULTILINE = 3 # signatures with < this many params stay one line
def _def_specs_by_line(tree: ast.AST) -> dict[int, tuple[int, bool]]:
"""Map each def keyword line to (param count, has-any-default).
``*`` / ``/`` markers aren't counted. A default exists if any positional default
is present or any keyword-only default is not ``None`` (``None`` in ``kw_defaults``
means a required keyword-only arg).
"""
out: dict[int, tuple[int, bool]] = {}
for node in ast.walk(tree):
if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
a = node.args
count = (
len(a.posonlyargs)
+ len(a.args)
+ len(a.kwonlyargs)
+ (1 if a.vararg else 0)
+ (1 if a.kwarg else 0)
)
has_default = bool(a.defaults) or any(d is not None for d in a.kw_defaults)
out[node.lineno] = (count, has_default)
return out
def normalize_def_trailing_comma(text: str) -> tuple[str, bool]:
"""Force a def signature one-per-line iff >= 3 params AND a default; else collapsible.
A qualifying signature gets a magic trailing comma added (ruff wraps it
one-per-line); every other signature has its trailing comma stripped so ruff
collapses it when it fits. Def parameter lists only, never call sites or
collection literals. Run BEFORE ruff format. Never changes the AST (re-checked).
"""
try:
tree = ast.parse(text)
toks = list(tokenize.generate_tokens(io.StringIO(text).readline))
except (tokenize.TokenError, IndentationError, SyntaxError):
return text, False
specs = _def_specs_by_line(tree)
n = len(toks)
edits: list[tuple[int, int, str]] = [] # (row, col, "del" | "ins")
i = 0
while i < n:
t = toks[i]
if t.type == tokenize.NAME and t.string == "def" and t.start[0] in specs:
cnt, has_default = specs[t.start[0]]
force_multiline = cnt >= _DEF_MIN_PARAMS_FOR_MULTILINE and has_default
j = i + 1
while j < n and not (toks[j].type == tokenize.OP and toks[j].string == "("):
if toks[j].type == tokenize.NEWLINE:
break
j += 1
if j < n and toks[j].type == tokenize.OP and toks[j].string == "(":
depth = 0
k = j
while k < n:
tk = toks[k]
if tk.type == tokenize.OP and tk.string == "(":
depth += 1
elif tk.type == tokenize.OP and tk.string == ")":
depth -= 1
if depth == 0:
m = k - 1
while m > j and toks[m].type in _STRING_TRIVIA:
m -= 1
last = toks[m]
has_comma = last.type == tokenize.OP and last.string == ","
empty = m == j # nothing between ( and )
if force_multiline and not has_comma and not empty:
edits.append((last.end[0], last.end[1], "ins"))
elif not force_multiline and has_comma:
edits.append((last.start[0], last.start[1], "del"))
break
k += 1
i = k + 1
continue
i += 1
if not edits:
return text, False
lines = text.splitlines(keepends=True)
for row, col, kind in sorted(edits, reverse=True):
ln = lines[row - 1]
if kind == "del":
if col < len(ln) and ln[col] == ",":
lines[row - 1] = ln[:col] + ln[col + 1 :]
else: # ins
lines[row - 1] = ln[:col] + "," + ln[col:]
out = "".join(lines)
try:
if ast.dump(ast.parse(out)) != ast.dump(ast.parse(text)):
return text, False
except SyntaxError:
return text, False
return out, True
def _split_string_token(s: str) -> tuple[str, str, str] | None:
"""Split a string literal source into (prefix, quote, body).
``prefix`` is the letters before the opening quote, ``quote`` the delimiter,
``body`` everything between. ``None`` if not a recognizable string literal.
"""
i = 0
while i < len(s) and s[i] not in ("'", '"'):
i += 1
if i >= len(s):
return None
prefix, rest = s[:i], s[i:]
for q in ('"""', "'''", '"', "'"):
if rest.startswith(q) and rest.endswith(q) and len(rest) >= 2 * len(q):
return prefix, q, rest[len(q) : len(rest) - len(q)]
return None
# A "piece" is one string literal in source: a plain STRING token, or a whole
# f-string spanning FSTRING_START..FSTRING_END. (kind, (row, col0), (row, col1), raw)
def _string_pieces(
toks: list[tokenize.TokenInfo], lines: list[str]
) -> list[tuple[str, tuple[int, int], tuple[int, int], str | None]]:
pieces: list[tuple[str, tuple[int, int], tuple[int, int], str | None]] = []
n = len(toks)
def raw_of(start: tuple[int, int], end: tuple[int, int]) -> str | None:
if start[0] != end[0]: # only single-physical-line pieces are mergeable
return None
return lines[start[0] - 1][start[1] : end[1]]
i = 0
while i < n:
t = toks[i]
if t.type == tokenize.STRING:
pieces.append(("str", t.start, t.end, raw_of(t.start, t.end)))
i += 1
elif t.type == tokenize.FSTRING_START:
depth = 0
j = i
while j < n: # walk to the matching FSTRING_END (f-strings can nest)
if toks[j].type == tokenize.FSTRING_START:
depth += 1
elif toks[j].type == tokenize.FSTRING_END:
depth -= 1
if depth == 0:
break
j += 1
end = toks[j].end
pieces.append(("f", t.start, end, raw_of(t.start, end)))
i = j + 1
else:
pieces.append(("other", t.start, t.end, None))
i += 1
return pieces
def _merge_string_run(pieces: list[tuple[str, str]]) -> str | None:
"""Merge a run of adjacent string pieces into one literal's source text.
``pieces`` is ``(kind, raw_source)`` with kind ``"str"`` or ``"f"``. Bytes are
left side-by-side (``None``); a run with no f-string merges plain/raw/unicode
sharing one prefix+quote by body concatenation; a run mixing an f-string with a
plain string (no bytes, no raw) folds into one f-string with plain braces escaped.
Runs of only f-strings are left alone. Caller re-checks the AST and drops a
differing change, so subtle cases are caught.
"""
parsed = []
for kind, raw in pieces:
pqb = _split_string_token(raw)
if pqb is None:
return None
prefix, quote, body = pqb
if "b" in prefix.lower():
return None # bytes: leave side-by-side
parsed.append((kind, prefix, quote, body))
if len({p[2] for p in parsed}) != 1:
return None # mixed quote style: not a safe textual merge
quote = parsed[0][2]
if not any(p[0] == "f" for p in parsed):
# No f-string: merge plain/raw/unicode sharing one prefix by concatenation.
if len({p[1].lower() for p in parsed}) != 1:
return None
return f"{parsed[0][1]}{quote}{''.join(p[3] for p in parsed)}{quote}"
# f-string fold only when a plain string is glued onto an f-string; a run of
# only f-strings is left side-by-side (folding long ones would force ruff to
# re-wrap the surrounding statement).
if all(p[0] == "f" for p in parsed):
return None
# raw mixed with f is too subtle (backslash + brace escaping) -> skip.
if any("r" in p[1].lower() for p in parsed):
return None
body = "".join(
b if kind == "f" else b.replace("{", "{{").replace("}", "}}")
for kind, _pfx, _q, b in parsed
)
return f"f{quote}{body}{quote}"
_LINE_LENGTH = 100 # ruff line-length; an f-fold must not push a statement past it
def _enclosing_stmt(tree: ast.AST, row: int) -> ast.stmt | None:
"""The innermost statement whose physical-line span contains ``row``."""
best: tuple[ast.stmt, int] | None = None
for node in ast.walk(tree):
if isinstance(node, ast.stmt):
lo = node.lineno
hi = node.end_lineno or lo
if lo <= row <= hi and (best is None or hi - lo < best[1]):
best = (node, hi - lo)
return best[0] if best else None
def _fold_collapses(
tree: ast.AST, lines: list[str], row: int, c0: int, c1: int, merged: str
) -> bool:
"""Whether an f-string fold at ``row[c0:c1]`` -> ``merged`` is safe to apply.
Only ``assert`` wraps awkwardly when a message folds (ruff parenthesizes the
condition once it no longer fits one line); every other construct wraps
acceptably so is always allowed. An ``assert`` fold is allowed only if already
one line, or its estimated folded one-line length fits the line length.
"""
stmt = _enclosing_stmt(tree, row)
if not isinstance(stmt, ast.Assert):
return True
lo, hi = stmt.lineno, stmt.end_lineno or stmt.lineno
if lo == hi:
return True
seg = []
for k in range(lo, hi + 1):
ln = lines[k - 1].rstrip("\n")
if k == row:
ln = ln[:c0] + merged + ln[c1:]
seg.append(ln)
indent = len(seg[0]) - len(seg[0].lstrip())
# Conservative over-estimate: join continuation lines with a single space
# (ruff joins bracketed wraps with none), so borderline cases skip the fold.
joined = " ".join(s.strip() for s in seg)
return indent + len(joined) <= _LINE_LENGTH
def merge_adjacent_string_literals(text: str) -> tuple[str, bool]:
"""Merge adjacent string literals on ONE physical line into a single literal.
Plain/raw/unicode runs merge by concatenation; an f-string + plain string folds
into one f-string (plain braces escaped) only while the statement still fits one
line. Runs of only f-strings, and bytes, are left side-by-side. The file AST is
re-checked and a differing change dropped, so meaning never changes.
"""
try:
toks = list(tokenize.generate_tokens(io.StringIO(text).readline))
tree = ast.parse(text)
except (tokenize.TokenError, IndentationError, SyntaxError):
return text, False
lines = text.splitlines(keepends=True)
pieces = _string_pieces(toks, lines)
# Group consecutive mergeable pieces (str/f, single line, same physical line).
runs: list[list[tuple[str, tuple[int, int], tuple[int, int], str]]] = []
cur: list[tuple[str, tuple[int, int], tuple[int, int], str]] = []
for kind, start, end, raw in pieces:
if kind in ("str", "f") and raw is not None:
if cur and cur[-1][2][0] != start[0]:
if len(cur) >= 2:
runs.append(cur)
cur = []
cur.append((kind, start, end, raw))
else:
if len(cur) >= 2:
runs.append(cur)
cur = []
if len(cur) >= 2:
runs.append(cur)
if not runs:
return text, False
edits = []
for run in runs:
merged = _merge_string_run([(kind, raw) for kind, _s, _e, raw in run])
if merged is None:
continue
row, c0, c1 = run[0][1][0], run[0][1][1], run[-1][2][1]
# An f-string fold must not push its statement onto extra lines; a plain
# concatenation always collapses cleanly so it skips this check.
if any(kind == "f" for kind, _s, _e, _r in run) and not _fold_collapses(
tree, lines, row, c0, c1, merged
):
continue
edits.append((row, c0, c1, merged))
if not edits:
return text, False
for row, c0, c1, repl in sorted(edits, key=lambda e: (e[0], e[1]), reverse=True):
ln = lines[row - 1]
lines[row - 1] = ln[:c0] + repl + ln[c1:]
out = "".join(lines)
try:
if ast.dump(ast.parse(text)) != ast.dump(ast.parse(out)):
return text, False
except SyntaxError:
return text, False
return out, True
def collapse_short_asserts(text: str) -> tuple[str, bool]:
"""Collapse a multi-line ``assert`` onto one line when it would fit.
When the statement's estimated one-line length fits, strip the magic trailing
commas (comma before a closer) holding it open so ruff rejoins it. Run BEFORE
ruff format. Skips asserts with a comment (would oscillate). Stripping is
non-semantic except for a one-element tuple; AST is re-checked and changing
asserts left alone.
"""
try:
tree = ast.parse(text)
toks = list(tokenize.generate_tokens(io.StringIO(text).readline))
except (tokenize.TokenError, IndentationError, SyntaxError):
return text, False
lines = text.splitlines(keepends=True)
multiline = [
(n.lineno, n.end_lineno)
for n in ast.walk(tree)
if isinstance(n, ast.Assert) and (n.end_lineno or n.lineno) > n.lineno
]
if not multiline:
return text, False
comment_rows = {t.start[0] for t in toks if t.type == tokenize.COMMENT}
targets = [] # (lo, hi) spans whose one-line form fits and have no comment
for lo, hi in multiline:
if any(lo <= r <= hi for r in comment_rows):
continue # a comment would keep ruff multi-line -> never collapses
seg = [lines[k].rstrip("\n") for k in range(lo - 1, hi)]
indent = len(seg[0]) - len(seg[0].lstrip())
# Over-estimate (join with a space; keep the comma) so a "fits" verdict
# is always at least as long as ruff's real one-line output -> no fight.
if indent + len(" ".join(s.strip() for s in seg)) <= _LINE_LENGTH:
targets.append((lo, hi))
if not targets:
return text, False
# Trailing commas (a ',' whose next significant token is a closer), grouped
# by the target assert they belong to.
sig = [t for t in toks if t.type not in _STRING_TRIVIA]
by_target: dict[tuple[int, int], list[tuple[int, int]]] = defaultdict(list)
for i, t in enumerate(sig):
if t.type == tokenize.OP and t.string == ",":
nxt = sig[i + 1] if i + 1 < len(sig) else None
if nxt and nxt.type == tokenize.OP and nxt.string in (")", "]", "}"):
for lo, hi in targets:
if lo <= t.start[0] <= hi:
by_target[(lo, hi)].append(t.start)
break
if not by_target:
return text, False
base_dump = ast.dump(tree)
working = lines[:]
changed = False
for positions in by_target.values(): # apply per assert; skip any that break AST
trial = working[:]
for row, col in sorted(positions, reverse=True):
ln = trial[row - 1]
if col < len(ln) and ln[col] == ",":
trial[row - 1] = ln[:col] + ln[col + 1 :]
try:
if ast.dump(ast.parse("".join(trial))) == base_dump:
working, changed = trial, True
except SyntaxError:
pass
return ("".join(working), True) if changed else (text, False)
def process_file(path: Path, pre: bool = False) -> bool:
try:
with tokenize.open(path) as handle:
original = handle.read()
@ -169,9 +601,23 @@ def process_file(path: Path) -> bool:
print(f"Failed to read {path}: {exc}", file=sys.stderr)
return False
if pre:
# Pre-ruff: normalize def-signature magic commas (>=3 params + a default
# add so ruff forces one-per-line; everything else strips so ruff
# collapses), and strip the magic trailing comma from a short multi-line
# assert so ruff joins it onto one line. Everything else runs post-ruff.
updated, normalized = normalize_def_trailing_comma(original)
updated, collapsed = collapse_short_asserts(updated)
if normalized or collapsed:
_atomic_write_text(path, updated, encoding)
return True
return False
updated, changed = enforce_spacing(original)
updated, blanked = remove_blank_after_short_import(updated)
updated, merged = merge_adjacent_string_literals(updated)
updated, removed = remove_redundant_passes(updated)
if changed or removed:
if changed or blanked or merged or removed:
_atomic_write_text(path, updated, encoding)
return True
return False
@ -180,6 +626,11 @@ def process_file(path: Path) -> bool:
def main(argv: list[str]) -> int:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("files", nargs="+", help="Python files to fix")
parser.add_argument(
"--pre",
action="store_true",
help="pre-ruff pass: normalize def-signature commas + collapse short multi-line asserts",
)
args = parser.parse_args(argv)
touched: list[Path] = []
@ -192,7 +643,7 @@ def main(argv: list[str]) -> int:
continue
if not path.exists() or path.is_dir():
continue
if process_file(path):
if process_file(path, pre=args.pre):
touched.append(path)
if touched:

2
scripts/install_gemma4_mlx.sh
View file

@ -1,4 +1,6 @@
#!/bin/bash
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
set -euo pipefail
# ============================================================

2
scripts/install_qwen3_6_mlx.sh
View file

@ -1,4 +1,6 @@
#!/bin/bash
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
set -euo pipefail
# ============================================================

289
scripts/install_rocm_wsl_strixhalo.sh Normal file
View file

@ -0,0 +1,289 @@
#!/usr/bin/env bash
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved.
#
# ──────────────────────────────────────────────────────────────────────────────
# Enable ROCm-on-WSL for AMD Strix Halo (Radeon 8060S / gfx1151)
# ──────────────────────────────────────────────────────────────────────────────
# install.sh already routes gfx1151 to the right ROCm wheels once a ROCm runtime
# is present; what it does NOT do is install AMD's ROCm userspace + the WSL DXG
# bridge. This helper automates that Linux-side prerequisite on Ubuntu 24.04
# WSL2 and is invoked by install.sh when it sees a Strix Halo APU in WSL (via
# /dev/dxg) but no ROCm runtime yet. Fully idempotent (re-run just re-verifies).
#
# Manual, admin-gated Windows prerequisite: an AMD Adrenalin driver with
# production ROCDXG/WSL support (26.2.2+). install.ps1 offers to update it. Once
# installed + rebooted, /dev/dxg is exposed to WSL and this script builds the rest.
#
# HOW ROCDXG WORKS (and why older /usr/lib/wsl/lib notes are wrong): librocdxg.so
# is AMD's user-mode bridge between the Linux HSA runtime and the Windows driver
# over /dev/dxg. The STANDARD hsa-rocr runtime (NOT the gone "roc4wsl" package)
# loads it when HSA_ENABLE_DXG_DETECTION=1. No hsa/rocm libs need injecting into
# /usr/lib/wsl/lib (it holds only d3d12/dxcore), yet rocminfo enumerates gfx1151
# fine -- so we gate on /dev/dxg, not on WSL lib injection.
#
# KNOWN CAVEAT (ROCm/ROCm#6022): librocdxg can cap usable ROCm VRAM at the WSL
# VM's RAM (.wslconfig [wsl2] memory=) on some BIOS UMA layouts, and amd-smi
# doesn't work in WSL. On OOM below capacity, raise memory= (then wsl --shutdown)
# and watch GPU use from Windows. Large-UMA BIOS exposes the full pool regardless.
#
# Verified on Ryzen AI Max+ PRO 395 / Radeon 8060S (gfx1151) with ROCm 7.2.1 +
# Ubuntu 24.04 + WSL2 + Adrenalin. These pins MOVE; bump + re-verify on newer ROCm.
# ──────────────────────────────────────────────────────────────────────────────
set -euo pipefail
# ── Tunables (override via env) ──────────────────────────────────────────────
ROCM_VER="${UNSLOTH_WSL_ROCM_VER:-7.2.1}" # ROCm release to install
GFX="gfx1151"
LIBROCDXG_REF="${UNSLOTH_LIBROCDXG_REF:-develop}" # ROCm/librocdxg git ref to build
# AMD's gfx1151 wheel index (same one install.sh uses); only for the smoke test.
TORCH_INDEX="${UNSLOTH_AMD_ROCM_MIRROR:-https://repo.amd.com/rocm/whl}/${GFX}/"
# Optional torch smoke test (throwaway venv). OFF by default: install.sh installs
# torch itself into the real venv right after, so a duplicate download is wasteful.
SMOKE_TEST="${UNSLOTH_WSL_SMOKE_TEST:-0}"
# REQUIRED constraint -- without it pip prefers PyPI's newer CUDA torch over the
# gfx1151 ROCm wheel. 2.11 carries AMD's real gfx1151 fix (matches install.sh).
TORCH_CONSTRAINT="${UNSLOTH_WSL_TORCH_CONSTRAINT:-torch>=2.11.0,<2.12.0}"
ROCM_DIR="" # resolved after install
say() { printf '\n\033[1;36m== %s\033[0m\n' "$*"; }
note() { printf ' %s\n' "$*"; }
die() { printf '\n\033[1;31m[BLOCKED] %s\033[0m\n' "$*" >&2; exit 1; }
# sudo only if not already root (WSL distros often run as root)
SUDO=""
if [ "$(id -u)" -ne 0 ]; then
command -v sudo >/dev/null 2>&1 || die "Need root or sudo to install ROCm."
SUDO="sudo"
fi
# ── Windows 11 SDK (headers for the librocdxg build) ─────────────────────────
# librocdxg's cmake build needs the Windows SDK 'shared' headers, which live on
# the Windows HOST under C:\Program Files (x86)\Windows Kits\10\Include\<ver>\.
_WIN_SDK_INC_BASE="/mnt/c/Program Files (x86)/Windows Kits/10/Include"
# Print the newest installed SDK include dir with 'shared' headers, or nothing.
# find + read loop (not `for ... in $(ls)`) since the base path has a space.
_find_win_sdk() {
[ -d "$_WIN_SDK_INC_BASE" ] || return 0
while IFS= read -r _inc; do
[ -n "$_inc" ] || continue
if [ -d "$_inc/shared" ]; then printf '%s' "$_inc"; return 0; fi
done < <(find "$_WIN_SDK_INC_BASE" -mindepth 1 -maxdepth 1 -type d 2>/dev/null | sort -Vr)
return 0
}
# Best-effort: install the Windows 11 SDK on the Windows HOST via winget so the
# build has its headers with no manual step. Elevates -> ONE UAC prompt; headers
# appear under /mnt/c immediately (no reboot). Never fatal -- failure falls
# through to a manual-install message. Opt out: UNSLOTH_SKIP_WIN_SDK_INSTALL=1.
_install_windows_sdk_via_winget() {
[ "${UNSLOTH_SKIP_WIN_SDK_INSTALL:-0}" = "1" ] && { note "Skipping Windows SDK auto-install (UNSLOTH_SKIP_WIN_SDK_INSTALL=1)."; return 0; }
command -v powershell.exe >/dev/null 2>&1 || return 0
# `command -v` succeeds even with WSL interop OFF (.exe on PATH but fails
# with "Exec format error"); verify it actually executes.
powershell.exe -NoProfile -Command "exit 0" >/dev/null 2>&1 || return 0
if ! powershell.exe -NoProfile -Command "if (Get-Command winget -ErrorAction SilentlyContinue) { exit 0 } else { exit 1 }" >/dev/null 2>&1; then
note "winget not available on the Windows host -- cannot auto-install the Windows SDK."
return 0
fi
say "Installing the Windows 11 SDK on the Windows host via winget"
note "librocdxg needs its headers. Approve the UAC prompt on the Windows desktop."
note "One-time (~1-3 GB download); opt out with UNSLOTH_SKIP_WIN_SDK_INSTALL=1."
# Newest SDK first, then a fallback. Header presence is the source of truth
# (re-check each attempt), not winget's exit code. </dev/null so winget never
# consumes a piped `curl | sh` stdin.
for _sdk_id in Microsoft.WindowsSDK.10.0.26100 Microsoft.WindowsSDK.10.0.22621; do
note "winget install ${_sdk_id} ..."
# --source winget: pin the community source so a broken default msstore
# source (the cert failure this PR fixes) can't abort SDK resolution.
powershell.exe -NoProfile -Command "winget install --id ${_sdk_id} -e --source winget --accept-source-agreements --accept-package-agreements --disable-interactivity" </dev/null || true
if [ -n "$(_find_win_sdk)" ]; then
note "Windows SDK headers present after install."
return 0
fi
done
note "Automatic Windows SDK install did not complete."
return 0
}
# ── PREFLIGHT ────────────────────────────────────────────────────────────────
say "Preflight checks"
# shellcheck disable=SC1091
. /etc/os-release 2>/dev/null || true
if [ "${VERSION_ID:-}" != "24.04" ]; then
die "This targets Ubuntu 24.04 (found '${VERSION_ID:-unknown}'). AMD's ROCm-on-WSL supports 24.04; create a dedicated distro: wsl --install Ubuntu-24.04 (do not run on 26.04 -- ROCm 7.2 does not target it yet)."
fi
if [ ! -e /dev/dxg ]; then
die "/dev/dxg missing -- WSL GPU paravirtualization not present. Ensure this is WSL2 (not WSL1) on a recent Windows build, and that an AMD GPU + ROCDXG-capable Adrenalin driver is installed on the Windows host (then reboot)."
fi
note "Ubuntu 24.04 + /dev/dxg present."
# Don't block on hsa/rocm libs in /usr/lib/wsl/lib: a working ROCDXG setup
# doesn't need them (only d3d12/dxcore). Real readiness is checked via rocminfo.
# ── Step 1: build/runtime prerequisites ──────────────────────────────────────
say "Installing build prerequisites"
export DEBIAN_FRONTEND=noninteractive
$SUDO apt-get update -y
# `make` is explicit: cmake shells out to it but Ubuntu only *recommends* it, so
# minimal images lack it and the librocdxg `make -j` build would fail.
$SUDO apt-get install -y cmake make gcc g++ git wget gpg ca-certificates python3-venv python3-pip
# ── Step 2: ROCm ${ROCM_VER} userspace (no DKMS -- WSL uses the Windows driver) ─
say "Installing ROCm ${ROCM_VER} userspace"
if ! command -v rocminfo >/dev/null 2>&1 && [ ! -x /opt/rocm/bin/rocminfo ]; then
# Direct apt-repo install (leaner than amdgpu-install; repo is indexed by
# ROCm version, e.g. .../apt/7.2.1).
$SUDO mkdir -p /etc/apt/keyrings
wget -qO- https://repo.radeon.com/rocm/rocm.gpg.key \
| gpg --dearmor | $SUDO tee /etc/apt/keyrings/rocm.gpg >/dev/null
echo "deb [arch=amd64 signed-by=/etc/apt/keyrings/rocm.gpg] https://repo.radeon.com/rocm/apt/${ROCM_VER} noble main" \
| $SUDO tee /etc/apt/sources.list.d/rocm.list >/dev/null
printf 'Package: *\nPin: release o=repo.radeon.com\nPin-Priority: 600\n' \
| $SUDO tee /etc/apt/preferences.d/rocm-pin-600 >/dev/null
$SUDO apt-get update -y
# rocm-libs pulls everything torch links at runtime (rocblas, hipblas,
# miopen-hip, rccl, ...); hsa-rocr + rocminfo come as deps. Large (~5 GB
# download / ~23 GB installed).
$SUDO apt-get install -y rocm-libs rocminfo hip-runtime-amd
else
note "ROCm already present -- skipping apt install."
fi
# Resolve the real ROCm dir and ensure the canonical /opt/rocm symlink. apt lays
# ROCm under /opt/rocm-<ver> and rocm-core symlinks /opt/rocm -> that; repair if
# an earlier partial run left /opt/rocm as a real dir blocking the symlink.
_real="$(ls -d /opt/rocm-* 2>/dev/null | sort -V | tail -1 || true)"
if [ -n "$_real" ] && [ ! -L /opt/rocm ] && [ -d /opt/rocm ]; then
# /opt/rocm is a real dir blocking the symlink. Only treat it as a removable
# stray stub if it's NOT a real ROCm install (a real one has bin/rocminfo /
# bin/hipcc / .info/version) -- this protects a user's pre-existing ROCm. Even
# then we MOVE IT ASIDE, never rm -rf, so a wrong guess can't lose data.
if [ -e /opt/rocm/bin/rocminfo ] || [ -e /opt/rocm/bin/hipcc ] || [ -e /opt/rocm/.info/version ]; then
note "/opt/rocm is a real ROCm install -- leaving it untouched (will install librocdxg into it)."
else
note "Moving stray /opt/rocm stub aside -> $_real (not deleting it)"
$SUDO cp -an /opt/rocm/. "$_real"/ 2>/dev/null || true
$SUDO mv /opt/rocm "/opt/rocm.unsloth-stub-bak.$(date +%s)" 2>/dev/null || true
[ -e /opt/rocm ] || $SUDO ln -s "$_real" /opt/rocm
fi
elif [ -n "$_real" ] && [ ! -e /opt/rocm ]; then
$SUDO ln -s "$_real" /opt/rocm
fi
if [ -L /opt/rocm ] || [ -d /opt/rocm ]; then ROCM_DIR="/opt/rocm"; else ROCM_DIR="$_real"; fi
{ [ -n "$ROCM_DIR" ] && [ -d "$ROCM_DIR" ]; } || die "ROCm not found under /opt after install."
note "ROCm at ${ROCM_DIR}"
# ── Step 3: build librocdxg (DXG <-> HSA bridge; not yet an apt package) ──────
say "Building librocdxg (${LIBROCDXG_REF})"
if [ -e "${ROCM_DIR}/lib/librocdxg.so" ]; then
note "librocdxg already installed -- skipping build."
else
# Discover the newest installed Win11 SDK (version differs per machine). If
# absent, auto-install via winget (one UAC prompt) and re-discover; only if
# that ALSO fails do we stop with manual instructions.
_win_sdk="$(_find_win_sdk)"
if [ -z "$_win_sdk" ]; then
note "Windows 11 SDK headers not found -- attempting automatic install..."
_install_windows_sdk_via_winget
_win_sdk="$(_find_win_sdk)"
fi
[ -n "$_win_sdk" ] || die "Windows 11 SDK headers not found under 'C:\\Program Files (x86)\\Windows Kits\\10\\Include\\*\\shared', and the automatic winget install did not complete. Install it on the Windows host (e.g. 'winget install Microsoft.WindowsSDK.10.0.26100') and re-run."
note "Windows SDK: ${_win_sdk}"
_src="${HOME}/.unsloth/librocdxg"
rm -rf "$_src"
git clone --depth 1 --branch "$LIBROCDXG_REF" https://github.com/ROCm/librocdxg.git "$_src" \
|| git clone "https://github.com/ROCm/librocdxg.git" "$_src"
(
cd "$_src"
git checkout "$LIBROCDXG_REF" 2>/dev/null || true
mkdir -p build && cd build
cmake .. -DWIN_SDK="${_win_sdk}/shared"
make -j"$(nproc)"
$SUDO make install
)
fi
# Ensure soname symlinks resolve to whatever version was built (e.g. 1.2.0).
_dxg_real="$(ls -1 "${ROCM_DIR}"/lib/librocdxg.so.*.* 2>/dev/null | sort -V | tail -1 || true)"
if [ -n "$_dxg_real" ]; then
_dxg_base="$(basename "$_dxg_real")" # librocdxg.so.1.2.0
_dxg_major="$(printf '%s' "$_dxg_base" | sed -E 's/librocdxg\.so\.([0-9]+).*/\1/')"
$SUDO ln -sf "$_dxg_base" "${ROCM_DIR}/lib/librocdxg.so.${_dxg_major}"
$SUDO ln -sf "librocdxg.so.${_dxg_major}" "${ROCM_DIR}/lib/librocdxg.so"
fi
echo "${ROCM_DIR}/lib" | $SUDO tee /etc/ld.so.conf.d/rocm.conf >/dev/null
$SUDO ldconfig
# ── Step 4: persist environment (system-wide so Studio's worker inherits it) ──
say "Persisting ROCm-on-WSL environment"
_envfile="/etc/profile.d/unsloth-rocm-wsl.sh"
$SUDO tee "$_envfile" >/dev/null <<EOF
# >>> Unsloth ROCm-on-WSL (gfx1151) >>>
export HSA_ENABLE_DXG_DETECTION=1
export TORCH_ROCM_AOTRITON_ENABLE_EXPERIMENTAL=1
export PATH="${ROCM_DIR}/bin:\${PATH}"
export LD_LIBRARY_PATH="${ROCM_DIR}/lib:\${LD_LIBRARY_PATH:-}"
# <<< Unsloth ROCm-on-WSL (gfx1151) <<<
EOF
# also drop into ~/.bashrc for interactive shells
if [ -n "${HOME:-}" ] && ! grep -q "Unsloth ROCm-on-WSL" "${HOME}/.bashrc" 2>/dev/null; then
cat "$_envfile" >> "${HOME}/.bashrc"
fi
# export into the current process so verification below works immediately
export HSA_ENABLE_DXG_DETECTION=1
export PATH="${ROCM_DIR}/bin:${PATH}"
export LD_LIBRARY_PATH="${ROCM_DIR}/lib:${LD_LIBRARY_PATH:-}"
# ── Step 5: verify the runtime enumerates the GPU ────────────────────────────
say "Verifying rocminfo sees ${GFX}"
# Capture rocminfo into a var BEFORE grepping: piping into `grep -q` SIGPIPEs
# rocminfo on first match, which under `set -o pipefail` turns a successful match
# into a pipeline failure. Match the gfx1151 ISA "Name:" agent exactly (not a
# broad gfx1[0-9]) so a generic fallback ISA or unrelated RDNA GPU can't pass.
_rocminfo_out="$(rocminfo 2>/dev/null || true)"
if ! printf '%s\n' "$_rocminfo_out" | grep -qE "Name:[[:space:]]*${GFX}([^0-9]|$)"; then
printf '%s\n' "$_rocminfo_out" | head -25 >&2 || true
die "rocminfo did not enumerate a ${GFX} GPU agent. Most common cause: the Windows AMD driver predates production ROCDXG -- update Adrenalin (install.ps1 offers this), reboot, and re-run."
fi
# Display-only summary: best-effort (|| true) so head's early pipe-close under
# `set -o pipefail` can't fail the bootstrap after verification already passed.
printf '%s\n' "$_rocminfo_out" | grep -E 'Marketing Name|Device Type|Compute Unit' | grep -iE "Radeon|GPU|Compute" | head -3 || true
note "ROCm-on-WSL runtime is live for ${GFX}."
# ── Step 6 (optional): torch smoke test from the gfx1151 index ───────────────
if [ "$SMOKE_TEST" = "1" ]; then
say "Smoke-testing PyTorch on ${GFX} (throwaway venv)"
_venv="${HOME}/.unsloth/rocm-smoketest"
rm -rf "$_venv"; python3 -m venv "$_venv"
"$_venv/bin/pip" install --quiet --upgrade pip
# gfx1151 index is primary (torch + triton); PyPI only an extra for pure-py
# deps. The constraint keeps pip on the ROCm wheel, not a newer PyPI CUDA torch.
"$_venv/bin/pip" install --index-url "$TORCH_INDEX" \
--extra-index-url https://pypi.org/simple "$TORCH_CONSTRAINT" || \
die "torch install from ${TORCH_INDEX} failed."
"$_venv/bin/python" - <<'PY'
import torch
ok = torch.cuda.is_available()
print("torch:", torch.__version__, "| cuda(rocm) available:", ok)
if ok:
print("device:", torch.cuda.get_device_name(0))
free, total = torch.cuda.mem_get_info(0)
print(f"mem: free={free/1e9:.1f} GB total={total/1e9:.1f} GB")
import time
a = torch.randn(4096, 4096, device="cuda", dtype=torch.float16)
b = torch.randn(4096, 4096, device="cuda", dtype=torch.float16)
torch.cuda.synchronize(); t0 = time.time()
for _ in range(10): c = a @ b
torch.cuda.synchronize()
print(f"matmul ok ({(time.time()-t0)/10*1e3:.1f} ms/iter)")
raise SystemExit(0 if ok else 1)
PY
rm -rf "$_venv"
fi
say "Done."
note "ROCm-on-WSL is ready for ${GFX}. If you ran this standalone, install Unsloth"
note "in THIS distro and it will detect the GPU automatically:"
note " curl -fsSL https://unsloth.ai/install.sh | sh"

48
scripts/lint_workflow_triggers.py
View file

@ -4,37 +4,19 @@
"""Refuse dangerous GitHub Actions trigger patterns at PR time.
Two patterns are banned outright, both of which powered the TanStack
GHSA-g7cv-rxg3-hmpx supply-chain compromise:
Bans patterns behind the TanStack GHSA-g7cv-rxg3-hmpx compromise:
1. `pull_request_target` -- runs a fork's workflow YAML against the
BASE repository's secrets and permissions. The fork can inject
arbitrary code into the base context. The TanStack worm used this
to land base-context execution from a fork PR. There is essentially
no safe use of this trigger for a public open-source project;
`pull_request` is the safe alternative.
1. `pull_request_target` -- runs a fork's workflow against the base
repo's secrets/permissions; use `pull_request` instead.
2. `workflow_run` chained to a PR-triggered workflow -- same trust
boundary problem one hop later (poisoned artifacts/caches run with
elevated permissions).
3. Cache keys shared between PR-triggered and publish/release/push
workflows -- a fork PR could poison a cache the publish workflow
restores. Partition the key namespaces.
2. `workflow_run` chained to a PR-triggered workflow -- carries the
same trust boundary problem one hop later. If a PR-triggered
workflow can poison artifacts/caches and a `workflow_run` trigger
fires off the result with elevated permissions, the attacker still
reaches the trusted context.
3. Shared cache keys between PR-triggered workflows and publish /
release / push-triggered workflows. The TanStack worm poisoned the
Actions cache from a fork PR and the legitimate release workflow
then restored the poisoned cache. Cache keys must be partitioned
so that nothing a PR can write is ever read by a workflow that
holds secrets.
Exit codes
==========
0 no findings
1 one or more findings; stderr lists each with file path
Run from repo root:
python3 scripts/lint_workflow_triggers.py
Exit codes: 0 = no findings, 1 = findings (listed on stderr).
Run from repo root: python3 scripts/lint_workflow_triggers.py
"""
from __future__ import annotations
@ -47,9 +29,7 @@ from pathlib import Path
try:
import yaml
except ImportError:
print(
"ERROR: PyYAML is required. Install with 'pip install pyyaml'", file = sys.stderr
)
print("ERROR: PyYAML is required. Install with 'pip install pyyaml'", file = sys.stderr)
sys.exit(2)
REPO_ROOT = Path(__file__).resolve().parents[1]
@ -153,9 +133,7 @@ def main() -> int:
)
if findings:
print(
"Workflow trigger lint failed with the following issues:", file = sys.stderr
)
print("Workflow trigger lint failed with the following issues:", file = sys.stderr)
for f in findings:
print(f" - {f}", file = sys.stderr)
return 1

220
scripts/lockfile_supply_chain_audit.py
View file

@ -5,64 +5,21 @@
"""Lockfile supply-chain audit for the Studio frontend and Tauri shell.
Runs BEFORE `npm ci` / `cargo fetch` in CI. Refuses to proceed when a
lockfile contains patterns that indicate the kind of supply-chain
injection seen in the npm Shai-Hulud waves and the cargo
crates.io brand-squat attempts.
lockfile contains patterns indicating supply-chain injection (npm
Shai-Hulud waves, cargo crates.io brand-squats).
What it checks
==============
Checks package-lock.json (lockfileVersion 2/3): `resolved` URL must be
the npm registry (direct git/github/file refs are the injection vector);
`integrity` SHA must be present; known IOC substrings grepped from the
body. Checks Cargo.lock: `source` must be the crates.io registry index;
known cargo IOC substrings.
studio/frontend/package-lock.json (lockfileVersion 2 or 3):
Exit codes: 0 = clean (or skip env var set to a justification >=5 chars,
not '1'/'true'); 1 = findings; 2 = internal error.
1. `resolved` URL origin. Every entry must resolve through
`https://registry.npmjs.org/`. Direct GitHub-hosted dependencies
(`git+ssh://`, `git+https://`, `github:owner/repo#sha`,
`file:`, `http://`) are refused -- npm's TanStack incident used
exactly this vector to land an unaudited GitHub commit hash as
an optional dependency.
2. `integrity` field presence. Every non-workspace entry must carry
an `integrity` SHA. A missing integrity means the registry can
swap the tarball after lockfile generation and CI will not
notice.
3. Known IOC strings. A hardcoded set of indicator-of-compromise
substrings is grepped across the entire lockfile body (file
names, dependency keys, URLs). The list is updated as new
campaigns surface. Catching one means the local install was
about to pull a publicly-known malicious release.
studio/src-tauri/Cargo.lock:
4. `source` field origin. Every entry with a `source` must point at
`registry+https://github.com/rust-lang/crates.io-index`. Direct
git sources (`git+https://...`) and `path+...` for cross-crate
paths warrant manual review and are flagged.
5. Known cargo IOC strings. Same idea as (3), separate list.
Exit codes
==========
0 no findings, or an opt-out env var (UNSLOTH_LOCKFILE_AUDIT_SKIP)
is set to a justification string (>=5 chars, not '1'/'true'/etc).
A value like '1' or 'true' is now REJECTED loudly and the audit
runs normally
1 one or more findings; stderr lists them with file path and line
number where derivable
2 internal error (missing dependency, malformed JSON, etc.)
Operational stance
==================
This scanner only PARSES the lockfiles -- it never executes anything
in them, never resolves anything against the network. Safe to run
ahead of every `npm ci`. The IOC list is short by design; this
complements (not replaces) `npm audit`, OSV-Scanner, and the
advisory-DB pipeline in `.github/workflows/security-audit.yml`. The
shape of the catch is "we refuse to proceed because the lockfile
itself is shaped wrong", which fires before any third-party install
script gets a chance to run on the runner.
Only PARSES the lockfiles, never executes or networks. Complements (not
replaces) `npm audit` / OSV-Scanner / the advisory-DB pipeline. Fires
before any third-party install script runs on the runner.
"""
from __future__ import annotations
@ -77,14 +34,9 @@ from pathlib import Path
REPO_ROOT = Path(__file__).resolve().parents[1]
# ─────────────────────────────────────────────────────────────────────
# Known IOC strings (case-sensitive substring match).
# ─────────────────────────────────────────────────────────────────────
#
# Keep these short and FACTUAL. Each entry is tied to a public advisory
# and is the literal string an attacker would have to embed for the
# attack to work. Adding speculative or generic patterns here would
# generate false positives on dependency upgrades.
# Known IOC strings (case-sensitive substring match). Each is tied to a
# public advisory; speculative/generic patterns would false-positive on
# upgrades.
NPM_IOC_STRINGS: tuple[str, ...] = (
# Shai-Hulud TanStack wave -- May 11, 2026 (GHSA-g7cv-rxg3-hmpx).
"router_init.js",
@ -328,36 +280,22 @@ BLOCKED_NPM_VERSIONS: dict[str, set[str]] = {
}
CARGO_IOC_STRINGS: tuple[str, ...] = (
# Reserved for future cargo-side incidents. Empty by default --
# `source` origin check below catches the structural pattern.
# Empty by default; the `source` origin check catches the structural
# pattern. Reserved for future cargo-side incidents.
)
# ─────────────────────────────────────────────────────────────────────
# Allowed lockfile origins.
# ─────────────────────────────────────────────────────────────────────
NPM_REGISTRY_PREFIX = "https://registry.npmjs.org/"
# Tarballs are also fetched from this mirror on some GH Actions cached
# runs (npm rewrites the resolved URL on cache hit). Allow either.
NPM_REGISTRY_PREFIXES_ALLOWED: tuple[str, ...] = (NPM_REGISTRY_PREFIX,)
CARGO_REGISTRY_SOURCE = "registry+https://github.com/rust-lang/crates.io-index"
# ─────────────────────────────────────────────────────────────────────
# Cargo non-registry source allowlist.
# ─────────────────────────────────────────────────────────────────────
#
# Each entry is `(crate_name, exact_source_string)`. The crate must
# match by name AND the source must match the full pinned-SHA string
# verbatim. Bumping the commit SHA forces a re-review here: the
# scanner fires until the new SHA is appended.
#
# Studio's Tauri shell pulls `fix-path-env` directly from
# tauri-apps/fix-path-env-rs because the crate is not published to
# crates.io. The pinned commit (c4c45d5) was reviewed at the time it
# landed; future bumps need explicit approval.
# Cargo non-registry source allowlist: `(crate_name, exact_source_string)`.
# Both must match verbatim; bumping the pinned SHA forces a re-review.
# Studio's Tauri shell pulls `fix-path-env` from git because it is not
# published to crates.io; commit c4c45d5 was reviewed when it landed.
CARGO_SOURCE_ALLOWLIST: tuple[tuple[str, str], ...] = (
(
"fix-path-env",
@ -367,11 +305,6 @@ CARGO_SOURCE_ALLOWLIST: tuple[tuple[str, str], ...] = (
)
# ─────────────────────────────────────────────────────────────────────
# Finding container.
# ─────────────────────────────────────────────────────────────────────
class Finding:
__slots__ = ("path", "package", "kind", "detail")
@ -390,27 +323,19 @@ class Finding:
def _gha_escape(text: str) -> str:
"""Escape a string for use in a GitHub Actions `::warning::` /
`::error::` workflow command message. GH Actions truncates
annotation messages at the first newline unless `\\n` is
escaped as `%0A`; carriage returns and the percent sign need
matching escapes per the workflow-commands spec. Order matters:
`%` must be replaced first so the subsequent `%0A` / `%0D`
sequences are not double-encoded.
"""Escape a string for a GH Actions `::warning::`/`::error::` message.
GH Actions truncates at the first newline unless `\\n`/`\\r` are
escaped as `%0A`/`%0D`. `%` must be replaced first to avoid
double-encoding the subsequent escapes.
"""
return text.replace("%", "%25").replace("\r", "%0D").replace("\n", "%0A")
# ─────────────────────────────────────────────────────────────────────
# package-lock.json audit.
# ─────────────────────────────────────────────────────────────────────
def audit_npm_lockfile(path: Path) -> list[Finding]:
findings: list[Finding] = []
if not path.exists():
# A missing requested lockfile is a config error, not a clean
# audit; surface it so a deleted default cannot pass silently.
# Missing lockfile is a config error, not a clean audit.
findings.append(
Finding(
path = str(path),
@ -427,8 +352,7 @@ def audit_npm_lockfile(path: Path) -> list[Finding]:
try:
raw = path.read_text(encoding = "utf-8")
except OSError as exc:
# Permission denied, is-a-directory, broken-pipe etc. -- surface
# as a finding instead of crashing CI with a raw traceback.
# Surface as a finding instead of crashing CI with a traceback.
findings.append(
Finding(
path = str(path),
@ -464,9 +388,7 @@ def audit_npm_lockfile(path: Path) -> list[Finding]:
packages = lock.get("packages") or {}
for key, entry in packages.items():
# The empty key "" is the project root; workspace entries use
# keys like "node_modules/foo" or "studio/frontend/sub-pkg".
# Skip the project root (it has no `resolved`).
# Empty key "" is the project root (no `resolved`); skip it.
if key == "":
continue
if entry.get("link"):
@ -474,12 +396,8 @@ def audit_npm_lockfile(path: Path) -> list[Finding]:
continue
resolved = entry.get("resolved")
# Entries living inside another package's `node_modules/`
# tree are bundled fold-ins -- the parent's tarball ships
# their source verbatim and the parent's `integrity` covers
# the whole subtree. npm represents them in lockfileVersion 3
# as nested entries with no `resolved` and no `integrity` of
# their own. Treat them as transparent to this audit.
# Entries nested in another package's node_modules are bundled
# fold-ins covered by the parent's integrity; treat as transparent.
nested = key.count("/node_modules/") >= 1
# 1. resolved-URL origin.
@ -541,18 +459,14 @@ def audit_npm_lockfile(path: Path) -> list[Finding]:
path = str(path),
package = key,
kind = "blocked-known-malicious",
detail = (
f"{pkg_name}@{version} is on the " "BLOCKED_NPM_VERSIONS list"
),
detail = (f"{pkg_name}@{version} is on the BLOCKED_NPM_VERSIONS list"),
)
)
# 4. Known IOC strings: scan the raw file body so we hit fields the
# structural pass above doesn't enumerate (scripts, optional
# dependencies, etc.). Cheap and complete.
# 4. Known IOC strings: scan the raw body to catch fields the
# structural pass doesn't enumerate (scripts, optional deps, etc.).
for ioc in NPM_IOC_STRINGS:
if ioc in raw:
# Best-effort line number lookup.
line_no = _first_line_containing(raw, ioc)
findings.append(
Finding(
@ -577,14 +491,7 @@ def _first_line_containing(text: str, needle: str) -> int | None:
return None
# ─────────────────────────────────────────────────────────────────────
# Cargo.lock audit.
# ─────────────────────────────────────────────────────────────────────
# Cargo.lock is TOML; parse with stdlib tomllib (Python 3.11+). The
# studio's Tauri shell already requires a modern toolchain so this is
# always available where CI runs.
# Cargo.lock is TOML; parsed with stdlib tomllib (Python 3.11+).
_PACKAGE_HEADER = re.compile(r"^\[\[package\]\]\s*$")
@ -708,24 +615,15 @@ def audit_cargo_lockfile(path: Path) -> list[Finding]:
return findings
# ─────────────────────────────────────────────────────────────────────
# CLI.
# ─────────────────────────────────────────────────────────────────────
# Finding kinds split into BLOCKING vs ADVISORY for the default run mode.
# Blocking findings come from public supply-chain attack indicators (a
# version we know is malicious, a string an attacker would have to embed
# for an attack to work). Advisory findings are structural lockfile
# anomalies (missing integrity, non-default registry, etc.) -- they
# WARN the maintainer but do not block merges. Pass --strict to make
# every finding blocking (PR-5479-style behavior for opt-in adopters).
# Blocking = public attack indicators (known-malicious version, IOC
# string). Advisory = structural anomalies that warn but don't block.
# --strict makes every finding blocking.
BLOCKING_KINDS: frozenset[str] = frozenset(
{
"blocked-known-malicious",
"known-ioc-string",
# Internal-failure kinds: a structurally broken lockfile MIGHT
# be hiding a real attack, so we keep these blocking too.
# A structurally broken lockfile might hide a real attack.
"malformed-lockfile",
"missing-lockfile",
"unreadable-lockfile",
@ -765,10 +663,7 @@ def main(argv: list[str] | None = None) -> int:
"--cargo-lockfile",
action = "append",
default = None,
help = (
"Path to a Cargo.lock (repeatable). "
"Default: studio/src-tauri/Cargo.lock."
),
help = ("Path to a Cargo.lock (repeatable). Default: studio/src-tauri/Cargo.lock."),
)
parser.add_argument(
"--strict",
@ -784,15 +679,9 @@ def main(argv: list[str] | None = None) -> int:
)
args = parser.parse_args(argv)
# SF4: require a real justification (e.g. JIRA ticket id) for the
# skip env var. Treat the trivially-set values ("1", "true", "yes",
# "on", empty) as INVALID -- they look like accidental flips and
# silently bypassed the supply-chain audit. A valid value is a
# non-empty string >=5 chars after stripping that does not match
# any of the boolean-shaped tokens above. An invalid value emits a
# loud GitHub Actions warning to stderr and FALLS THROUGH to run
# the audit normally (fail-safe). A valid value emits a warning
# naming the reason and skips with rc=0 (compat).
# Require a real justification (>=5 chars, not a boolean-shaped token)
# for the skip env var. An invalid value warns and falls through to
# run the audit (fail-safe); a valid one warns and skips with rc=0.
_skip_raw = os.environ.get("UNSLOTH_LOCKFILE_AUDIT_SKIP")
if _skip_raw is not None:
_skip = _skip_raw.strip()
@ -814,8 +703,7 @@ def main(argv: list[str] | None = None) -> int:
return 0
root = Path(args.root).resolve()
# Explicit --npm-lockfile/--cargo-lockfile scopes the scan to those
# paths; defaults apply only to the no-args CI invocation.
# Explicit flags scope the scan; defaults apply only to no-args CI.
_user_explicit = args.npm_lockfile is not None or args.cargo_lockfile is not None
if _user_explicit:
npm_paths = [root / p for p in (args.npm_lockfile or ())]
@ -840,11 +728,9 @@ def main(argv: list[str] | None = None) -> int:
)
return 0
# Split findings into blocking (known-malicious / IOC / structurally
# broken) and advisory (everything else, e.g. missing integrity on a
# registry-published tarball). In default mode advisory findings are
# printed but do not change the exit code; --strict treats every
# finding as blocking.
# Split into blocking (known-malicious / IOC / structurally broken)
# and advisory (everything else). Default mode prints advisories
# without changing the exit code; --strict makes all blocking.
blocking = [f for f in all_findings if f.kind in BLOCKING_KINDS]
advisory = [f for f in all_findings if f.kind not in BLOCKING_KINDS]
@ -859,12 +745,8 @@ def main(argv: list[str] | None = None) -> int:
file = sys.stderr,
)
for f in advisory:
# Surface in GitHub Actions UI as a warning annotation when run
# under Actions; harmless prefix elsewhere. GH Actions
# truncates annotation messages at the first newline unless
# newlines are escaped as `%0A`, so the full multi-line
# Finding (kind + path + package + detail) only renders in
# the UI after _gha_escape collapses it onto one line.
# GH Actions warning annotation; _gha_escape collapses the
# multi-line Finding onto one line so it renders fully in the UI.
print(f"::warning::{_gha_escape(str(f))}", file = sys.stderr)
print(file = sys.stderr)
@ -881,9 +763,7 @@ def main(argv: list[str] | None = None) -> int:
file = sys.stderr,
)
for f in blocking:
# Same %-encoding rationale as the advisory branch above: the
# GH Actions annotation is truncated at the first newline
# unless the message is escaped.
# Same %-encoding rationale as the advisory branch above.
print(f"::error::{_gha_escape(str(f))}", file = sys.stderr)
print(file = sys.stderr)
print(

89
scripts/notebook_to_python.py
View file

@ -22,19 +22,14 @@ import urllib.parse
from pathlib import Path
# Hosts we are willing to fetch raw notebook JSON from. Anything else
# is rejected before `urlopen` so a typoed / hostile URL cannot pull
# code from arbitrary infrastructure.
# Allowlist of hosts for raw notebook fetches; anything else rejected before urlopen.
_ALLOWED_NOTEBOOK_HOSTS = {
"raw.githubusercontent.com",
"gist.githubusercontent.com",
}
# Shell metacharacters that imply the cell's `!cmd` line cannot be
# parsed as a flat argv. If any of these appears, `shlex.split` would
# either fail or, worse, silently strip the operator -- so we keep
# `shell=True` for that command and emit a review marker.
# Metacharacters that mean a `!cmd` line can't be a flat argv -> keep shell=True + review marker.
_SHELL_METACHARS_RE = re.compile(r"\$\(|`|\|\||\||&&|>>?|<<?|\*|\?|;")
@ -46,12 +41,8 @@ def needs_fstring(cmd: str) -> bool:
def github_blob_to_raw(url: str) -> str:
"""Convert GitHub blob URL to raw URL."""
# https://github.com/user/repo/blob/branch/path
# -> https://raw.githubusercontent.com/user/repo/branch/path
# Compare the parsed host exactly (not as a substring) so a URL
# like https://attacker.example.com/github.com/blob/... does NOT
# get rewritten to a github raw URL. Closes CodeQL alert
# py/incomplete-url-substring-sanitization.
# github.com/user/repo/blob/branch/path -> raw.githubusercontent.com/user/repo/branch/path
# Exact host match (not substring) so attacker.example.com/github.com/blob/... is not rewritten.
parsed = urllib.parse.urlparse(url)
if parsed.netloc != "github.com" or "/blob/" not in parsed.path:
return url
@ -63,18 +54,12 @@ def github_blob_to_raw(url: str) -> str:
def download_notebook(url: str) -> tuple[str, str]:
"""Download notebook from URL. Returns (content, filename)."""
# Convert blob URL to raw if needed
raw_url = github_blob_to_raw(url)
# Extract filename from URL
parsed = urllib.parse.urlparse(raw_url)
filename = os.path.basename(urllib.parse.unquote(parsed.path))
# Host allowlist. Refuse to fetch from anywhere the campaign IOC
# tables flag (or just anywhere we don't recognise). The blob->raw
# conversion above only emits `raw.githubusercontent.com`, so a
# rejection here means the caller hand-typed a URL pointing
# somewhere we don't trust.
# Host allowlist: refuse to fetch from anything we don't recognise.
host = parsed.hostname
if host not in _ALLOWED_NOTEBOOK_HOSTS:
raise ValueError(
@ -82,7 +67,6 @@ def download_notebook(url: str) -> tuple[str, str]:
f"{sorted(_ALLOWED_NOTEBOOK_HOSTS)}"
)
# Download
print(f"Downloading {url}...")
with urllib.request.urlopen(raw_url, timeout = 60) as response:
content = response.read().decode("utf-8")
@ -97,29 +81,18 @@ def is_url(path: str) -> bool:
def replace_colab_paths(source: str) -> str:
"""Replace Colab-specific /content/ paths with current working directory."""
# Replace /content/ with f-string using _WORKING_DIR
source = source.replace('"/content/', 'f"{_WORKING_DIR}/')
source = source.replace("'/content/", "f'{_WORKING_DIR}/")
return source
def _emit_shell_command(indent: str, full_cmd: str, *, allow_shell: bool) -> list[str]:
"""Render a `!cmd` notebook line as one or more Python statements.
"""Render a `!cmd` notebook line as Python statements.
When the command body is f-string-interpolated, contains shell
metacharacters, or spans multiple lines, falling back to
`shell=True` is the only correct option -- `shlex.split` would
either drop operators or fail outright. We surface that with a
`# WARNING: shell=True; reviewed for hostile input` comment so a
reviewer cannot miss it.
Otherwise we emit `subprocess.run(shlex.split(cmd), shell=False)`
so the converted script is not a re-injection vector if the
notebook ever interpolates user-controlled data.
`allow_shell` defaults to True at the CLI for backwards
compatibility. Setting it to False makes `shell=True` emission a
hard error (no surprise behaviour).
f-string interpolation, shell metacharacters, or multiline force
shell=True (shlex.split would drop operators), flagged with a
WARNING comment. Otherwise emit shell=False argv form. allow_shell
False makes shell=True emission a hard error.
"""
needs_f = needs_fstring(full_cmd)
has_meta = bool(_SHELL_METACHARS_RE.search(full_cmd))
@ -144,7 +117,6 @@ def _emit_shell_command(indent: str, full_cmd: str, *, allow_shell: bool) -> lis
stmt = f"{indent}subprocess.run({f_prefix}{full_cmd!r}, shell=True)"
return [warn, stmt]
# Shell-safe argv form.
return [f"{indent}subprocess.run(shlex.split({full_cmd!r}), shell=False)"]
@ -159,12 +131,10 @@ def convert_cell_to_python(source: str, *, allow_shell: bool = True) -> str:
stripped = line.strip()
indent = line[: len(line) - len(line.lstrip())]
# Skip %%capture
if stripped.startswith("%%capture"):
i += 1
continue
# Handle %%file magic
if stripped.startswith("%%file "):
filename = stripped[7:].strip()
file_lines = []
@ -178,7 +148,6 @@ def convert_cell_to_python(source: str, *, allow_shell: bool = True) -> str:
result.append(f'{indent} _f.write("""{file_content}""")')
continue
# Handle ! shell commands
if stripped.startswith("!"):
cmd_lines = [stripped[1:]]
while cmd_lines[-1].rstrip().endswith("\\") and i + 1 < len(lines):
@ -186,9 +155,7 @@ def convert_cell_to_python(source: str, *, allow_shell: bool = True) -> str:
cmd_lines.append(lines[i].strip())
full_cmd = "\n".join(cmd_lines)
result.extend(
_emit_shell_command(indent, full_cmd, allow_shell = allow_shell)
)
result.extend(_emit_shell_command(indent, full_cmd, allow_shell = allow_shell))
# %cd path -> os.chdir(path)
elif stripped.startswith("%cd "):
@ -310,23 +277,16 @@ def convert_notebook_to_script(
content = f.read()
source_name = source
# Generate output filename
output_filename = filename.replace(".ipynb", ".py")
# Clean up filename
output_filename = (
output_filename.replace("(", "").replace(")", "").replace("-", "_")
)
output_filename = output_filename.replace("(", "").replace(")", "").replace("-", "_")
# Add output directory if specified
if output_dir:
output_path = os.path.join(output_dir, output_filename)
else:
output_path = output_filename
# Convert
script = convert_notebook(content, source_name, allow_shell = allow_shell)
# Write output
with open(output_path, "w", encoding = "utf-8") as f:
f.write(script)
@ -337,9 +297,7 @@ def convert_notebook_to_script(
def main():
import argparse
class Formatter(
argparse.ArgumentDefaultsHelpFormatter, argparse.RawDescriptionHelpFormatter
):
class Formatter(argparse.ArgumentDefaultsHelpFormatter, argparse.RawDescriptionHelpFormatter):
pass
parser = argparse.ArgumentParser(
@ -353,17 +311,9 @@ Examples:
python notebook_to_python.py https://github.com/unslothai/notebooks/blob/main/nb/Oute_TTS_(1B).ipynb
""",
)
parser.add_argument(
"notebooks", nargs = "+", help = "Notebook files or URLs to convert."
)
parser.add_argument(
"-o", "--output", dest = "output_dir", default = ".", help = "Output directory."
)
# Default True for backwards compatibility: existing Colab notebooks
# routinely use pipes / redirection / interpolation in `!cmd` lines
# and the converted script needs to keep working. Operators who
# convert untrusted notebooks should pass --no-allow-shell to force
# a hard error on every metacharacter-bearing cell.
parser.add_argument("notebooks", nargs = "+", help = "Notebook files or URLs to convert.")
parser.add_argument("-o", "--output", dest = "output_dir", default = ".", help = "Output directory.")
# Default True for backwards compat; pass --no-allow-shell for untrusted notebooks.
parser.add_argument(
"--allow-shell",
dest = "allow_shell",
@ -381,14 +331,9 @@ Examples:
args = parser.parse_args()
# Create output directory if needed
os.makedirs(args.output_dir, exist_ok = True)
# SF2: track per-notebook failures so a CI invocation that converts
# 10 notebooks but silently fails on 3 is no longer reported as
# success. Each failure is collected and the loop continues so the
# caller sees the full set; final exit status is 1 if anything
# failed.
# Track per-notebook failures; continue the loop and exit 1 if any failed.
failures: list[tuple[str, str]] = []
ok = 0
total = len(args.notebooks)

199
scripts/notebook_validator.py
View file

@ -49,12 +49,9 @@ from typing import Any, Iterable, Iterator
def _atomic_write_bytes(path: pathlib.Path, data: bytes) -> None:
"""Atomic write helper. See `scripts/scan_packages.py::update_req_file`.
A crash between `mkstemp` and `os.replace` leaves the prior file
untouched, so a half-downloaded PyPI metadata cache file cannot
poison subsequent runs of the validator.
"""
"""Atomic write (see scripts/scan_packages.py::update_req_file). A crash
between mkstemp and os.replace leaves the prior file intact, so a
half-downloaded cache file can't poison later runs."""
path.parent.mkdir(parents = True, exist_ok = True)
dirpath = str(path.parent) or "."
fd, tmp_path = tempfile.mkstemp(prefix = ".nb_val.", dir = dirpath)
@ -81,20 +78,16 @@ COLAB_PIP_FREEZE_URL = (
)
COLAB_FALLBACK_FILE = DATA_DIR / "colab_pip_freeze.gpu.txt"
# Oracle files we snapshot from googlecolab/backend-info. The diff
# subcommand fetches each, compares against the committed snapshot,
# and surfaces NEW / REMOVED / CHANGED entries so upstream Colab base
# image rotations land in CI within ~24h instead of when a notebook
# breaks. Every rule in this validator that resolves against the
# Colab preinstall (R-INST-002/003/004/005) gets earlier signal.
# Oracle files snapshotted from googlecolab/backend-info. The colab-diff
# subcommand surfaces NEW/REMOVED/CHANGED entries so upstream Colab base
# image rotations land in CI within ~24h, giving R-INST-002/003/004/005
# earlier signal.
COLAB_ORACLE_FILES: dict[str, str] = {
"pip-freeze.gpu.txt": "colab_pip_freeze.gpu.txt",
"apt-list-gpu.txt": "colab_apt_list.gpu.txt",
"os-info-gpu.txt": "colab_os_info.gpu.txt",
}
COLAB_ORACLE_BASE_URL = (
"https://raw.githubusercontent.com/googlecolab/backend-info/main/"
)
COLAB_ORACLE_BASE_URL = "https://raw.githubusercontent.com/googlecolab/backend-info/main/"
# ----- Compat tables. PRs add rows as new releases land. ----- #
@ -147,9 +140,8 @@ class Finding:
def iter_notebooks(
notebooks_dir: pathlib.Path, include_templates: bool = False
) -> Iterator[pathlib.Path]:
"""Yield user-facing .ipynb files under nb/ and kaggle/. Pass
include_templates=True to also walk original_template/ (used by the
convert subcommand which doesn't lint install cells)."""
"""Yield user-facing .ipynb files under nb/ and kaggle/.
include_templates=True also walks original_template/ (for convert)."""
subs = ("nb", "kaggle")
if include_templates:
subs = ("nb", "kaggle", "original_template")
@ -195,18 +187,13 @@ def install_cells(nb: dict[str, Any]) -> list[tuple[int, str]]:
if first and first[0].strip().startswith("%%capture"):
out.append((i, src))
continue
if re.search(
r"^[ \t]*!\s*(uv\s+)?pip\s+(install|uninstall)\b", src, re.MULTILINE
):
if re.search(r"^[ \t]*!\s*(uv\s+)?pip\s+(install|uninstall)\b", src, re.MULTILINE):
out.append((i, src))
return out
# Notebook target environment. The Colab oracle (pip-freeze.gpu.txt) only
# applies to notebooks that actually run on Colab; AMD-Dev-Cloud,
# Kaggle, HuggingFace-Course, and DGX-Spark notebooks have their own
# preinstalled environments and the Colab-vs-cell rules are not
# applicable to them.
# Colab oracle only applies to notebooks that run on Colab; AMD, Kaggle,
# DGX-Spark have their own preinstalls and the Colab-vs-cell rules don't apply.
def target_environment(notebook_name: str) -> str:
parts = pathlib.PurePath(notebook_name).parts
base = parts[-1] if parts else notebook_name
@ -331,9 +318,7 @@ def parse_pip_line(line: str, line_no: int = 0) -> PipInvocation | None:
if t in ("install", "uninstall"):
continue
packages.append(t)
return PipInvocation(
tool = tool, flags = flags, packages = packages, raw = line, line_no = line_no
)
return PipInvocation(tool = tool, flags = flags, packages = packages, raw = line, line_no = line_no)
def _glue_line_continuations(text: str) -> list[tuple[int, str]]:
@ -418,9 +403,7 @@ def pypi_metadata(name: str, version: str) -> dict[str, Any] | None:
return data
def transitive_constraint(
name: str, version: str, target: str
) -> tuple[str | None, list[str]]:
def transitive_constraint(name: str, version: str, target: str) -> tuple[str | None, list[str]]:
"""Return (raw_specifier_string_or_None, list_of_(op,version) tuples)
for the constraint that `name==version` places on `target`.
"""
@ -474,19 +457,12 @@ def constraint_satisfied(version: str, ops: list[tuple[str, str]]) -> bool:
def resolved_set(install_cell: str, colab: dict[str, str]) -> dict[str, str]:
"""Merge install-cell explicit constraints with Colab pip-freeze. Cell
wins.
"""Merge install-cell constraints with Colab pip-freeze (cell wins).
Resolution order per package, when more than one form is present:
1. Exact `==V` pin in any install line (definitive).
2. Upper-bound `<=V` constraint (pip picks the highest
allowed; that's V).
3. Colab pip-freeze fallback.
The lower-bound `>=V` is intentionally NOT reflected here — a `>=V`
by itself doesn't change the resolved version when a higher
Colab-preinstalled version is already in scope. (R-INST-003 calls
`_install_cell_lower_bound` separately to model that case.)
Resolution order per package: (1) exact `==V` pin, (2) upper-bound `<=V`
(pip picks the highest allowed = V), (3) Colab fallback. Lower-bound `>=V`
is intentionally NOT reflected (it doesn't lower an already-higher Colab
version); R-INST-003 models that via `_install_cell_lower_bound`.
"""
out = dict(colab)
pinned: set[str] = set()
@ -501,10 +477,7 @@ def resolved_set(install_cell: str, colab: dict[str, str]) -> dict[str, str]:
out[sp.name] = ver
pinned.add(sp.name)
elif op == "<=" and sp.name not in pinned:
if (
sp.name not in upper_bounds
or cmp_versions(ver, upper_bounds[sp.name]) < 0
):
if sp.name not in upper_bounds or cmp_versions(ver, upper_bounds[sp.name]) < 0:
upper_bounds[sp.name] = ver
# Apply upper bounds where Colab's preinstall violates them.
for name, ub in upper_bounds.items():
@ -519,9 +492,7 @@ def resolved_set(install_cell: str, colab: dict[str, str]) -> dict[str, str]:
# ----- Rules ----- #
def rule_inst_001_git_plus(
install_cell: str, file: str, cell_idx: int
) -> list[Finding]:
def rule_inst_001_git_plus(install_cell: str, file: str, cell_idx: int) -> list[Finding]:
findings: list[Finding] = []
for inv in iter_pip_invocations(install_cell):
if any("git+" in p for p in inv.packages) or "git+" in inv.raw:
@ -556,8 +527,7 @@ def rule_inst_002_no_deps_transitive(
v = explicit_pin(sp)
if v is None:
continue
# Check transitive constraints on a curated short list of pkgs we
# care about (transformers/peft/trl/accelerate/torchao/torchcodec).
# Check transitive constraints on a curated short list of pkgs.
for target in (
"tokenizers",
"torchao",
@ -588,10 +558,9 @@ def rule_inst_002_no_deps_transitive(
def _install_cell_lower_bound(install_cell: str, target: str) -> str | None:
"""Return the highest LOWER bound that any install line places on `target`,
or None if no constraint is present. Treats `==V` as both lower and upper.
Used by R-INST-003: a `pip install torchao>=0.16.0` line is enough to
satisfy a `torchao>=0.16.0` floor even though it's not a `==` pin."""
"""Return the highest lower bound any install line places on `target`
(treating `==V` as both bounds), or None. Used by R-INST-003 so a
`torchao>=0.16.0` line satisfies the floor without a `==` pin."""
best: str | None = None
for inv in iter_pip_invocations(install_cell):
for raw in inv.packages:
@ -665,20 +634,17 @@ def rule_inst_004_torchcodec_torch(
def rule_inst_005_transformers_tokenizers(
install_cell: str, colab: dict[str, str], file: str, cell_idx: int
) -> list[Finding]:
"""Fires only when transformers is installed with `--no-deps`. Without
`--no-deps`, pip resolves the correct tokenizers transitively, so the
rule would be a false positive (this is the case for older notebooks
that pin `transformers==4.51.3` but rely on pip's transitive resolver).
The rule targets the exact pattern PR #261b / #264 fixed:
`pip install --no-deps transformers==X` next to a Colab preinstall
`tokenizers` outside transformers's window."""
"""Fires only when transformers is installed with `--no-deps` (otherwise
pip resolves tokenizers transitively and flagging would be a false
positive). Targets the PR #261b/#264 pattern: `--no-deps transformers==X`
next to a Colab `tokenizers` outside transformers's window."""
findings: list[Finding] = []
res = resolved_set(install_cell, colab)
tf = res.get("transformers")
tok = res.get("tokenizers")
if not tf or tok is None:
return findings
# Find the install line that pins transformers and check for --no-deps.
# Find the transformers pin and check for --no-deps.
transformers_line_no_deps = False
for inv in iter_pip_invocations(install_cell):
for raw in inv.packages:
@ -714,9 +680,7 @@ def rule_inst_005_transformers_tokenizers(
_RE_DOUBLE_BANG = re.compile(r"^[ \t]*!{2,}\s*pip\b", re.MULTILINE)
def rule_inst_006_double_bang(
install_cell: str, file: str, cell_idx: int
) -> list[Finding]:
def rule_inst_006_double_bang(install_cell: str, file: str, cell_idx: int) -> list[Finding]:
findings: list[Finding] = []
for m in _RE_DOUBLE_BANG.finditer(install_cell):
line_no = install_cell.count("\n", 0, m.start()) + 1
@ -739,11 +703,9 @@ def rule_inst_006_double_bang(
class _APIScanner(ast.NodeVisitor):
"""Scan user-facing code cells for known deprecated patterns. R-API-001
(`for_training`/`for_inference`) is intentionally absent: those helpers
are still part of the live unsloth surface as of 2026-05; PR #221 removed
the calls cosmetically from Vision notebooks but did not deprecate the
methods. R-API-004 (live API surface diff) catches actual removals
dynamically without us hand-coding them."""
(`for_training`/`for_inference`) is intentionally absent: those helpers are
still live as of 2026-05 (PR #221 removed them cosmetically, not as a
deprecation). R-API-004 catches actual removals dynamically."""
def __init__(self, file: str, cell_idx: int):
self.file = file
@ -751,14 +713,10 @@ class _APIScanner(ast.NodeVisitor):
self.findings: list[Finding] = []
def visit_Call(self, node: ast.Call) -> None:
# SFTConfig with suboptimal optim (R-API-003).
# NOTE: PR #221 also stripped `gradient_checkpointing` /
# `gradient_checkpointing_kwargs` from a handful of vision notebooks,
# but those kwargs are still accepted by live TRL (verified against
# trl==0.25.1 in the unsloth workspace) so removing them was
# cosmetic, not a deprecation. We do NOT flag them. R-API-004 (live
# API surface diff in the api subcommand) is the right way to catch
# actual TRL signature drift.
# SFTConfig with suboptimal optim (R-API-003).
# NOTE: PR #221 also stripped gradient_checkpointing kwargs from some
# vision notebooks, but they're still accepted by live TRL (trl==0.25.1)
# so that was cosmetic. We don't flag them; R-API-004 catches real drift.
if isinstance(node.func, ast.Name) and node.func.id == "SFTConfig":
for kw in node.keywords:
if (
@ -813,16 +771,10 @@ POLICY_CLAUSES_DEFAULT = [
]
def extract_policy_clauses(
update_script: pathlib.Path,
) -> list[tuple[str, re.Pattern[str], Any]]:
"""Best-effort: scan update_all_notebooks.py for canonical phrases used by
multiple templates. Falls back to POLICY_CLAUSES_DEFAULT.
Today we use POLICY_CLAUSES_DEFAULT directly; the regex form is
intentionally permissive so a template-side reword (e.g. comment changes)
doesn't cause false positives. New clauses become 1-line PRs to this list.
"""
def extract_policy_clauses(update_script: pathlib.Path) -> list[tuple[str, re.Pattern[str], Any]]:
"""Best-effort scan of update_all_notebooks.py for canonical phrases;
falls back to POLICY_CLAUSES_DEFAULT (which we use directly today). The
permissive regexes avoid false positives on template rewords."""
return list(POLICY_CLAUSES_DEFAULT)
@ -879,24 +831,15 @@ def cmd_drift(args: argparse.Namespace) -> int:
print(f"FAIL: {update_script} not found", file = sys.stderr)
return 2
# Stash any pre-existing dirty state, run the updater, diff, restore.
head = (
subprocess.check_output(["git", "rev-parse", "HEAD"], cwd = nbdir)
.decode()
.strip()
)
head = subprocess.check_output(["git", "rev-parse", "HEAD"], cwd = nbdir).decode().strip()
subprocess.run(
["git", "-C", str(nbdir), "stash", "--include-untracked"],
check = False,
capture_output = True,
)
# SF3: the restore MUST run even on SystemExit / KeyboardInterrupt /
# segfault-propagated exception, otherwise the user's working tree
# silently stays rolled back into the stash. A bare try/finally
# (NOT try/except/finally) preserves the original exception and
# still runs the cleanup. The pre-existing try/except around
# `subprocess.run` of the updater is folded inside the new outer
# try so its early returns still happen, but the stash pop is
# protected.
# The restore MUST run even on SystemExit/KeyboardInterrupt, else the
# working tree stays rolled back into the stash. A bare try/finally keeps
# the original exception while still running the cleanup (stash pop).
findings: list[Finding] = []
rc: int
try:
@ -941,8 +884,7 @@ def cmd_drift(args: argparse.Namespace) -> int:
)
rc = 0 if not findings else 1
finally:
# Restore the working tree. Both commands MUST run regardless of
# how the try block exited (including SystemExit/KeyboardInterrupt).
# Restore the working tree (both commands run regardless of exit path).
subprocess.run(
["git", "-C", str(nbdir), "checkout", "."],
check = False,
@ -990,9 +932,7 @@ def cmd_convert(args: argparse.Namespace) -> int:
hint = proc.stderr[-200:].strip(),
)
)
print(
f"converted {len(notebooks) - len(failed)}/{len(notebooks)} notebooks to {out}"
)
print(f"converted {len(notebooks) - len(failed)}/{len(notebooks)} notebooks to {out}")
_emit(failed)
return 0 if not failed else 1
@ -1002,11 +942,7 @@ def cmd_convert(args: argparse.Namespace) -> int:
def cmd_lint(args: argparse.Namespace) -> int:
nbdir = pathlib.Path(args.notebooks_dir).resolve()
colab_path = (
pathlib.Path(args.colab_pin).resolve()
if args.colab_pin
else COLAB_FALLBACK_FILE
)
colab_path = pathlib.Path(args.colab_pin).resolve() if args.colab_pin else COLAB_FALLBACK_FILE
colab = parse_pip_freeze(colab_path)
if not colab:
print(
@ -1031,31 +967,24 @@ def cmd_lint(args: argparse.Namespace) -> int:
continue
rel = str(path.relative_to(nbdir))
env = target_environment(rel)
# The Colab oracle is the source of truth ONLY for Colab notebooks.
# Other targets (amd / kaggle / dgx_spark) have their own runtime
# preinstall sets that aren't tracked here yet, so we apply the
# environment-agnostic rules and skip the Colab-specific ones.
# Colab oracle applies only to Colab notebooks; other targets get the
# environment-agnostic rules only (their preinstalls aren't tracked).
oracle = colab if env == "colab" else {}
cells = install_cells(nb)
# Per-cell rules: forbid-pattern checks scoped to a single line.
# Per-cell forbid-pattern checks.
for idx, cell in cells:
findings += rule_inst_001_git_plus(cell, rel, idx)
findings += rule_inst_006_double_bang(cell, rel, idx)
# Whole-notebook rules: a notebook's install steps are sometimes split
# across multiple cells (initial install + post-install bumps). Merge
# all install cells before resolving compat against Colab.
# Whole-notebook rules: install steps may span multiple cells, so merge
# before resolving compat against Colab.
merged = "\n".join(c for _, c in cells)
if env == "colab" and merged:
first_cell = cells[0][0] if cells else None
findings += rule_inst_003_peft_torchao(merged, oracle, rel, first_cell)
findings += rule_inst_004_torchcodec_torch(merged, oracle, rel, first_cell)
findings += rule_inst_005_transformers_tokenizers(
merged, oracle, rel, first_cell
)
findings += rule_inst_005_transformers_tokenizers(merged, oracle, rel, first_cell)
if not args.no_pypi:
findings += rule_inst_002_no_deps_transitive(
merged, oracle, rel, first_cell
)
findings += rule_inst_002_no_deps_transitive(merged, oracle, rel, first_cell)
findings += scan_user_cells(nb, rel)
_emit(findings)
return 0 if not any(f.severity == "error" for f in findings) else 1
@ -1178,8 +1107,7 @@ def _parse_apt_lines(text: str) -> dict[str, str]:
def _parse_os_lines(text: str) -> dict[str, str]:
"""Free-form `<tool> <version>` lines. Skip comments. The key is the
first token lower-cased; the value is the rest of the line."""
"""Free-form `<tool> <version>` lines -> {tool_lower: rest}."""
out: dict[str, str] = {}
for line in text.splitlines():
line = line.strip()
@ -1216,10 +1144,9 @@ def _diff_oracle(
def cmd_colab_diff(args: argparse.Namespace) -> int:
"""Fetch every Colab oracle file in COLAB_ORACLE_FILES, diff against
the committed snapshot, and print NEW / REMOVED / CHANGED. Advisory
by default (rc=0); --strict promotes any diff to rc=1 so the daily
cron can fail loudly when upstream rotates."""
"""Diff each Colab oracle file against its committed snapshot and print
NEW/REMOVED/CHANGED. Advisory (rc=0) by default; --strict makes any diff
rc=1 so the daily cron fails loudly on upstream rotation."""
snapshot_dir = pathlib.Path(args.snapshot_dir).resolve()
any_diff = False
for upstream_name, snapshot_name in COLAB_ORACLE_FILES.items():
@ -1232,9 +1159,7 @@ def cmd_colab_diff(args: argparse.Namespace) -> int:
print(f"::warning::colab-diff: could not fetch {url}: {e}")
continue
if not snap_path.exists():
print(
f"::warning::colab-diff: no committed snapshot at {snap_path}; skipping"
)
print(f"::warning::colab-diff: no committed snapshot at {snap_path}; skipping")
continue
snapshot_text = snap_path.read_text(encoding = "utf-8", errors = "replace")
parser = _COLAB_ORACLE_PARSERS[upstream_name]

14
scripts/run_ruff_format.py
View file

@ -1,5 +1,7 @@
#!/usr/bin/env python3
"""Run `ruff format` followed by kwarg spacing enforcement."""
"""Run a pre-pass (normalize def-signature magic commas + collapse short
multi-line asserts), then `ruff format`, then the kwarg-spacing / import /
string-merge post-pass."""
from __future__ import annotations
@ -15,12 +17,20 @@ def main(argv: list[str]) -> int:
if not files:
return 0
spacing_script = HERE / "enforce_kwargs_spacing.py"
# Pre-ruff: normalize def-signature magic commas and strip the magic comma
# from short multi-line asserts so ruff wraps/joins accordingly.
pre_cmd = [sys.executable, str(spacing_script), "--pre", *files]
pre_proc = subprocess.run(pre_cmd)
if pre_proc.returncode != 0:
return pre_proc.returncode
ruff_cmd = [sys.executable, "-m", "ruff", "format", *files]
ruff_proc = subprocess.run(ruff_cmd)
if ruff_proc.returncode != 0:
return ruff_proc.returncode
spacing_script = HERE / "enforce_kwargs_spacing.py"
spacing_cmd = [sys.executable, str(spacing_script), *files]
spacing_proc = subprocess.run(spacing_cmd)
return spacing_proc.returncode

264
scripts/scan_npm_packages.py
View file

@ -7,73 +7,32 @@
"""scan_npm_packages.py -- npm-side content scanner.
Counterpart to scripts/scan_packages.py for the pip ecosystem. Reads
npm counterpart to scripts/scan_packages.py. Reads
studio/frontend/package-lock.json, downloads each resolved tarball
DIRECTLY from registry.npmjs.org (never via `npm install` -- no
lifecycle scripts ever run), verifies the lockfile integrity hash,
unpacks each tarball into a sandboxed temp dir behind size / count /
path-escape / symlink guards, and pattern-scans the extracted file
contents for the signatures common to npm supply-chain attacks:
lifecycle scripts run), verifies the lockfile integrity hash, unpacks
each into a sandboxed temp dir behind size/count/path-escape/symlink
guards, and pattern-scans extracted contents for npm supply-chain
attack signatures: malicious lifecycle scripts, C2 / exfil hosts,
credential-stealing references, known IOC filenames, and obfuscation
shapes.
- Lifecycle (preinstall / install / postinstall / prepare) scripts
in any package.json that fetch + execute external code.
- C2 / exfiltration hosts (getsession.org, AWS IMDS endpoints,
Kubernetes ServiceAccount token paths, GitHub Actions OIDC,
HashiCorp Vault endpoints).
- Credential-stealing references (~/.npmrc, ~/.aws/credentials,
GITHUB_TOKEN / NPM_TOKEN in JS sources).
- Known IOC filenames from public advisories
(router_init.js, tanstack_runner.js, router_runtime.js).
- Obfuscation shapes (large single JS in package root with a low
whitespace ratio + Function/eval against a base64-decoded blob).
Safety stance
=============
This script ingests attacker-controlled archives. Every parse path
assumes the worst:
1. Downloads ONLY from `registry.npmjs.org`. Any tarball URL with a
different hostname is refused without fetching.
2. Tarball download is size-capped (HARD_MAX_TARBALL_BYTES default
64 MiB). HEAD-style probe via the Content-Length response header
plus a chunked read that aborts on overflow.
Safety stance (ingests attacker-controlled archives; assumes worst):
1. Downloads ONLY from registry.npmjs.org; other hosts refused.
2. Tarball download size-capped via Content-Length probe + chunked
read that aborts on overflow.
3. SHA-512 integrity verified against the lockfile entry BEFORE the
tarball is even opened. A mismatch aborts that package -- the
scanner does not "fall back" to the registry-published hash.
4. tar extraction goes through `safe_extract`:
- rejects symbolic links (`SYMTYPE`, `LNKTYPE`)
- rejects absolute paths, `..` traversal, paths outside the
extract root after resolution
- rejects character / block / FIFO devices
- per-file uncompressed size cap (HARD_MAX_FILE_BYTES, default
8 MiB) AND cumulative cap (HARD_MAX_TOTAL_BYTES, default
128 MiB) AND member-count cap (HARD_MAX_MEMBERS, default
50_000)
- tar reads happen via `tarfile.open(mode='r|gz')` streaming
so an oversized file is detected before write
5. NOTHING from the extracted tree is ever executed. Files are read
as raw bytes, decoded with `errors='replace'`, and grepped. We
never call `node`, `eval`, `compile`, `subprocess.run`,
`os.system`, or anything that would touch the tarball's
declared scripts.
6. Tempdir is created with `tempfile.mkdtemp(prefix='npm-scan-')`,
fully resolved with .resolve(), and registered with atexit to be
wiped on every termination path.
7. Stdlib only. No third-party deps -- adding one would itself be a
supply-chain liability.
tarball is opened; mismatch aborts that package (no fallback).
4. tar extraction via `safe_extract`: rejects symlinks, absolute /
`..` paths, device files; enforces per-file, cumulative, and
member-count caps; streams (`r|gz`) so oversize is caught early.
5. NOTHING extracted is executed -- files are read as bytes and
grepped only.
6. Tempdir resolved and atexit-wiped on every termination path.
7. Stdlib only (a dep would be a supply-chain liability itself).
Exit codes
==========
0 no findings of severity HIGH or higher
1 one or more HIGH/CRITICAL findings (or pre-scan structural
anomalies -- non-registry resolved URL, missing integrity)
2 internal error (lockfile missing, integrity mismatch on
download, malformed tarball, etc.)
The script is meant to be run in CI on every PR that touches
package-lock.json and on a nightly schedule.
Exit codes: 0 = no HIGH+ findings; 1 = HIGH/CRITICAL or pre-scan
structural anomaly; 2 = internal error. Run in CI per-PR and nightly.
"""
from __future__ import annotations
@ -565,12 +524,9 @@ CRED_HOST_NEEDS_CONTEXT: tuple[tuple[str, str], ...] = (
),
)
# Credentials a frontend package should NEVER need to read. Bare
# substring match is too noisy (object-treeify ships a `docker` dev
# script that mounts ~/.npmrc -- legitimate dev tooling, never run
# at install time). We instead surface these only when they appear
# inside a LIFECYCLE script (preinstall / install / postinstall /
# prepare), which is the only path that runs automatically on
# Credentials a frontend package should never read. Bare substring
# match is too noisy (legit dev tooling mounts ~/.npmrc), so we flag
# these only inside lifecycle scripts -- the only auto-run path on
# `npm ci`. See `scan_package_json` below.
CRED_PATH_SUBSTRINGS: tuple[tuple[str, str], ...] = (
("/.npmrc", "npm credentials file"),
@ -603,9 +559,8 @@ _JS_FETCH_EVAL = re.compile(
""",
)
# `process.env.GITHUB_TOKEN` / `NPM_TOKEN` / `AWS_*` access in
# top-level / install-time code is suspicious. We also catch
# `os.environ["GITHUB_TOKEN"]` for the rare Python-in-npm postinstall.
# Token env access in install-time code; also catches os.environ[...]
# for the rare Python-in-npm postinstall.
_JS_ENV_TOKEN = re.compile(
r"""(process\.env\.|os\.environ\[?['"])(?:
GITHUB_TOKEN | GH_TOKEN | NPM_TOKEN | NODE_AUTH_TOKEN
@ -616,11 +571,9 @@ _JS_ENV_TOKEN = re.compile(
re.VERBOSE,
)
# Suspicious lifecycle-script payloads. Anything in a package.json
# `scripts` field that wgets/curls an external resource and executes
# it. We do NOT block ALL curl/wget in scripts (some legit packages
# fetch test fixtures into devDependencies), but we DO block the
# fetch+exec chain.
# Lifecycle-script fetch+exec chain: curl/wget an external resource
# and run it. Bare curl/wget is allowed (legit fixture fetches); only
# the fetch+exec chain is blocked.
_LIFECYCLE_FETCH_EXEC = re.compile(
r"""(?xs)
(?:curl|wget|fetch|http\.get|axios\.get)\s+ # fetch verb
@ -634,9 +587,8 @@ _LIFECYCLE_FETCH_EXEC = re.compile(
""",
)
# Obfuscation: large JS file that is mostly one line of base64-ish
# blob with a Function() / eval() bookend. Tuned against the
# router_init.js shape (2.3 MB obfuscated single-blob).
# Obfuscation: large single-line base64-ish blob behind Function()/
# eval(). Tuned against the router_init.js shape (2.3 MB blob).
_OBFUSC_BLOB = re.compile(
r"""(?xs)
(?:Function|eval)\s*\(\s*['"`]?
@ -653,11 +605,9 @@ _OBFUSC_BLOB = re.compile(
def parse_lockfile(path: Path) -> tuple[list[PackageEntry], list[Finding]]:
"""Return (entries, structural_findings).
Structural findings here are HIGH-severity refusals that should
short-circuit the scan -- a lockfile with non-registry resolved
URLs is itself a finding (covered by scripts/lockfile_supply_chain
_audit.py in detail; we surface a summary here so this scanner is
standalone-runnable).
Structural findings are HIGH-severity refusals that short-circuit
the scan (e.g. non-registry resolved URLs). A summary is surfaced
here so this scanner is standalone-runnable.
"""
entries: list[PackageEntry] = []
findings: list[Finding] = []
@ -701,9 +651,8 @@ def parse_lockfile(path: Path) -> tuple[list[PackageEntry], list[Finding]]:
resolved = entry.get("resolved")
if not resolved:
continue
# Strict registry origin check. lockfile_supply_chain_audit
# already catches this; double-defend here so this scanner
# cannot be tricked into fetching from an attacker-chosen URL.
# Strict registry origin check so this scanner can't be tricked
# into fetching from an attacker-chosen URL.
parsed = urllib.parse.urlparse(resolved)
if parsed.scheme != "https" or parsed.hostname != ALLOWED_DOWNLOAD_HOST:
findings.append(
@ -775,16 +724,12 @@ def download_tarball(
timeout: float = HARD_HTTP_TIMEOUT_S,
max_bytes: int = HARD_MAX_TARBALL_BYTES,
) -> tuple[Path, str | None]:
"""Stream-download entry.resolved to dest. Verify SRI integrity.
"""Stream-download entry.resolved to dest and verify SRI integrity.
Returns (downloaded_path, error_or_none). On any error the
returned path may not exist. Network access is restricted to
https://{ALLOWED_DOWNLOAD_HOST}/ -- the caller passes a Request
we already validated.
Returns (downloaded_path, error_or_none); on error the path may not
exist. Network access is restricted to ALLOWED_DOWNLOAD_HOST.
"""
# Re-assert hostname; the entry was validated at parse time but a
# defence-in-depth check here means a future refactor cannot
# accidentally bypass it.
# Re-assert hostname (defence-in-depth against a future refactor).
parsed = urllib.parse.urlparse(entry.resolved)
if parsed.scheme != "https" or parsed.hostname != ALLOWED_DOWNLOAD_HOST:
return dest, (f"refused download from non-allowlisted URL {entry.resolved!r}")
@ -823,8 +768,7 @@ def download_tarball(
written += len(chunk)
if written > max_bytes:
return dest, (
f"download exceeded cap {max_bytes} bytes "
f"after {written} bytes"
f"download exceeded cap {max_bytes} bytes " f"after {written} bytes"
)
h.update(chunk)
out.write(chunk)
@ -874,8 +818,7 @@ def safe_extract(
total = 0
count = 0
try:
# Open in streaming mode so we never seek backwards in the
# input. `r|gz` rejects malformed gzip frames immediately.
# Streaming mode (no backward seeks); `r|gz` rejects bad gzip.
with tarfile.open(tarball_path, mode = "r|gz") as tf:
for member in tf:
count += 1
@ -890,9 +833,8 @@ def safe_extract(
return f"refused link member {name!r} (sym/lnk)"
if member.isdev() or member.isfifo():
return f"refused special member {name!r}"
# Cumulative cap is checked against DECLARED size up
# front to short-circuit obvious bombs without reading
# the body.
# Check declared size up front to short-circuit bombs
# without reading the body.
declared = max(member.size, 0)
if declared > HARD_MAX_BINARY_FILE_BYTES:
return (
@ -904,9 +846,8 @@ def safe_extract(
f"cumulative bytes {total + declared} > cap "
f"{max_total_bytes} at {name!r}"
)
# Strip leading "package/" -- the npm convention. We do
# NOT trust npm to be right, so we explicitly resolve
# the destination and refuse anything that escapes.
# Resolve destination and refuse anything escaping root
# (don't trust the npm "package/" convention).
dest = extract_root / name
if not _is_within(extract_root, dest):
return f"refused escape: {name!r} resolved outside root"
@ -920,17 +861,11 @@ def safe_extract(
src = tf.extractfile(member)
if src is None:
continue
# Sniff first 16 bytes to classify text vs binary.
# Text-cap members get the tight 16 MiB limit; binary
# members (executables, .node, .wasm, native libs)
# get the generous binary cap. We bound BOTH cases.
# Sniff first 16 bytes to classify text vs binary;
# each gets its own cap (both are bounded).
header = src.read(16)
is_binary = _looks_binary(name, header)
file_cap = (
HARD_MAX_BINARY_FILE_BYTES
if is_binary
else HARD_MAX_TEXT_FILE_BYTES
)
file_cap = HARD_MAX_BINARY_FILE_BYTES if is_binary else HARD_MAX_TEXT_FILE_BYTES
if declared > file_cap:
return (
f"member {name!r} declared size {declared} > "
@ -946,8 +881,7 @@ def safe_extract(
f"({'binary' if is_binary else 'text'})"
)
total += len(data)
# Write with restrictive mode (rw-r--r--) so even if
# someone runs the extract dir nothing is executable.
# Restrictive mode (rw-r--r--): nothing executable.
with open(dest, "wb") as out:
out.write(data)
os.chmod(dest, 0o644)
@ -963,7 +897,11 @@ def safe_extract(
# ─────────────────────────────────────────────────────────────────────
def _evidence(text: str, pat: re.Pattern, max_chars: int = 200) -> str:
def _evidence(
text: str,
pat: re.Pattern,
max_chars: int = 200,
) -> str:
m = pat.search(text)
if not m:
return ""
@ -978,11 +916,7 @@ def _evidence(text: str, pat: re.Pattern, max_chars: int = 200) -> str:
LIFECYCLE_HOOKS = ("preinstall", "install", "postinstall", "prepare")
def scan_package_json(
pkg: PackageEntry,
rel: str,
text: str,
) -> list[Finding]:
def scan_package_json(pkg: PackageEntry, rel: str, text: str) -> list[Finding]:
findings: list[Finding] = []
try:
meta = json.loads(text)
@ -1013,10 +947,8 @@ def scan_package_json(
),
)
)
# Credential file paths inside a lifecycle script are
# exfiltration prep -- npm runs these scripts automatically
# on `npm ci`. Manual `scripts.*` entries (like a `docker`
# dev script) are out of scope: npm does not run them.
# Cred file paths in a lifecycle script are exfil prep (npm
# auto-runs these on `npm ci`); manual scripts are out of scope.
for path_substr, why in CRED_PATH_SUBSTRINGS:
if path_substr in body:
findings.append(
@ -1056,9 +988,7 @@ def scan_package_json(
if isinstance(opt, dict):
for k, v in opt.items():
if isinstance(v, str) and (
v.startswith("github:")
or v.startswith("git+")
or v.startswith("git://")
v.startswith("github:") or v.startswith("git+") or v.startswith("git://")
):
findings.append(
Finding(
@ -1078,20 +1008,12 @@ def scan_package_json(
def _host_in_outbound_context(text: str, host: str) -> bool:
"""True if `host` appears in a way consistent with an outbound call.
"""True if `host` appears consistent with an outbound call.
A bare `"169.254.169.254"` array literal (defensive blocklist) is
safe; a `fetch("http://169.254.169.254/...")` is not. The signal
is co-occurrence with either an HTTP URL scheme or a fetch verb
within a short window.
A defensive blocklist looks like:
const CLOUD_METADATA_IPS = ["169.254.169.254", "169.254.170.2"];
An exfil call looks like:
fetch("http://169.254.169.254/latest/meta-data/...")
http.request({ host: "169.254.169.254", path: "/..." })
A bare array literal (defensive blocklist) is safe; co-occurrence
with an HTTP URL scheme or a fetch verb in a short window is not.
"""
# Esc for use in a regex (IPs contain dots).
# Escape for regex (IPs contain dots).
host_re = re.escape(host)
# 1. URL form: http://host or https://host or //host/ or //host"
url_form = re.compile(
@ -1117,11 +1039,7 @@ def _host_in_outbound_context(text: str, host: str) -> bool:
return False
def scan_text_blob(
pkg: PackageEntry,
rel: str,
text: str,
) -> list[Finding]:
def scan_text_blob(pkg: PackageEntry, rel: str, text: str) -> list[Finding]:
findings: list[Finding] = []
# IOC substrings (literal, case-sensitive).
@ -1138,8 +1056,7 @@ def scan_text_blob(
)
)
# Credential surfaces. Tier 1: hosts with no legitimate use,
# bare substring is enough.
# Cred surfaces, tier 1: hosts with no legit use; bare substring.
for needle, why in CRED_HOST_ALWAYS_BAD:
if needle in text:
findings.append(
@ -1156,8 +1073,8 @@ def scan_text_blob(
)
)
# Credential surfaces. Tier 2: hosts that do appear in defensive
# code; require co-occurrence with a fetch verb or URL prefix.
# Cred surfaces, tier 2: hosts that appear in defensive code too;
# require co-occurrence with a fetch verb or URL prefix.
for needle, why in CRED_HOST_NEEDS_CONTEXT:
if needle in text and _host_in_outbound_context(text, needle):
findings.append(
@ -1175,11 +1092,8 @@ def scan_text_blob(
)
)
# Credential PATHS are deliberately not scanned here; they have
# too high a false-positive rate at file scope (defensive code,
# docker mounts, AWS SDK docs strings). `scan_package_json`
# catches the malicious case -- credential paths inside a
# lifecycle script run automatically on `npm ci`.
# Credential PATHS aren't scanned here (too many FPs at file
# scope); scan_package_json catches them inside lifecycle scripts.
# JS-specific regex.
if _JS_FETCH_EVAL.search(text):
@ -1190,10 +1104,7 @@ def scan_text_blob(
filename = rel,
pattern = "js-fetch-eval",
evidence = _evidence(text, _JS_FETCH_EVAL),
detail = (
"Function/eval against base64-decoded payload "
"(obfuscated dropper shape)"
),
detail = ("Function/eval against base64-decoded payload (obfuscated dropper shape)"),
)
)
if _JS_ENV_TOKEN.search(text):
@ -1225,9 +1136,8 @@ def scan_text_blob(
return findings
# Filename suffix decides which scanners run. We deliberately treat
# *.cjs/*.mjs/*.ts the same as *.js -- attackers use whichever
# extension the consumer's bundler / loader resolves.
# Filename suffix decides which scanners run; .cjs/.mjs/.ts are
# treated like .js (attackers use whichever the loader resolves).
_TEXT_SUFFIXES = (
".js",
".mjs",
@ -1247,10 +1157,7 @@ _TEXT_SUFFIXES = (
)
def scan_extracted_tree(
pkg: PackageEntry,
root: Path,
) -> list[Finding]:
def scan_extracted_tree(pkg: PackageEntry, root: Path) -> list[Finding]:
findings: list[Finding] = []
for path in sorted(root.rglob("*")):
if not path.is_file():
@ -1258,12 +1165,9 @@ def scan_extracted_tree(
rel = path.relative_to(root).as_posix()
lower = rel.lower()
if not lower.endswith(_TEXT_SUFFIXES):
# Skip native binaries entirely -- regex over compiled
# machine code is just noise (false positives in WASM
# opcodes, .node BSS segments, image pixel data). Use
# content-magic detection so extensionless executables
# (eg `package/biome`) and versioned shared libraries
# are also skipped.
# Skip native binaries (regex over machine code is noise);
# content-magic detection also skips extensionless
# executables and versioned shared libraries.
try:
if path.stat().st_size > HARD_MAX_TEXT_FILE_BYTES:
continue
@ -1304,16 +1208,13 @@ def scan_extracted_tree(
# ─────────────────────────────────────────────────────────────────────
def scan_one(
pkg: PackageEntry,
workspace: Path,
) -> tuple[list[Finding], str | None]:
"""Download + extract + scan a single package. Cleans up its dir.
def scan_one(pkg: PackageEntry, workspace: Path) -> tuple[list[Finding], str | None]:
"""Download + extract + scan a single package; cleans up its dir.
Returns (findings, error). `error` is non-None only on hard
failures (download error, integrity mismatch, malformed tarball);
on a clean run with findings the error is None and the caller
decides exit code based on severity.
failures (download, integrity mismatch, malformed tarball); on a
clean run with findings, error is None and the caller decides the
exit code from severity.
"""
pkg_dir = workspace / f"{pkg.name.replace('/', '_')}-{pkg.version}"
pkg_dir.mkdir(parents = True, exist_ok = True)
@ -1444,8 +1345,7 @@ def main(argv: list[str] | None = None) -> int:
if hard_errors or blocking:
if blocking:
print(
f"\n[scan-npm] FAIL: {len(blocking)} finding(s) "
f"at or above {threshold}",
f"\n[scan-npm] FAIL: {len(blocking)} finding(s) " f"at or above {threshold}",
file = sys.stderr,
)
return 1

325
scripts/scan_packages.py
View file

@ -56,27 +56,22 @@ from dataclasses import dataclass, field
from pathlib import Path
# ---------------------------------------------------------------------------
# Severity
# ---------------------------------------------------------------------------
CRITICAL = "CRITICAL"
HIGH = "HIGH"
MEDIUM = "MEDIUM"
SEVERITY_ORDER = {CRITICAL: 0, HIGH: 1, MEDIUM: 2}
# Hard pin-blocks for publicly confirmed malicious PyPI versions.
# Source: Socket.dev 2026-05-12 disclosure (Mini Shai-Hulud May-12 wave) and
# earlier Semgrep / Endor reports for the `lightning` entries.
# Hard pin-blocks for confirmed malicious PyPI versions (Socket.dev 2026-05-12
# Mini Shai-Hulud wave; earlier Semgrep/Endor reports for `lightning`).
BLOCKED_PYPI_VERSIONS: dict[str, set[str]] = {
"guardrails-ai": {"0.10.1"},
"mistralai": {"2.4.6"},
"lightning": {"2.6.2", "2.6.3"},
}
# ---------------------------------------------------------------------------
# Pattern definitions
# ---------------------------------------------------------------------------
# Subprocess / OS exec patterns
RE_SUBPROCESS = re.compile(
@ -86,8 +81,7 @@ RE_SUBPROCESS = re.compile(
# Encoding / obfuscation
RE_BASE64 = re.compile(
r"\bbase64\s*\.\s*(b64decode|decodebytes|b32decode|b16decode)\b"
r"|\bcodecs\s*\.\s*decode\b",
r"\bbase64\s*\.\s*(b64decode|decodebytes|b32decode|b16decode)\b|\bcodecs\s*\.\s*decode\b",
)
# exec / eval
@ -299,9 +293,7 @@ RE_CRYPTO_THEFT = re.compile(
RE_PTH_IMPORT = re.compile(r"^\s*import\s+", re.MULTILINE)
# openssl CLI invocations via subprocess (encrypted exfiltration)
RE_OPENSSL_CLI = re.compile(
r"\bopenssl\s+(enc|rand|rsautl|pkeyutl|genrsa|dgst|s_client)\b"
)
RE_OPENSSL_CLI = re.compile(r"\bopenssl\s+(enc|rand|rsautl|pkeyutl|genrsa|dgst|s_client)\b")
# Write to /tmp then execute (staged dropper)
RE_TEMP_EXEC = re.compile(
@ -315,10 +307,8 @@ RE_C2_POLLING = re.compile(
re.DOTALL,
)
# Developer-tool persistence hooks. The PyTorch Lightning 2.6.x compromise
# planted SessionStart hooks into Claude Code, VS Code tasks, and Cursor
# settings so the payload re-attached on every editor open. Catches any
# package writing into a known dev-tool config that supports auto-run.
# Developer-tool persistence hooks. Lightning 2.6.x planted SessionStart hooks
# into Claude Code / VS Code / Cursor so the payload re-attached on editor open.
RE_DEV_TOOL_HIJACK = re.compile(
r"\.claude/settings\.json"
r"|\.cursor/.*hooks"
@ -329,9 +319,8 @@ RE_DEV_TOOL_HIJACK = re.compile(
r"|\bautomator\b.*\.workflow\b",
)
# Hard-coded credential / API-token regexes embedded in source. Packages
# that ship regexes for OTHER people's secrets are nearly always
# stealers (litellm 1.82.7, elementary-data 0.23.3, Shai-Hulud).
# Hard-coded credential / API-token regexes embedded in source. Packages that
# ship regexes for OTHER people's secrets are nearly always stealers.
RE_TOKEN_REGEX = re.compile(
r"\bgh[psoru]_[A-Za-z0-9_]{20,}" # GitHub PAT/OAuth/etc.
r"|\bgithub_pat_[A-Za-z0-9_]{20,}"
@ -345,20 +334,16 @@ RE_TOKEN_REGEX = re.compile(
r"|\bglpat-[0-9A-Za-z_-]{20,}", # GitLab PAT
)
# Mini Shai-Hulud May-12 2026 wave indicators. The dropper artifact name
# `transformers.pyz` is high-confidence (no legit PyPI package ships a `.pyz`
# named after `transformers`); the host + slogans are CRITICAL.
# Mini Shai-Hulud May-12 2026 wave indicators. `transformers.pyz` dropper name
# is high-confidence; the host + slogans are CRITICAL.
RE_MAY12_IOC = re.compile(
r"(git-tanstack\.com|/tmp/transformers\.pyz|transformers\.pyz"
r"|With Love TeamPCP|We've been online over 2 hours)",
re.IGNORECASE,
)
# JavaScript-side obfuscation. The npm chalk/debug compromise and the
# Lightning router_runtime.js use the same minifier-style hex-var name
# pattern; a bundle full of `_0x1f2e3d` identifiers is a near-universal
# tell for a malicious npm payload (and very rare in legit minified code
# that ships in PyPI wheels).
# JavaScript-side obfuscation. A bundle full of `_0x1f2e3d` hex-var identifiers
# is a near-universal tell for a malicious npm payload, rare in legit wheels.
RE_JS_OBFUSCATION = re.compile(
r"_0x[a-f0-9]{4,6}\s*=\s*function"
r"|var\s+_0x[a-f0-9]{4,6}\b"
@ -366,9 +351,8 @@ RE_JS_OBFUSCATION = re.compile(
r"|String\.fromCharCode\s*\(\s*\d+\s*(?:,\s*\d+\s*){10,}\)",
)
# Web3 / wallet-hijack pattern. The Qix npm phish overrode fetch /
# XMLHttpRequest and attached a `window.ethereum` listener that
# Levenshtein-swapped recipient addresses on the way to the network.
# Web3 / wallet-hijack pattern. The Qix npm phish overrode fetch/XMLHttpRequest
# and swapped recipient addresses via a `window.ethereum` listener.
RE_WEB3_HIJACK = re.compile(
r"\bwindow\.ethereum\b"
r"|\bweb3\.eth\.\w+\s*\("
@ -377,11 +361,9 @@ RE_WEB3_HIJACK = re.compile(
r"|TronWeb|solanaWeb3",
)
# Self-propagating supply-chain worms (Shai-Hulud, ForceMemo) plant
# their own GitHub workflow in every repo they can reach, and lean on
# trufflehog/gitleaks for credential discovery. The combo of any of
# these strings inside a *package payload* is overwhelming evidence of
# repo-takeover intent.
# Self-propagating worms (Shai-Hulud, ForceMemo) plant their own GitHub workflow
# in every repo they reach and use trufflehog/gitleaks for credential discovery.
# Any of these strings in a package payload is strong repo-takeover evidence.
RE_WORKFLOW_INJECT = re.compile(
r"\.github/workflows/[^\"\']*\.ya?ml"
r"|\btrufflehog\b|\bgitleaks\b"
@ -391,9 +373,8 @@ RE_WORKFLOW_INJECT = re.compile(
re.IGNORECASE | re.DOTALL,
)
# Shell-side patterns specific to install.sh / postinstall scripts that
# pipe remote code into a shell. `curl ... | sh` and friends are the
# canonical npm postinstall dropper.
# install.sh / postinstall scripts piping remote code into a shell.
# `curl ... | sh` is the canonical npm postinstall dropper.
RE_SHELL_DROPPER = re.compile(
r"\bcurl\b[^\n|]*\|\s*(?:sh|bash|zsh)\b"
r"|\bwget\b[^\n|]*-O-\s*\|\s*(?:sh|bash|zsh)\b"
@ -403,9 +384,6 @@ RE_SHELL_DROPPER = re.compile(
)
# ---------------------------------------------------------------------------
# Finding dataclass
# ---------------------------------------------------------------------------
@dataclass
class Finding:
severity: str
@ -415,9 +393,7 @@ class Finding:
evidence: str = ""
# ---------------------------------------------------------------------------
# Checkers
# ---------------------------------------------------------------------------
def check_pth_file(content: str, filename: str, package: str) -> list[Finding]:
@ -428,7 +404,7 @@ def check_pth_file(content: str, filename: str, package: str) -> list[Finding]:
"""
findings = []
# Only care about .pth files that have import lines (executable)
# Only .pth files with import lines are executable
import_lines = [line for line in content.splitlines() if RE_PTH_IMPORT.match(line)]
if not import_lines:
return findings # Pure path entries, inert
@ -473,7 +449,7 @@ def check_pth_file(content: str, filename: str, package: str) -> list[Finding]:
)
)
# Large base64 blob (special handling for blob size)
# Large base64 blob
if RE_LARGE_BLOB.search(content):
blob = RE_LARGE_BLOB.search(content).group()
findings.append(
@ -486,7 +462,7 @@ def check_pth_file(content: str, filename: str, package: str) -> list[Finding]:
)
)
# Catch-all: any import line at all in .pth (if nothing else triggered)
# Catch-all: any import line in .pth if nothing else triggered
if not findings and import_lines:
evidence = "\n".join(import_lines[:5])
if len(import_lines) > 5:
@ -524,7 +500,7 @@ def check_py_file(content: str, filename: str, package: str) -> list[Finding]:
is_setup = basename in ("setup.py", "setup.cfg")
is_init = basename == "__init__.py"
# Pre-compute all pattern matches
# Pre-compute pattern matches
has_network = bool(RE_NETWORK.search(content))
has_subprocess = bool(RE_SUBPROCESS.search(content))
has_base64 = bool(RE_BASE64.search(content))
@ -549,9 +525,7 @@ def check_py_file(content: str, filename: str, package: str) -> list[Finding]:
has_c2_polling = bool(RE_C2_POLLING.search(content))
has_may12_ioc = bool(RE_MAY12_IOC.search(content))
# ---------------------------------------------------------------
# CRITICAL: combination patterns that strongly indicate malice
# ---------------------------------------------------------------
# base64 decode + subprocess execution (staged payload)
if has_base64 and has_subprocess:
@ -743,9 +717,7 @@ def check_py_file(content: str, filename: str, package: str) -> list[Finding]:
)
)
# ---------------------------------------------------------------
# HIGH: single strong signals or weaker combinations
# ---------------------------------------------------------------
# Obfuscated payload: base64 + exec/eval + large blob
if has_base64 and has_exec_eval and has_blob:
@ -882,9 +854,7 @@ def check_py_file(content: str, filename: str, package: str) -> list[Finding]:
)
)
# ---------------------------------------------------------------
# MEDIUM: standalone signals (informational, may be legitimate)
# ---------------------------------------------------------------
# base64 + exec/eval without blob
if has_base64 and has_exec_eval and not has_blob:
@ -962,7 +932,11 @@ def check_py_file(content: str, filename: str, package: str) -> list[Finding]:
return findings
def _extract_evidence(content: str, pattern: re.Pattern, max_matches: int = 3) -> str:
def _extract_evidence(
content: str,
pattern: re.Pattern,
max_matches: int = 3,
) -> str:
"""Pull matching lines as evidence snippets."""
lines = content.splitlines()
matches = []
@ -977,23 +951,18 @@ def _extract_evidence(content: str, pattern: re.Pattern, max_matches: int = 3) -
return " | ".join(matches) if matches else ""
# ---------------------------------------------------------------------------
# Non-Python checkers
# ---------------------------------------------------------------------------
# Several recent PyPI compromises (PyTorch Lightning 2.6.x, ForceMemo)
# carried the active payload in a bundled .js / .sh / workflow yaml so
# the Python imports looked clean on first glance. These checkers scan
# those file types when they appear inside a Python wheel/sdist.
# Recent PyPI compromises (Lightning 2.6.x, ForceMemo) carried the payload in a
# bundled .js / .sh / workflow yaml so the Python imports looked clean. These
# checkers scan those file types when they appear inside a wheel/sdist.
def check_js_file(content: str, filename: str, package: str) -> list[Finding]:
"""Run JS-side checks. Triggered by .js / .mjs / .cjs / .ts."""
findings = []
# A JS file *inside a Python wheel* that's larger than 100 KB is
# itself anomalous (legit Python packages don't ship hand-written
# JS bundles). Combined with ANY of the other JS heuristics it is
# CRITICAL; standalone it is HIGH.
# A >100 KB JS file inside a Python wheel is anomalous: CRITICAL combined
# with any other JS heuristic, HIGH standalone.
is_large = len(content) > 100 * 1024
has_obf = bool(RE_JS_OBFUSCATION.search(content))
has_web3 = bool(RE_WEB3_HIJACK.search(content))
@ -1118,10 +1087,8 @@ def check_shell_file(content: str, filename: str, package: str) -> list[Finding]
def check_workflow_file(content: str, filename: str, package: str) -> list[Finding]:
"""Run GitHub-Actions workflow checks. Triggered by .github/workflows/*.yml."""
findings = []
# A GitHub workflow file inside a *PyPI package* is itself
# suspicious (Shai-Hulud's whole MO is to plant `shai-hulud.yml`
# in every repo it can write to). Anything matching the workflow
# injection signature gets flagged CRITICAL.
# A workflow file inside a PyPI package is suspicious (Shai-Hulud plants
# `shai-hulud.yml` everywhere); injection-signature matches are CRITICAL.
if RE_WORKFLOW_INJECT.search(content):
findings.append(
Finding(
@ -1165,15 +1132,11 @@ def check_workflow_file(content: str, filename: str, package: str) -> list[Findi
return findings
# ---------------------------------------------------------------------------
# Archive handling
# ---------------------------------------------------------------------------
# Tarbomb caps, mirrored from scripts/scan_npm_packages.py::safe_extract.
# Refuses zip-of-death / tar-of-death archives so a hostile sdist or
# wheel cannot exhaust memory or fill the temp dir before content
# scanning even starts. Keep these constants in sync with the npm side;
# we duplicate rather than import to keep `scan_packages.py` standalone.
# Refuses zip/tar-of-death so a hostile archive cannot exhaust memory before
# scanning. Keep in sync with the npm side; duplicated to stay standalone.
HARD_MAX_FILE_BYTES = 64 * 1024 * 1024 # 64 MiB per member
HARD_MAX_TOTAL_BYTES = 512 * 1024 * 1024 # 512 MiB cumulative
HARD_MAX_MEMBERS = 50_000 # entries per archive
@ -1182,11 +1145,8 @@ HARD_MAX_MEMBERS = 50_000 # entries per archive
def _refuse_unsafe_member_name(name: str) -> str | None:
"""Return a refusal reason for a member name, or None if safe.
Mirrors `scan_npm_packages.py::safe_extract` semantics: no absolute
paths, no `..` traversal segments. The caller is responsible for
checking the resolved path lands inside the extract root, but for
iter_archive_files we never write to disk so the name-shape check
plus the in-memory size cap is sufficient.
Mirrors `safe_extract`: no absolute paths, no `..` traversal. We never write
to disk, so the name-shape check plus the in-memory size cap is sufficient.
"""
if name.startswith("/") or ".." in Path(name).parts:
return f"unsafe member name {name!r}"
@ -1196,9 +1156,8 @@ def _refuse_unsafe_member_name(name: str) -> str | None:
def iter_archive_files(archive_path: str):
"""Yield (filename, text_content) for every file in a wheel/sdist.
Streams members with size + count caps applied at the member level
so a tarbomb / zipbomb cannot blow up the scanner's memory budget.
On cap breach we emit a `[WARN]` log and short-circuit the archive.
Streams members with per-member size + count caps so a tarbomb/zipbomb can't
blow the memory budget. On cap breach, emits a `[WARN]` and short-circuits.
"""
path = Path(archive_path)
@ -1224,7 +1183,7 @@ def iter_archive_files(archive_path: str):
file = sys.stderr,
)
continue
# Declared (uncompressed) size cap.
# Declared (uncompressed) size cap
if info.file_size > HARD_MAX_FILE_BYTES:
print(
f" [WARN] {path.name}: skipped {info.filename!r} "
@ -1261,20 +1220,17 @@ def iter_archive_files(archive_path: str):
file = sys.stderr,
)
return
# Refuse symlinks / hardlinks / devices outright -- the
# scanner never writes them anyway, but tar parsers
# have historically dereferenced them on extract.
# Refuse symlinks/hardlinks/devices: tar parsers have
# historically dereferenced them on extract.
if member.issym() or member.islnk():
print(
f" [WARN] {path.name}: refused link member "
f"{member.name!r}",
f" [WARN] {path.name}: refused link member " f"{member.name!r}",
file = sys.stderr,
)
continue
if member.isdev() or member.isfifo():
print(
f" [WARN] {path.name}: refused special member "
f"{member.name!r}",
f" [WARN] {path.name}: refused special member " f"{member.name!r}",
file = sys.stderr,
)
continue
@ -1306,8 +1262,7 @@ def iter_archive_files(archive_path: str):
f = tf.extractfile(member)
if f is None:
continue
# Bound the read so a tar header that lies about
# size cannot OOM us.
# Bound the read: a tar header may lie about size
data = f.read(HARD_MAX_FILE_BYTES + 1)
if len(data) > HARD_MAX_FILE_BYTES:
print(
@ -1328,11 +1283,9 @@ def iter_archive_files(archive_path: str):
def scan_archive(archive_path: str, package: str) -> list[Finding]:
"""Scan all files in an archive for malicious patterns.
A corrupted archive container (truncated wheel, bad gzip header,
etc.) used to be silently skipped by an ``except Exception: continue``
inside ``iter_archive_files``. Per the silent-failure hardening
(SF1) it now emits a CRITICAL ``archive_corrupted`` finding so the
main loop counts and surfaces it rather than reporting "0 findings".
A corrupted archive container (truncated wheel, bad gzip header, etc.) emits
a CRITICAL ``archive_corrupted`` finding rather than being silently skipped
and reported as "0 findings" (silent-failure hardening SF1).
"""
findings: list[Finding] = []
try:
@ -1343,22 +1296,17 @@ def scan_archive(archive_path: str, package: str) -> list[Finding]:
elif lower.endswith(".py"):
findings.extend(check_py_file(content, filename, package))
elif lower.endswith((".js", ".mjs", ".cjs", ".ts")):
# Lightning 2.6.x hid its real payload in a 14.8 MB
# router_runtime.js inside a Python wheel. Without this
# branch we'd have only seen the small Python loader.
# Lightning 2.6.x hid its payload in a 14.8 MB router_runtime.js;
# without this branch we'd only see the small Python loader.
findings.extend(check_js_file(content, filename, package))
elif lower.endswith((".sh", ".bash")):
findings.extend(check_shell_file(content, filename, package))
elif "/.github/workflows/" in lower and lower.endswith((".yml", ".yaml")):
# Shai-Hulud / ForceMemo plant their own GHA workflow.
# A workflow file inside a *PyPI package* is on its own
# already a yellow flag; pattern-match the worm signatures.
# Shai-Hulud/ForceMemo plant their own GHA workflow
findings.extend(check_workflow_file(content, filename, package))
except (zipfile.BadZipFile, tarfile.TarError, EOFError, OSError) as exc:
# The archive cannot be opened or is structurally broken. A
# benign wheel/sdist always opens; a malformed one is either a
# transport corruption (treat as scan failure) or a deliberate
# attempt to bypass scanners that swallow archive errors.
# Archive cannot be opened / is structurally broken: either transport
# corruption or a deliberate attempt to bypass error-swallowing scanners.
findings.append(
Finding(
CRITICAL,
@ -1371,24 +1319,18 @@ def scan_archive(archive_path: str, package: str) -> list[Finding]:
return findings
# ---------------------------------------------------------------------------
# Download packages
# ---------------------------------------------------------------------------
_RE_PYPI_SPEC_VERSION = re.compile(r"==\s*([A-Za-z0-9_.\-+!]+)")
def _check_blocked_pypi_versions(
specs: list[str],
) -> tuple[list[str], list[Finding]]:
def _check_blocked_pypi_versions(specs: list[str]) -> tuple[list[str], list[Finding]]:
"""Filter ``specs`` against ``BLOCKED_PYPI_VERSIONS``.
Returns ``(safe_specs, findings)``. Each blocked spec emits a CRITICAL
``Finding`` and is removed from the returned spec list so the caller
never fetches the malicious tarball. Specs without an ``==X.Y.Z`` pin
pass through unchanged -- pip will resolve them at download time and
the existing scanners will catch the payload via the IOC regexes.
``Finding`` and is dropped so the malicious tarball is never fetched. Specs
without an ``==X.Y.Z`` pin pass through; the IOC regexes catch them later.
"""
safe: list[str] = []
findings: list[Finding] = []
@ -1410,7 +1352,7 @@ def _check_blocked_pypi_versions(
f"{name}=={version} is on the BLOCKED_PYPI_VERSIONS list",
)
)
# Drop the spec; do not download.
# Drop the spec; do not download
continue
safe.append(spec)
return safe, findings
@ -1419,14 +1361,11 @@ def _check_blocked_pypi_versions(
def _pip_download_env() -> dict[str, str]:
"""Return a scrubbed environment for invoking `pip download`.
Hostile shells / CI configs can override the index with PIP_INDEX_URL,
PIP_EXTRA_INDEX_URL, or a user `pip.conf`. We strip every PIP_*
override and route the resolver explicitly at PyPI. PIP_CONFIG_FILE
is forced to /dev/null so a stray ~/.pip/pip.conf with an
extra-index-url cannot bypass the pin.
Strips every PIP_* override and forces the resolver at PyPI; PIP_CONFIG_FILE
is /dev/null so a stray pip.conf extra-index-url cannot bypass the pin.
"""
env = {**os.environ}
# Drop any user override.
# Drop any user override
for key in [k for k in env if k.startswith("PIP_")]:
env.pop(key, None)
env["PIP_INDEX_URL"] = "https://pypi.org/simple"
@ -1436,10 +1375,8 @@ def _pip_download_env() -> dict[str, str]:
return env
# Pip resolver flags shared by both download branches. Pinning the
# index URL on the CLI is belt + braces with the env scrub above.
# `--no-build-isolation` is deliberately NOT set; we never invoke
# setup.py at all because of `--only-binary :all:`.
# Pip resolver flags shared by both download branches. CLI index-URL pin is
# belt + braces with the env scrub; `--only-binary :all:` avoids running setup.py.
_PIP_DOWNLOAD_PIN_FLAGS = [
"--index-url",
"https://pypi.org/simple",
@ -1448,9 +1385,8 @@ _PIP_DOWNLOAD_PIN_FLAGS = [
]
# Strip any character that could escape `dest` via `os.path.join`. This
# is the last line of defence before `pkg_dir = os.path.join(dest, ...)`
# so a spec like `../../etc/foo==1.0` cannot land outside the temp tree.
# Strip characters that could escape `dest` via `os.path.join`, so a spec like
# `../../etc/foo==1.0` cannot land outside the temp tree.
_RE_PKG_NAME_SANITIZE = re.compile(r"[^A-Za-z0-9._-]")
@ -1462,27 +1398,21 @@ def download_packages(
) -> tuple[list[tuple[str, str]], list[str]]:
"""Download packages to dest using pip download. NEVER installs.
Returns ``(results, download_errors)`` where ``results`` is a list of
``(spec_or_name, filepath)`` for every downloaded archive and
``download_errors`` is a list of one-line transport-failure summaries.
A non-empty ``download_errors`` MUST cause the caller to exit non-zero
even if no findings were produced; a silent ``0 findings, scan
incomplete`` is the bug class this return-shape was widened to fix.
Returns ``(results, download_errors)``: ``results`` is ``(spec_or_name,
filepath)`` per archive; ``download_errors`` is one-line transport-failure
summaries. A non-empty ``download_errors`` MUST make the caller exit
non-zero so a partial scan can't masquerade as "0 findings, all clean".
When with_deps=True, downloads the full transitive dependency tree
in a single pip invocation (all archives land in one flat dir).
When with_deps=False (default), downloads each spec individually
with --no-deps.
with_deps=True downloads the full transitive tree in one pip call (flat dir);
with_deps=False (default) downloads each spec individually with --no-deps.
"""
results: list[tuple[str, str]] = []
download_errors: list[str] = []
env = _pip_download_env()
if with_deps:
# Single pip download call for all specs + their transitive deps.
# `--only-binary :all:` refuses sdists so we never execute a
# setup.py just to learn dependency metadata; combined with the
# scrubbed env, pip is wired hard at pypi.org.
# Single pip download for all specs + transitive deps. `--only-binary
# :all:` refuses sdists so we never execute setup.py for metadata.
os.makedirs(dest, exist_ok = True)
cmd = [
sys.executable,
@ -1498,13 +1428,11 @@ def download_packages(
cmd,
capture_output = True,
text = True,
timeout = 600, # transitive resolution can be slow
timeout = 600, # transitive resolution is slow
env = env,
)
if proc.returncode != 0:
msg = (
f"pip download (with deps) failed: " f"{proc.stderr.strip()[:500]}"
)
msg = f"pip download (with deps) failed: " f"{proc.stderr.strip()[:500]}"
print(f" [ERROR] {msg}", file = sys.stderr)
download_errors.append(msg)
except subprocess.TimeoutExpired:
@ -1516,14 +1444,12 @@ def download_packages(
for fname in sorted(os.listdir(dest)):
fpath = os.path.join(dest, fname)
if os.path.isfile(fpath):
# Derive package name from filename
pkg_name = fname.split("-")[0].replace("_", "-").lower()
results.append((pkg_name, fpath))
else:
for spec in specs:
raw_name = _extract_pkg_name(spec)
# Sanitize before joining into `dest` so a hostile spec
# cannot path-traverse out of the destination directory.
# Sanitize before joining into `dest` to prevent path traversal
safe_name = _RE_PKG_NAME_SANITIZE.sub("_", raw_name) or "_pkg"
pkg_dir = os.path.join(dest, safe_name)
os.makedirs(pkg_dir, exist_ok = True)
@ -1547,10 +1473,7 @@ def download_packages(
env = env,
)
if proc.returncode != 0:
msg = (
f"pip download failed for {spec}: "
f"{proc.stderr.strip()[:500]}"
)
msg = f"pip download failed for {spec}: " f"{proc.stderr.strip()[:500]}"
print(f" [ERROR] {msg}", file = sys.stderr)
download_errors.append(msg)
continue
@ -1560,7 +1483,6 @@ def download_packages(
download_errors.append(msg)
continue
# Find downloaded file(s)
for fname in os.listdir(pkg_dir):
fpath = os.path.join(pkg_dir, fname)
if os.path.isfile(fpath):
@ -1568,9 +1490,7 @@ def download_packages(
return results, download_errors
# ---------------------------------------------------------------------------
# Parse requirements files
# ---------------------------------------------------------------------------
_RE_NAME = re.compile(r"^([A-Za-z0-9]([A-Za-z0-9._-]*[A-Za-z0-9])?)")
@ -1579,9 +1499,7 @@ def _extract_pkg_name(spec: str) -> str:
"""Extract the package name from a pip spec string."""
m = _RE_NAME.match(spec)
return (
m.group(1)
if m
else spec.split("==")[0].split(">=")[0].split("<=")[0].split("[")[0].strip()
m.group(1) if m else spec.split("==")[0].split(">=")[0].split("<=")[0].split("[")[0].strip()
)
@ -1601,7 +1519,7 @@ def parse_requirements(req_files: list[str]) -> list[dict]:
if not line or line.startswith("#") or line.startswith("-"):
continue
is_git = line.startswith("git+") or "git+" in line.split("#")[0]
# Strip inline comments and environment markers for spec
# Strip inline comments and env markers
spec = line.split("#")[0].strip()
spec = spec.split(";")[0].strip()
if not spec:
@ -1634,7 +1552,7 @@ def get_downloaded_version(archive_path: str) -> str | None:
parts = basename[:-4].split("-")
if len(parts) >= 2:
return parts[1]
# Sdist: name-version.tar.gz / .tar.bz2 / .zip
# Sdist: name-version.<ext>
for ext in (".tar.gz", ".tar.bz2", ".tar.xz", ".tar", ".zip"):
if basename.endswith(ext):
stem = basename[: -len(ext)]
@ -1644,9 +1562,7 @@ def get_downloaded_version(archive_path: str) -> str | None:
return None
# ---------------------------------------------------------------------------
# Display
# ---------------------------------------------------------------------------
def severity_color(sev: str) -> str:
@ -1662,7 +1578,6 @@ def print_findings(findings: list[Finding]) -> None:
print("\n All clean. No suspicious patterns found.")
return
# Sort by severity
findings.sort(key = lambda f: SEVERITY_ORDER.get(f.severity, 99))
print(f"\n {'=' * 72}")
@ -1692,9 +1607,7 @@ def print_findings(findings: list[Finding]) -> None:
print(f" Summary: {', '.join(parts)}")
# ---------------------------------------------------------------------------
# PyPI version queries and --fix logic
# ---------------------------------------------------------------------------
def version_sort_key(v: str) -> tuple:
@ -1718,15 +1631,13 @@ def version_sort_key(v: str) -> tuple:
base = v_clean[0]
suffix = v[len(base) :]
# Parse numeric parts
parts = []
for seg in base.split("."):
try:
parts.append(int(seg))
except ValueError:
parts.append(0)
# Pad to at least 3 parts
while len(parts) < 3:
while len(parts) < 3: # pad to at least 3 parts
parts.append(0)
# Suffix ordering: dev < alpha < beta < rc < (none) < post
@ -1742,7 +1653,7 @@ def version_sort_key(v: str) -> tuple:
elif suffix_lower.startswith("post"):
suffix_rank = 1
else:
suffix_rank = 0 # stable release
suffix_rank = 0 # stable
return (epoch, tuple(parts), suffix_rank, suffix)
@ -1782,11 +1693,10 @@ def find_safe_version(
print(f" [WARN] No versions found on PyPI for {name}", file = sys.stderr)
return None
# Find index of bad version
try:
bad_idx = versions.index(bad_ver)
except ValueError:
# bad_ver might have been resolved to a different string; search by sort key
# bad_ver may resolve to a different string; search by sort key
bad_key = version_sort_key(bad_ver)
bad_idx = None
for i, v in enumerate(versions):
@ -1830,7 +1740,6 @@ def find_safe_version(
print(f" {ver} -- CRITICAL finding(s), skipping")
break
# Clean up scan dir for this version
shutil.rmtree(scan_dir, ignore_errors = True)
if clean:
@ -1860,8 +1769,7 @@ def update_req_line(raw_line: str, safe_ver: str, old_ver: str | None) -> str:
code_part, marker = code_part.split(";", 1)
marker = ";" + marker
# Replace version specifier
# Match patterns like ==1.2.3, >=1.2, ~=1.0, <=2.0, !=1.1, or bare name
# Replace version specifier (==1.2.3, >=1.2, ~=1.0, !=1.1, or bare name)
rewritten = re.sub(
r"([A-Za-z0-9._-]+)\s*(?:[><=!~]=?[^;#,\s]*(?:\s*,\s*[><=!~]=?[^;#,\s]*)*)?",
lambda m: f"{m.group(1)}=={safe_ver}",
@ -1880,12 +1788,8 @@ def update_req_file(filepath: str, updates: dict[int, str]) -> None:
updates: {line_num (1-indexed): new_line_text}
Writes atomically: stage in a sibling tmp file on the same
filesystem, fsync, then `os.replace` over the original. A SIGKILL
or power loss mid-write therefore either leaves the original
intact or leaves the fully new file -- never a half-written
requirements file (which would silently re-introduce a malicious
pin).
Writes atomically (sibling tmp file, fsync, os.replace) so a crash mid-write
never leaves a half-written file that re-introduces a malicious pin.
"""
with open(filepath) as f:
lines = f.readlines()
@ -1893,8 +1797,7 @@ def update_req_file(filepath: str, updates: dict[int, str]) -> None:
for line_num, new_text in updates.items():
idx = line_num - 1
if 0 <= idx < len(lines):
# Preserve original line ending
ending = "\n" if lines[idx].endswith("\n") else ""
ending = "\n" if lines[idx].endswith("\n") else "" # preserve line ending
lines[idx] = new_text + ending
dirpath = os.path.dirname(os.path.abspath(filepath)) or "."
@ -1917,13 +1820,9 @@ def update_req_file(filepath: str, updates: dict[int, str]) -> None:
raise
def _run_fix(
critical_pkgs: set[str],
entries: list[dict],
max_search: int,
) -> None:
def _run_fix(critical_pkgs: set[str], entries: list[dict], max_search: int) -> None:
"""Run the --fix flow: find safe versions, update requirements files."""
# Map package names to their entries for source tracking
# Map package names to entries for source tracking
pkg_entries: dict[str, list[dict]] = {}
for e in entries:
norm = e["name"].lower().replace("-", "_").replace(".", "_")
@ -1941,14 +1840,11 @@ def _run_fix(
if git_entries:
for e in git_entries:
src = e["source_file"] or "CLI"
print(
f" [SKIP] {pkg_name} is a git URL dep in {src}, cannot auto-update"
)
print(f" [SKIP] {pkg_name} is a git URL dep in {src}, cannot auto-update")
changes_summary.append(f" SKIP {pkg_name} (git URL)")
continue
# Get the currently resolved version
# Try to extract from the spec (e.g. name==1.2.3)
# Resolved version: try to extract from the spec (name==1.2.3)
current_ver = None
for e in related:
spec = e["spec"]
@ -1963,13 +1859,10 @@ def _run_fix(
downloaded = download_packages([pkg_name], dl_dir)
if downloaded:
current_ver = get_downloaded_version(downloaded[0][1])
# Delete resolution download immediately
shutil.rmtree(dl_dir, ignore_errors = True)
if not current_ver:
print(
f" [WARN] Cannot determine current version of {pkg_name}, skipping fix"
)
print(f" [WARN] Cannot determine current version of {pkg_name}, skipping fix")
changes_summary.append(f" SKIP {pkg_name} (version unknown)")
continue
@ -1986,9 +1879,7 @@ def _run_fix(
continue
print(f" [OK] {pkg_name}: {current_ver} -> {safe_ver}")
changes_summary.append(
f" FIX {pkg_name}=={current_ver} -> {pkg_name}=={safe_ver}"
)
changes_summary.append(f" FIX {pkg_name}=={current_ver} -> {pkg_name}=={safe_ver}")
# Update all occurrences in requirements files
file_updates: dict[str, dict[int, str]] = {}
@ -2006,7 +1897,6 @@ def _run_fix(
for filepath, updates in file_updates.items():
update_req_file(filepath, updates)
# Print summary
print(f"\n {'=' * 72}")
print(f" FIX SUMMARY")
print(f" {'=' * 72}")
@ -2015,9 +1905,7 @@ def _run_fix(
print(f"\n Re-run without --fix to verify the scan is clean.")
# ---------------------------------------------------------------------------
# Directory scanning
# ---------------------------------------------------------------------------
def _find_requirements_files(root: str) -> list[str]:
@ -2034,30 +1922,25 @@ def _find_requirements_files(root: str) -> list[str]:
skip_dirs = {"__pycache__", "node_modules", "venv", ".venv", "site-packages"}
results = []
for dirpath, dirnames, filenames in os.walk(root):
# Skip hidden dirs and known non-requirement dirs
# Skip hidden and known non-requirement dirs
dirnames[:] = [
d
for d in dirnames
if not d.startswith(".")
and d not in skip_dirs
and not d.endswith(".egg-info")
if not d.startswith(".") and d not in skip_dirs and not d.endswith(".egg-info")
]
dirname = os.path.basename(dirpath)
for fname in sorted(filenames):
if not fname.endswith(".txt"):
continue
# Match requirements*.txt anywhere
if fnmatch.fnmatch(fname.lower(), "requirements*.txt"):
results.append(os.path.join(dirpath, fname))
# Match *.txt inside a directory named "requirements"
# *.txt inside a directory named "requirements"
elif dirname == "requirements":
results.append(os.path.join(dirpath, fname))
return sorted(results)
# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------
def main() -> int:
@ -2115,9 +1998,7 @@ def main() -> int:
print(f" {f}")
req_files.extend(found)
else:
print(
f" [WARN] No requirements files found in {scan_dir}/", file = sys.stderr
)
print(f" [WARN] No requirements files found in {scan_dir}/", file = sys.stderr)
# Build unified entry list: list of dicts with source tracking
entries: list[dict] = []
@ -2158,7 +2039,7 @@ def main() -> int:
all_findings: list[Finding] = []
# Hard pin-block: refuse to download known-malicious PyPI versions.
# Hard pin-block: refuse to download known-malicious PyPI versions
specs, blocked_findings = _check_blocked_pypi_versions(specs)
all_findings.extend(blocked_findings)
@ -2196,10 +2077,8 @@ def main() -> int:
)
_run_fix(critical_pkgs, entries, args.max_search)
# Surface any pip-download failures BEFORE the scan-result exit code so
# an empty / partial download cannot mask itself as "0 findings, all
# clean". This is item (4) of the silent-failure hardening: an
# unresolvable spec or PyPI timeout used to print to stderr and exit 0.
# Surface pip-download failures BEFORE the exit code so a partial download
# can't masquerade as "0 findings, all clean" (silent-failure hardening 4).
if download_errors:
print(
f"\n {'=' * 72}\n"
@ -2211,7 +2090,7 @@ def main() -> int:
for err in download_errors:
print(f" [ERROR] {err}", file = sys.stderr)
print(
" Refusing to report 'all clean' on a partial scan; " "exiting 2.",
" Refusing to report 'all clean' on a partial scan; exiting 2.",
file = sys.stderr,
)
return 2

17
scripts/stamp_studio_release.py
View file

@ -17,12 +17,13 @@ import zipfile
from pathlib import Path
def _atomic_write_text(path: Path, data: str, encoding: str = "utf-8") -> None:
"""Atomic version of ``Path.write_text``.
A crash or signal mid-write leaves the prior file intact; the
Studio build never reads a partial ``_studio_release_build.py``.
"""
def _atomic_write_text(
path: Path,
data: str,
encoding: str = "utf-8",
) -> None:
"""Atomic ``Path.write_text``: a crash mid-write leaves the prior file
intact, so the build never reads a partial ``_studio_release_build.py``."""
dirpath = str(path.parent) or "."
path.parent.mkdir(parents = True, exist_ok = True)
fd, tmp_path = tempfile.mkstemp(prefix = ".stamp_studio.", dir = dirpath)
@ -41,9 +42,7 @@ def _atomic_write_text(path: Path, data: str, encoding: str = "utf-8") -> None:
REPO_ROOT = Path(__file__).resolve().parents[1]
BUILD_INFO_PATH = (
REPO_ROOT / "studio" / "backend" / "utils" / "_studio_release_build.py"
)
BUILD_INFO_PATH = REPO_ROOT / "studio" / "backend" / "utils" / "_studio_release_build.py"
BUILD_INFO_SUFFIX = "studio/backend/utils/_studio_release_build.py"
VERSION_RE = re.compile(r"^v\d+\.\d+\.\d+(?:-[0-9A-Za-z.][0-9A-Za-z.-]*)?$")
GIT_DESCRIBE_SUFFIX_RE = re.compile(r"-\d+-g[0-9A-Fa-f]+(?:-dirty)?$")

143
scripts/sync_allow_scripts_pins.py Normal file
View file

@ -0,0 +1,143 @@
#!/usr/bin/env python3
"""Keep `allowScripts` pins in studio/frontend/package.json in sync with
package-lock.json.
`npm approve-scripts` writes version-pinned entries ("pkg@1.2.3": true).
A dependency bump strands the pin, so the approval (or denial) silently
stops matching and the package's install scripts fall back to
"unreviewed". This tool re-pins existing entries to the versions the
lockfile actually resolves; it never adds or removes entries, so
approving a brand-new script-bearing package stays a human decision.
Usage:
python scripts/sync_allow_scripts_pins.py --check # CI: exit 1 on drift
python scripts/sync_allow_scripts_pins.py --fix # rewrite package.json
Pinned keys follow npm's allowScripts grammar: "name@1.2.3" or
"name@1.2.3 || 1.2.4". Bare names (no version) match every version and
are left alone. Entries whose range is not an exact-version disjunction
(wildcards, tags) are left alone too.
"""
from __future__ import annotations
import argparse
import json
import re
import sys
from pathlib import Path
REPO_ROOT = Path(__file__).resolve().parents[1]
DEFAULT_DIR = REPO_ROOT / "studio" / "frontend"
EXACT_VERSION_RE = re.compile(r"^\d+\.\d+\.\d+(?:[-+][0-9A-Za-z.+-]+)?$")
def split_spec(key: str) -> tuple[str, str | None]:
"""'@scope/name@1.2.3' -> ('@scope/name', '1.2.3'); bare names -> (key, None)."""
if key.startswith("@"):
rest = key[1:]
if "@" not in rest:
return key, None
name, rng = rest.split("@", 1)
return "@" + name, rng
if "@" not in key:
return key, None
name, rng = key.split("@", 1)
return name, rng
def is_exact_disjunction(rng: str) -> bool:
parts = [p.strip() for p in rng.split("||")]
return all(EXACT_VERSION_RE.match(p) for p in parts) and bool(parts)
def version_sort_key(version: str) -> tuple:
release = version.split("-", 1)[0].split("+", 1)[0]
return tuple(int(x) for x in release.split(".")), version
def script_versions_from_lock(lock: dict) -> dict[str, list[str]]:
"""Map package name -> sorted versions that carry install scripts."""
out: dict[str, set[str]] = {}
for path, meta in (lock.get("packages") or {}).items():
if not path or not meta.get("hasInstallScript"):
continue
name = path.rsplit("node_modules/", 1)[-1]
version = meta.get("version")
if name and version:
out.setdefault(name, set()).add(version)
return {n: sorted(vs, key = version_sort_key) for n, vs in out.items()}
def desired_key(name: str, versions: list[str]) -> str:
return f"{name}@{' || '.join(versions)}"
def compute_renames(policy: dict, lock_versions: dict[str, list[str]]) -> dict[str, str]:
renames: dict[str, str] = {}
for key in policy:
name, rng = split_spec(key)
if rng is None or not is_exact_disjunction(rng):
continue # bare name or non-exact spec: matches by name, never stale
versions = lock_versions.get(name)
if not versions:
continue # package gone or script-free now: stale pin is inert
want = desired_key(name, versions)
if key != want:
renames[key] = want
return renames
def main(argv: list[str] | None = None) -> int:
ap = argparse.ArgumentParser(description = __doc__)
mode = ap.add_mutually_exclusive_group(required = True)
mode.add_argument("--check", action = "store_true", help = "exit 1 if pins are stale")
mode.add_argument("--fix", action = "store_true", help = "rewrite package.json in place")
ap.add_argument(
"--dir",
type = Path,
default = DEFAULT_DIR,
help = "directory holding package.json + package-lock.json",
)
args = ap.parse_args(argv)
pkg_path = args.dir / "package.json"
lock_path = args.dir / "package-lock.json"
if not pkg_path.exists() or not lock_path.exists():
print(f"sync-allow-scripts: nothing to do ({args.dir} has no package.json + lockfile)")
return 0
pkg = json.loads(pkg_path.read_text(encoding = "utf-8"))
policy = pkg.get("allowScripts")
if not isinstance(policy, dict) or not policy:
print("sync-allow-scripts: no allowScripts policy in package.json, nothing to do")
return 0
lock = json.loads(lock_path.read_text(encoding = "utf-8"))
renames = compute_renames(policy, script_versions_from_lock(lock))
if not renames:
print(f"sync-allow-scripts: {len(policy)} allowScripts entries in sync with the lockfile")
return 0
for old, new in renames.items():
print(f' stale pin: "{old}" -> "{new}"')
if args.check:
print(
"sync-allow-scripts: pins are stale; run "
"`python scripts/sync_allow_scripts_pins.py --fix` and commit the result"
)
return 1
pkg["allowScripts"] = {renames.get(k, k): v for k, v in policy.items()}
pkg_path.write_text(json.dumps(pkg, indent = 2, ensure_ascii = False) + "\n", encoding = "utf-8")
print(
f"sync-allow-scripts: re-pinned {len(renames)} entr{'y' if len(renames) == 1 else 'ies'} in {pkg_path}"
)
return 0
if __name__ == "__main__":
sys.exit(main())

85
scripts/uninstall.ps1
View file

@ -1,3 +1,6 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
#
# Unsloth Studio uninstaller for Windows PowerShell.
# Stops running servers and removes install dir, launcher data, CLI shim,
# desktop and Start Menu shortcuts, the user PATH entry, and the PathBackup
@ -13,15 +16,19 @@ function Uninstall-UnslothStudio {
function _Step { param([string]$Msg) Write-Host $Msg }
function _Substep { param([string]$Msg, [string]$Color = "Gray") Write-Host " $Msg" -ForegroundColor $Color }
# Remove a file/dir/symlink only if it exists. Idempotent.
# Remove a file/dir/symlink if present. Idempotent; retries since a just-killed
# process can briefly hold a handle (Windows refuses the delete until released).
function _RemovePath {
param([string]$Path)
if ([string]::IsNullOrWhiteSpace($Path)) { return }
if (Test-Path -LiteralPath $Path) {
if (-not (Test-Path -LiteralPath $Path)) { return }
for ($attempt = 1; $attempt -le 3; $attempt++) {
try {
Remove-Item -LiteralPath $Path -Recurse -Force -ErrorAction Stop
_Substep "removed: $Path" "Green"
return
} catch {
if ($attempt -lt 3) { Start-Sleep -Milliseconds 700; continue }
_Substep "could not remove: $Path ($($_.Exception.Message))" "Yellow"
}
}
@ -233,9 +240,58 @@ function Uninstall-UnslothStudio {
} catch { }
}
# Stop processes that would block deleting the paths we remove. Unlike
# _StopStudioProcesses (venv exe only), this also catches llama-server/llama-cli,
# the unsloth.exe shim, and orphaned mp workers under SYSTEM python holding a
# venv DLL (an open DLL handle blocks the dir delete) -- found by scanning each
# candidate's loaded modules, not just its image path.
function _StopProcessesLockingRoots {
param([string[]]$Roots)
$clean = @($Roots | Where-Object { $_ } | ForEach-Object { $_.TrimEnd('\','/') })
if ($clean.Count -eq 0) { return }
$underRoot = {
param($p)
if (-not $p) { return $false }
foreach ($r in $clean) { if ($p -ieq $r -or $p -ilike "$r\*") { return $true } }
return $false
}
# 1. Image path under a target root (venv python, shim, llama-server).
try {
foreach ($proc in (Get-CimInstance Win32_Process -ErrorAction SilentlyContinue)) {
if ((& $underRoot $proc.ExecutablePath)) {
try { Stop-Process -Id $proc.ProcessId -Force -ErrorAction SilentlyContinue } catch { }
}
}
} catch { }
# 2. A loaded module under a target root (orphaned mp-fork python holding a
# venv DLL). Scoped to names that load our DLLs to keep the scan fast.
try {
$cands = Get-Process -Name python, pythonw, unsloth, llama-server, llama-cli -ErrorAction SilentlyContinue
foreach ($proc in $cands) {
$hit = $false
try {
foreach ($m in $proc.Modules) { if ((& $underRoot $m.FileName)) { $hit = $true; break } }
} catch { } # access denied enumerating modules -> skip
if ($hit) { try { Stop-Process -Id $proc.Id -Force -ErrorAction SilentlyContinue } catch { } }
}
} catch { }
}
# Default install root + default data dir.
$defaultStudioHome = if ($env:USERPROFILE) { Join-Path $env:USERPROFILE ".unsloth\studio" } else { $null }
$defaultDataDir = if ($env:LOCALAPPDATA) { Join-Path $env:LOCALAPPDATA "Unsloth Studio" } else { $null }
# Default-mode ~/.unsloth holds a SHARED llama.cpp build + .cache that are
# siblings of studio (not under it), so deleting <studio> misses them -- handle
# explicitly. No-op in env/custom mode (nested under the custom root, removed
# with it). A user-set UNSLOTH_LLAMA_CPP_PATH is left alone.
$defaultUnslothHome = if ($env:USERPROFILE) { Join-Path $env:USERPROFILE ".unsloth" } else { $null }
$defaultLlamaCpp = if ($defaultUnslothHome) { Join-Path $defaultUnslothHome "llama.cpp" } else { $null }
$defaultCache = if ($defaultUnslothHome) { Join-Path $defaultUnslothHome ".cache" } else { $null }
# llama.cpp atomic-install staging root (install_llama_prebuilt.py .staging,
# sibling of the install dir). Usually pruned after activate, but an interrupted
# build can leave a "<name>.staging-XXXX" tree; removing it lets the empty-dir
# cleanup of ~/.unsloth below succeed. No-op in env/custom mode and when absent.
$defaultStaging = if ($defaultUnslothHome) { Join-Path $defaultUnslothHome ".staging" } else { $null }
# Build known-root list FIRST so the port-file kill can verify ownership.
$customRoots = @(_CustomStudioRoots)
@ -252,6 +308,9 @@ function Uninstall-UnslothStudio {
_StopByPortFile -PortFile (Join-Path $r "share\studio.port") -KnownRoots $knownRoots
}
_StopStudioProcesses -KnownRoots $knownRoots
# Also stop anything holding a handle on the exact paths we delete (llama-server,
# the CLI shim, an mp-fork python with a venv DLL) so the dir delete isn't refused.
_StopProcessesLockingRoots -Roots (@($knownRoots) + @($defaultDataDir, $defaultLlamaCpp, $defaultCache))
# ── Remove custom-root install trees ──
_Step "Removing data and install directories..."
@ -270,6 +329,16 @@ function Uninstall-UnslothStudio {
if ($defaultStudioHome) { _RemovePath $defaultStudioHome }
# Default data dir.
if ($defaultDataDir) { _RemovePath $defaultDataDir }
# Default-mode shared llama.cpp build + cache (siblings of studio under
# ~/.unsloth). No-op in env/custom mode and when absent.
if ($defaultLlamaCpp) { _RemovePath $defaultLlamaCpp }
if ($defaultCache) { _RemovePath $defaultCache }
if ($defaultStaging) { _RemovePath $defaultStaging }
# Drop ~/.unsloth itself, but ONLY if now empty -- never nuke unrelated content.
if ($defaultUnslothHome -and (Test-Path -LiteralPath $defaultUnslothHome) -and
-not (Get-ChildItem -LiteralPath $defaultUnslothHome -Force -ErrorAction SilentlyContinue)) {
_RemovePath $defaultUnslothHome
}
# ── Remove desktop and Start Menu shortcuts ──
_Step "Removing desktop and Start Menu shortcuts..."
@ -280,6 +349,18 @@ function Uninstall-UnslothStudio {
if ($env:APPDATA) {
_RemovePath (Join-Path $env:APPDATA "Microsoft\Windows\Start Menu\Programs\Unsloth Studio.lnk")
}
# Invalidate the Win11 Start Menu tile cache so the removed shortcut's tile
# disappears promptly instead of lingering stale (mirrors install.ps1's
# New-StudioShortcuts). Preserves start2.bin (the pin layout).
try {
$smehTemp = Join-Path $env:LOCALAPPDATA "Packages\Microsoft.Windows.StartMenuExperienceHost_cw5n1h2txyewy\TempState"
if (Test-Path -LiteralPath $smehTemp) {
Get-ChildItem -LiteralPath $smehTemp -Filter "TileCache_*" -ErrorAction SilentlyContinue |
Remove-Item -Force -ErrorAction SilentlyContinue
Remove-Item -LiteralPath (Join-Path $smehTemp "StartUnifiedTileModelCache.dat") -Force -ErrorAction SilentlyContinue
Stop-Process -Name StartMenuExperienceHost -Force -ErrorAction SilentlyContinue
}
} catch { }
# ── Clean user PATH and registry backup ──
_Step "Cleaning user PATH and registry..."

115
scripts/uninstall.sh
View file

@ -1,4 +1,7 @@
#!/usr/bin/env sh
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
#
# Unsloth Studio uninstaller (macOS / Linux / WSL).
# Stops running servers and removes install dir, launcher data,
# CLI shim, desktop shortcut, .app bundle, and Launch Services entry.
@ -209,6 +212,21 @@ _custom_studio_roots | while IFS= read -r _custom_root; do
_remove_path "$_custom_root"
done
_remove_path "$HOME/.unsloth/studio"
# Default-mode shared llama.cpp build + cache are siblings of studio (not removed
# by deleting it). No-op in env/custom mode (they nest under the custom root) and
# when absent. A user-set UNSLOTH_LLAMA_CPP_PATH is intentionally kept.
_remove_path "$HOME/.unsloth/llama.cpp"
_remove_path "$HOME/.unsloth/.cache"
# llama.cpp atomic-install staging root (install_llama_prebuilt.py .staging).
# Normally pruned after activate, but an interrupted build can leave it behind;
# removing it lets the rmdir below succeed. No-op in env/custom mode and absent.
_remove_path "$HOME/.unsloth/.staging"
# ROCm-on-WSL helper artifacts (librocdxg build clone + smoke-test venv). No-op
# where they don't exist; removing them lets the rmdir below succeed.
_remove_path "$HOME/.unsloth/librocdxg"
_remove_path "$HOME/.unsloth/rocm-smoketest"
# Drop ~/.unsloth only if now empty (rmdir refuses non-empty, so user content is kept).
rmdir "$HOME/.unsloth" 2>/dev/null || true
_remove_path "$HOME/.local/share/unsloth"
# CLI shim: only the symlink Studio created, never a pip-installed file.
_remove_cli_shim
@ -241,23 +259,102 @@ case "$_os" in
Linux)
if [ "$_is_wsl" = "1" ]; then
echo "Removing WSL Windows-side shortcuts..."
# install.sh creates 'Unsloth Studio.lnk' on the Windows Desktop and
# Start Menu Programs folder via powershell.exe; mirror that path.
if command -v powershell.exe >/dev/null 2>&1; then
# install.sh creates per-distro 'Unsloth Studio (WSL - <distro>).lnk'
# on the Windows Desktop + Start Menu via powershell.exe. Scope removal
# to THIS distro (passed as $args[0]) so a multi-distro install keeps the
# other distros' launchers; the TARGET=wsl.exe check still spares a
# native install's "Unsloth Studio.lnk". Prefer powershell.exe; test it
# can EXECUTE (`command -v` succeeds even with interop OFF -- .exe then
# fails "Exec format error", common on systemd-enabled distros).
_wsl_distro="${WSL_DISTRO_NAME:-}"
_ps_ran=0
if command -v powershell.exe >/dev/null 2>&1 && \
powershell.exe -NoProfile -Command "exit 0" >/dev/null 2>&1; then
_ps_ran=1
# Inject the distro into the command: a -Command string does not
# receive trailing tokens as $args. WSL distro names are safe to
# embed (no quotes/$/backtick).
# shellcheck disable=SC2016
# $env:APPDATA is a PowerShell expansion; intentionally literal at shell level.
powershell.exe -NoProfile -Command '
$names = @("Desktop","StartMenu");
powershell.exe -NoProfile -Command '$distro = "'"$_wsl_distro"'";
$dirs = @(
[Environment]::GetFolderPath("Desktop"),
(Join-Path $env:APPDATA "Microsoft\Windows\Start Menu\Programs")
);
$ws = New-Object -ComObject WScript.Shell;
foreach ($d in $dirs) {
if (-not $d) { continue }
$p = Join-Path $d "Unsloth Studio.lnk";
if (Test-Path -LiteralPath $p) { Remove-Item -LiteralPath $p -Force }
if (-not $d -or -not (Test-Path -LiteralPath $d)) { continue }
Get-ChildItem -LiteralPath $d -Filter "Unsloth Studio*.lnk" -ErrorAction SilentlyContinue | ForEach-Object {
try {
$sc = $ws.CreateShortcut($_.FullName);
if ("$($sc.TargetPath) $($sc.Arguments)" -notmatch "wsl\.exe") { return }
# When the distro is known, require the per-distro
# name for this distro or its -d "<distro>" argument
# so launchers for other distros are not removed.
if ($distro) {
$nameMatch = ($_.Name -eq "Unsloth Studio (WSL - $distro).lnk");
$argMatch = ($sc.Arguments -match ("-d\s+`"?" + [regex]::Escape($distro) + "`"?"));
if (-not ($nameMatch -or $argMatch)) { return }
}
Remove-Item -LiteralPath $_.FullName -Force -ErrorAction SilentlyContinue
} catch { }
}
}' >/dev/null 2>&1 || true
fi
# Fallback when powershell.exe can't run (interop disabled): remove the
# WSL .lnk files via drvfs. The "Unsloth Studio (WSL..." name is
# WSL-specific, so a native install's "Unsloth Studio.lnk" never matches.
if [ "$_ps_ran" = "0" ]; then
for _drive in /mnt/c /mnt/d /mnt/e; do
[ -d "$_drive/Users" ] || continue
for _udir in "$_drive"/Users/*; do
[ -d "$_udir" ] || continue
for _scdir in \
"$_udir/Desktop" \
"$_udir/OneDrive/Desktop" \
"$_udir"/OneDrive*/Desktop \
"$_udir/AppData/Roaming/Microsoft/Windows/Start Menu/Programs"; do
[ -d "$_scdir" ] || continue
if [ -n "$_wsl_distro" ]; then
# Exact per-distro name (no glob) so other distros survive.
_lnk="$_scdir/Unsloth Studio (WSL - ${_wsl_distro}).lnk"
[ -e "$_lnk" ] && rm -f "$_lnk" 2>/dev/null && echo " removed: $_lnk" || true
else
# Distro unknown: fall back to the broad WSL prefix.
for _lnk in "$_scdir"/"Unsloth Studio (WSL"*.lnk; do
[ -e "$_lnk" ] && rm -f "$_lnk" 2>/dev/null && echo " removed: $_lnk" || true
done
fi
done
done
done
fi
# ── ROCm-on-WSL config (install_rocm_wsl_strixhalo.sh) ──
# Remove Unsloth's own ROCDXG config (the env it persisted). The system
# ROCm userspace is a shared prereq (like CUDA) and is LEFT IN PLACE by
# default; set UNSLOTH_UNINSTALL_ROCM=1 to remove it too.
echo "Removing ROCm-on-WSL config..."
_sudo=""
if [ "$_uid" != "0" ] && command -v sudo >/dev/null 2>&1; then _sudo="sudo"; fi
$_sudo rm -f /etc/profile.d/unsloth-rocm-wsl.sh 2>/dev/null || true
if [ -f "$HOME/.bashrc" ] && grep -q "Unsloth ROCm-on-WSL" "$HOME/.bashrc" 2>/dev/null; then
_bk=$(mktemp 2>/dev/null || echo "$HOME/.bashrc.unsloth.tmp")
if sed '/# >>> Unsloth ROCm-on-WSL/,/# <<< Unsloth ROCm-on-WSL/d' "$HOME/.bashrc" > "$_bk" 2>/dev/null; then
cat "$_bk" > "$HOME/.bashrc" 2>/dev/null || true
echo " cleaned ROCm-on-WSL block from ~/.bashrc"
fi
rm -f "$_bk" 2>/dev/null || true
fi
if [ "${UNSLOTH_UNINSTALL_ROCM:-0}" = "1" ]; then
echo " removing system ROCm (UNSLOTH_UNINSTALL_ROCM=1)..."
$_sudo rm -f /etc/apt/sources.list.d/rocm.list /etc/apt/preferences.d/rocm-pin-600 \
/etc/apt/keyrings/rocm.gpg /etc/ld.so.conf.d/rocm.conf 2>/dev/null || true
$_sudo sh -c 'rm -rf /opt/rocm /opt/rocm-*' 2>/dev/null || true
if command -v ldconfig >/dev/null 2>&1; then $_sudo ldconfig 2>/dev/null || true; fi
elif [ -d /opt/rocm ]; then
echo " Note: ROCm userspace (/opt/rocm*) left in place (shared prereq)."
echo " Remove it by re-running with UNSLOTH_UNINSTALL_ROCM=1, or manually:"
echo " sudo rm -rf /opt/rocm /opt/rocm-* && sudo ldconfig"
fi
fi
echo "Removing Linux .desktop entry..."
_remove_path "$HOME/.local/share/applications/unsloth-studio.desktop"

22
scripts/verify_comment_only_diff.py
View file

@ -57,9 +57,7 @@ def _git_show(rev: str, path: str) -> str:
def _strip_docstrings(tree: ast.AST) -> ast.AST:
"""Remove every string-literal docstring (Module / FunctionDef /
AsyncFunctionDef / ClassDef). Empty body becomes ``pass`` so
ast.unparse stays valid."""
"""Remove docstrings; empty bodies become ``pass`` so unparse stays valid."""
for node in ast.walk(tree):
if isinstance(
node,
@ -87,9 +85,8 @@ def _normalize_py(src: str) -> str:
def _strip_shell_comments(s: str) -> str:
"""Strip pure-comment lines and inline trailing comments from a shell
snippet, then collapse runs of blank lines. Heuristic only: leaves a
line untouched if it has an odd quote count (open string)."""
"""Strip shell comments and collapse blank lines. Heuristic: skips lines
with an odd quote count (open string)."""
out = []
for line in s.splitlines():
stripped = line.lstrip()
@ -116,9 +113,7 @@ def _strip_shell_comments(s: str) -> str:
def _normalize_yaml_run_strings(obj: Any) -> Any:
"""Walk the parsed YAML object; for any multi-line string (i.e. a
``run: |`` script body), strip shell comments. Returns a normalised
copy."""
"""Strip shell comments from any multi-line string (``run: |`` body)."""
if isinstance(obj, dict):
return {k: _normalize_yaml_run_strings(v) for k, v in obj.items()}
if isinstance(obj, list):
@ -128,12 +123,15 @@ def _normalize_yaml_run_strings(obj: Any) -> Any:
return obj
def _walk_yaml_diff(b: Any, a: Any, prefix: str = "") -> None:
def _walk_yaml_diff(
b: Any,
a: Any,
prefix: str = "",
) -> None:
"""Print a path-keyed summary of the first structural / scalar diff."""
if type(b) is not type(a):
print(
f" type-diff at {prefix or '/'}: "
f"{type(b).__name__} -> {type(a).__name__}",
f" type-diff at {prefix or '/'}: " f"{type(b).__name__} -> {type(a).__name__}",
)
return
if isinstance(b, dict):

812
scripts/verify_import_hoist.py Normal file
View file

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

23
studio/Unsloth_Studio_Colab.ipynb
View file

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

27
studio/backend/_platform_compat.py
View file

@ -4,17 +4,16 @@
"""
Compatibility shim for Anaconda/conda-forge Python builds.
Anaconda modifies sys.version to include distributor metadata between pipe
characters, e.g. '3.12.4 | packaged by Anaconda, Inc. | (main, ...) [MSC ...]'.
Python's platform._sys_version() has a hardcoded regex that cannot parse this,
raising ValueError. CPython closed this as "not planned" (cpython#102396).
Anaconda puts distributor metadata between pipes in sys.version, e.g.
'3.12.4 | packaged by Anaconda, Inc. | (main, ...) [MSC ...]'. The regex in
platform._sys_version() can't parse this and raises ValueError (cpython#102396,
closed as "not planned").
This module seeds platform._sys_version_cache so the stdlib parser never sees
the problematic string, fixing the import chain:
We seed platform._sys_version_cache so the stdlib parser never sees the bad
string, fixing the import chain:
structlog -> rich.pretty -> attrs._compat -> platform.python_implementation()
Import this module before any library imports that may trigger the above chain.
Safe to import multiple times (no-op if cache is already seeded or no pipes).
Import before any library that may trigger that chain. Idempotent.
"""
import platform
@ -23,18 +22,17 @@ import sys
def _seed_sys_version_cache() -> None:
"""One-shot cache prime: parse a cleaned sys.version and seed the cache."""
"""Parse a cleaned sys.version and seed the cache once."""
raw = sys.version
# Strip paired |...| segments (Anaconda, conda-forge metadata)
cleaned = re.sub(r"\s*\|[^|]*\|\s*", " ", raw).strip()
# Format B: "ver (build) | label | (build_dup) \n[compiler]"
# After pipe-strip, two consecutive (...) groups remain; drop the second.
# Pipe-strip can leave two consecutive (...) groups; drop the second.
cleaned = re.sub(r"(\([^)]*\))\s+\([^)]*\)", r"\1", cleaned)
if "|" in cleaned:
# Unpaired pipe remaining -- keep version + everything from "(" onward
# Unpaired pipe left: keep version + everything from "(" onward
m = re.match(r"([\w.+]+)\s*", cleaned)
p = cleaned.find("(")
if m and p > 0:
@ -43,13 +41,12 @@ def _seed_sys_version_cache() -> None:
if cleaned == raw:
return # Nothing to fix
# Parse the cleaned string through the real stdlib parser
try:
result = platform._sys_version(cleaned)
except ValueError:
return # Cleaning didn't produce a parseable string; don't make things worse
return # Still unparsable; don't make things worse
# Seed the cache so future calls with the raw string skip parsing entirely
# Seed the cache so future calls with the raw string skip parsing
cache = getattr(platform, "_sys_version_cache", None)
if isinstance(cache, dict):
cache[raw] = result

397
studio/backend/assets/chat_templates/gemma-4-edge.jinja Normal file
View file

@ -0,0 +1,397 @@
{#-
Gemma 4 chat template (E2B / E4B edge variant), vendored for Unsloth Studio.
Source: google/gemma-4-31B-it HF discussion/PR #118 (adds the preserve_thinking
flag plus null-rendering, string-arguments validation, balanced turn tags, empty
messages handling, and OpenAI image_url/input_audio aliases).
Studio-local changes vs PR #118:
1. preserve_thinking defaults to false (see SETUP block below).
2. The empty "<|channel>thought\n<channel|>" block on enable_thinking=false is
NOT emitted. Google ships a distinct template for E2B/E4B (google/gemma-4-E2B-it,
google/gemma-4-E4B-it) that omits it; only the 12b/26B-A4B/31B family emits it.
This file matches the E2B/E4B behavior; gemma-4.jinja keeps the larger-model one.
Applied to unsloth/gemma-4-E2B-it-GGUF and unsloth/gemma-4-E4B-it-GGUF so the
embedded GGUF template does not need re-downloading.
-#}
{%- macro format_parameters(properties, required, filter_keys=false) -%}
{%- set standard_keys = ['description', 'type', 'properties', 'required', 'nullable'] -%}
{%- set ns = namespace(found_first=false) -%}
{%- for key, value in properties | dictsort -%}
{%- set add_comma = false -%}
{%- if not filter_keys or key not in standard_keys -%}
{%- if ns.found_first %},{% endif -%}
{%- set ns.found_first = true -%}
{{ key }}:{
{%- if value['description'] -%}
description:<|"|>{{ value['description'] }}<|"|>
{%- set add_comma = true -%}
{%- endif -%}
{%- if value['type'] | upper == 'STRING' -%}
{%- if value['enum'] -%}
{%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
enum:{{ format_argument(value['enum']) }}
{%- endif -%}
{%- elif value['type'] | upper == 'ARRAY' -%}
{%- if value['items'] is mapping and value['items'] -%}
{%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
items:{
{%- set ns_items = namespace(found_first=false) -%}
{%- for item_key, item_value in value['items'] | dictsort -%}
{%- if item_value is not none -%}
{%- if ns_items.found_first %},{% endif -%}
{%- set ns_items.found_first = true -%}
{%- if item_key == 'properties' -%}
properties:{
{%- if item_value is mapping -%}
{{- format_parameters(item_value, value['items']['required'] | default([])) -}}
{%- endif -%}
}
{%- elif item_key == 'required' -%}
required:[
{%- for req_item in item_value -%}
<|"|>{{- req_item -}}<|"|>
{%- if not loop.last %},{% endif -%}
{%- endfor -%}
]
{%- elif item_key == 'type' -%}
{%- if item_value is string -%}
type:{{ format_argument(item_value | upper) }}
{%- else -%}
type:{{ format_argument(item_value | map('upper') | list) }}
{%- endif -%}
{%- else -%}
{{ item_key }}:{{ format_argument(item_value) }}
{%- endif -%}
{%- endif -%}
{%- endfor -%}
}
{%- endif -%}
{%- endif -%}
{%- if value['nullable'] %}
{%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
nullable:true
{%- endif -%}
{%- if value['type'] | upper == 'OBJECT' -%}
{%- if value['properties'] is defined and value['properties'] is mapping -%}
{%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
properties:{
{{- format_parameters(value['properties'], value['required'] | default([])) -}}
}
{%- elif value is mapping -%}
{%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
properties:{
{{- format_parameters(value, value['required'] | default([]), filter_keys=true) -}}
}
{%- endif -%}
{%- if value['required'] -%}
{%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
required:[
{%- for item in value['required'] | default([]) -%}
<|"|>{{- item -}}<|"|>
{%- if not loop.last %},{% endif -%}
{%- endfor -%}
]
{%- endif -%}
{%- endif -%}
{%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
type:<|"|>{{ value['type'] | upper }}<|"|>}
{%- endif -%}
{%- endfor -%}
{%- endmacro -%}
{%- macro format_function_declaration(tool_data) -%}
declaration:{{- tool_data['function']['name'] -}}{description:<|"|>{{- tool_data['function']['description'] -}}<|"|>
{%- set params = tool_data['function']['parameters'] -%}
{%- if params -%}
,parameters:{
{%- if params['properties'] -%}
properties:{ {{- format_parameters(params['properties'], params['required']) -}} },
{%- endif -%}
{%- if params['required'] -%}
required:[
{%- for item in params['required'] -%}
<|"|>{{- item -}}<|"|>
{{- ',' if not loop.last -}}
{%- endfor -%}
],
{%- endif -%}
{%- if params['type'] -%}
type:<|"|>{{- params['type'] | upper -}}<|"|>}
{%- endif -%}
{%- endif -%}
{%- if 'response' in tool_data['function'] -%}
{%- set response_declaration = tool_data['function']['response'] -%}
,response:{
{%- if response_declaration['description'] -%}
description:<|"|>{{- response_declaration['description'] -}}<|"|>,
{%- endif -%}
{%- if response_declaration['type'] | upper == 'OBJECT' -%}
type:<|"|>{{- response_declaration['type'] | upper -}}<|"|>}
{%- endif -%}
{%- endif -%}
}
{%- endmacro -%}
{%- macro format_argument(argument, escape_keys=True) -%}
{%- if argument is none -%}
{{- 'null' -}}
{%- elif argument is string -%}
{{- '<|"|>' + argument + '<|"|>' -}}
{%- elif argument is boolean -%}
{{- 'true' if argument else 'false' -}}
{%- elif argument is mapping -%}
{{- '{' -}}
{%- set ns = namespace(found_first=false) -%}
{%- for key, value in argument | dictsort -%}
{%- if ns.found_first %},{% endif -%}
{%- set ns.found_first = true -%}
{%- if escape_keys -%}
{{- '<|"|>' + key + '<|"|>' -}}
{%- else -%}
{{- key -}}
{%- endif -%}
:{{- format_argument(value, escape_keys=escape_keys) -}}
{%- endfor -%}
{{- '}' -}}
{%- elif argument is sequence -%}
{{- '[' -}}
{%- for item in argument -%}
{{- format_argument(item, escape_keys=escape_keys) -}}
{%- if not loop.last %},{% endif -%}
{%- endfor -%}
{{- ']' -}}
{%- else -%}
{{- argument -}}
{%- endif -%}
{%- endmacro -%}
{%- macro strip_thinking(text) -%}
{%- set ns = namespace(result='') -%}
{%- for part in text.split('<channel|>') -%}
{%- if '<|channel>' in part -%}
{%- set ns.result = ns.result + part.split('<|channel>')[0] -%}
{%- else -%}
{%- set ns.result = ns.result + part -%}
{%- endif -%}
{%- endfor -%}
{{- ns.result | trim -}}
{%- endmacro -%}
{%- macro format_tool_response_block(tool_name, response) -%}
{{- '<|tool_response>' -}}
{%- if response is mapping -%}
{{- 'response:' + tool_name + '{' -}}
{%- for key, value in response | dictsort -%}
{{- key -}}:{{- format_argument(value, escape_keys=False) -}}
{%- if not loop.last %},{% endif -%}
{%- endfor -%}
{{- '}' -}}
{%- else -%}
{{- 'response:' + tool_name + '{value:' + format_argument(response, escape_keys=False) + '}' -}}
{%- endif -%}
{{- '<tool_response|>' -}}
{%- endmacro -%}
{#- ===== SETUP ===== -#}
{%- set ns = namespace(prev_message_type=None, prev_non_tool_role=None) -%}
{%- set loop_messages = messages -%}
{%- set enable_thinking = enable_thinking | default(false) -%}
{#- Unsloth Studio: preserve_thinking defaults OFF (upstream PR #118 defaults true). -#}
{%- set preserve_thinking = preserve_thinking | default(false) -%}
{{- bos_token -}}
{#- Handle System/Tool Definitions Block -#}
{%- if enable_thinking or tools or (messages and messages[0]['role'] in ['system', 'developer']) -%}
{{- '<|turn>system\n' -}}
{#- Inject Thinking token at the very top of the FIRST system turn -#}
{%- if enable_thinking -%}
{{- '<|think|>\n' -}}
{%- set ns.prev_message_type = 'think' -%}
{%- endif -%}
{%- if messages and messages[0]['role'] in ['system', 'developer'] -%}
{%- if messages[0]['content'] is string -%}
{{- messages[0]['content'] | trim -}}
{%- elif messages[0]['content'] is sequence -%}
{%- for item in messages[0]['content'] -%}
{{- item['text'] | trim + ' '-}}
{%- endfor -%}
{%- endif -%}
{%- set loop_messages = messages[1:] -%}
{%- endif -%}
{%- if tools -%}
{%- for tool in tools %}
{{- '<|tool>' -}}
{{- format_function_declaration(tool) | trim -}}
{{- '<tool|>' -}}
{%- endfor %}
{%- set ns.prev_message_type = 'tool' -%}
{%- endif -%}
{{- '<turn|>\n' -}}
{%- endif %}
{#- Pre-scan: find last user message index for reasoning guard -#}
{%- set ns_turn = namespace(last_user_idx=-1) -%}
{%- for i in range(loop_messages | length) -%}
{%- if loop_messages[i]['role'] == 'user' -%}
{%- set ns_turn.last_user_idx = i -%}
{%- endif -%}
{%- endfor -%}
{#- Loop through messages -#}
{%- for message in loop_messages -%}
{%- if message['role'] != 'tool' -%}
{%- set ns.prev_message_type = None -%}
{%- set role = 'model' if message['role'] == 'assistant' else message['role'] -%}
{#- Detect continuation using tracked state - O(1) instead of O(n) backward scan -#}
{%- set continue_same_model_turn = (role == 'model' and ns.prev_non_tool_role == 'assistant') -%}
{%- if not continue_same_model_turn -%}
{{- '<|turn>' + role + '\n' }}
{%- endif -%}
{#- Render reasoning/reasoning_content as thinking channel (tool-call turns only) -#}
{%- set thinking_text = message.get('reasoning') or message.get('reasoning_content') -%}
{%- set thinking_gate = (loop.index0 > ns_turn.last_user_idx) or preserve_thinking -%}
{%- if thinking_text and thinking_gate and message.get('tool_calls') -%}
{{- '<|channel>thought\n' + thinking_text + '\n<channel|>' -}}
{%- endif -%}
{%- if message.get('tool_calls') -%}
{%- for tool_call in message.get('tool_calls') -%}
{%- set function = tool_call['function'] -%}
{{- '<|tool_call>call:' + function['name'] + '{' -}}
{%- if function['arguments'] is mapping -%}
{%- set ns_args = namespace(found_first=false) -%}
{%- for key, value in function['arguments'] | dictsort -%}
{%- if ns_args.found_first %},{% endif -%}
{%- set ns_args.found_first = true -%}
{{- key -}}:{{- format_argument(value, escape_keys=False) -}}
{%- endfor -%}
{%- elif function['arguments'] is none -%}
{%- else -%}
{{- raise_exception(
"chat_template: tool_calls[].function.arguments must be a "
"JSON object (mapping), not a string. Deserialize arguments "
"before passing to the template."
) -}}
{%- endif -%}
{{- '}<tool_call|>' -}}
{%- endfor -%}
{%- set ns.prev_message_type = 'tool_call' -%}
{%- endif -%}
{%- set ns_tr_out = namespace(flag=false) -%}
{%- if message.get('tool_responses') -%}
{#- Legacy: tool_responses embedded on the assistant message (Google/Gemma native) -#}
{%- for tool_response in message.get('tool_responses') -%}
{{- format_tool_response_block(tool_response['name'] | default('unknown', true), tool_response['response']) -}}
{%- set ns_tr_out.flag = true -%}
{%- set ns.prev_message_type = 'tool_response' -%}
{%- endfor -%}
{%- elif message.get('tool_calls') -%}
{#- OpenAI Chat Completions: forward-scan consecutive role:tool messages -#}
{%- set ns_tool_scan = namespace(stopped=false) -%}
{%- for k in range(loop.index0 + 1, loop_messages | length) -%}
{%- if ns_tool_scan.stopped -%}
{%- elif loop_messages[k]['role'] != 'tool' -%}
{%- set ns_tool_scan.stopped = true -%}
{%- else -%}
{%- set follow = loop_messages[k] -%}
{#- Resolve tool_call_id to function name -#}
{%- set ns_tname = namespace(name=follow.get('name') or 'unknown') -%}
{%- for tc in message.get('tool_calls') -%}
{%- if tc.get('id') == follow.get('tool_call_id') -%}
{%- set ns_tname.name = tc['function']['name'] -%}
{%- endif -%}
{%- endfor -%}
{#- Handle content as string or content-parts array -#}
{%- set tool_body = follow.get('content') -%}
{%- if tool_body is string -%}
{{- format_tool_response_block(ns_tname.name, tool_body) -}}
{%- elif tool_body is sequence and tool_body is not string -%}
{%- set ns_txt = namespace(s='') -%}
{%- for part in tool_body -%}
{%- if part.get('type') == 'text' -%}
{%- set ns_txt.s = ns_txt.s + (part.get('text') | default('')) -%}
{%- endif -%}
{%- endfor -%}
{{- format_tool_response_block(ns_tname.name, ns_txt.s) -}}
{%- for part in tool_body -%}
{%- if part.get('type') in ['image', 'image_url'] -%}
{{- '<|image|>' -}}
{%- elif part.get('type') in ['audio', 'input_audio'] -%}
{{- '<|audio|>' -}}
{%- elif part.get('type') == 'video' -%}
{{- '<|video|>' -}}
{%- endif -%}
{%- endfor -%}
{%- else -%}
{{- format_tool_response_block(ns_tname.name, tool_body) -}}
{%- endif -%}
{%- set ns_tr_out.flag = true -%}
{%- set ns.prev_message_type = 'tool_response' -%}
{%- endif -%}
{%- endfor -%}
{%- endif -%}
{%- set captured_content -%}
{%- if message.get('content') is string -%}
{%- if role == 'model' -%}
{{- strip_thinking(message['content']) -}}
{%- else -%}
{{- message['content'] | trim -}}
{%- endif -%}
{%- elif message.get('content') is sequence -%}
{%- for item in message['content'] -%}
{%- if item.get('type') == 'text' -%}
{%- if role == 'model' -%}
{{- strip_thinking(item['text']) -}}
{%- else -%}
{{- item['text'] | trim -}}
{%- endif -%}
{%- elif item.get('type') in ['image', 'image_url'] -%}
{{- '<|image|>' -}}
{%- elif item.get('type') in ['audio', 'input_audio'] -%}
{{- '<|audio|>' -}}
{%- elif item.get('type') == 'video' -%}
{{- '<|video|>' -}}
{%- endif -%}
{%- endfor -%}
{%- endif -%}
{%- endset -%}
{{- captured_content -}}
{%- set has_content = captured_content | trim | length > 0 -%}
{#- Forward-scan: find next non-tool message role for continuation detection -#}
{%- set next_nt = namespace(role=None, found=false) -%}
{%- for j in range(loop.index0 + 1, loop_messages | length) -%}
{%- if not next_nt.found -%}
{%- if loop_messages[j]['role'] != 'tool' -%}
{%- set next_nt.role = loop_messages[j]['role'] -%}
{%- set next_nt.found = true -%}
{%- endif -%}
{%- endif -%}
{%- endfor -%}
{%- set continues_into_next = (
role == 'model'
and next_nt.role == 'assistant'
and not message.get('tool_calls')
and not ns_tr_out.flag
) -%}
{%- if ns.prev_message_type == 'tool_call' and not ns_tr_out.flag -%}
{{- '<|tool_response>' -}}
{%- elif continues_into_next -%}
{{- '\n' -}}
{%- elif not (ns_tr_out.flag and not has_content) -%}
{{- '<turn|>\n' -}}
{%- endif -%}
{#- Track previous non-tool role for next iteration (avoids O(n) backward scan) -#}
{%- set ns.prev_non_tool_role = message['role'] -%}
{%- endif -%}
{%- endfor -%}
{%- if add_generation_prompt -%}
{%- if ns.prev_message_type != 'tool_response' and ns.prev_message_type != 'tool_call' -%}
{{- '<|turn>model\n' -}}
{%- endif -%}
{#- E2B/E4B do NOT emit an empty thought block when enable_thinking is false
(unlike the 12b/26B-A4B/31B family); see header. -#}
{%- endif -%}

397
studio/backend/assets/chat_templates/gemma-4.jinja Normal file
View file

@ -0,0 +1,397 @@
{#-
Gemma 4 chat template, vendored for Unsloth Studio.
Source: google/gemma-4-31B-it HF discussion/PR #118 (adds the preserve_thinking
flag plus null-rendering, string-arguments validation, balanced turn tags, empty
messages handling, and OpenAI image_url/input_audio aliases).
Studio-local change: preserve_thinking defaults to false (see SETUP block below).
Applied to unsloth/gemma-4-*-GGUF models so the embedded GGUF template does not
need re-downloading. Keep in sync with upstream if PR #118 changes.
-#}
{%- macro format_parameters(properties, required, filter_keys=false) -%}
{%- set standard_keys = ['description', 'type', 'properties', 'required', 'nullable'] -%}
{%- set ns = namespace(found_first=false) -%}
{%- for key, value in properties | dictsort -%}
{%- set add_comma = false -%}
{%- if not filter_keys or key not in standard_keys -%}
{%- if ns.found_first %},{% endif -%}
{%- set ns.found_first = true -%}
{{ key }}:{
{%- if value['description'] -%}
description:<|"|>{{ value['description'] }}<|"|>
{%- set add_comma = true -%}
{%- endif -%}
{%- if value['type'] | upper == 'STRING' -%}
{%- if value['enum'] -%}
{%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
enum:{{ format_argument(value['enum']) }}
{%- endif -%}
{%- elif value['type'] | upper == 'ARRAY' -%}
{%- if value['items'] is mapping and value['items'] -%}
{%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
items:{
{%- set ns_items = namespace(found_first=false) -%}
{%- for item_key, item_value in value['items'] | dictsort -%}
{%- if item_value is not none -%}
{%- if ns_items.found_first %},{% endif -%}
{%- set ns_items.found_first = true -%}
{%- if item_key == 'properties' -%}
properties:{
{%- if item_value is mapping -%}
{{- format_parameters(item_value, value['items']['required'] | default([])) -}}
{%- endif -%}
}
{%- elif item_key == 'required' -%}
required:[
{%- for req_item in item_value -%}
<|"|>{{- req_item -}}<|"|>
{%- if not loop.last %},{% endif -%}
{%- endfor -%}
]
{%- elif item_key == 'type' -%}
{%- if item_value is string -%}
type:{{ format_argument(item_value | upper) }}
{%- else -%}
type:{{ format_argument(item_value | map('upper') | list) }}
{%- endif -%}
{%- else -%}
{{ item_key }}:{{ format_argument(item_value) }}
{%- endif -%}
{%- endif -%}
{%- endfor -%}
}
{%- endif -%}
{%- endif -%}
{%- if value['nullable'] %}
{%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
nullable:true
{%- endif -%}
{%- if value['type'] | upper == 'OBJECT' -%}
{%- if value['properties'] is defined and value['properties'] is mapping -%}
{%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
properties:{
{{- format_parameters(value['properties'], value['required'] | default([])) -}}
}
{%- elif value is mapping -%}
{%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
properties:{
{{- format_parameters(value, value['required'] | default([]), filter_keys=true) -}}
}
{%- endif -%}
{%- if value['required'] -%}
{%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
required:[
{%- for item in value['required'] | default([]) -%}
<|"|>{{- item -}}<|"|>
{%- if not loop.last %},{% endif -%}
{%- endfor -%}
]
{%- endif -%}
{%- endif -%}
{%- if add_comma %},{%- else -%} {%- set add_comma = true -%} {% endif -%}
type:<|"|>{{ value['type'] | upper }}<|"|>}
{%- endif -%}
{%- endfor -%}
{%- endmacro -%}
{%- macro format_function_declaration(tool_data) -%}
declaration:{{- tool_data['function']['name'] -}}{description:<|"|>{{- tool_data['function']['description'] -}}<|"|>
{%- set params = tool_data['function']['parameters'] -%}
{%- if params -%}
,parameters:{
{%- if params['properties'] -%}
properties:{ {{- format_parameters(params['properties'], params['required']) -}} },
{%- endif -%}
{%- if params['required'] -%}
required:[
{%- for item in params['required'] -%}
<|"|>{{- item -}}<|"|>
{{- ',' if not loop.last -}}
{%- endfor -%}
],
{%- endif -%}
{%- if params['type'] -%}
type:<|"|>{{- params['type'] | upper -}}<|"|>}
{%- endif -%}
{%- endif -%}
{%- if 'response' in tool_data['function'] -%}
{%- set response_declaration = tool_data['function']['response'] -%}
,response:{
{%- if response_declaration['description'] -%}
description:<|"|>{{- response_declaration['description'] -}}<|"|>,
{%- endif -%}
{%- if response_declaration['type'] | upper == 'OBJECT' -%}
type:<|"|>{{- response_declaration['type'] | upper -}}<|"|>}
{%- endif -%}
{%- endif -%}
}
{%- endmacro -%}
{%- macro format_argument(argument, escape_keys=True) -%}
{%- if argument is none -%}
{{- 'null' -}}
{%- elif argument is string -%}
{{- '<|"|>' + argument + '<|"|>' -}}
{%- elif argument is boolean -%}
{{- 'true' if argument else 'false' -}}
{%- elif argument is mapping -%}
{{- '{' -}}
{%- set ns = namespace(found_first=false) -%}
{%- for key, value in argument | dictsort -%}
{%- if ns.found_first %},{% endif -%}
{%- set ns.found_first = true -%}
{%- if escape_keys -%}
{{- '<|"|>' + key + '<|"|>' -}}
{%- else -%}
{{- key -}}
{%- endif -%}
:{{- format_argument(value, escape_keys=escape_keys) -}}
{%- endfor -%}
{{- '}' -}}
{%- elif argument is sequence -%}
{{- '[' -}}
{%- for item in argument -%}
{{- format_argument(item, escape_keys=escape_keys) -}}
{%- if not loop.last %},{% endif -%}
{%- endfor -%}
{{- ']' -}}
{%- else -%}
{{- argument -}}
{%- endif -%}
{%- endmacro -%}
{%- macro strip_thinking(text) -%}
{%- set ns = namespace(result='') -%}
{%- for part in text.split('<channel|>') -%}
{%- if '<|channel>' in part -%}
{%- set ns.result = ns.result + part.split('<|channel>')[0] -%}
{%- else -%}
{%- set ns.result = ns.result + part -%}
{%- endif -%}
{%- endfor -%}
{{- ns.result | trim -}}
{%- endmacro -%}
{%- macro format_tool_response_block(tool_name, response) -%}
{{- '<|tool_response>' -}}
{%- if response is mapping -%}
{{- 'response:' + tool_name + '{' -}}
{%- for key, value in response | dictsort -%}
{{- key -}}:{{- format_argument(value, escape_keys=False) -}}
{%- if not loop.last %},{% endif -%}
{%- endfor -%}
{{- '}' -}}
{%- else -%}
{{- 'response:' + tool_name + '{value:' + format_argument(response, escape_keys=False) + '}' -}}
{%- endif -%}
{{- '<tool_response|>' -}}
{%- endmacro -%}
{#- ===== SETUP ===== -#}
{%- set ns = namespace(prev_message_type=None, prev_non_tool_role=None) -%}
{%- set loop_messages = messages -%}
{%- set enable_thinking = enable_thinking | default(false) -%}
{#- Unsloth Studio: preserve_thinking defaults OFF (upstream PR #118 defaults true). -#}
{%- set preserve_thinking = preserve_thinking | default(false) -%}
{{- bos_token -}}
{#- Handle System/Tool Definitions Block -#}
{%- if enable_thinking or tools or (messages and messages[0]['role'] in ['system', 'developer']) -%}
{{- '<|turn>system\n' -}}
{#- Inject Thinking token at the very top of the FIRST system turn -#}
{%- if enable_thinking -%}
{{- '<|think|>\n' -}}
{%- set ns.prev_message_type = 'think' -%}
{%- endif -%}
{%- if messages and messages[0]['role'] in ['system', 'developer'] -%}
{%- if messages[0]['content'] is string -%}
{{- messages[0]['content'] | trim -}}
{%- elif messages[0]['content'] is sequence -%}
{%- for item in messages[0]['content'] -%}
{{- item['text'] | trim + ' '-}}
{%- endfor -%}
{%- endif -%}
{%- set loop_messages = messages[1:] -%}
{%- endif -%}
{%- if tools -%}
{%- for tool in tools %}
{{- '<|tool>' -}}
{{- format_function_declaration(tool) | trim -}}
{{- '<tool|>' -}}
{%- endfor %}
{%- set ns.prev_message_type = 'tool' -%}
{%- endif -%}
{{- '<turn|>\n' -}}
{%- endif %}
{#- Pre-scan: find last user message index for reasoning guard -#}
{%- set ns_turn = namespace(last_user_idx=-1) -%}
{%- for i in range(loop_messages | length) -%}
{%- if loop_messages[i]['role'] == 'user' -%}
{%- set ns_turn.last_user_idx = i -%}
{%- endif -%}
{%- endfor -%}
{#- Loop through messages -#}
{%- for message in loop_messages -%}
{%- if message['role'] != 'tool' -%}
{%- set ns.prev_message_type = None -%}
{%- set role = 'model' if message['role'] == 'assistant' else message['role'] -%}
{#- Detect continuation using tracked state - O(1) instead of O(n) backward scan -#}
{%- set continue_same_model_turn = (role == 'model' and ns.prev_non_tool_role == 'assistant') -%}
{%- if not continue_same_model_turn -%}
{{- '<|turn>' + role + '\n' }}
{%- endif -%}
{#- Render reasoning/reasoning_content as thinking channel (tool-call turns only) -#}
{%- set thinking_text = message.get('reasoning') or message.get('reasoning_content') -%}
{%- set thinking_gate = (loop.index0 > ns_turn.last_user_idx) or preserve_thinking -%}
{%- if thinking_text and thinking_gate and message.get('tool_calls') -%}
{{- '<|channel>thought\n' + thinking_text + '\n<channel|>' -}}
{%- endif -%}
{%- if message.get('tool_calls') -%}
{%- for tool_call in message.get('tool_calls') -%}
{%- set function = tool_call['function'] -%}
{{- '<|tool_call>call:' + function['name'] + '{' -}}
{%- if function['arguments'] is mapping -%}
{%- set ns_args = namespace(found_first=false) -%}
{%- for key, value in function['arguments'] | dictsort -%}
{%- if ns_args.found_first %},{% endif -%}
{%- set ns_args.found_first = true -%}
{{- key -}}:{{- format_argument(value, escape_keys=False) -}}
{%- endfor -%}
{%- elif function['arguments'] is none -%}
{%- else -%}
{{- raise_exception(
"chat_template: tool_calls[].function.arguments must be a "
"JSON object (mapping), not a string. Deserialize arguments "
"before passing to the template."
) -}}
{%- endif -%}
{{- '}<tool_call|>' -}}
{%- endfor -%}
{%- set ns.prev_message_type = 'tool_call' -%}
{%- endif -%}
{%- set ns_tr_out = namespace(flag=false) -%}
{%- if message.get('tool_responses') -%}
{#- Legacy: tool_responses embedded on the assistant message (Google/Gemma native) -#}
{%- for tool_response in message.get('tool_responses') -%}
{{- format_tool_response_block(tool_response['name'] | default('unknown', true), tool_response['response']) -}}
{%- set ns_tr_out.flag = true -%}
{%- set ns.prev_message_type = 'tool_response' -%}
{%- endfor -%}
{%- elif message.get('tool_calls') -%}
{#- OpenAI Chat Completions: forward-scan consecutive role:tool messages -#}
{%- set ns_tool_scan = namespace(stopped=false) -%}
{%- for k in range(loop.index0 + 1, loop_messages | length) -%}
{%- if ns_tool_scan.stopped -%}
{%- elif loop_messages[k]['role'] != 'tool' -%}
{%- set ns_tool_scan.stopped = true -%}
{%- else -%}
{%- set follow = loop_messages[k] -%}
{#- Resolve tool_call_id to function name -#}
{%- set ns_tname = namespace(name=follow.get('name') or 'unknown') -%}
{%- for tc in message.get('tool_calls') -%}
{%- if tc.get('id') == follow.get('tool_call_id') -%}
{%- set ns_tname.name = tc['function']['name'] -%}
{%- endif -%}
{%- endfor -%}
{#- Handle content as string or content-parts array -#}
{%- set tool_body = follow.get('content') -%}
{%- if tool_body is string -%}
{{- format_tool_response_block(ns_tname.name, tool_body) -}}
{%- elif tool_body is sequence and tool_body is not string -%}
{%- set ns_txt = namespace(s='') -%}
{%- for part in tool_body -%}
{%- if part.get('type') == 'text' -%}
{%- set ns_txt.s = ns_txt.s + (part.get('text') | default('')) -%}
{%- endif -%}
{%- endfor -%}
{{- format_tool_response_block(ns_tname.name, ns_txt.s) -}}
{%- for part in tool_body -%}
{%- if part.get('type') in ['image', 'image_url'] -%}
{{- '<|image|>' -}}
{%- elif part.get('type') in ['audio', 'input_audio'] -%}
{{- '<|audio|>' -}}
{%- elif part.get('type') == 'video' -%}
{{- '<|video|>' -}}
{%- endif -%}
{%- endfor -%}
{%- else -%}
{{- format_tool_response_block(ns_tname.name, tool_body) -}}
{%- endif -%}
{%- set ns_tr_out.flag = true -%}
{%- set ns.prev_message_type = 'tool_response' -%}
{%- endif -%}
{%- endfor -%}
{%- endif -%}
{%- set captured_content -%}
{%- if message.get('content') is string -%}
{%- if role == 'model' -%}
{{- strip_thinking(message['content']) -}}
{%- else -%}
{{- message['content'] | trim -}}
{%- endif -%}
{%- elif message.get('content') is sequence -%}
{%- for item in message['content'] -%}
{%- if item.get('type') == 'text' -%}
{%- if role == 'model' -%}
{{- strip_thinking(item['text']) -}}
{%- else -%}
{{- item['text'] | trim -}}
{%- endif -%}
{%- elif item.get('type') in ['image', 'image_url'] -%}
{{- '<|image|>' -}}
{%- elif item.get('type') in ['audio', 'input_audio'] -%}
{{- '<|audio|>' -}}
{%- elif item.get('type') == 'video' -%}
{{- '<|video|>' -}}
{%- endif -%}
{%- endfor -%}
{%- endif -%}
{%- endset -%}
{{- captured_content -}}
{%- set has_content = captured_content | trim | length > 0 -%}
{#- Forward-scan: find next non-tool message role for continuation detection -#}
{%- set next_nt = namespace(role=None, found=false) -%}
{%- for j in range(loop.index0 + 1, loop_messages | length) -%}
{%- if not next_nt.found -%}
{%- if loop_messages[j]['role'] != 'tool' -%}
{%- set next_nt.role = loop_messages[j]['role'] -%}
{%- set next_nt.found = true -%}
{%- endif -%}
{%- endif -%}
{%- endfor -%}
{%- set continues_into_next = (
role == 'model'
and next_nt.role == 'assistant'
and not message.get('tool_calls')
and not ns_tr_out.flag
) -%}
{%- if ns.prev_message_type == 'tool_call' and not ns_tr_out.flag -%}
{{- '<|tool_response>' -}}
{%- elif continues_into_next -%}
{{- '\n' -}}
{%- elif not (ns_tr_out.flag and not has_content) -%}
{{- '<turn|>\n' -}}
{%- endif -%}
{#- Track previous non-tool role for next iteration (avoids O(n) backward scan) -#}
{%- set ns.prev_non_tool_role = message['role'] -%}
{%- endif -%}
{%- endfor -%}
{%- if add_generation_prompt -%}
{%- if ns.prev_message_type != 'tool_response' and ns.prev_message_type != 'tool_call' -%}
{{- '<|turn>model\n' -}}
{%- endif -%}
{%- if not enable_thinking -%}
{#- Suppress thinking - but not when awaiting tool responses -#}
{%- if ns.prev_message_type != 'tool_call' -%}
{{- '<|channel>thought\n<channel|>' -}}
{%- endif -%}
{%- endif -%}
{%- endif -%}

9
studio/backend/assets/configs/inference_defaults.json
View file

@ -277,6 +277,13 @@
"min_p": 0.01,
"repetition_penalty": 1.0
},
"minimax-m2.7": {
"temperature": 1.0,
"top_p": 0.95,
"top_k": 40,
"min_p": 0.01,
"repetition_penalty": 1.0
},
"minimax-m2.5": {
"temperature": 1.0,
"top_p": 0.95,
@ -390,7 +397,7 @@
"deepseek-r1", "deepseek-v3", "deepseek-ocr",
"glm-5", "glm-4",
"nemotron",
"minimax-m2.5", "minimax",
"minimax-m2.7", "minimax-m2.5", "minimax",
"gpt-oss", "granite-4",
"kimi-k2", "kimi",
"lfm2", "smollm", "olmo", "falcon", "ernie", "seed", "grok", "mimo"

4
studio/backend/auth/__init__.py
View file

@ -1,9 +1,7 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""
Authentication module for JWT-based auth with SQLite storage.
"""
"""Authentication module for JWT-based auth with SQLite storage."""
from .authentication import (
create_access_token,

32
studio/backend/auth/authentication.py
View file

@ -58,7 +58,7 @@ def create_access_token(
"""
Create a signed JWT for the given subject (e.g. username).
Tokens are valid across restarts because the signing secret is stored in SQLite.
Valid across restarts: the signing secret is stored in SQLite.
"""
to_encode = {"sub": subject}
if desktop:
@ -100,7 +100,7 @@ def create_refresh_token(subject: str, *, desktop: bool = False) -> str:
"""
Create a random refresh token, store its hash in SQLite, and return it.
Refresh tokens are opaque (not JWTs) and expire after REFRESH_TOKEN_EXPIRE_DAYS.
Refresh tokens are opaque (not JWTs); expire after REFRESH_TOKEN_EXPIRE_DAYS.
"""
token = secrets.token_urlsafe(48)
expires_at = datetime.now(timezone.utc) + timedelta(days = REFRESH_TOKEN_EXPIRE_DAYS)
@ -108,14 +108,12 @@ def create_refresh_token(subject: str, *, desktop: bool = False) -> str:
return token
def refresh_access_token(
refresh_token: str,
) -> Tuple[Optional[str], Optional[str], bool]:
def refresh_access_token(refresh_token: str) -> Tuple[Optional[str], Optional[str], bool]:
"""
Validate a refresh token and issue a new access token.
The refresh token itself is NOT consumed — it stays valid until expiry.
Returns a new access_token or None if the refresh token is invalid/expired.
The refresh token is NOT consumed; it stays valid until expiry.
Returns a new access_token, or None if the refresh token is invalid/expired.
"""
verified = verify_refresh_token(refresh_token)
if verified is None:
@ -130,16 +128,14 @@ def refresh_access_token(
def reload_secret() -> None:
"""
Keep legacy API compatibility for callers expecting auth storage init.
Legacy API compat for callers expecting auth storage init.
Auth now resolves the current signing secret directly from SQLite.
"""
load_jwt_secret()
async def get_current_subject(
credentials: HTTPAuthorizationCredentials = Depends(security),
) -> str:
async def get_current_subject(credentials: HTTPAuthorizationCredentials = Depends(security)) -> str:
"""Validate JWT and require the password-change flow to be completed."""
return await _get_current_subject(
credentials,
@ -158,19 +154,9 @@ async def get_current_subject_allow_password_change(
async def _get_current_subject(
credentials: HTTPAuthorizationCredentials,
*,
allow_password_change: bool,
credentials: HTTPAuthorizationCredentials, *, allow_password_change: bool
) -> str:
"""
FastAPI dependency to validate the JWT and return the subject.
Use this as a dependency on routes that should be protected, e.g.:
@router.get("/secure")
async def secure_endpoint(current_subject: str = Depends(get_current_subject)):
...
"""
"""FastAPI dependency: validate the JWT and return the subject. Use on protected routes."""
token = credentials.credentials
# --- API key path (sk-unsloth-...) ---

2
studio/backend/auth/hashing.py
View file

@ -23,7 +23,7 @@ def hash_password(password: str, salt: str | None = None) -> Tuple[str, str]:
"sha256",
password.encode("utf-8"),
salt.encode("utf-8"),
100_000, # 100k iterations
100_000,
)
return salt, dk.hex()

133
studio/backend/auth/storage.py
View file

@ -1,14 +1,14 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""
SQLite storage for authentication data (user credentials + JWT secret).
"""
"""SQLite storage for auth data (user credentials + JWT secret)."""
import hashlib
import hmac
import os
import secrets
import sqlite3
import threading
from datetime import datetime, timezone
from typing import Optional, Tuple
@ -17,42 +17,40 @@ from utils.paths import auth_db_path, ensure_dir
DB_PATH = auth_db_path()
DEFAULT_ADMIN_USERNAME = "unsloth"
# Plaintext bootstrap password file — lives beside auth.db, deleted on
# first password change so the credential never lingers on disk.
# Plaintext bootstrap password file beside auth.db, deleted on first password
# change so the credential never lingers on disk.
_BOOTSTRAP_PW_PATH = DB_PATH.parent / ".bootstrap_password"
# In-process cache so we don't re-read the file on every HTML serve.
# In-process cache to avoid re-reading the file on every HTML serve.
_bootstrap_password: Optional[str] = None
def generate_bootstrap_password() -> str:
"""Generate a 4-word diceware passphrase and persist it to disk.
The passphrase is written to ``_BOOTSTRAP_PW_PATH`` so that it
survives server restarts (the DB only stores the *hash*). On
subsequent calls / restarts, the persisted value is returned.
Persisted (the DB stores only the hash) so it survives restarts; later
calls return the persisted value.
"""
global _bootstrap_password
# 1. Already cached in this process?
# Cached in this process?
if _bootstrap_password is not None:
return _bootstrap_password
# 2. Already persisted from a previous run?
# Persisted from a previous run?
if _BOOTSTRAP_PW_PATH.is_file():
_bootstrap_password = _BOOTSTRAP_PW_PATH.read_text().strip()
if _bootstrap_password:
return _bootstrap_password
# 3. First-ever startup — generate a fresh passphrase.
# First startup: generate a fresh passphrase.
import diceware
_bootstrap_password = diceware.get_passphrase(
options = diceware.handle_options(args = ["-n", "4", "-d", "", "-c"])
)
# Persist so the *same* passphrase is used if the server restarts
# before the user changes the password.
# Persist so the same passphrase survives restarts until password change.
ensure_dir(_BOOTSTRAP_PW_PATH.parent)
_BOOTSTRAP_PW_PATH.write_text(_bootstrap_password)
try:
@ -88,21 +86,12 @@ def clear_bootstrap_password() -> None:
def _hash_token(token: str) -> str:
"""SHA-256 hash helper used for refresh token storage.
"""SHA-256 hash helper for refresh token storage.
Plain SHA-256 is intentional here: refresh tokens are high-entropy
random strings from ``secrets.token_urlsafe(48)`` (384 bits of
entropy), so a slow KDF (Argon2 / bcrypt / PBKDF2) provides zero
additional security — no attacker can brute-force 2^384 regardless
of hash speed — while adding tens of ms of CPU to every refresh.
See the OWASP Password Storage Cheat Sheet on fast-vs-slow hashing
of high-entropy inputs.
API keys use the separate ``_pbkdf2_api_key`` helper below, which
runs PBKDF2-HMAC-SHA256 with a persistent server-side salt — not
for cryptographic reasons (128-bit random tokens don't need slow
hashing), but because CodeQL's ``py/weak-sensitive-data-hashing``
query mislabels API keys as passwords and demands a KDF.
Plain SHA-256 is intentional: refresh tokens are 384-bit random strings, so
a slow KDF adds no security while costing per-refresh latency. API keys use
the separate ``_pbkdf2_api_key`` helper, only to satisfy CodeQL's
``py/weak-sensitive-data-hashing`` query, not for crypto reasons.
"""
return hashlib.sha256(token.encode("utf-8")).hexdigest()
@ -151,13 +140,9 @@ def get_connection() -> sqlite3.Connection:
);
"""
)
api_key_columns = {
row["name"] for row in conn.execute("PRAGMA table_info(api_keys)")
}
api_key_columns = {row["name"] for row in conn.execute("PRAGMA table_info(api_keys)")}
if "is_internal" not in api_key_columns:
conn.execute(
"ALTER TABLE api_keys ADD COLUMN is_internal INTEGER NOT NULL DEFAULT 0"
)
conn.execute("ALTER TABLE api_keys ADD COLUMN is_internal INTEGER NOT NULL DEFAULT 0")
conn.execute(
"""
CREATE TABLE IF NOT EXISTS app_secrets (
@ -171,35 +156,28 @@ def get_connection() -> sqlite3.Connection:
conn.execute(
"ALTER TABLE auth_user ADD COLUMN must_change_password INTEGER NOT NULL DEFAULT 0"
)
refresh_columns = {
row["name"] for row in conn.execute("PRAGMA table_info(refresh_tokens)")
}
refresh_columns = {row["name"] for row in conn.execute("PRAGMA table_info(refresh_tokens)")}
if "is_desktop" not in refresh_columns:
conn.execute(
"ALTER TABLE refresh_tokens ADD COLUMN is_desktop INTEGER NOT NULL DEFAULT 0"
)
conn.execute("ALTER TABLE refresh_tokens ADD COLUMN is_desktop INTEGER NOT NULL DEFAULT 0")
conn.commit()
return conn
# ── API-key PBKDF2 salt ────────────────────────────────────────────────
#
# Module-level cache for the persistent API-key PBKDF2 salt. Populated
# lazily on first use via ``_get_or_create_api_key_pbkdf2_salt``. Not
# protected by a lock because (a) the ``INSERT OR IGNORE`` provides
# atomicity at the SQLite layer and (b) concurrent populations converge
# on the same value, so the worst case is a harmless duplicate read on
# startup.
# Module-level cache for the persistent API-key PBKDF2 salt, populated lazily
# via ``_get_or_create_api_key_pbkdf2_salt``. No lock needed: (a) ``INSERT OR
# IGNORE`` is atomic at the SQLite layer and (b) concurrent populations
# converge on the same value, so the worst case is a harmless duplicate read
# on startup.
_api_key_pbkdf2_salt_cache: Optional[bytes] = None
def _get_or_create_api_key_pbkdf2_salt() -> bytes:
"""Return the persistent API-key PBKDF2 salt, generating it once if missing.
Stored as a hex-encoded 32-byte random value in the ``app_secrets``
table under key ``"api_key_pbkdf2_salt"``. Regenerated only if the row
is missing (i.e. fresh install, or operator manually deleted the row
and accepts invalidating existing API keys).
Hex-encoded 32-byte random value in ``app_secrets``. Regenerated only when
the row is missing (fresh install, or operator deleted it).
"""
global _api_key_pbkdf2_salt_cache
if _api_key_pbkdf2_salt_cache is not None:
@ -241,22 +219,10 @@ _DESKTOP_SECRET_CREATED_AT_KEY = "desktop_secret_created_at"
def _pbkdf2_api_key(raw_key: str) -> str:
"""PBKDF2-HMAC-SHA256 an API key with a persistent server-side salt.
Used for API-key storage ONLY, not refresh tokens. Matches the
PBKDF2 algorithm + iteration count used by the password hasher in
``auth/hashing.py`` so the codebase is consistent on which KDF it
uses for credential storage.
Notes on why a slow KDF here is *only* a CodeQL appeasement and
*not* a cryptographic requirement: API keys are cryptographically
random 128-bit tokens (via ``secrets.token_hex``), so brute force
against 2^128 is infeasible regardless of hash speed. CodeQL's
``py/weak-sensitive-data-hashing`` query mislabels these tokens as
"password" sensitive data and then demands a KDF from its
allowlist (Argon2 / scrypt / bcrypt / PBKDF2). Per the query's
own recommendation page we use PBKDF2. The persistent salt is
still loaded from ``app_secrets`` so an attacker dumping the
``api_keys`` table alone cannot derive hashes for candidate
tokens without also obtaining the salt row.
For API-key storage ONLY, not refresh tokens. The slow KDF is only to
appease CodeQL's ``py/weak-sensitive-data-hashing`` query, not a crypto
requirement (API keys are random 128-bit tokens). The salt lives in
``app_secrets`` so dumping ``api_keys`` alone can't derive hashes.
"""
salt = _get_or_create_api_key_pbkdf2_salt()
dk = hashlib.pbkdf2_hmac(
@ -272,6 +238,29 @@ def _pbkdf2_desktop_secret(raw_secret: str) -> str:
return _pbkdf2_api_key(raw_secret)
# Memoize the deterministic raw-key -> PBKDF2-hash derivation so the 100k-round
# KDF runs once per key instead of on every authenticated request. Keyed by a
# salted HMAC of the key (not the key itself); revocation/expiry are still
# enforced by the SQLite read on every call, so a cache hit only skips the KDF.
# Only keys present in the DB are cached, so unknown-key spam can't grow it.
_api_key_hash_cache: dict[str, str] = {}
_API_KEY_HASH_CACHE_MAX = 4096
_api_key_hash_cache_lock = threading.Lock()
def _api_key_cache_id(raw_key: str) -> str:
"""Cache id for a raw key: salted HMAC-SHA256 (not the key itself)."""
return hmac.new(
_get_or_create_api_key_pbkdf2_salt(), raw_key.encode("utf-8"), hashlib.sha256
).hexdigest()
def _reset_api_key_hash_cache() -> None:
"""Drop memoized derivations (tests / salt change)."""
with _api_key_hash_cache_lock:
_api_key_hash_cache.clear()
def is_initialized() -> bool:
"""Check if auth is ready for login (at least one user exists in DB)."""
conn = get_connection()
@ -521,7 +510,7 @@ def verify_refresh_token(token: str) -> Optional[Tuple[str, bool]]:
token_hash = _hash_token(token)
conn = get_connection()
try:
# Clean up any expired tokens while we're here
# Opportunistically clean up expired tokens
conn.execute(
"DELETE FROM refresh_tokens WHERE expires_at < ?",
(datetime.now(timezone.utc).isoformat(),),
@ -740,7 +729,9 @@ def validate_api_key(raw_key: str) -> Optional[str]:
Also updates ``last_used_at`` on success.
"""
key_hash = _pbkdf2_api_key(raw_key)
cache_id = _api_key_cache_id(raw_key)
cached_hash = _api_key_hash_cache.get(cache_id)
key_hash = cached_hash if cached_hash is not None else _pbkdf2_api_key(raw_key)
conn = get_connection()
try:
cur = conn.execute(
@ -750,6 +741,12 @@ def validate_api_key(raw_key: str) -> Optional[str]:
row = cur.fetchone()
if row is None:
return None
# Real key: memoize so later requests skip the KDF. Bounded; clear on overflow.
if cached_hash is None:
with _api_key_hash_cache_lock:
if len(_api_key_hash_cache) >= _API_KEY_HASH_CACHE_MAX:
_api_key_hash_cache.clear()
_api_key_hash_cache[cache_id] = key_hash
if not row["is_active"]:
return None
if row["expires_at"] is not None:

375
studio/backend/cloudflare_tunnel.py Normal file
View file

@ -0,0 +1,375 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Free Cloudflare quick tunnel for Studio's 0.0.0.0 launches.
The raw http://<ip>:<port> is often unreachable (https-vs-http, blocked ports,
closed security groups); a cloudflared quick tunnel gives a free
https://*.trycloudflare.com URL that works anywhere, with no account or domain.
Best-effort throughout: any failure collapses to "no URL" and Studio keeps
running. Stdlib only (back-end imports are lazy) so it is safe to import early.
"""
from __future__ import annotations
import os
import platform
import re
import shutil
import subprocess
import sys
import threading
from pathlib import Path
from typing import Optional, Tuple
# cloudflared logs the quick-tunnel URL; match only the URL so we do not depend
# on the surrounding wording, which Cloudflare may change. The negative lookahead
# drops cloudflared's own API host, which appears in failure lines such as
# failed to request quick Tunnel: Post "https://api.trycloudflare.com/tunnel"
# and must never be mistaken for a usable tunnel URL.
_URL_RE = re.compile(r"https://(?!api\.)[A-Za-z0-9-]+\.trycloudflare\.com")
# cloudflared logs this once per edge connection it establishes. Until at least
# one appears the quick-tunnel URL returns Cloudflare error 1033 (HTTP 530), so
# we wait for it before advertising the URL.
_REGISTERED_MARKER = "Registered tunnel connection"
_RELEASE_BASE = "https://github.com/cloudflare/cloudflared/releases/latest/download"
_READY_TIMEOUT = 15.0 # seconds to wait for the URL + a registered edge connection
_DOWNLOAD_TIMEOUT = 60 # urlopen timeout for the one-time binary download
def _windows_hidden_kwargs() -> dict:
"""Suppress a child console window on Windows; no-op elsewhere."""
if sys.platform != "win32":
return {}
flags = getattr(subprocess, "CREATE_NO_WINDOW", 0)
return {"creationflags": flags} if flags else {}
def _asset_name() -> Optional[Tuple[str, bool]]:
"""(release asset filename, is_tgz) for this OS/arch, or None if unsupported."""
system = platform.system().lower()
machine = platform.machine().lower()
is_x64 = machine in ("x86_64", "amd64", "x64")
is_arm64 = machine in ("aarch64", "arm64")
is_x86 = machine in ("i386", "i686", "x86")
if system == "linux":
if is_x64:
return ("cloudflared-linux-amd64", False)
if is_arm64:
return ("cloudflared-linux-arm64", False)
elif system == "darwin":
if is_arm64:
return ("cloudflared-darwin-arm64.tgz", True)
if is_x64:
return ("cloudflared-darwin-amd64.tgz", True)
elif system == "windows":
if is_x64:
return ("cloudflared-windows-amd64.exe", False)
if is_x86:
return ("cloudflared-windows-386.exe", False)
return None
def _cache_path() -> Optional[Path]:
"""studio_bin_root()/cloudflared(.exe), or None if the studio home is unresolvable."""
try:
from utils.paths.storage_roots import studio_bin_root # lazy: backend-only import
except Exception:
return None
name = "cloudflared.exe" if sys.platform == "win32" else "cloudflared"
return studio_bin_root() / name
def find_cloudflared() -> Optional[str]:
"""Locate an existing cloudflared: PATH first, then the Studio bin cache."""
on_path = shutil.which("cloudflared")
if on_path:
return on_path
cached = _cache_path()
if cached is not None and cached.is_file() and os.access(cached, os.X_OK):
return str(cached)
return None
def _download(url: str, dest: Path) -> bool:
"""Download url to dest via urllib (temp file + atomic rename). Best-effort -> bool."""
import tempfile
import urllib.request
tmp_path: Optional[Path] = None
try:
dest.parent.mkdir(parents = True, exist_ok = True)
with tempfile.NamedTemporaryFile(
prefix = dest.name + ".tmp-", dir = dest.parent, delete = False
) as handle:
tmp_path = Path(handle.name)
# GitHub's CDN 403s the default Python-urllib User-Agent.
req = urllib.request.Request(url, headers = {"User-Agent": "unsloth-studio"})
with urllib.request.urlopen(req, timeout = _DOWNLOAD_TIMEOUT) as response:
shutil.copyfileobj(response, handle)
if tmp_path.stat().st_size == 0:
raise RuntimeError("empty download")
os.replace(tmp_path, dest)
return True
except Exception:
if tmp_path is not None:
try:
tmp_path.unlink(missing_ok = True)
except Exception:
pass
return False
def _extract_tgz_member(tgz_path: Path, dest: Path) -> bool:
"""Extract just the `cloudflared` member from a darwin .tgz to dest.
Rejects absolute paths and `..` traversal so a hostile archive cannot write
outside dest. Best-effort -> bool.
"""
import tarfile
try:
with tarfile.open(tgz_path, "r:gz") as tar:
member = None
for m in tar.getmembers():
if not m.isfile() or os.path.basename(m.name) != "cloudflared":
continue
if m.name.startswith("/") or ".." in Path(m.name).parts:
continue
member = m
break
if member is None:
return False
src = tar.extractfile(member)
if src is None:
return False
with src, open(dest, "wb") as out:
shutil.copyfileobj(src, out)
return True
except Exception:
return False
def ensure_cloudflared() -> Optional[str]:
"""Return a cloudflared path, downloading + caching the binary once if missing."""
existing = find_cloudflared()
if existing:
return existing
asset = _asset_name()
cached = _cache_path()
if asset is None or cached is None:
return None
name, is_tgz = asset
url = f"{_RELEASE_BASE}/{name}"
try:
cached.parent.mkdir(parents = True, exist_ok = True)
if is_tgz:
tgz = cached.with_suffix(".tgz")
if not _download(url, tgz) or not _extract_tgz_member(tgz, cached):
tgz.unlink(missing_ok = True)
return None
tgz.unlink(missing_ok = True)
elif not _download(url, cached):
return None
if sys.platform != "win32":
os.chmod(cached, 0o755)
return str(cached)
except Exception:
return None
class CloudflareTunnel:
"""A cloudflared quick tunnel to http://localhost:<port>. Best-effort throughout.
Use localhost (not the wildcard bind) as the tunnel origin so cloudflared's
upstream stays local-only.
"""
def __init__(
self,
port: int,
binary: str,
protocol: Optional[str] = None,
):
self.port = port
self.binary = binary
# None lets cloudflared pick its default (quic, with its own http2
# fallback); set to "http2" to force it when quic is blocked.
self.protocol = protocol
self._proc: Optional[subprocess.Popen] = None
self._lock = threading.Lock()
self._stopped = False
self._url_event = threading.Event()
self._ready_event = threading.Event()
self.url: Optional[str] = None
self.ready = False
self.error: Optional[str] = None
def start(self) -> None:
cmd = [
self.binary,
"tunnel",
"--url",
f"http://localhost:{self.port}",
"--no-autoupdate",
]
if self.protocol:
cmd += ["--protocol", self.protocol]
with self._lock:
# A stop() that landed before us (e.g. a shutdown in the caller's
# register->start window) marks the tunnel stopped; spawning now would
# orphan a process nobody owns, so refuse.
if self._stopped:
return
proc = subprocess.Popen(
cmd,
stdout = subprocess.PIPE,
stderr = subprocess.STDOUT,
stdin = subprocess.DEVNULL,
text = True,
errors = "replace",
bufsize = 1,
**_windows_hidden_kwargs(),
)
self._proc = proc
threading.Thread(
target = self._reader, args = (proc,), name = "cloudflared-reader", daemon = True
).start()
def _reader(self, proc: subprocess.Popen) -> None:
# Drain cloudflared's output: capture the first trycloudflare URL and the
# first edge-connection registration, and keep draining so it never
# blocks on a full pipe.
try:
if proc.stdout is not None:
for line in proc.stdout:
if self.url is None:
match = _URL_RE.search(line)
if match:
self.url = match.group(0)
self._url_event.set()
if not self.ready and _REGISTERED_MARKER in line:
self.ready = True
self._ready_event.set()
except Exception:
pass
finally:
# stdout closed -> cloudflared has exited. Record why, and unblock any
# waiters at once instead of letting them wait out the full timeout.
if self.url is None:
self.error = "cloudflared exited before emitting a tunnel URL"
elif not self.ready:
self.error = "cloudflared exited before the tunnel connection registered"
self._url_event.set()
self._ready_event.set()
def wait_for_ready(self, timeout: float = _READY_TIMEOUT) -> Optional[str]:
"""Block until the tunnel is actually serving -- the URL has been minted
*and* at least one edge connection has registered -- or until timeout.
Returns the URL only when ready, so callers never advertise a URL that
would return Cloudflare error 1033 (HTTP 530)."""
self._ready_event.wait(timeout)
return self.url if self.ready else None
def stop(self) -> None:
"""Terminate the tunnel. Idempotent and safe to call from a signal handler."""
with self._lock:
# Mark stopped so a start() racing behind us refuses to spawn.
self._stopped = True
proc, self._proc = self._proc, None
if proc is None:
return
try:
if proc.poll() is None:
proc.terminate()
try:
proc.wait(timeout = 5)
except subprocess.TimeoutExpired:
proc.kill()
try:
proc.wait(timeout = 5)
except Exception:
pass
except Exception:
pass
# Single serving process per Studio launch, so one module-level tunnel handle is
# enough; the lock guards the start/stop/shutdown races.
_active_tunnel: Optional[CloudflareTunnel] = None
_active_lock = threading.Lock()
# Latched by stop_studio_tunnel so a shutdown landing *between* a start's retry
# attempts aborts the loop instead of starting a tunnel nobody will ever stop.
_shutdown_requested = False
def start_studio_tunnel(port: int, timeout: float = _READY_TIMEOUT) -> Optional[str]:
"""Start a quick tunnel and return its public URL once it is actually
serving, or None (best-effort).
Waits for cloudflared to both mint the URL and register an edge connection
before returning, so the caller never advertises a URL that yields Cloudflare
error 1033 (HTTP 530). If a URL is minted but no connection registers within
the window (e.g. quic is blocked on this network), retries once forcing the
http2 protocol. On any failure the tunnel is stopped and None is returned.
"""
global _active_tunnel, _shutdown_requested
binary = ensure_cloudflared()
if not binary:
return None
with _active_lock:
_shutdown_requested = False # fresh session
# Default protocol first (quic, with cloudflared's own http2 fallback); if a
# URL appears but no connection registers, quic is likely blocked -> retry
# once forcing http2.
for protocol in (None, "http2"):
# Create + register under the lock, and bail if a stop already landed
# (e.g. between this and the previous attempt) so we never start a tunnel
# after shutdown has run.
with _active_lock:
if _shutdown_requested:
_active_tunnel = None
return None
tunnel = CloudflareTunnel(port, binary, protocol = protocol)
prior, _active_tunnel = _active_tunnel, tunnel
if prior is not None:
prior.stop()
try:
tunnel.start()
url = tunnel.wait_for_ready(timeout)
except Exception:
url = None
if url:
return url
saw_url = tunnel.url is not None
# Not ready: drop it, but only if we are still the active tunnel.
with _active_lock:
was_active = _active_tunnel is tunnel
if was_active:
_active_tunnel = None
tunnel.stop()
# A concurrent shutdown or start took over while we waited; retrying would
# spawn a tunnel nobody owns (orphaned after shutdown), so bail instead.
if not was_active:
return None
# No URL at all is an API/network failure, not a protocol one; forcing
# http2 will not help, so do not burn another window on it.
if not saw_url:
return None
return None
def stop_studio_tunnel() -> None:
"""Terminate the active tunnel, if any. Idempotent."""
global _active_tunnel, _shutdown_requested
with _active_lock:
# Latch so an in-flight start_studio_tunnel won't start a fresh tunnel
# (e.g. its http2 retry) after we have already torn down.
_shutdown_requested = True
tunnel, _active_tunnel = _active_tunnel, None
if tunnel is not None:
tunnel.stop()

196
studio/backend/colab.py
View file

@ -2,15 +2,13 @@
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""
Colab-specific helpers for running Unsloth Studio.
Uses Colab's built-in proxy - no external tunneling needed!
Colab helpers for Unsloth Studio. Uses Colab's built-in proxy.
"""
from pathlib import Path
import sys
# Fix for Anaconda/conda-forge Python: seed platform._sys_version_cache before
# any library imports that trigger attrs -> rich -> structlog -> platform crash.
# Seed platform._sys_version_cache before attrs->rich->structlog->platform crash on conda Python.
# See: https://github.com/python/cpython/issues/102396
_backend_dir = str(Path(__file__).parent)
if _backend_dir not in sys.path:
@ -25,31 +23,59 @@ logger = get_logger(__name__)
def get_colab_url(port: int = 8888) -> str:
"""
Get the actual Colab proxy URL for a port.
Get the Colab proxy URL for a port.
Retries up to 3 times, validating the result is a real HTTPS Colab URL.
Falls back to http://localhost:{port} only when all attempts fail.
"""
import time as _time
fallback = f"http://localhost:{port}"
try:
from google.colab.output import eval_js
except ImportError:
return fallback
# Use Colab's proxy mechanism
url = eval_js(f"google.colab.kernel.proxyPort({port})", timeout_sec = 5)
return url if url else f"http://localhost:{port}"
except Exception as e:
logger.info(f"Note: Could not get Colab URL ({e})")
return f"http://localhost:{port}"
for attempt in range(3):
try:
url = eval_js(f"google.colab.kernel.proxyPort({port})", timeout_sec = 10)
# Valid proxy URL is https:// and embeds the port.
if url and isinstance(url, str) and url.startswith("https://") and str(port) in url:
return url.rstrip("/")
except Exception as e:
logger.info(f"Note: Could not get Colab URL (attempt {attempt + 1}/3: {e})")
if attempt < 2:
_time.sleep(1)
logger.warning(
f"Could not get a valid Colab proxy URL after 3 attempts — using localhost fallback. "
f"The link/iframe may not work from outside the runtime."
)
return fallback
def show_link(port: int = 8888):
"""Display a styled clickable link to the UI."""
def show_link(port: int = 8888, *, _url: "str | None" = None):
"""Display a styled clickable link to the UI.
*_url* is an optional pre-fetched proxy URL; pass it to avoid a second eval_js round-trip.
"""
from IPython.display import display, HTML
# Get real Colab proxy URL
url = get_colab_url(port)
url = _url if _url is not None else get_colab_url(port)
# Truncated display URL; try/except so an odd URL shape still renders the link.
try:
port_prefix = f"{port}-"
idx = url.index(port_prefix)
next_dash = url.index("-", idx + len(port_prefix))
short_url = url[: next_dash + 1] + "..."
except (ValueError, IndexError):
short_url = url
# Plain-text line so the URL shows even if HTML display fails.
logger.info(f"🌐 Unsloth Studio URL: {url}")
short_url = (
url[: url.index("-", url.index(f"{port}-") + len(str(port)) + 1) + 1] + "..."
if f"{port}-" in url
else url
)
html = f"""
<div style="display: inline-block; padding: 20px; background: #ffffff; border: 2px solid #000000;
border-radius: 12px; margin: 10px 0; font-family: system-ui, -apple-system, sans-serif;">
@ -59,10 +85,10 @@ def show_link(port: int = 8888):
height="48" style="display:block;">
Unsloth Studio is Ready!
</h2>
<a href="{url}" target="_blank"
<a href="{url}" onclick="var w=window.open(this.href,'_blank');if(!w){{return true;}}return false;"
style="display: inline-flex; align-items: center; gap: 10px; padding: 14px 28px;
background: #000000; color: white; text-decoration: none; border-radius: 8px;
font-weight: 800; font-size: 16px;">
font-weight: 800; font-size: 16px; cursor: pointer;">
<svg xmlns="http://www.w3.org/2000/svg" width="18" height="18" viewBox="0 0 24 24" fill="white"><polygon points="5,3 19,12 5,21"/></svg>
Open Unsloth Studio
</a>
@ -77,6 +103,67 @@ def show_link(port: int = 8888):
display(HTML(html))
def _is_studio_healthy(port: int, timeout: float = 2.0) -> bool:
"""Return True if a Studio backend is already answering health checks on *port*."""
import urllib.request
try:
with urllib.request.urlopen(f"http://localhost:{port}/api/health", timeout = timeout):
return True
except Exception:
return False
def _show_and_embed(port: int):
"""Embed the Studio inline for *port* with a branded header bar.
Fetches the proxy URL once (registering the port), then renders header bar +
iframe. Falls back to serve_kernel_port_as_iframe if IPython HTML is unavailable.
"""
url = get_colab_url(port)
logger.info(f"🌐 Unsloth Studio URL: {url}")
try:
from IPython.display import HTML, display
iframe_id = f"unsloth-studio-{port}"
# Truncated header URL — best-effort, falls back to full URL.
try:
port_prefix = f"{port}-"
idx = url.index(port_prefix)
next_dash = url.index("-", idx + len(port_prefix))
short_url = url[: next_dash + 1] + "..."
except (ValueError, IndexError):
short_url = url
display(
HTML(f"""
<div style="font-family:system-ui,-apple-system,sans-serif;margin:8px 0;
border-radius:12px;overflow:hidden;box-shadow:0 2px 16px rgba(0,0,0,0.18);">
<div style="display:flex;align-items:center;gap:10px;padding:10px 16px;background:#000;">
<img src="https://github.com/unslothai/unsloth/raw/main/studio/frontend/public/unsloth-gem.png"
height="26" style="display:block;">
<span style="color:#fff;font-weight:700;font-size:15px;letter-spacing:-0.2px;">Unsloth Studio</span>
<span style="margin-left:auto;color:#666;font-size:11px;font-family:monospace;">{short_url}</span>
</div>
<iframe
id="{iframe_id}"
src="{url}"
style="width:100%;height:82vh;min-height:600px;max-height:1100px;border:none;display:block;box-sizing:border-box;"
allow="clipboard-read; clipboard-write"
></iframe>
</div>
""")
)
except Exception:
# Fallback: Colab's built-in helper.
try:
from google.colab import output as colab_output
colab_output.serve_kernel_port_as_iframe(port, height = 900, width = "100%")
except ImportError:
pass
def start(port: int = 8888):
"""
Start Unsloth Studio server in Colab and display the URL.
@ -85,10 +172,23 @@ def start(port: int = 8888):
from colab import start
start()
"""
import sys
import time
logger.info("🦥 Starting Unsloth Studio...")
# Fast path: Studio already running (cell re-run). Re-launching would collide on
# the port, so just re-show the link and iframe.
if _is_studio_healthy(port):
logger.info(f" Studio is already running on port {port} — reusing existing server.")
_show_and_embed(port)
try:
for _ in range(10000):
time.sleep(300)
print("=", end = "", flush = True)
except KeyboardInterrupt:
logger.info("\nUnsloth Studio keepalive stopped.")
return
logger.info(" Loading backend...")
from run import run_server
@ -96,18 +196,56 @@ def start(port: int = 8888):
repo_root = Path(__file__).parent.parent
frontend_path = repo_root / "frontend" / "dist"
if not frontend_path.exists():
if not (frontend_path / "index.html").exists():
logger.info("❌ Frontend not built! Please run the setup cell first.")
return
logger.info(" Starting server...")
# Start server silently
run_server(host = "0.0.0.0", port = port, frontend_path = frontend_path, silent = True)
try:
app = run_server(host = "0.0.0.0", port = port, frontend_path = frontend_path, silent = True)
except SystemExit as exc:
logger.error(f"❌ Unsloth Studio failed to start: {exc}")
return
except Exception as exc:
logger.error(f"❌ Unsloth Studio failed to start: {exc}")
return
logger.info(" Server started!")
# run_server auto-increments the port if in use; read back the bound port so the
# proxy URL and iframe point at the right place.
actual_port: int = getattr(getattr(app, "state", None), "server_port", None) or port
# Show the clickable link with real URL
show_link(port)
logger.info(f" Server started on port {actual_port}!")
# Poll health endpoint before showing the link — avoids the race where ready_event
# fires but the process hasn't finished binding.
import urllib.request
server_ready = False
for _ in range(40):
try:
with urllib.request.urlopen(f"http://localhost:{actual_port}/api/health", timeout = 1):
server_ready = True
break
except Exception:
time.sleep(0.5)
if not server_ready:
logger.error(
f"❌ Unsloth Studio did not become healthy on port {actual_port}. "
"Check for errors above."
)
return
_show_and_embed(actual_port)
# Keep kernel alive so the daemon server thread runs; handle KeyboardInterrupt
# cleanly so interrupting the cell gives a readable message.
try:
for _ in range(10000):
time.sleep(300)
print("=", end = "", flush = True)
except KeyboardInterrupt:
logger.info("\nUnsloth Studio keepalive stopped.")
if __name__ == "__main__":

15
studio/backend/core/__init__.py
View file

@ -4,18 +4,16 @@
"""
Unified core module for Unsloth backend
Imports are LAZY (via __getattr__) so that training subprocesses can
import core.training.worker without pulling in heavy ML dependencies
like unsloth, transformers, or torch before the version activation
code has a chance to run.
Imports are LAZY (via __getattr__) so training subprocesses can import
core.training.worker without pulling in heavy ML deps (unsloth, transformers,
torch) before the version-activation code runs.
"""
import sys
from pathlib import Path
# Ensure the backend directory is on sys.path so that bare "from utils.*"
# imports used throughout the backend work when core is imported as a package
# (e.g. from the CLI: "from studio.backend.core import ModelConfig").
# Add backend dir to sys.path so bare "from utils.*" imports work when core
# is imported as a package.
_backend_dir = str(Path(__file__).resolve().parent.parent)
if _backend_dir not in sys.path:
sys.path.insert(0, _backend_dir)
@ -69,7 +67,7 @@ def __getattr__(name):
globals()["TrainingProgress"] = TrainingProgress
return globals()[name]
# Config (from utils.models)
# Config (utils.models)
if name in (
"is_vision_model",
"ModelConfig",
@ -140,7 +138,6 @@ def __getattr__(name):
# Datasets
if name == "format_and_template_dataset":
from utils.datasets import format_and_template_dataset
globals()["format_and_template_dataset"] = format_and_template_dataset
return format_and_template_dataset

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

@ -0,0 +1,135 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Shared torchao Windows-ROCm import stub.
torchao (pulled in by transformers.quantizers) imports distributed_c10d.py
unconditionally, which crashes on Windows ROCm because the RCCL backend
(torch._C._distributed_c10d) is absent. Stubbing torchao short-circuits its
import chain; _StubSubpackageFinder handles any depth of torchao.xxx.yyy.
Worker subprocesses call install_torchao_windows_rocm_stub() before importing
transformers / unsloth_zoo.
"""
from __future__ import annotations
import sys
import types
import importlib.abc
import importlib.machinery
_STUB_SENTINEL = object()
# Metaclass for stub types so isinstance(x, StubClass) returns False instead of
# raising TypeError -- peft's lora/torchao.py does isinstance() against torchao
# types, which fails if those names resolve to stub modules rather than types.
class _StubTypeMeta(type):
def __instancecheck__(cls, instance):
return False
def __subclasscheck__(cls, subclass):
return False
def __getattr__(cls, attr):
if attr.startswith("__"):
raise AttributeError(attr)
child = _StubTypeMeta(attr, (), {})
setattr(cls, attr, child)
return child
def __call__(cls, *args, **kwargs):
return None
def _make_stub_type(name):
"""Stub class: accepted by isinstance() (always False), supports attr access."""
return _StubTypeMeta(name, (), {})
def _make_mod_stub(mod_name):
m = types.ModuleType(mod_name)
m.__path__ = []
m.__package__ = mod_name
m._unsloth_stub = _STUB_SENTINEL
m.__spec__ = importlib.machinery.ModuleSpec(mod_name, loader = None, is_package = True)
def _ga(
attr,
_m = m,
_n = mod_name,
):
if attr.startswith("__"):
raise AttributeError(attr)
# Return a stub CLASS (not module) so isinstance() returns False, not TypeError.
child = _make_stub_type(f"{_n}.{attr}")
setattr(_m, attr, child)
return child
m.__getattr__ = _ga
return m
class _StubSubpackageLoader(importlib.abc.Loader):
def __init__(self, mod_name):
self._mod_name = mod_name
def create_module(self, spec):
return _make_mod_stub(self._mod_name)
def exec_module(self, module):
pass
class _StubSubpackageFinder(importlib.abc.MetaPathFinder):
def find_spec(
self,
fullname,
path,
target = None,
):
if "." not in fullname:
return None
parent = sys.modules.get(fullname.rsplit(".", 1)[0])
if parent is None:
return None
if getattr(parent, "_unsloth_stub", None) is not _STUB_SENTINEL:
return None
return importlib.machinery.ModuleSpec(
fullname, _StubSubpackageLoader(fullname), is_package = True
)
def install_torchao_windows_rocm_stub() -> None:
"""Pre-stub torchao on Windows ROCm so transformers/peft imports don't crash.
No-op elsewhere (incl. Windows CUDA, where torchao is real). Must run before
importing transformers / unsloth_zoo. Safe to call once per worker.
"""
# Gate on the active torch runtime, not env-var presence -- HIP_PATH/ROCM_PATH
# persist after reverting to a CUDA wheel. Some ROCm wheels lack
# torch.version.hip but still encode "rocm" in __version__, so accept either.
_is_win32_rocm = False
if sys.platform == "win32":
try:
import torch as _torch_probe
_is_win32_rocm = bool(
getattr(getattr(_torch_probe, "version", None), "hip", None)
or "rocm" in getattr(_torch_probe, "__version__", "").lower()
)
del _torch_probe
except Exception:
pass
if _is_win32_rocm:
# Register the finder only on Windows ROCm.
sys.meta_path.append(_StubSubpackageFinder())
# Seed torchao top-level + key submodules; the finder handles the rest.
for _tao_name in (
"torchao",
"torchao.quantization",
"torchao.dtypes",
"torchao.float8",
"torchao.utils",
):
if _tao_name not in sys.modules:
sys.modules[_tao_name] = _make_mod_stub(_tao_name)

4
studio/backend/core/data_recipe/__init__.py
View file

@ -1,9 +1,7 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""
Data Recipe core (DataDesigner wrapper + job runner).
"""
"""Data Recipe core (DataDesigner wrapper + job runner)."""
from .jobs import JobManager, get_job_manager

7
studio/backend/core/data_recipe/huggingface.py
View file

@ -36,9 +36,7 @@ def _resolve_recipe_artifact_path(artifact_path: str) -> Path:
if not resolved.exists():
raise RecipeDatasetPublishError("Execution artifacts are no longer available.")
if not resolved.is_dir():
raise RecipeDatasetPublishError(
"Execution artifact path is not a dataset folder."
)
raise RecipeDatasetPublishError("Execution artifact path is not a dataset folder.")
return resolved
@ -97,8 +95,7 @@ def publish_recipe_dataset(
tags = None,
)
card.text = card.text.replace(_DATA_DESIGNER_FOOTER, _UNSLOTH_STUDIO_FOOTER)
# Data Designer currently drops the explicit token when pushing the
# dataset card. Push it ourselves so auth stays request-local.
# Data Designer drops the explicit token, so push the card ourselves to keep auth request-local.
card.push_to_hub(repo_id, token = hf_token, repo_type = "dataset")
client._upload_main_dataset_files(

63
studio/backend/core/data_recipe/jobs/manager.py
View file

@ -108,9 +108,7 @@ class Subscription:
event_id = self._next_id
body = json.dumps(event, separators = (",", ":"), ensure_ascii = False)
event_type = event.get("type") or "message"
return (
f"id: {event_id}\n" f"event: {event_type}\n" f"data: {body}\n\n"
).encode("utf-8")
return (f"id: {event_id}\n" f"event: {event_type}\n" f"data: {body}\n\n").encode("utf-8")
class JobManager:
@ -134,10 +132,9 @@ class JobManager:
) -> str:
"""Spawn the job subprocess (one at a time, no cap).
``internal_api_key_id`` is the row id of a workflow-scoped
sk-unsloth-* key minted by the route layer for local providers.
JobManager revokes it when the job reaches a terminal state so the
key's live window is no longer than the run.
``internal_api_key_id`` is a workflow-scoped sk-unsloth-* key row id
minted by the route layer; revoked on terminal state so the key's
live window is no longer than the run.
"""
llm_columns = recipe.get("columns") or []
llm_column_count = 0
@ -158,9 +155,7 @@ class JobManager:
job_id = uuid.uuid4().hex
self._job = Job(job_id = job_id, status = "pending", started_at = time.time())
self._job.progress_columns_total = llm_column_count
self._job.source_progress_estimated_total = _github_source_estimated_total(
recipe
)
self._job.source_progress_estimated_total = _github_source_estimated_total(recipe)
self._job.internal_api_key_id = internal_api_key_id
self._events.clear()
self._seq = 0
@ -187,9 +182,7 @@ class JobManager:
self._pump_thread = threading.Thread(target = self._pump_loop, daemon = True)
self._pump_thread.start()
self._emit(
{"type": EVENT_JOB_ENQUEUED, "ts": time.time(), "job_id": job_id}
)
self._emit({"type": EVENT_JOB_ENQUEUED, "ts": time.time(), "job_id": job_id})
return job_id
def cancel(self, job_id: str) -> bool:
@ -200,9 +193,7 @@ class JobManager:
if self._proc is None or not self._proc.is_alive():
return True
self._job.status = "cancelling"
self._emit(
{"type": EVENT_JOB_CANCELLING, "ts": time.time(), "job_id": job_id}
)
self._emit({"type": EVENT_JOB_CANCELLING, "ts": time.time(), "job_id": job_id})
try:
self._proc.terminate()
except (AttributeError, OSError):
@ -210,7 +201,7 @@ class JobManager:
return True
def get_status(self, job_id: str) -> dict | None:
"""UI friendly snapshot that we need. Alternative to sse kinda of and structured"""
"""UI-friendly structured snapshot; an alternative to SSE."""
with self._lock:
if self._job is None or self._job.job_id != job_id:
return None
@ -319,19 +310,12 @@ class JobManager:
if not parquet_dir.exists():
return {"error": f"dataset path missing: {parquet_dir}"}
return self._load_dataset_page(
parquet_dir = parquet_dir, limit = limit, offset = offset
)
return self._load_dataset_page(parquet_dir = parquet_dir, limit = limit, offset = offset)
except Exception as exc:
return {"error": f"dataset load failed: {exc}"}
@staticmethod
def _load_dataset_page(
*,
parquet_dir: Path,
limit: int,
offset: int,
) -> dict[str, Any]:
def _load_dataset_page(*, parquet_dir: Path, limit: int, offset: int) -> dict[str, Any]:
dataset_page = JobManager._load_dataset_page_with_duckdb(
parquet_dir = parquet_dir,
limit = limit,
@ -347,10 +331,7 @@ class JobManager:
@staticmethod
def _load_dataset_page_with_duckdb(
*,
parquet_dir: Path,
limit: int,
offset: int,
*, parquet_dir: Path, limit: int, offset: int
) -> dict[str, Any] | None:
parquet_glob = str((parquet_dir / "*.parquet").resolve())
try:
@ -389,10 +370,7 @@ class JobManager:
@staticmethod
def _load_dataset_page_with_data_designer(
*,
parquet_dir: Path,
limit: int,
offset: int,
*, parquet_dir: Path, limit: int, offset: int
) -> dict[str, Any]:
from data_designer.config.utils.io_helpers import read_parquet_dataset
@ -402,7 +380,10 @@ class JobManager:
return {"dataset": to_preview_jsonable(rows), "total": total}
def subscribe(
self, job_id: str, *, after_seq: int | None = None
self,
job_id: str,
*,
after_seq: int | None = None,
) -> Subscription | None:
"""SSE subscribe: get replay buffer + live events stream."""
with self._lock:
@ -497,9 +478,7 @@ class JobManager:
self._job.error = self._job.error or "process exited"
self._job.finished_at = time.time()
event_type = (
EVENT_JOB_CANCELLED
if self._job.status == "cancelled"
else EVENT_JOB_ERROR
EVENT_JOB_CANCELLED if self._job.status == "cancelled" else EVENT_JOB_ERROR
)
self._emit(
{
@ -557,16 +536,14 @@ class JobManager:
def _retire_workflow_key(self, job: Job) -> None:
"""Revoke the workflow-scoped sk-unsloth-* key, if one was minted.
Best-effort: revocation failures are swallowed. The key would
expire on its own after 24h, so a missed revoke is a latency
concern, not a correctness one.
Best-effort: failures are swallowed. The key expires after 24h, so a
missed revoke is a latency, not correctness, concern.
"""
key_id = getattr(job, "internal_api_key_id", None)
if not key_id:
return
try:
from auth import storage # deferred: avoids circular import
from auth import storage # deferred: avoid circular import
storage.revoke_internal_api_key(int(key_id))
except Exception:
pass

42
studio/backend/core/data_recipe/jobs/parse.py
View file

@ -46,7 +46,7 @@ class ParsedUpdate:
source_progress: SourceProgress | None = None
# kinda of a bummber but currently only option, Best effort parser from data-designer logs -> structured status for UI.
# Best-effort parser from data-designer logs -> structured status for UI.
_RE_SAMPLERS = re.compile(
r"Preparing samplers to generate (?P<rows>\d+) records across (?P<cols>\d+) columns"
)
@ -119,8 +119,7 @@ def parse_log_message(msg: str) -> ParsedUpdate | None:
page_items = page_items,
rate_remaining = int(m.group("remaining")),
message = (
f"Scraping GitHub source: {repo} "
f"{resource} page {page} (+{page_items})"
f"Scraping GitHub source: {repo} " f"{resource} page {page} (+{page_items})"
),
),
)
@ -134,10 +133,7 @@ def parse_log_message(msg: str) -> ParsedUpdate | None:
source = "github",
status = "rate_limited",
retry_after_sec = seconds,
message = (
"Waiting for GitHub rate limit. "
"Studio will resume automatically."
),
message = ("Waiting for GitHub rate limit. Studio will resume automatically."),
),
)
@ -151,8 +147,7 @@ def parse_log_message(msg: str) -> ParsedUpdate | None:
status = "rate_limited",
retry_after_sec = seconds,
message = (
"Waiting for GitHub secondary rate limit. "
"Studio will resume automatically."
"Waiting for GitHub secondary rate limit. Studio will resume automatically."
),
),
)
@ -166,10 +161,7 @@ def parse_log_message(msg: str) -> ParsedUpdate | None:
source = "github",
status = "rate_limited",
retry_after_sec = seconds,
message = (
"Waiting for GitHub rate limit. "
"Studio will resume automatically."
),
message = ("Waiting for GitHub rate limit. Studio will resume automatically."),
),
)
@ -335,7 +327,7 @@ def apply_update(job: Job, update: ParsedUpdate) -> None:
_apply_source_progress(job, update.source_progress)
if update.stage in USAGE_RESET_STAGES:
# usage summary is a short block so we reset once we move into the next stage.
# Usage summary is a short block; reset on the next stage.
job._in_usage_summary = False
if update.usage_section_start is not None:
@ -387,15 +379,13 @@ def _apply_source_progress(job: Job, progress: SourceProgress) -> None:
count_key = f"{progress.repo}:{progress.resource}"
if page_key not in job._source_seen_pages:
job._source_seen_pages.add(page_key)
job._source_counts[count_key] = int(
job._source_counts.get(count_key, 0)
) + int(page_items or 0)
job._source_counts[count_key] = int(job._source_counts.get(count_key, 0)) + int(
page_items or 0
)
fetched_items = sum(job._source_counts.values())
if fetched_items <= 0:
fetched_items = progress.fetched_items or (
previous.fetched_items if previous else None
)
fetched_items = progress.fetched_items or (previous.fetched_items if previous else None)
estimated_total = (
progress.estimated_total
@ -415,14 +405,10 @@ def _apply_source_progress(job: Job, progress: SourceProgress) -> None:
repo = progress.repo or (previous.repo if previous else None),
resource = progress.resource or (previous.resource if previous else None),
page = (
progress.page
if progress.page is not None
else (previous.page if previous else None)
progress.page if progress.page is not None else (previous.page if previous else None)
),
page_items = (
page_items
if page_items is not None
else (previous.page_items if previous else None)
page_items if page_items is not None else (previous.page_items if previous else None)
),
fetched_items = fetched_items,
estimated_total = estimated_total,
@ -453,9 +439,7 @@ def _compute_overall_progress(job: Job, column_progress: Progress) -> Progress:
if len(job._column_done) == 0:
done = current_done
else:
sum_done = sum(
max(0, min(value, total_rows)) for value in job._column_done.values()
)
sum_done = sum(max(0, min(value, total_rows)) for value in job._column_done.values())
done = int(sum_done / total_columns)
prev_done = int(job.progress.done or 0)

5
studio/backend/core/data_recipe/jobs/types.py
View file

@ -90,9 +90,8 @@ class Job:
progress_columns_total: int | None = None
source_progress_estimated_total: int | None = None
completed_columns: list[str] = field(default_factory = list)
# Id of the internal sk-unsloth-* API key minted for a local-model
# workflow. Revoked when the job terminates so the key's live window
# matches the run rather than its 24h TTL.
# Id of the internal sk-unsloth-* API key minted for a local-model workflow.
# Revoked when the job ends so the key's window matches the run, not its 24h TTL.
internal_api_key_id: int | None = None
_current_usage_model: str | None = None
_in_usage_summary: bool = False

32
studio/backend/core/data_recipe/jobs/worker.py
View file

@ -60,9 +60,7 @@ def _slugify_run_name(value: str) -> str:
return slug[:80].strip("-")
def _build_dataset_name(
*, run_name: str | None, job_id: str, artifact_root: Path
) -> str:
def _build_dataset_name(*, run_name: str | None, job_id: str, artifact_root: Path) -> str:
fallback = f"recipe_{job_id}"
slug = _slugify_run_name(run_name or "")
base_name = f"recipe_{slug}" if slug else fallback
@ -74,21 +72,11 @@ def _build_dataset_name(
return candidate
def run_job_process(
*,
event_queue,
recipe: dict[str, Any],
run: dict[str, Any],
) -> None:
"""
Subprocess entrypoint.
Sends events to `event_queue`.
"""
def run_job_process(*, event_queue, recipe: dict[str, Any], run: dict[str, Any]) -> None:
"""Subprocess entrypoint. Sends events to `event_queue`."""
import os
os.environ["PYTHONWARNINGS"] = (
"ignore" # Suppress warnings at C-level before imports
)
os.environ["PYTHONWARNINGS"] = "ignore" # suppress C-level warnings before imports
import warnings
from loggers.config import LogConfig
@ -124,8 +112,8 @@ def run_job_process(
builder = build_config_builder(recipe)
designer = create_data_designer(recipe, artifact_path = str(_ARTIFACT_ROOT))
# DataDesigner configures root logging in DataDesigner.__init__.
# Attach queue logger directly to `data_designer` so parser events survive root resets.
# DataDesigner resets root logging in __init__; attach the queue handler
# to the named loggers directly so parser events survive.
handler = _QueueLogHandler(event_queue)
handler.setLevel(logging.INFO)
for logger_name in (
@ -172,14 +160,10 @@ def run_job_process(
}
)
else:
results = designer.create(
builder, num_records = rows, dataset_name = dataset_name
)
results = designer.create(builder, num_records = rows, dataset_name = dataset_name)
analysis = to_jsonable(results.load_analysis().model_dump(mode = "json"))
if merge_batches:
_merge_batches_to_single_parquet(
results.artifact_storage.base_dataset_path
)
_merge_batches_to_single_parquet(results.artifact_storage.base_dataset_path)
artifact_path = str(results.artifact_storage.base_dataset_path)
event_queue.put(
{

2
studio/backend/core/data_recipe/jsonable.py
View file

@ -23,7 +23,6 @@ def _pil_to_preview_payload(image: Any) -> dict[str, Any]:
def _open_pil_image_from_bytes(raw_bytes: bytes):
from PIL import Image # type: ignore
with Image.open(io.BytesIO(raw_bytes)) as image:
return image.copy()
@ -52,7 +51,6 @@ def _to_pil_from_hf_image_dict(value: Any) -> Any | None:
if isinstance(path_value, str) and path_value.strip():
try:
from PIL import Image # type: ignore
with Image.open(Path(path_value)) as image:
return image.copy()
except (OSError, ValueError, TypeError):

52
studio/backend/core/data_recipe/local_callable_validators.py
View file

@ -81,9 +81,7 @@ def split_oxc_local_callable_validators(
def register_oxc_local_callable_validators(
*,
builder,
specs: list[OxcLocalCallableValidatorSpec],
*, builder, specs: list[OxcLocalCallableValidatorSpec]
) -> None:
if not specs:
return
@ -114,10 +112,7 @@ def register_oxc_local_callable_validators(
)
def _parse_oxc_spec(
*,
column: dict[str, Any],
) -> OxcLocalCallableValidatorSpec | None:
def _parse_oxc_spec(*, column: dict[str, Any]) -> OxcLocalCallableValidatorSpec | None:
if str(column.get("column_type") or "").strip() != "validation":
return None
if str(column.get("validator_type") or "").strip() != "local_callable":
@ -138,11 +133,7 @@ def _parse_oxc_spec(
target_columns_raw = column.get("target_columns")
target_columns = (
[
value.strip()
for value in target_columns_raw
if isinstance(value, str) and value.strip()
]
[value.strip() for value in target_columns_raw if isinstance(value, str) and value.strip()]
if isinstance(target_columns_raw, list)
else []
)
@ -182,9 +173,7 @@ def _parse_oxc_validation_marker(fn_name: str) -> tuple[str, str, str]:
return "javascript", "syntax", "auto"
code_lang = parts[0] if parts[0] in _OXC_LANG_TO_NODE_LANG else "javascript"
mode = parts[1] if parts[1] in _OXC_VALIDATION_MODES else "syntax"
code_shape = (
parts[2] if len(parts) >= 3 and parts[2] in _OXC_CODE_SHAPES else "auto"
)
code_shape = parts[2] if len(parts) >= 3 and parts[2] in _OXC_CODE_SHAPES else "auto"
return code_lang, mode, code_shape
@ -195,7 +184,7 @@ def _build_oxc_validation_function(lang: str, validation_mode: str, code_shape:
normalized_code_shape = code_shape if code_shape in _OXC_CODE_SHAPES else "auto"
def _validator(df):
import pandas as pd # imported lazily for local callable runtime
import pandas as pd # lazy import for local callable runtime
row_count = int(len(df.index))
if row_count == 0:
@ -205,10 +194,7 @@ def _build_oxc_validation_function(lang: str, validation_mode: str, code_shape:
code_values = (
["" for _ in range(row_count)]
if not code_column
else [
"" if value is None else str(value)
for value in df[code_column].tolist()
]
else ["" if value is None else str(value) for value in df[code_column].tolist()]
)
results = _run_oxc_batch(
@ -224,16 +210,14 @@ def _build_oxc_validation_function(lang: str, validation_mode: str, code_shape:
)
return pd.DataFrame(results)
_validator.__name__ = f"{OXC_VALIDATION_FN_MARKER}_{node_lang}_{mode.replace('+', '_')}_{normalized_code_shape}"
_validator.__name__ = (
f"{OXC_VALIDATION_FN_MARKER}_{node_lang}_{mode.replace('+', '_')}_{normalized_code_shape}"
)
return _validator
def _run_oxc_batch(
*,
node_lang: str,
validation_mode: str,
code_shape: str,
code_values: list[str],
*, node_lang: str, validation_mode: str, code_shape: str, code_values: list[str]
) -> list[dict[str, Any]]:
if not _OXC_RUNNER_PATH.exists():
return _fallback_results(
@ -308,21 +292,13 @@ def _run_oxc_batch(
warning_count_raw = item.get("warning_count")
out.append(
{
"is_valid": bool(is_valid_raw)
if isinstance(is_valid_raw, bool)
else False,
"error_count": int(error_count_raw)
if isinstance(error_count_raw, int)
else 0,
"is_valid": bool(is_valid_raw) if isinstance(is_valid_raw, bool) else False,
"error_count": int(error_count_raw) if isinstance(error_count_raw, int) else 0,
"error_message": str(message_raw or ""),
"severity": str(severity_raw)
if isinstance(severity_raw, str)
else None,
"severity": str(severity_raw) if isinstance(severity_raw, str) else None,
"code": str(code_raw) if isinstance(code_raw, str) else None,
"labels": labels_raw if isinstance(labels_raw, list) else [],
"codeframe": str(codeframe_raw)
if isinstance(codeframe_raw, str)
else None,
"codeframe": str(codeframe_raw) if isinstance(codeframe_raw, str) else None,
"warning_count": int(warning_count_raw)
if isinstance(warning_count_raw, int)
else 0,

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

@ -22,9 +22,7 @@ def _encode_bytes_to_base64(value: bytes | bytearray) -> str:
return base64.b64encode(bytes(value)).decode("utf-8")
def _load_image_file_to_base64(
path_value: str, *, base_path: str | None = None
) -> str | None:
def _load_image_file_to_base64(path_value: str, *, base_path: str | None = None) -> str | None:
try:
path = Path(path_value)
candidates: list[Path] = []
@ -109,7 +107,7 @@ def _apply_data_designer_image_context_patch() -> None:
return
try:
from data_designer.config.models import ImageContext
from data_designer.config.models import ImageContext # pyright: ignore[reportMissingImports]
except ImportError:
return
@ -119,9 +117,7 @@ def _apply_data_designer_image_context_patch() -> None:
original_auto_resolve = ImageContext._auto_resolve_context_value
def _patched_auto_resolve(
self: Any, context_value: Any, base_path: str | None
) -> Any:
def _patched_auto_resolve(self: Any, context_value: Any, base_path: str | None) -> Any:
normalized = _normalize_image_context_value(context_value, base_path = base_path)
return original_auto_resolve(self, normalized, base_path)
@ -131,7 +127,7 @@ def _apply_data_designer_image_context_patch() -> None:
def build_model_providers(recipe: dict[str, Any]):
from data_designer.config.models import ModelProvider
from data_designer.config.models import ModelProvider # pyright: ignore[reportMissingImports]
providers: list[ModelProvider] = []
for provider in recipe.get("model_providers", []):
@ -163,18 +159,19 @@ def _recipe_has_llm_columns(recipe: dict[str, Any]) -> bool:
return False
def _validate_recipe_runtime_support(
recipe: dict[str, Any],
model_providers: list[Any],
) -> None:
def _validate_recipe_runtime_support(recipe: dict[str, Any], model_providers: list[Any]) -> None:
if _recipe_has_llm_columns(recipe) and not model_providers:
raise ValueError("Add a Provider connection block before running this recipe.")
def build_mcp_providers(
recipe: dict[str, Any],
) -> list:
from data_designer.config.mcp import LocalStdioMCPProvider, MCPProvider
def build_mcp_providers(recipe: dict[str, Any]) -> list:
from data_designer.config.mcp import LocalStdioMCPProvider, MCPProvider # pyright: ignore[reportMissingImports]
# Same gate as the chat MCP path: stdio providers spawn a local subprocess,
# so build them only when this host allows it (desktop / explicit opt-in).
from core.inference.mcp_client import stdio_mcp_enabled
stdio_allowed = stdio_mcp_enabled()
providers: list[MCPProvider | LocalStdioMCPProvider] = []
for provider in recipe.get("mcp_providers", []):
@ -182,6 +179,8 @@ def build_mcp_providers(
continue
provider_type = provider.get("provider_type")
if provider_type == "stdio":
if not stdio_allowed:
continue
env = provider.get("env")
if not isinstance(env, dict):
env = {}
@ -214,19 +213,43 @@ def build_mcp_providers(
return providers
def _strip_frontend_model_config_metadata(recipe: dict[str, Any]) -> dict[str, Any]:
model_configs = recipe.get("model_configs")
if not isinstance(model_configs, list):
return recipe
changed = False
next_model_configs: list[Any] = []
for model_config in model_configs:
if isinstance(model_config, dict) and "gguf_variant" in model_config:
next_model_config = dict(model_config)
next_model_config.pop("gguf_variant", None)
next_model_configs.append(next_model_config)
changed = True
continue
next_model_configs.append(model_config)
if not changed:
return recipe
return {
**recipe,
"model_configs": next_model_configs,
}
def build_config_builder(recipe: dict[str, Any]):
_apply_data_designer_image_context_patch()
from data_designer.config import DataDesignerConfigBuilder
from data_designer.config.processors import ProcessorType
from data_designer.config import DataDesignerConfigBuilder # pyright: ignore[reportMissingImports]
from data_designer.config.processors import ProcessorType # pyright: ignore[reportMissingImports]
recipe_core = {
key: value
for key, value in recipe.items()
if key not in {"model_providers", "mcp_providers"}
}
recipe_core, oxc_local_callable_specs = split_oxc_local_callable_validators(
recipe_core
)
recipe_core = _strip_frontend_model_config_metadata(recipe_core)
recipe_core, oxc_local_callable_specs = split_oxc_local_callable_validators(recipe_core)
builder = DataDesignerConfigBuilder.from_config({"data_designer": recipe_core})
register_oxc_local_callable_validators(
builder = builder,
@ -234,7 +257,7 @@ def build_config_builder(recipe: dict[str, Any]):
)
# DataDesignerConfigBuilder.from_config currently skips processors.
# Re-attach explicitly so drop_columns/schema_transform survive API payload.
# Re-attach so drop_columns/schema_transform survive the API payload.
for processor in recipe_core.get("processors") or []:
if not isinstance(processor, dict):
continue
@ -250,23 +273,18 @@ def build_config_builder(recipe: dict[str, Any]):
return builder
def create_data_designer(
recipe: dict[str, Any],
*,
artifact_path: str | None = None,
):
def create_data_designer(recipe: dict[str, Any], *, artifact_path: str | None = None):
_apply_data_designer_image_context_patch()
from data_designer.interface.data_designer import DataDesigner
from data_designer.interface.data_designer import DataDesigner # pyright: ignore[reportMissingImports]
recipe = _strip_frontend_model_config_metadata(recipe)
model_providers = build_model_providers(recipe)
_validate_recipe_runtime_support(recipe, model_providers)
# DataDesigner requires at least one model provider in its registry even
# when the pipeline contains no LLM columns. Supply a lightweight stub
# so sampler/expression-only recipes can run without a real provider.
# DataDesigner requires >=1 model provider even with no LLM columns; stub
# one so sampler/expression-only recipes run without a real provider.
if not model_providers:
from data_designer.config.models import ModelProvider
from data_designer.config.models import ModelProvider # pyright: ignore[reportMissingImports]
model_providers = [
ModelProvider(
name = "_unused",
@ -290,8 +308,7 @@ def validate_recipe(recipe: dict[str, Any]) -> None:
def preview_recipe(
recipe: dict[str, Any],
num_records: int,
recipe: dict[str, Any], num_records: int
) -> tuple[list[dict[str, Any]], dict[str, Any] | None, dict[str, Any] | None]:
builder = build_config_builder(recipe)
designer = create_data_designer(recipe)
@ -303,14 +320,10 @@ def preview_recipe(
dataset = [to_jsonable(row) for row in raw_rows]
artifacts = (
None
if results.processor_artifacts is None
else to_jsonable(results.processor_artifacts)
None if results.processor_artifacts is None else to_jsonable(results.processor_artifacts)
)
analysis = (
None
if results.analysis is None
else to_jsonable(results.analysis.model_dump(mode = "json"))
None if results.analysis is None else to_jsonable(results.analysis.model_dump(mode = "json"))
)
return dataset, artifacts, analysis

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

@ -1,12 +1,11 @@
# 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 submodule - Model export operations
"""Export submodule - model export operations.
The default get_export_backend() returns an ExportOrchestrator that
delegates to a subprocess. The original ExportBackend runs inside
the subprocess and can be imported directly from .export when needed.
get_export_backend() returns an ExportOrchestrator that delegates to a
subprocess. The original ExportBackend runs inside the subprocess and can be
imported directly from .export when needed.
"""
from .orchestrator import ExportOrchestrator, get_export_backend

163
studio/backend/core/export/export.py
View file

@ -1,10 +1,7 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
# backend/export.py
"""
Export backend - handles model exporting in various formats
"""
"""Export backend - exports models in various formats."""
import glob
import json
@ -21,7 +18,12 @@ from utils.hardware import clear_gpu_cache
from utils.models import is_vision_model, get_base_model_from_lora
from utils.models.model_config import detect_audio_type
from utils.paths import ensure_dir, outputs_root, resolve_export_dir, resolve_output_dir
from utils.paths import (
ensure_dir,
outputs_root,
resolve_export_write_dir,
resolve_output_dir,
)
from core.inference import get_inference_backend
# GPU-only imports — guarded for Apple Silicon where these aren't needed
@ -46,10 +48,8 @@ def _is_wsl():
def _apply_wsl_sudo_patch():
"""On WSL, monkey-patch do_we_need_sudo() to return False.
WSL doesn't have passwordless sudo, and do_we_need_sudo() runs
`sudo apt-get update` which hangs waiting for a stdin password
inside a non-interactive subprocess. setup.sh pre-installs the
build dependencies on WSL, so sudo is not needed at runtime.
WSL lacks passwordless sudo and do_we_need_sudo()'s `sudo apt-get update`
hangs on a stdin password; setup.sh pre-installs the build deps anyway.
"""
if not _is_wsl():
return
@ -58,16 +58,11 @@ def _apply_wsl_sudo_patch():
import unsloth_zoo.llama_cpp as llama_cpp_module
def _wsl_do_we_need_sudo(system_type = "debian"):
logger.info(
"WSL detected — skipping sudo check "
"(build deps pre-installed by setup.sh)"
)
logger.info("WSL detected — skipping sudo check (build deps pre-installed by setup.sh)")
return False
llama_cpp_module.do_we_need_sudo = _wsl_do_we_need_sudo
logger.info(
"Applied WSL sudo patch to " "unsloth_zoo.llama_cpp.do_we_need_sudo"
)
logger.info("Applied WSL sudo patch to unsloth_zoo.llama_cpp.do_we_need_sudo")
except Exception as e:
logger.warning(f"Could not apply WSL sudo patch: {e}")
@ -115,18 +110,15 @@ class ExportBackend:
try:
logger.info("Starting memory cleanup...")
# Unload all models from inference backend
model_names = list(self.inference_backend.models.keys())
for model_name in model_names:
self.inference_backend.unload_model(model_name)
# Clear current export state
self.current_model = None
self.current_tokenizer = None
self.current_checkpoint = None
self._audio_type = None
# Clear GPU memory cache (handles gc + backend-specific cleanup)
clear_gpu_cache()
logger.info("Memory cleanup completed successfully")
@ -142,11 +134,9 @@ class ExportBackend:
"""
Scan outputs folder for training runs and their checkpoints.
Returns:
List of tuples: [(model_name, [(display_name, checkpoint_path), ...]), ...]
Returns: [(model_name, [(display_name, checkpoint_path), ...]), ...]
"""
from utils.models.checkpoints import scan_checkpoints
return scan_checkpoints(outputs_dir = outputs_dir)
def load_checkpoint(
@ -165,12 +155,11 @@ class ExportBackend:
try:
logger.info(f"Loading checkpoint: {checkpoint_path}")
# First, cleanup existing models
self.cleanup_memory()
checkpoint_path_obj = Path(checkpoint_path)
# Determine the model identity for type detection
# Model identity for type detection
adapter_config = checkpoint_path_obj / "adapter_config.json"
base_model = None
if adapter_config.exists():
@ -180,11 +169,9 @@ class ExportBackend:
model_id = base_model or checkpoint_path
# Detect audio type and vision
self._audio_type = detect_audio_type(model_id)
self.is_vision = not self._audio_type and is_vision_model(model_id)
# Load model based on type
if self._audio_type == "csm":
from unsloth import FastModel
from transformers import CsmForConditionalGeneration
@ -224,7 +211,6 @@ class ExportBackend:
elif self._audio_type == "bicodec":
from unsloth import FastModel
logger.info("Loading as BiCodec (Spark-TTS) audio model...")
model, tokenizer = FastModel.from_pretrained(
model_name = checkpoint_path,
@ -236,7 +222,6 @@ class ExportBackend:
elif self._audio_type == "dac":
from unsloth import FastModel
logger.info("Loading as DAC (OuteTTS) audio model...")
model, tokenizer = FastModel.from_pretrained(
model_name = checkpoint_path,
@ -254,7 +239,7 @@ class ExportBackend:
load_in_4bit = load_in_4bit,
trust_remote_code = trust_remote_code,
)
tokenizer = processor # For vision models, processor acts as tokenizer
tokenizer = processor # vision: processor acts as tokenizer
else:
logger.info("Loading as text model...")
@ -266,14 +251,12 @@ class ExportBackend:
trust_remote_code = trust_remote_code,
)
# Check if PEFT / LoRA model
if _IS_MLX:
# MLX doesn't use PeftModel — detect LoRA via adapter_config.json
self.is_peft = adapter_config.exists()
else:
self.is_peft = isinstance(model, (PeftModel, PeftModelForCausalLM))
# Store loaded model
self.current_model = model
self.current_tokenizer = tokenizer
self.current_checkpoint = checkpoint_path
@ -348,9 +331,7 @@ class ExportBackend:
output_path: Optional[str] = None
try:
if _IS_MLX:
mlx_save_method = (
"merged_4bit" if format_type == "4-bit (FP4)" else "merged_16bit"
)
mlx_save_method = "merged_4bit" if format_type == "4-bit (FP4)" else "merged_16bit"
else:
if format_type == "4-bit (FP4)":
save_method = "merged_4bit_forced"
@ -359,9 +340,8 @@ class ExportBackend:
else:
save_method = "merged_16bit"
# Save locally if requested
if save_directory:
save_directory = str(resolve_export_dir(save_directory))
save_directory = str(resolve_export_write_dir(save_directory))
logger.info(f"Saving merged model locally to: {save_directory}")
ensure_dir(Path(save_directory))
@ -380,7 +360,6 @@ class ExportBackend:
logger.info(f"Model saved successfully to {save_directory}")
output_path = str(Path(save_directory).resolve())
# Push to hub if requested
if push_to_hub:
if not repo_id or not hf_token:
return (
@ -415,9 +394,7 @@ class ExportBackend:
private = private,
)
else:
hub_save_method = (
save_method if save_method is not None else "merged_16bit"
)
hub_save_method = save_method if save_method is not None else "merged_16bit"
self.current_model.push_to_hub_merged(
repo_id,
self.current_tokenizer,
@ -463,15 +440,14 @@ class ExportBackend:
output_path: Optional[str] = None
try:
# Save locally if requested
if save_directory:
save_directory = str(resolve_export_dir(save_directory))
save_directory = str(resolve_export_write_dir(save_directory))
logger.info(f"Saving base model locally to: {save_directory}")
ensure_dir(Path(save_directory))
if _IS_MLX:
# MLX: save_pretrained_merged handles non-LoRA models too
# (fuse() is a no-op when there are no LoRA layers)
# (fuse() is a no-op without LoRA layers)
self.current_model.save_pretrained_merged(
save_directory,
self.current_tokenizer,
@ -486,7 +462,6 @@ class ExportBackend:
logger.info(f"Model saved successfully to {save_directory}")
output_path = str(Path(save_directory).resolve())
# Push to hub if requested
if push_to_hub:
if not repo_id or not hf_token:
return (
@ -521,14 +496,11 @@ class ExportBackend:
private = private,
)
else:
# Get base model name from request or model config
# Base model name from request or model config
base_model = (
base_model_id
or self.current_model.config._name_or_path
or "unknown"
base_model_id or self.current_model.config._name_or_path or "unknown"
)
# Create repo
hf_api = HfApi(token = hf_token)
repo_id = PushToHubMixin._create_repo(
PushToHubMixin,
@ -538,7 +510,6 @@ class ExportBackend:
)
username = repo_id.split("/")[0]
# Create and push model card
content = MODEL_CARD.format(
username = username,
base_model = base_model,
@ -547,11 +518,8 @@ class ExportBackend:
extra = "unsloth",
)
card = ModelCard(content)
card.push_to_hub(
repo_id, token = hf_token, commit_message = "Unsloth Model Card"
)
card.push_to_hub(repo_id, token = hf_token, commit_message = "Unsloth Model Card")
# Upload model files
if save_directory:
hf_api.upload_folder(
folder_path = save_directory,
@ -600,24 +568,20 @@ class ExportBackend:
return False, "No model loaded. Please select a checkpoint first.", None
output_path: Optional[str] = None
model_tmp_to_cleanup: Optional[str] = None
try:
# Convert quantization method to lowercase for unsloth
# unsloth expects lowercase quant method
quant_method = quantization_method.lower()
# Pin convert_hf_to_gguf.py to the same llama.cpp ref as the
# llama-quantize binary (Studio installs at a tagged ref via
# setup.sh) so it can't drift past the pinned binary's gguf API.
# Set before both branches; hub-only export has save_directory == "".
# Pin convert_hf_to_gguf.py to setup.sh's tagged llama.cpp ref so it
# can't drift past the pinned llama-quantize binary's gguf API.
global _LLAMA_CPP_SCRIPTS_WARNING_EMITTED
try:
from unsloth_zoo.llama_cpp import (
LLAMA_CPP_DEFAULT_DIR,
_resolve_local_convert_script, # noqa: F401
)
os.environ.setdefault(
"UNSLOTH_LLAMA_CPP_SCRIPTS_DIR", LLAMA_CPP_DEFAULT_DIR
)
os.environ.setdefault("UNSLOTH_LLAMA_CPP_SCRIPTS_DIR", LLAMA_CPP_DEFAULT_DIR)
except ImportError:
if not _LLAMA_CPP_SCRIPTS_WARNING_EMITTED:
logger.warning(
@ -629,71 +593,61 @@ class ExportBackend:
)
_LLAMA_CPP_SCRIPTS_WARNING_EMITTED = True
# Save locally if requested
if save_directory:
save_directory = str(resolve_export_dir(save_directory))
# Resolve to absolute path so unsloth's relative-path internals
# (check_llama_cpp, use_local_gguf, _download_convert_hf_to_gguf)
# all resolve against the repo root cwd, NOT the export directory.
save_directory = str(resolve_export_write_dir(save_directory))
# Keep unsloth relative-path internals anchored to the repo cwd.
abs_save_dir = os.path.abspath(save_directory)
logger.info(f"Saving GGUF model locally to: {abs_save_dir}")
# Create the directory if it doesn't exist
ensure_dir(Path(abs_save_dir))
# On WSL, patch out sudo check before llama.cpp build
_apply_wsl_sudo_patch()
# Snapshot existing .gguf files in cwd before conversion.
# unsloth's convert_to_gguf writes output files relative to
# cwd (repo root), so we diff afterwards and relocate them.
# convert_to_gguf writes output relative to cwd (repo root);
# snapshot existing .gguf so we can diff and relocate afterwards.
cwd = os.getcwd()
pre_existing_ggufs = set(glob.glob(os.path.join(cwd, "*.gguf")))
# Pass absolute path — no os.chdir needed.
# unsloth saves intermediate HF model files into model_save_path.
# unsloth-zoo's check_llama_cpp() uses ~/.unsloth/llama.cpp by default.
model_save_path = os.path.join(abs_save_dir, "model")
pre_existing_subs = {d.name for d in Path(abs_save_dir).iterdir() if d.is_dir()}
# Avoid clobbering an existing user-owned model/ directory.
import uuid
_model_tmp = os.path.join(abs_save_dir, f"_tmp_model_{uuid.uuid4().hex[:8]}")
model_tmp_to_cleanup = _model_tmp
self.current_model.save_pretrained_gguf(
model_save_path,
_model_tmp,
self.current_tokenizer,
quantization_method = quant_method,
)
# Relocate GGUF artifacts into the export directory.
# convert_to_gguf writes .gguf files to cwd (repo root)
# because --outfile is a relative path like "model.Q4_K_M.gguf".
new_ggufs = (
set(glob.glob(os.path.join(cwd, "*.gguf"))) - pre_existing_ggufs
)
# Relocate the .gguf that convert_to_gguf wrote to cwd (repo root).
new_ggufs = set(glob.glob(os.path.join(cwd, "*.gguf"))) - pre_existing_ggufs
for src in sorted(new_ggufs):
dest = os.path.join(abs_save_dir, os.path.basename(src))
shutil.move(src, dest)
logger.info(
f"Relocated GGUF: {os.path.basename(src)} → {abs_save_dir}/"
)
logger.info(f"Relocated GGUF: {os.path.basename(src)} → {abs_save_dir}/")
# Flatten any .gguf files from subdirectories into abs_save_dir.
# save_pretrained_gguf may create subdirs (e.g. model_gguf/)
# with a name different from model_save_path.
# Flatten GGUF files from subdirs created during this export.
for sub in list(Path(abs_save_dir).iterdir()):
if not sub.is_dir():
continue
if sub.name in pre_existing_subs:
continue
for src in sub.glob("*.gguf"):
dest = os.path.join(abs_save_dir, src.name)
shutil.move(str(src), dest)
logger.info(f"Relocated GGUF: {src.name} → {abs_save_dir}/")
# Clean up the subdirectory (intermediate HF files, etc.)
shutil.rmtree(str(sub), ignore_errors = True)
logger.info(f"Cleaned up subdirectory: {sub.name}")
# For non-PEFT models, save_pretrained_gguf redirects to the
# checkpoint path, leaving a *_gguf directory in outputs/.
# Relocate any GGUFs from there and clean it up.
# For non-PEFT models, save_pretrained_gguf leaves a *_gguf dir at
# the checkpoint path; relocate its GGUFs and clean it up.
if self.current_checkpoint:
ckpt = Path(self.current_checkpoint)
gguf_dir = ckpt.parent / f"{ckpt.name}_gguf"
if gguf_dir.is_dir():
if gguf_dir.is_dir() and gguf_dir.resolve() != Path(abs_save_dir).resolve():
for src in gguf_dir.glob("*.gguf"):
dest = os.path.join(abs_save_dir, src.name)
shutil.move(str(src), dest)
@ -701,9 +655,7 @@ class ExportBackend:
# Also relocate Ollama Modelfile if present
modelfile = gguf_dir / "Modelfile"
if modelfile.is_file():
shutil.move(
str(modelfile), os.path.join(abs_save_dir, "Modelfile")
)
shutil.move(str(modelfile), os.path.join(abs_save_dir, "Modelfile"))
logger.info(f"Relocated Modelfile → {abs_save_dir}/")
shutil.rmtree(str(gguf_dir), ignore_errors = True)
logger.info(f"Cleaned up intermediate GGUF dir: {gguf_dir}")
@ -711,8 +663,6 @@ class ExportBackend:
# Write export metadata so the Chat page can identify the base model
self._write_export_metadata(abs_save_dir)
# Log final file locations (after relocation) so it's clear
# where the GGUF files actually ended up.
final_ggufs = sorted(glob.glob(os.path.join(abs_save_dir, "*.gguf")))
logger.info(
"GGUF export complete. Final files in %s:\n %s",
@ -721,7 +671,6 @@ class ExportBackend:
)
output_path = str(Path(abs_save_dir).resolve())
# Push to hub if requested
if push_to_hub:
if not repo_id or not hf_token:
return (
@ -747,6 +696,8 @@ class ExportBackend:
)
except Exception as e:
if model_tmp_to_cleanup:
shutil.rmtree(model_tmp_to_cleanup, ignore_errors = True)
logger.error(f"Error exporting GGUF model: {e}")
import traceback
@ -775,9 +726,8 @@ class ExportBackend:
output_path: Optional[str] = None
try:
# Save locally if requested
if save_directory:
save_directory = str(resolve_export_dir(save_directory))
save_directory = str(resolve_export_write_dir(save_directory))
logger.info(f"Saving LoRA adapter locally to: {save_directory}")
ensure_dir(Path(save_directory))
@ -791,7 +741,6 @@ class ExportBackend:
logger.info(f"Adapter saved successfully to {save_directory}")
output_path = str(Path(save_directory).resolve())
# Push to hub if requested
if push_to_hub:
if not repo_id or not hf_token:
return (
@ -814,12 +763,8 @@ class ExportBackend:
repo_type = "model",
)
else:
self.current_model.push_to_hub(
repo_id, token = hf_token, private = private
)
self.current_tokenizer.push_to_hub(
repo_id, token = hf_token, private = private
)
self.current_model.push_to_hub(repo_id, token = hf_token, private = private)
self.current_tokenizer.push_to_hub(repo_id, token = hf_token, private = private)
logger.info(f"Adapter pushed successfully to {repo_id}")
return True, "LoRA adapter exported successfully", output_path

141
studio/backend/core/export/orchestrator.py
View file

@ -1,15 +1,13 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""
Export orchestrator — subprocess-based.
"""Export orchestrator — subprocess-based.
Provides the same API as ExportBackend, but delegates all ML work
to a persistent subprocess. The subprocess is spawned on first checkpoint
load and stays alive for subsequent export operations.
Same API as ExportBackend, but delegates all ML work to a persistent
subprocess spawned on first checkpoint load and reused for later exports.
When switching between checkpoints that need different transformers versions,
the old subprocess is killed and a new one is spawned with the correct version.
When switching between checkpoints needing different transformers
versions, the old subprocess is killed and a new one spawned.
Pattern follows core/inference/orchestrator.py.
"""
@ -30,9 +28,7 @@ logger = get_logger(__name__)
_CTX = mp.get_context("spawn")
# Maximum number of captured log lines kept in memory per export
# orchestrator. Acts as scrollback for the live export log panel in the
# UI. 4000 lines is ~1 MB worst-case at 256 chars/line.
# Max log lines kept per orchestrator (live log panel scrollback); ~1 MB worst-case.
_LOG_BUFFER_MAXLEN = 4000
@ -41,46 +37,31 @@ class ExportOrchestrator:
Export backend orchestrator — subprocess-based.
Exposes the same API surface as ExportBackend so routes/export.py
needs minimal changes. Internally, all heavy ML operations happen in
a persistent subprocess.
needs minimal changes. All heavy ML work happens in a persistent
subprocess.
"""
def __init__(self):
# Subprocess state
self._proc: Optional[mp.Process] = None
self._cmd_queue: Any = None
self._resp_queue: Any = None
# Serializes export operations (load_checkpoint, export_*,
# cleanup) so concurrent HTTP requests can never interleave
# commands on the subprocess queue. Previously unused.
# Serializes export ops so concurrent HTTP requests can't interleave commands.
self._lock = threading.Lock()
# Local state mirrors (updated from subprocess responses)
# Local state mirrors (updated from subprocess responses).
self.current_checkpoint: Optional[str] = None
self.is_vision: bool = False
self.is_peft: bool = False
# ── Live log capture ─────────────────────────────────────
# Thread-safe ring buffer of log lines forwarded from the
# worker subprocess. Powers the GET /api/export/logs/stream
# SSE endpoint that the export dialog consumes.
# Thread-safe ring buffer of worker log lines; powers the export logs SSE endpoint.
self._log_buffer: Deque[Dict[str, Any]] = deque(maxlen = _LOG_BUFFER_MAXLEN)
self._log_lock = threading.Lock()
# Monotonically increasing sequence number. Never reset across
# operations, so SSE clients can use it as a stable cursor even
# if clear_logs() is called mid-session.
# Monotonic seq, never reset, so SSE clients have a stable cursor across clear_logs().
self._log_seq: int = 0
# Snapshot of _log_seq captured at the start of the current run
# (updated by clear_logs()). The SSE endpoint defaults its
# cursor to this value so a client that connects AFTER the
# worker has already emitted its first lines still sees the
# full run. Every line appended during the current run has seq
# strictly greater than _run_start_seq, and every line from
# prior runs has seq less than or equal to it.
# _log_seq snapshot at the current run's start; SSE defaults its cursor here so a
# late-connecting client still sees the full run. Current run has seq > this.
self._run_start_seq: int = 0
# True while an export operation (load/export/cleanup) is
# running. The SSE endpoint ends the stream 1 second after
# this flips back to False to drain any trailing log lines.
# True while an export op runs; SSE ends the stream 1s after this flips False.
self._export_active: bool = False
atexit.register(self._cleanup)
@ -91,13 +72,7 @@ class ExportOrchestrator:
# ------------------------------------------------------------------
def _append_log(self, entry: Dict[str, Any]) -> None:
"""Append a log line from the worker subprocess to the buffer.
Entries look like {"type": "log", "stream": "stdout"|"stderr",
"line": "...", "ts": ...}. Each is stamped with a monotonic
seq number before it lands in the buffer so SSE clients can
cursor through new lines.
"""
"""Append a worker log line to the buffer, stamped with a monotonic seq."""
line = entry.get("line")
if not line:
return
@ -113,18 +88,10 @@ class ExportOrchestrator:
)
def clear_logs(self) -> None:
"""Drop any buffered log lines from a previous operation.
"""Drop buffered log lines from a previous op so the UI shows only this run.
Called at the start of each export op so the UI shows only the
output of the current run. The seq counter is NOT reset, so an
SSE client that captured the cursor before clear_logs() will
still see new lines (with strictly greater seq numbers).
Also snapshots the current seq into ``_run_start_seq`` so the
SSE endpoint can anchor its default cursor at the start of
this run. Anything appended after this call has seq strictly
greater than the snapshot and is reachable via
``get_logs_since(get_run_start_seq())``.
The seq counter is NOT reset (clients keep a stable cursor); the current seq
is snapshotted into ``_run_start_seq`` to anchor the SSE default cursor.
"""
with self._log_lock:
self._log_buffer.clear()
@ -144,12 +111,7 @@ class ExportOrchestrator:
return self._log_seq
def get_run_start_seq(self) -> int:
"""Return the seq value captured at the start of the current run.
The SSE endpoint uses this as the default cursor so a client
that connects AFTER the worker has already started emitting
output still sees every line from the current run.
"""
"""Return the seq captured at the current run's start (SSE default cursor)."""
with self._log_lock:
return self._run_start_seq
@ -193,22 +155,19 @@ class ExportOrchestrator:
self._proc = None
return
# 1. Drain stale responses
self._drain_queue()
# 2. Send shutdown command
try:
self._cmd_queue.put({"type": "shutdown"})
except (OSError, ValueError):
pass
# 3. Wait for graceful shutdown
try:
self._proc.join(timeout = timeout)
except Exception:
pass
# 4. Force kill if still alive
# Force kill if still alive.
if self._proc is not None and self._proc.is_alive():
logger.warning("Export subprocess did not exit gracefully, terminating")
try:
@ -261,12 +220,15 @@ class ExportOrchestrator:
except (EOFError, OSError, ValueError):
return None
def _wait_response(self, expected_type: str, timeout: float = 3600.0) -> dict:
def _wait_response(
self,
expected_type: str,
timeout: float = 3600.0,
) -> dict:
"""Block until a response of the expected type arrives.
Export operations can take a very long time — GGUF conversion for
large models (30B+) easily takes 20-30 minutes. Default timeout
is 1 hour.
Export ops can take a long time — GGUF conversion for large
models (30B+) easily takes 20-30 minutes. Default timeout 1 hour.
"""
deadline = time.monotonic() + timeout
@ -275,7 +237,6 @@ class ExportOrchestrator:
resp = self._read_resp(timeout = min(remaining, 2.0))
if resp is None:
# Check subprocess health
if not self._ensure_subprocess_alive():
raise RuntimeError("Export subprocess crashed during wait")
continue
@ -290,17 +251,14 @@ class ExportOrchestrator:
raise RuntimeError(f"Subprocess error: {error_msg}")
if rtype == "log":
# Forwarded stdout/stderr line from the worker process.
# Forwarded stdout/stderr line from the worker.
self._append_log(resp)
continue
if rtype == "status":
message = resp.get("message", "")
logger.info("Export subprocess status: %s", message)
# Surface status messages in the live log panel too so
# users see high level progress (e.g. "Importing
# Unsloth...", "Loading checkpoint: ...") alongside
# subprocess output.
# Surface status in the live log panel for high-level progress.
if message:
self._append_log(
{
@ -311,16 +269,14 @@ class ExportOrchestrator:
)
continue
# Other response types during wait — skip
# Other response types during wait — skip.
logger.debug(
"Skipping response type '%s' while waiting for '%s'",
rtype,
expected_type,
)
raise RuntimeError(
f"Timeout waiting for '{expected_type}' response after {timeout}s"
)
raise RuntimeError(f"Timeout waiting for '{expected_type}' response after {timeout}s")
def _drain_queue(self) -> list:
"""Drain all pending responses."""
@ -360,20 +316,17 @@ class ExportOrchestrator:
}
with self._lock:
# Start a fresh log buffer for this operation so the UI
# sees only the current run's output.
# Fresh log buffer so the UI sees only this run's output.
self.clear_logs()
self._export_active = True
try:
# Always kill existing subprocess and spawn fresh.
# Always kill any existing subprocess and spawn fresh.
if self._ensure_subprocess_alive():
self._shutdown_subprocess()
elif self._proc is not None:
self._shutdown_subprocess(timeout = 2)
logger.info(
"Spawning fresh export subprocess for '%s'", checkpoint_path
)
logger.info("Spawning fresh export subprocess for '%s'", checkpoint_path)
self._spawn_subprocess(sub_config)
try:
@ -485,17 +438,11 @@ class ExportOrchestrator:
},
)
def _run_export(
self, export_type: str, params: dict
) -> Tuple[bool, str, Optional[str]]:
"""Send an export command to the subprocess and wait for result.
def _run_export(self, export_type: str, params: dict) -> Tuple[bool, str, Optional[str]]:
"""Send an export command and wait for the result.
Returns ``(success, message, output_path)``. ``output_path`` is the
resolved on-disk directory the worker actually wrote to (None when
the export only pushed to Hub or failed before any file was
written). Surfaced via the export route's ``details.output_path``
so the dialog's success screen can show the user where the model
landed.
Returns ``(success, message, output_path)``. ``output_path`` is the on-disk
dir the worker wrote to (None if it only pushed to Hub or failed pre-write).
"""
with self._lock:
if not self._ensure_subprocess_alive():
@ -529,7 +476,6 @@ class ExportOrchestrator:
"""Cleanup export-related models from memory."""
with self._lock:
if not self._ensure_subprocess_alive():
# No subprocess — just clear local state
self.current_checkpoint = None
self.is_vision = False
self.is_peft = False
@ -544,7 +490,7 @@ class ExportOrchestrator:
except RuntimeError:
success = False
# Shut down subprocess after cleanup — no model loaded
# Shut down subprocess after cleanup — no model loaded.
self._shutdown_subprocess()
self.current_checkpoint = None
@ -554,12 +500,9 @@ class ExportOrchestrator:
finally:
self._export_active = False
def scan_checkpoints(
self, outputs_dir: str = str(outputs_root())
) -> List[Tuple[str, list]]:
"""Scan for checkpoints — no ML imports needed, runs locally."""
def scan_checkpoints(self, outputs_dir: str = str(outputs_root())) -> List[Tuple[str, list]]:
"""Scan for checkpoints — runs locally, no ML imports."""
from utils.models.checkpoints import scan_checkpoints
return scan_checkpoints(outputs_dir = outputs_dir)

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

@ -1,16 +1,12 @@
# 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 subprocess entry point.
"""Export subprocess entry point.
Each export session runs in a persistent subprocess (mp.get_context("spawn")).
This gives us a clean Python interpreter with no stale module state —
solving the transformers version-switching problem completely.
The subprocess stays alive while a model is loaded, accepting commands
(load, export_merged, export_base, export_gguf, export_lora, cleanup,
shutdown) via mp.Queue.
Each export session runs in a persistent subprocess (mp spawn), giving a clean
interpreter with no stale module state, which solves transformers version
switching. The subprocess stays alive while a model is loaded, accepting commands
(load, export_*, cleanup, shutdown) via mp.Queue.
Pattern follows core/inference/worker.py and core/training/worker.py.
"""
@ -31,42 +27,31 @@ from typing import Any
logger = get_logger(__name__)
# Gate that controls whether captured stdout/stderr lines are forwarded
# to the parent's resp_queue (and from there to the export-dialog SSE
# stream). Closed by default so the noisy bootstrap phase -- transformers
# venv activation, Unsloth/torch imports, base-model resolution, "Top
# GGUF/hub models" lists, vision detection, weight loading bars -- is
# suppressed in the UI. _handle_export() opens the gate at the start of
# the actual export work and leaves it open; the orchestrator always
# spawns a fresh subprocess for the next checkpoint load (see
# orchestrator._spawn_subprocess) which resets this state.
#
# Lines dropped while the gate is closed are still echoed to the saved
# original stdout/stderr fds so the server console / log file keeps the
# full output for debugging.
# Gate controlling whether captured stdout/stderr lines are forwarded to the
# parent's resp_queue (and on to the export-dialog SSE stream). Closed by default
# so the noisy bootstrap phase (imports, model resolution, loading bars) is
# suppressed in the UI; _handle_export() opens it when export work starts. The
# orchestrator spawns a fresh subprocess per checkpoint load, resetting this.
# Dropped lines are still echoed to the saved fds so the server log keeps them.
_log_forward_gate = threading.Event()
def _setup_log_capture(resp_queue: Any) -> None:
"""Redirect fds 1 and 2 through pipes so every line printed by this
worker process and any child process it spawns is forwarded to the
parent process via resp_queue as {"type": "log", ...} messages.
"""Redirect fds 1 and 2 through pipes so every line printed by this worker
and any child it spawns is forwarded to the parent via resp_queue as
{"type": "log", ...} messages.
Must be called BEFORE LogConfig.setup_logging and BEFORE any ML
imports, otherwise library handlers may capture the original stderr
reference and bypass the pipe.
Lines are also echoed back to the original stdout/stderr so the
server console keeps receiving the full subprocess output, even
while ``_log_forward_gate`` is closed.
Must run BEFORE LogConfig.setup_logging and any ML imports, else library
handlers may capture the original stderr reference and bypass the pipe.
Lines are also echoed back to the original fds so the server console keeps
the full output even while ``_log_forward_gate`` is closed.
"""
try:
saved_out_fd = os.dup(1)
saved_err_fd = os.dup(2)
except OSError:
# dup failed (exotic platforms) - give up quietly, export still
# works, just no live log streaming.
# dup failed; give up quietly (export still works, no live streaming).
return
try:
@ -88,13 +73,11 @@ def _setup_log_capture(resp_queue: Any) -> None:
pass
return
# Close the write ends we just dup2'd (fds 1 and 2 are the real
# write ends now).
# Close the write ends we just dup2'd (fds 1 and 2 are the real write ends).
os.close(w_out)
os.close(w_err)
# Replace Python's sys.stdout/sys.stderr with line-buffered writers
# bound to the (now-redirected) fds 1 and 2.
# Replace sys.stdout/sys.stderr with line-buffered writers on fds 1 and 2.
try:
sys.stdout = os.fdopen(1, "w", buffering = 1, encoding = "utf-8", errors = "replace")
sys.stderr = os.fdopen(2, "w", buffering = 1, encoding = "utf-8", errors = "replace")
@ -112,8 +95,7 @@ def _setup_log_capture(resp_queue: Any) -> None:
continue
if not chunk:
break
# Echo to the original fd so the server console still sees
# the full output.
# Echo to the original fd so the server console keeps the full output.
try:
os.write(echo_fd, chunk)
except OSError:
@ -133,9 +115,8 @@ def _setup_log_capture(resp_queue: Any) -> None:
if not line:
continue
if not _log_forward_gate.is_set():
# Gate closed (bootstrap phase) -- already echoed to
# the saved console fd above; drop the line so the
# export dialog doesn't see import / vendoring noise.
# Gate closed (bootstrap): already echoed above; drop the
# line so the export dialog skips import noise.
continue
try:
resp_queue.put_nowait(
@ -147,8 +128,7 @@ def _setup_log_capture(resp_queue: Any) -> None:
}
)
except Exception:
# Queue put failed (full, closed, etc.) - drop the
# line rather than crash the reader thread.
# Queue put failed; drop the line rather than crash the thread.
pass
if buf and _log_forward_gate.is_set():
try:
@ -181,7 +161,7 @@ def _setup_log_capture(resp_queue: Any) -> None:
def _activate_transformers_version(model_name: str) -> None:
"""Activate the correct transformers version BEFORE any ML imports."""
# Ensure backend is on path for utils imports
# Ensure backend is on sys.path for utils imports.
backend_path = str(Path(__file__).resolve().parent.parent.parent)
if backend_path not in sys.path:
sys.path.insert(0, backend_path)
@ -267,11 +247,9 @@ def _handle_export(backend, cmd: dict, resp_queue: Any) -> None:
export_type = cmd["export_type"] # "merged", "base", "gguf", "lora"
response_type = f"export_{export_type}_done"
# Open the log forwarding gate so the user sees the actual export
# progress (Unsloth merge bars, file copies, GGUF conversion, etc.)
# in the live log panel. The gate stays open for the rest of this
# subprocess's life; the orchestrator spawns a fresh subprocess for
# the next checkpoint load, which resets the gate to closed.
# Open the log forwarding gate so the user sees export progress in the live
# log panel. Stays open for the rest of this subprocess's life; the
# orchestrator spawns a fresh subprocess per checkpoint load, resetting it.
_log_forward_gate.set()
output_path: Any = None
@ -362,12 +340,7 @@ def _handle_cleanup(backend, resp_queue: Any) -> None:
)
def run_export_process(
*,
cmd_queue: Any,
resp_queue: Any,
config: dict,
) -> None:
def run_export_process(*, cmd_queue: Any, resp_queue: Any, config: dict) -> None:
"""Subprocess entrypoint. Persistent — runs command loop until shutdown.
Args:
@ -377,25 +350,16 @@ def run_export_process(
"""
import queue as _queue
# Install fd-level stdout/stderr capture FIRST so every subsequent
# print and every child process inherits the redirected fds. This
# is what powers the live export log stream in the UI.
# Install fd-level stdout/stderr capture FIRST so every subsequent print and
# every child process inherits the redirected fds (powers the live log stream).
_setup_log_capture(resp_queue)
os.environ["TOKENIZERS_PARALLELISM"] = "false"
os.environ["PYTHONWARNINGS"] = (
"ignore" # Suppress warnings at C-level before imports
)
# Force unbuffered output from any child Python process (e.g. the
# GGUF converter) so their prints surface in the log stream as they
# happen rather than at the end.
os.environ["PYTHONWARNINGS"] = "ignore" # suppress C-level warnings before imports
# Unbuffered output from child Python (e.g. GGUF converter) so prints surface live.
os.environ["PYTHONUNBUFFERED"] = "1"
# tqdm defaults to a 10-second mininterval when stdout is not a tty
# (which it isn't here -- we redirected fd 1/2 to a pipe). That makes
# multi-step progress bars look frozen in the export log panel. Force
# frequent flushes so the user sees movement during merge / GGUF
# conversion. Has no effect on single-step bars (e.g. "Copying 1
# files") which only emit start/end events regardless.
# tqdm defaults to a 10s mininterval when stdout isn't a tty (we redirected
# fd 1/2 to a pipe), making multi-step bars look frozen; force frequent flushes.
os.environ.setdefault("TQDM_MININTERVAL", "0.5")
import warnings
@ -426,11 +390,10 @@ def run_export_process(
)
return
# ── 1b. On Windows, check Triton availability (must be before import torch) ──
# ── 1b. Check Triton on Windows (must precede import torch) ──
if sys.platform == "win32":
try:
import triton # noqa: F401
logger.info("Triton available — torch.compile enabled")
except ImportError:
os.environ["TORCHDYNAMO_DISABLE"] = "1"
@ -439,6 +402,13 @@ def run_export_process(
'Install for better performance: pip install "triton-windows<3.7"'
)
# ── 1c. Stub torchao on Windows ROCm ──
# See core/_torchao_stub.py: torchao crashes on Windows ROCm (RCCL absent).
# No-op off Windows ROCm. Must run before importing transformers / unsloth_zoo.
from core._torchao_stub import install_torchao_windows_rocm_stub
install_torchao_windows_rocm_stub()
# ── 2. Import ML libraries (fresh in this clean process) ──
try:
_send_response(
@ -458,9 +428,7 @@ def run_export_process(
import transformers
logger.info(
"Export subprocess loaded transformers %s", transformers.__version__
)
logger.info("Export subprocess loaded transformers %s", transformers.__version__)
except Exception as exc:
_send_response(
@ -512,7 +480,7 @@ def run_export_process(
try:
if cmd_type == "load":
# Load a new checkpoint (reusing this subprocess)
# Load a new checkpoint, reusing this subprocess.
backend.cleanup_memory()
_handle_load(backend, cmd, resp_queue)
@ -561,9 +529,7 @@ def run_export_process(
)
except Exception as exc:
logger.error(
"Error handling command '%s': %s", cmd_type, exc, exc_info = True
)
logger.error("Error handling command '%s': %s", cmd_type, exc, exc_info = True)
_send_response(
resp_queue,
{

8
studio/backend/core/inference/__init__.py
View file

@ -2,17 +2,17 @@
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""
Inference submodule - Inference backend for model loading and generation
Inference submodule - backend for model loading and generation.
The default get_inference_backend() returns an InferenceOrchestrator that
delegates to a subprocess. The original InferenceBackend runs inside
the subprocess and can be imported directly from .inference when needed.
delegates to a subprocess. The original InferenceBackend runs inside the
subprocess and can be imported directly from .inference when needed.
"""
from .orchestrator import InferenceOrchestrator, get_inference_backend
from .llama_cpp import LlamaCppBackend
# Expose InferenceOrchestrator as InferenceBackend for backward compat
# Expose InferenceOrchestrator as InferenceBackend for backward compat.
InferenceBackend = InferenceOrchestrator
__all__ = [

68
studio/backend/core/inference/_html_to_md.py
View file

@ -5,7 +5,7 @@
Minimal HTML-to-Markdown converter using only the standard library.
Replaces the external ``html2text`` (GPL-3.0) dependency with a ~250-line
``html.parser.HTMLParser`` subclass. Covers headings, links, bold/italic,
``html.parser.HTMLParser`` subclass. Covers headings, links, bold/italic,
lists, tables, blockquotes, code blocks, and entity decoding.
"""
@ -81,9 +81,7 @@ class _MarkdownRenderer(HTMLParser):
self._pre_parts: list[str] = []
self._in_inline_code: bool = False
# Blockquote state -- stack of output buffers so nested
# blockquotes each collect their own content and get prefixed
# with the correct number of ">" markers on close.
# Blockquote state: stack of buffers so nested blockquotes get the right ">" depth.
self._bq_stack: list[list[str]] = []
# ------------------------------------------------------------------
@ -102,7 +100,7 @@ class _MarkdownRenderer(HTMLParser):
# ------------------------------------------------------------------
def _prefix_blockquote(self, content: str) -> str:
"""Prefix every line of *content* with ``> ``."""
# Strip trailing whitespace first, then collapse blank lines
# Strip trailing whitespace, then collapse blank lines.
content = re.sub(r"[ \t]+$", "", content, flags = re.MULTILINE)
content = re.sub(r"\n{3,}", "\n\n", content).strip()
if not content:
@ -116,10 +114,7 @@ class _MarkdownRenderer(HTMLParser):
prefixed.append(">")
return "\n".join(prefixed)
# ------------------------------------------------------------------
# Table helpers -- flush open cells and rows so that HTML with
# omitted optional end tags (</td>, </tr>) does not lose data.
# ------------------------------------------------------------------
# Table helpers: flush open cells/rows so omitted </td>/</tr> don't lose data.
def _finish_cell(self) -> None:
if not self._in_cell:
return
@ -142,10 +137,7 @@ class _MarkdownRenderer(HTMLParser):
self._current_row = []
self._row_has_th = False
# ------------------------------------------------------------------
# Link text helper -- normalize whitespace so block-level content
# inside an <a> does not produce multiline Markdown link labels.
# ------------------------------------------------------------------
# Link text helper: normalize whitespace so block content in <a> stays single-line.
def _finish_link(self) -> None:
text = re.sub(r"\s+", " ", "".join(self._link_text_parts)).strip()
href = self._link_href or ""
@ -234,22 +226,19 @@ class _MarkdownRenderer(HTMLParser):
self._emit("\n\n")
elif tag == "tr":
# Flush any open cell/row from a previous row that may
# have omitted its optional </td> or </tr> end tags.
# Flush open cell/row from a prior row that omitted </td>/</tr>.
self._finish_cell()
self._finish_row()
elif tag in ("th", "td"):
# Flush any open cell (handles omitted </td>/<th>)
self._finish_cell()
self._finish_cell() # handles omitted </td>/</th>
self._cell_parts = []
self._in_cell = True
if tag == "th":
self._row_has_th = True
elif tag == "img":
# Skip images -- keeps fetched page text focused on readable
# content and avoids data-URI amplification.
# Skip images: keeps text readable, avoids data-URI amplification.
return
def handle_endtag(self, tag: str) -> None:
@ -310,7 +299,7 @@ class _MarkdownRenderer(HTMLParser):
self._finish_row()
elif tag == "table":
# Flush any remaining row (handles omitted </tr>)
# Flush remaining row (handles omitted </tr>).
self._finish_cell()
self._finish_row()
self._in_table = False
@ -325,15 +314,13 @@ class _MarkdownRenderer(HTMLParser):
if self._in_pre:
self._pre_parts.append(data)
return
# Preserve literal whitespace inside inline <code> spans
# Preserve literal whitespace inside inline <code> spans.
if self._in_inline_code:
self._emit(data)
return
# Collapse all whitespace (including newlines) per HTML rules
# Collapse all whitespace (including newlines) per HTML rules.
text = re.sub(r"\s+", " ", data)
# Suppress whitespace-only text nodes between table structural
# elements (indentation from source HTML) to prevent leading
# spaces from breaking Markdown table row alignment.
# Suppress whitespace-only nodes between table elements (source indentation).
if self._in_table and not self._in_cell and not text.strip():
return
self._emit(text)
@ -348,18 +335,10 @@ class _MarkdownRenderer(HTMLParser):
return
self._emit(html.unescape(f"&#{name};"))
# ------------------------------------------------------------------
# Flush pending buffers (handles truncated HTML from capped fetches)
# ------------------------------------------------------------------
def flush_pending(self) -> None:
"""Flush any open side-buffers into ``_out``.
Called after ``close()`` to recover content from truncated HTML
where closing tags were never seen (common when ``_fetch_page_text``
caps the download by byte count).
"""
"""Flush open side-buffers into ``_out`` after close(), recovering truncated HTML."""
# Flush innermost buffers first so their content propagates outward.
if self._in_link:
self._finish_link()
@ -376,7 +355,7 @@ class _MarkdownRenderer(HTMLParser):
block = "```\n" + raw + "\n```"
self._emit("\n\n" + block + "\n\n")
# Flatten any open blockquote buffers (innermost first)
# Flatten any open blockquote buffers (innermost first).
while self._bq_stack:
content = "".join(self._bq_stack.pop())
prefixed = self._prefix_blockquote(content)
@ -388,15 +367,9 @@ class _MarkdownRenderer(HTMLParser):
self._out.append("\n\n" + prefixed + "\n\n")
# ------------------------------------------------------------------
# Post-processing
# ------------------------------------------------------------------
def _cleanup(text: str) -> str:
"""Normalize whitespace and blank lines in the final output.
Preserves content inside fenced code blocks verbatim so that
intentional blank lines in ``<pre>`` content are not collapsed.
"""
"""Normalize whitespace and blank lines, preserving fenced code blocks verbatim."""
lines = text.split("\n")
out: list[str] = []
in_fence = False
@ -411,7 +384,6 @@ def _cleanup(text: str) -> str:
continue
if in_fence:
# Preserve code block content exactly as-is
out.append(line)
continue
@ -427,17 +399,13 @@ def _cleanup(text: str) -> str:
return "\n".join(out).strip()
# ------------------------------------------------------------------
# Public API
# ------------------------------------------------------------------
def html_to_markdown(source_html: str) -> str:
"""Convert an HTML string to Markdown.
"""Convert HTML to Markdown (headings, links, emphasis, lists, tables, blockquotes, code, entities).
Handles headings, links, bold/italic, lists (ordered and unordered),
tables, blockquotes, code blocks, and HTML entities. ``<script>``,
``<style>``, and ``<head>`` sections are stripped entirely.
``<script>``, ``<style>``, and ``<head>`` are stripped entirely.
"""
# Normalize line endings before parsing
# Normalize line endings before parsing.
source_html = source_html.replace("\r\n", "\n").replace("\r", "\n")
renderer = _MarkdownRenderer()
renderer.feed(source_html)

219
studio/backend/core/inference/anthropic_compat.py
View file

@ -4,15 +4,43 @@
"""
Anthropic Messages API ↔ OpenAI format translation utilities.
Pure functions and a stateful stream emitter — no FastAPI, no I/O.
Pure functions plus stateful stream emitters; no FastAPI, no I/O.
"""
from __future__ import annotations
import json
import uuid
from typing import Any, Optional, Union
def openai_finish_to_anthropic_stop(finish_reason, had_tool_calls = False) -> str:
"""Map an OpenAI finish_reason to an Anthropic stop_reason.
'length' -> 'max_tokens' (truncation wins even mid tool call, so a cut-off
tool call isn't mislabeled tool_use); tool_calls / had_tool_calls -> 'tool_use';
'stop_sequence' -> 'stop_sequence'; 'stop'/None/unknown -> 'end_turn'."""
# Truncation takes precedence: a tool call cut off at max_tokens has possibly
# incomplete arguments, so report max_tokens rather than telling the client to
# run the tool.
if finish_reason == "length":
return "max_tokens"
if finish_reason == "tool_calls" or had_tool_calls:
return "tool_use"
if finish_reason == "stop_sequence":
return "stop_sequence"
# "stop", None, and any unknown value collapse to end_turn.
return "end_turn"
def anthropic_tool_use_id(upstream_id = None) -> str:
"""Return an Anthropic-style tool_use id (prefix 'toolu_'). Reuses an
upstream id only if it already starts with 'toolu_'; otherwise mints a fresh
'toolu_<24 hex>'."""
if upstream_id and isinstance(upstream_id, str) and upstream_id.startswith("toolu_"):
return upstream_id
return f"toolu_{uuid.uuid4().hex[:24]}"
def _anthropic_image_block_to_openai_part(block: dict) -> Optional[dict]:
"""Translate one Anthropic ``image`` block to an OpenAI ``image_url`` part.
@ -42,14 +70,13 @@ def _anthropic_image_block_to_openai_part(block: dict) -> Optional[dict]:
def anthropic_messages_to_openai(
messages: list[dict],
system: Optional[Union[str, list]] = None,
messages: list[dict], system: Optional[Union[str, list]] = None
) -> list[dict]:
"""Convert Anthropic messages + system to OpenAI-format message dicts.
User messages that carry ``image`` blocks are emitted as OpenAI
multimodal content arrays (``[{type: "text", ...}, {type: "image_url", ...}]``)
so they flow through llama-server's native vision pathway.
User messages with ``image`` blocks are emitted as OpenAI multimodal
content arrays (``[{type: "text", ...}, {type: "image_url", ...}]``) so
they flow through llama-server's native vision pathway.
"""
result: list[dict] = []
@ -76,8 +103,7 @@ def anthropic_messages_to_openai(
continue
if role == "assistant":
# Assistant content carries text + tool_use; images aren't
# part of Anthropic's assistant content model.
# Assistant content: text + tool_use only (no images in Anthropic's model).
text_parts: list[str] = []
tool_calls: list[dict] = []
for block in content:
@ -105,9 +131,7 @@ def anthropic_messages_to_openai(
continue
if role == "user":
# Build an ordered part list so text/image interleaving is
# preserved (e.g. [text, image, text, image]). tool_result
# blocks become their own OpenAI "tool" role messages.
# Ordered parts preserve text/image interleaving; tool_result -> own "tool" messages.
user_parts: list[dict] = []
has_image = False
tool_results: list[dict] = []
@ -125,9 +149,7 @@ def anthropic_messages_to_openai(
tc = b.get("content", "")
if isinstance(tc, list):
tc = " ".join(
p["text"]
for p in tc
if isinstance(p, dict) and p.get("type") == "text"
p["text"] for p in tc if isinstance(p, dict) and p.get("type") == "text"
)
tool_results.append(
{
@ -140,8 +162,7 @@ def anthropic_messages_to_openai(
if has_image:
result.append({"role": "user", "content": user_parts})
else:
# No images — collapse text parts to a plain string so
# existing text-only callers keep their simple shape.
# No images: collapse text parts to a plain string.
text = "\n".join(p["text"] for p in user_parts)
if text:
result.append({"role": "user", "content": text})
@ -184,8 +205,8 @@ def anthropic_tool_choice_to_openai(tc: Any) -> Any:
- ``{"type": "tool", "name": "get_weather"}``
→ ``{"type": "function", "function": {"name": "get_weather"}}``
Returns ``None`` for ``None`` or any unrecognized shape (caller may
then fall back to its own default, typically ``"auto"``).
Returns ``None`` for ``None`` or any unrecognized shape (caller falls
back to its own default, typically ``"auto"``).
"""
if tc is None:
return None
@ -211,17 +232,40 @@ def build_anthropic_sse_event(event_type: str, data: dict) -> str:
return f"event: {event_type}\ndata: {json.dumps(data)}\n\n"
def _message_delta_usage(usage: Optional[dict]) -> dict:
"""Usage block for a message_delta event (cumulative token counts). Cache
fields are always 0 — no prompt caching backend. ``usage`` may be None when a
metadata event carried usage=None (e.g. only finish_reason set)."""
usage = usage or {}
return {
"input_tokens": usage.get("prompt_tokens", 0),
"cache_creation_input_tokens": 0,
"cache_read_input_tokens": 0,
"output_tokens": usage.get("completion_tokens", 0),
}
class AnthropicStreamEmitter:
"""Converts generator events from generate_chat_completion_with_tools()
into Anthropic Messages SSE strings."""
"""Converts generate_chat_completion_with_tools() events into Anthropic
Messages SSE strings."""
def __init__(self) -> None:
self.block_index: int = 0
self._text_block_open: bool = False
self._open_tool_call_id: Optional[str] = None
# The mapped Anthropic ``toolu_*`` id published in content_block_start,
# reused for the paired tool_result so consumers can correlate them.
self._open_tool_use_id: Optional[str] = None
self._open_tool_args_sent: bool = False
self._prev_text: str = ""
self._usage: dict = {}
def start(self, message_id: str, model: str) -> list[str]:
def start(
self,
message_id: str,
model: str,
input_tokens: int = 0,
) -> list[str]:
"""Emit message_start and open the first text content block."""
events = []
events.append(
@ -237,7 +281,12 @@ class AnthropicStreamEmitter:
"model": model,
"stop_reason": None,
"stop_sequence": None,
"usage": {"input_tokens": 0, "output_tokens": 0},
"usage": {
"input_tokens": input_tokens,
"output_tokens": 0,
"cache_creation_input_tokens": 0,
"cache_read_input_tokens": 0,
},
},
},
)
@ -260,20 +309,28 @@ class AnthropicStreamEmitter:
# status events — no Anthropic equivalent
return []
def finish(self, stop_reason: str = "end_turn") -> list[str]:
def finish(
self,
stop_reason: str = "end_turn",
stop_sequence = None,
) -> list[str]:
"""Close any open block and emit message_delta + message_stop."""
events = []
if self._text_block_open:
if self._text_block_open or self._open_tool_call_id is not None:
events.append(self._close_block())
self._open_tool_call_id = None
self._open_tool_use_id = None
self._open_tool_args_sent = False
events.append(
build_anthropic_sse_event(
"message_delta",
{
"type": "message_delta",
"delta": {"stop_reason": stop_reason, "stop_sequence": None},
"usage": {
"output_tokens": self._usage.get("completion_tokens", 0),
"delta": {
"stop_reason": stop_reason,
"stop_sequence": stop_sequence,
},
"usage": _message_delta_usage(self._usage),
},
)
)
@ -310,12 +367,26 @@ class AnthropicStreamEmitter:
return events
def _handle_tool_start(self, event: dict) -> list[str]:
tool_call_id = event.get("tool_call_id", "")
args = event.get("arguments", {})
if tool_call_id and self._open_tool_call_id == tool_call_id:
return self._tool_arguments_delta(args)
events = []
# Close current text block if open
if self._text_block_open:
events.append(self._close_block())
# Open a tool_use block
# Defensive: close a stale open tool_use block before starting another.
elif self._open_tool_call_id is not None:
events.append(self._close_block())
self._open_tool_call_id = None
self._open_tool_use_id = None
self._open_tool_args_sent = False
# Open a tool_use block.
self.block_index += 1
self._open_tool_call_id = tool_call_id
self._open_tool_use_id = anthropic_tool_use_id(tool_call_id)
self._open_tool_args_sent = False
events.append(
build_anthropic_sse_event(
"content_block_start",
@ -324,42 +395,54 @@ class AnthropicStreamEmitter:
"index": self.block_index,
"content_block": {
"type": "tool_use",
"id": event.get("tool_call_id", ""),
"id": self._open_tool_use_id,
"name": event.get("tool_name", ""),
"input": {},
},
},
)
)
# Emit the arguments as input_json_delta
args = event.get("arguments", {})
if args:
events.append(
build_anthropic_sse_event(
"content_block_delta",
{
"type": "content_block_delta",
"index": self.block_index,
"delta": {
"type": "input_json_delta",
"partial_json": json.dumps(args),
},
},
)
)
events.extend(self._tool_arguments_delta(args))
return events
def _tool_arguments_delta(self, args: dict) -> list[str]:
if not args:
return []
if self._open_tool_args_sent:
return []
self._open_tool_args_sent = True
return [
build_anthropic_sse_event(
"content_block_delta",
{
"type": "content_block_delta",
"index": self.block_index,
"delta": {
"type": "input_json_delta",
"partial_json": json.dumps(args),
},
},
)
]
def _handle_tool_end(self, event: dict) -> list[str]:
events = []
# Close the tool_use block
events.append(self._close_block())
# Close the tool_use block.
if self._open_tool_call_id is not None or self._text_block_open:
events.append(self._close_block())
# Reuse the id published in content_block_start; fall back to mapping
# the raw id only if no tool_start preceded this end.
tool_use_id = self._open_tool_use_id or anthropic_tool_use_id(event.get("tool_call_id", ""))
self._open_tool_call_id = None
self._open_tool_use_id = None
self._open_tool_args_sent = False
# Emit custom tool_result event (non-standard, ignored by SDKs)
events.append(
build_anthropic_sse_event(
"tool_result",
{
"type": "tool_result",
"tool_use_id": event.get("tool_call_id", ""),
"tool_use_id": tool_use_id,
"content": event.get("result", ""),
},
)
@ -398,10 +481,10 @@ class AnthropicStreamEmitter:
class AnthropicPassthroughEmitter:
"""Converts llama-server's OpenAI-format streaming chunks into Anthropic SSE.
Used for the client-side tool-use pass-through path: the client (e.g. Claude
Code) sends its own tool definitions in the ``tools`` field and expects to
execute them itself. We forward them to llama-server and translate the
streaming response back to Anthropic format without executing anything.
Used for the client-side tool-use pass-through path: the client (e.g.
Claude Code) sends its own tool definitions in ``tools`` and executes
them itself. We forward them to llama-server and translate the streaming
response back to Anthropic format without executing anything.
"""
def __init__(self) -> None:
@ -410,8 +493,14 @@ class AnthropicPassthroughEmitter:
self._tool_call_states: dict = {} # delta index -> {block_index, id, name}
self._usage: dict = {}
self._stop_reason: str = "end_turn"
self._stop_sequence: Optional[str] = None
def start(self, message_id: str, model: str) -> list[str]:
def start(
self,
message_id: str,
model: str,
input_tokens: int = 0,
) -> list[str]:
return [
build_anthropic_sse_event(
"message_start",
@ -425,7 +514,12 @@ class AnthropicPassthroughEmitter:
"model": model,
"stop_reason": None,
"stop_sequence": None,
"usage": {"input_tokens": 0, "output_tokens": 0},
"usage": {
"input_tokens": input_tokens,
"output_tokens": 0,
"cache_creation_input_tokens": 0,
"cache_read_input_tokens": 0,
},
},
},
)
@ -475,7 +569,7 @@ class AnthropicPassthroughEmitter:
# New tool call — close prior block, open tool_use block
if self._current_block_type is not None:
events.append(self._close_current_block())
tc_id = tc.get("id", "")
tc_id = anthropic_tool_use_id(tc.get("id", ""))
tc_name = fn.get("name", "")
self.block_index += 1
self._current_block_type = "tool_use"
@ -518,12 +612,7 @@ class AnthropicPassthroughEmitter:
# ── Finish reason ──
if finish_reason:
if finish_reason == "tool_calls":
self._stop_reason = "tool_use"
elif finish_reason == "length":
self._stop_reason = "max_tokens"
else:
self._stop_reason = "end_turn"
self._stop_reason = openai_finish_to_anthropic_stop(finish_reason)
return events
@ -538,11 +627,9 @@ class AnthropicPassthroughEmitter:
"type": "message_delta",
"delta": {
"stop_reason": self._stop_reason,
"stop_sequence": None,
},
"usage": {
"output_tokens": self._usage.get("completion_tokens", 0),
"stop_sequence": self._stop_sequence,
},
"usage": _message_delta_usage(self._usage),
},
)
)

73
studio/backend/core/inference/audio_codecs.py
View file

@ -77,22 +77,22 @@ class AudioCodecManager:
return
from snac import SNAC
self._snac_model = (
SNAC.from_pretrained("hubertsiuzdak/snac_24khz").to(device).eval()
)
self._snac_model = SNAC.from_pretrained("hubertsiuzdak/snac_24khz").to(device).eval()
logger.info("Loaded SNAC codec (24kHz)")
def _load_bicodec(self, device: str, model_repo_path: Optional[str] = None) -> None:
def _load_bicodec(
self,
device: str,
model_repo_path: Optional[str] = None,
) -> None:
if self._bicodec_tokenizer is not None:
return
import os
import sys
# Clone SparkAudio/Spark-TTS GitHub repo for the sparktts Python package
# (same approach as training — the HF model repos don't contain the package)
spark_code_dir = os.path.join(
os.path.dirname(model_repo_path or "."), "Spark-TTS"
)
# Clone SparkAudio/Spark-TTS for the sparktts package (HF model repos
# don't contain it)
spark_code_dir = os.path.join(os.path.dirname(model_repo_path or "."), "Spark-TTS")
sparktts_pkg = os.path.join(spark_code_dir, "sparktts")
if not os.path.isdir(sparktts_pkg):
logger.info(f"Cloning SparkAudio/Spark-TTS to {spark_code_dir}...")
@ -115,7 +115,7 @@ class AudioCodecManager:
from sparktts.models.audio_tokenizer import BiCodecTokenizer
# BiCodecTokenizer needs the MODEL repo path (contains BiCodec/ weights)
# BiCodecTokenizer needs the MODEL repo path (has BiCodec/ weights)
tokenizer_path = model_repo_path or spark_code_dir
self._bicodec_repo_path = tokenizer_path
self._bicodec_tokenizer = BiCodecTokenizer(tokenizer_path, device)
@ -127,9 +127,8 @@ class AudioCodecManager:
import os
import sys
# Clone OuteTTS repo (same pattern as Spark-TTS / BiCodec)
# The pip package has problematic dependencies; the notebook clones and
# removes gguf_model.py, interface.py, __init__.py before importing.
# Clone OuteTTS (the pip package has problematic deps; we remove
# gguf_model.py, interface.py, __init__.py before importing).
base_dir = os.path.dirname(os.path.abspath(__file__))
outetts_code_dir = os.path.join(base_dir, "OuteTTS")
outetts_pkg = os.path.join(outetts_code_dir, "outetts")
@ -148,8 +147,7 @@ class AudioCodecManager:
env = child_env_without_native_path_secret(),
**_windows_hidden_subprocess_kwargs(),
)
# Remove files that pull in heavy / incompatible dependencies
# (matches notebook: gguf_model.py is under models/, others under outetts/)
# Remove files pulling in heavy / incompatible deps
remove_paths = [
os.path.join(outetts_pkg, "models", "gguf_model.py"),
os.path.join(outetts_pkg, "interface.py"),
@ -177,16 +175,11 @@ class AudioCodecManager:
# ── Decoders ─────────────────────────────────────────────────
def decode_snac(
self, generated_ids: torch.Tensor, device: str
) -> Tuple[bytes, int]:
"""
Decode SNAC tokens (Orpheus) into WAV bytes.
def decode_snac(self, generated_ids: torch.Tensor, device: str) -> Tuple[bytes, int]:
"""Decode SNAC tokens (Orpheus) into WAV bytes.
generated_ids: full model output including prompt tokens.
Looks for START_OF_SPEECH (128257) marker, extracts codes after it,
Finds the START_OF_SPEECH (128257) marker, extracts codes after it,
strips EOS (128258), redistributes 7-per-frame codes into 3 SNAC layers.
Returns (wav_bytes, 24000).
"""
# Find START_OF_SPEECH token (128257)
@ -194,10 +187,8 @@ class AudioCodecManager:
if len(token_indices[1]) > 0:
cropped = generated_ids[:, token_indices[1][-1] + 1 :]
else:
# Gracefully fall back to using entire output if marker not found
logger.warning(
"No START_OF_SPEECH token (128257) found — using full generated output"
)
# Fall back to the entire output if the marker is missing
logger.warning("No START_OF_SPEECH token (128257) found — using full generated output")
cropped = generated_ids
row = cropped[0]
@ -223,8 +214,7 @@ class AudioCodecManager:
layer_3.append(codes[7 * i + 6] - 24576)
snac_codes = [
torch.tensor(layer).unsqueeze(0).to(device)
for layer in [layer_1, layer_2, layer_3]
torch.tensor(layer).unsqueeze(0).to(device) for layer in [layer_1, layer_2, layer_3]
]
with torch.no_grad():
@ -234,16 +224,13 @@ class AudioCodecManager:
return _numpy_to_wav_bytes(waveform, 24000), 24000
def decode_csm(self, audio_values: torch.Tensor) -> Tuple[bytes, int]:
"""
Decode CSM output (already a waveform from model.generate(output_audio=True)).
Returns (wav_bytes, 24000).
"""
"""Decode CSM output (already a waveform). Returns (wav_bytes, 24000)."""
waveform = audio_values[0].to(torch.float32).cpu().numpy()
return _numpy_to_wav_bytes(waveform, 24000), 24000
def decode_bicodec(self, generated_text: str, device: str) -> Tuple[bytes, int]:
"""
Decode BiCodec tokens (Spark-TTS) from generated text.
"""Decode BiCodec tokens (Spark-TTS) from generated text.
Extracts bicodec_semantic_N and bicodec_global_N tokens via regex.
Returns (wav_bytes, sample_rate).
"""
@ -254,19 +241,15 @@ class AudioCodecManager:
f"BiCodec decode: {len(global_matches)} global tokens, {len(semantic_matches)} semantic tokens"
)
if len(global_matches) < 10:
logger.info(
f"BiCodec generated text (first 500 chars): {generated_text[:500]}"
)
logger.info(f"BiCodec generated text (first 500 chars): {generated_text[:500]}")
if not semantic_matches:
raise ValueError("No bicodec_semantic tokens found in generated output")
semantic_ids = (
torch.tensor([int(t) for t in semantic_matches]).long().unsqueeze(0)
)
semantic_ids = torch.tensor([int(t) for t in semantic_matches]).long().unsqueeze(0)
# Speaker encoder expects exactly 32 global tokens (token_num=32 in BiCodec config).
# Pad with zeros or truncate to 32.
# Speaker encoder expects exactly 32 global tokens (token_num=32);
# pad with zeros or truncate.
GLOBAL_TOKEN_NUM = 32
if global_matches:
raw = [int(t) for t in global_matches]
@ -288,8 +271,8 @@ class AudioCodecManager:
return _numpy_to_wav_bytes(wav_np, sr), sr
def decode_dac(self, generated_text: str, device: str) -> Tuple[bytes, int]:
"""
Decode DAC tokens (OuteTTS) from generated text.
"""Decode DAC tokens (OuteTTS) from generated text.
Extracts c1_N and c2_N codec code tokens via regex.
Returns (wav_bytes, 24000).
"""

8
studio/backend/core/inference/chat_template_helpers.py
View file

@ -2,8 +2,8 @@
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""
Dependency-light wrapper around tokenizer.apply_chat_template with a
kwarg fallback for templates that reject reasoning/tools args.
Dependency-light wrapper around tokenizer.apply_chat_template with a kwarg
fallback for templates that reject reasoning/tools args.
"""
from typing import Optional
@ -55,6 +55,4 @@ def apply_chat_template_for_generation(
break
if last_exc is not None:
raise last_exc
raise RuntimeError(
"apply_chat_template_for_generation: no attempt produced a result"
)
raise RuntimeError("apply_chat_template_for_generation: no attempt produced a result")

109
studio/backend/core/inference/chat_templates.py Normal file
View file

@ -0,0 +1,109 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Bundled chat-template selection for GGUF inference.
Some shipped GGUF quants embed an older chat template. Rather than re-cutting and
asking users to re-download every quant, Studio can override the embedded template
at llama-server launch time with a bundled, up-to-date Jinja template for known
model families. The override is wired through the existing ``chat_template_override``
-> ``--chat-template-file`` path in ``LlamaCppBackend.load_model``.
Currently this covers ``unsloth/gemma-4-*-GGUF``, which gains the upstream PR #118
``preserve_thinking`` flag (defaulted OFF here) so the Studio "Preserve thinking"
toggle appears while staying disabled by default.
"""
import re
from functools import lru_cache
from pathlib import Path
from typing import Optional
# assets live at <backend>/assets/chat_templates/. This module is at
# <backend>/core/inference/chat_templates.py, so walk up three parents to
# <backend> (mirrors utils/inference/inference_config.py).
_ASSETS_DIR = Path(__file__).parent.parent.parent / "assets" / "chat_templates"
# unsloth/gemma-4-<variant>-GGUF (case-insensitive). The "-GGUF" suffix is retained
# on ModelConfig.identifier for HF GGUF repos, so this matches E2B / E4B / 31B /
# 26B-A4B and any future unsloth/gemma-4-*-GGUF, while excluding gemma-3,
# non-Unsloth, and non-GGUF identifiers (e.g. the bf16 "unsloth/gemma-4-E2B-it").
_GEMMA4_GGUF_RE = re.compile(r"^unsloth/gemma-4-.+-gguf$", re.IGNORECASE)
# Google ships two distinct gemma-4 chat templates: E2B/E4B omit the empty
# "<|channel>thought<channel|>" block on enable_thinking=false, while the
# 12b/26B-A4B/31B family emits it. Route the two GGUF families to the matching
# bundled template so each keeps its model's intended behavior.
_GEMMA4_EDGE_GGUF_RE = re.compile(r"^unsloth/gemma-4-e[24]b-it-gguf$", re.IGNORECASE)
_GEMMA4_TEMPLATE_FILE = "gemma-4.jinja" # 12b / 26B-A4B / 31B
_GEMMA4_EDGE_TEMPLATE_FILE = "gemma-4-edge.jinja" # E2B / E4B
def _canonical_repo_id(model_identifier: str) -> str:
"""Mirror ``ModelConfig.from_identifier``: a bare HF shorthand with no owner
(e.g. ``gemma-4-E2B-it-GGUF``) defaults to the ``unsloth/`` org. The resolver
runs on the raw ``request.model_path`` (before that canonicalization), so apply
the same rule here, otherwise shorthand loads would skip the override.
"""
mid = model_identifier.strip()
if mid and "/" not in mid:
mid = f"unsloth/{mid}"
return mid
def is_unsloth_gemma4_gguf(model_identifier: Optional[str]) -> bool:
"""True for canonical ``unsloth/gemma-4-*-GGUF`` repo identifiers (and the
owner-less shorthand that resolves to the same Unsloth repo)."""
if not model_identifier:
return False
return bool(_GEMMA4_GGUF_RE.match(_canonical_repo_id(model_identifier)))
def is_unsloth_gemma4_edge_gguf(model_identifier: Optional[str]) -> bool:
"""True for the E2B / E4B GGUF repos, which use the edge-variant template."""
if not model_identifier:
return False
return bool(_GEMMA4_EDGE_GGUF_RE.match(_canonical_repo_id(model_identifier)))
def _gemma4_template_file(model_identifier: Optional[str]) -> Optional[str]:
"""Return the bundled template filename for a gemma-4 GGUF id, else None."""
if is_unsloth_gemma4_edge_gguf(model_identifier):
return _GEMMA4_EDGE_TEMPLATE_FILE
if is_unsloth_gemma4_gguf(model_identifier):
return _GEMMA4_TEMPLATE_FILE
return None
@lru_cache(maxsize=8)
def load_bundled_chat_template(name: str) -> str:
"""Read a bundled chat-template asset by filename (cached for the process)."""
return (_ASSETS_DIR / name).read_text(encoding="utf-8")
def resolve_effective_chat_template_override(
*,
model_identifier: Optional[str],
user_override: Optional[str],
) -> Optional[str]:
"""Resolve which chat-template text to launch llama-server with.
Precedence:
1. An explicit, non-empty user override always wins (advanced users).
2. For ``unsloth/gemma-4-*-GGUF``, return the bundled gemma-4 template
(adds ``preserve_thinking``, default off) so the embedded GGUF template
is overridden without re-downloading quants. E2B/E4B get the edge
variant; 12b/26B-A4B/31B get the standard one.
3. Otherwise ``None`` -> llama-server renders the GGUF's embedded template.
The result is fed to ``LlamaCppBackend.load_model(chat_template_override=...)``
and must be computed before the route-level reload-dedup check so the live
backend state and the incoming request compare consistently.
"""
if user_override and user_override.strip():
return user_override
template_file = _gemma4_template_file(model_identifier)
if template_file is not None:
return load_bundled_chat_template(template_file)
return None

2
studio/backend/core/inference/defaults.py
View file

@ -49,7 +49,7 @@ DEFAULT_MODELS_STANDARD = [
def get_default_models() -> list[str]:
hw.get_device() # ensure detect_hardware() has run
hw.get_device() # ensures detect_hardware() has run
if hw.CHAT_ONLY:
return list(DEFAULT_MODELS_GGUF)
return list(DEFAULT_MODELS_STANDARD)

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

File diff suppressed because it is too large Load diff

592
studio/backend/core/inference/inference.py
View file

File diff suppressed because it is too large Load diff

21
studio/backend/core/inference/key_exchange.py
View file

@ -4,13 +4,13 @@
"""
RSA key pair for encrypting API keys in transit.
The frontend encrypts API keys with the server's public key before
including them in requests. The backend decrypts with its private key
before forwarding to external providers.
The frontend encrypts API keys with the server's public key before sending
them; the backend decrypts with its private key before forwarding to external
providers.
The key pair is generated at server startup and lives only in memory —
it is regenerated on each restart. The frontend fetches the public key
via GET /api/providers/public-key on load.
The key pair is generated at server startup, lives only in memory, and is
regenerated on each restart. The frontend fetches the public key via
GET /api/providers/public-key on load.
"""
import base64
@ -36,9 +36,7 @@ def init_key_pair() -> None:
"""Generate an RSA-2048 key pair. Called once at server startup."""
global _private_key, _public_key_pem, _public_key_fingerprint
if _private_key is not None:
# Re-entry is suspicious — every fresh keypair invalidates all
# in-flight ciphertext encrypted against the previous public key.
# Log loudly so a regression that calls init twice is visible.
# Re-entry invalidates in-flight ciphertext from the old public key; log loudly.
logger.warning(
"init_key_pair called again — replacing existing RSA keypair "
"(previous fingerprint=%s). Any frontend that cached the old "
@ -111,9 +109,8 @@ def decrypt_api_key(encrypted_b64: str) -> str:
),
)
except Exception as exc:
# Surface enough state to distinguish key mismatch (wrong public key
# used on encrypt) from a padding/algo mismatch or corrupted bytes.
# Expected ciphertext length for RSA-2048 is exactly 256 bytes.
# Log state to distinguish key mismatch from padding/algo mismatch or
# corrupted bytes. RSA-2048 ciphertext is exactly 256 bytes.
logger.warning(
"decrypt_api_key: RSA decrypt failed (ciphertext_len=%d, expected=256, "
"fingerprint=%s, exc=%s): %s",

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

File diff suppressed because it is too large Load diff

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

@ -1,46 +1,28 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Validator for user-supplied llama-server pass-through args.
"""Boundary validator for user-supplied llama-server pass-through args.
Studio runs llama-server as a managed subprocess and lets callers pass
extra flags directly (CLI: ``unsloth run ... --top-k 20``; HTTP:
``LoadRequest.llama_extra_args``). This module is the boundary that
rejects only flags Studio fundamentally cannot share with the user --
model identity, the auth key, and the network endpoint Studio's HTTP
proxy targets. Anything else passes through.
Reject only flags Studio manages (model identity, auth, network, parallel
slots). Everything else (sampling, ``-c``, ``-ngl``, ``--flash-attn``,
``--cache-type-*``, ``--spec-*``, ``--jinja``, ...) is appended after
Studio's auto-set flags so llama.cpp's last-wins parser lets the user override.
User-supplied args are appended to ``cmd`` after Studio's auto-set
flags, so llama.cpp's last-wins CLI parsing makes the user's value
override the auto-set one. That covers tunable knobs the user might
reasonably want to override -- ``-c``/``--ctx-size``,
``-np``/``--parallel``, ``-fa``/``--flash-attn``,
``-ngl``/``--gpu-layers``, ``-t``/``--threads``, ``-fit``/``--fit*``,
``--cache-type-k/v``, ``--chat-template-file/-kwargs``,
``--spec-*``, ``--jinja``/``--no-jinja``,
``--no-context-shift``/``--context-shift``, sampling params, etc.
Reference: https://github.com/ggml-org/llama.cpp/blob/master/tools/server/README.md
Ref: https://github.com/ggml-org/llama.cpp/blob/master/tools/server/README.md
"""
from __future__ import annotations
from typing import Iterable, Optional
# Each group is the full set of aliases (short + long) for one
# hard-denied flag, taken from the llama-server README. If llama.cpp
# adds a new alias for an existing denied flag, extend the relevant
# group.
#
# Flags NOT in this list (e.g. -c, --parallel, --flash-attn, -ngl,
# -t/--threads, --jinja, --no-context-shift, --fit*, --cache-type-*,
# --chat-template-*, --spec-*) pass through and override Studio's
# auto-set version via llama.cpp's last-wins CLI parsing.
# Each group = every alias (short + long) of one hard-denied flag.
# Extend the matching group when llama.cpp adds a new alias.
_DENYLIST_GROUPS: tuple[frozenset[str], ...] = (
# Model identity -- Studio resolves the model from LoadRequest and
# passes -m / mmproj after downloading from HF if needed. A second
# -m would point at a different model than the one Studio thinks
# is loaded.
# Parallel slots: owned by typer --parallel; a pass-through would desync
# app.state.llama_parallel_slots from llama-server.
frozenset({"-np", "--parallel", "--n-parallel"}),
# Model identity: Studio resolves it from LoadRequest; a second -m would
# load a different model than Studio thinks it loaded.
frozenset({"-m", "--model"}),
frozenset({"-mu", "--model-url"}),
frozenset({"-dr", "--docker-repo"}),
@ -51,28 +33,21 @@ _DENYLIST_GROUPS: tuple[frozenset[str], ...] = (
frozenset({"-hft", "--hf-token"}),
frozenset({"-mm", "--mmproj"}),
frozenset({"-mmu", "--mmproj-url"}),
# Networking -- Studio binds llama-server's port and reverse-proxies
# HTTP traffic to it. Retargeting host/port/path/prefix would
# orphan Studio's proxy and the UI would lose the server.
# Networking: Studio binds + proxies; retargeting orphans the proxy.
frozenset({"--host"}),
frozenset({"--port"}),
frozenset({"--path"}),
frozenset({"--api-prefix"}),
frozenset({"--reuse-port"}),
# Auth / TLS -- Studio terminates auth at its own layer; an
# upstream --api-key would shadow Studio's UNSLOTH_DIRECT_STREAM
# key, and TLS on llama-server would break the local proxy hop.
# Auth / TLS: Studio terminates auth; upstream --api-key / TLS shadows
# Studio's key and breaks the proxy hop.
frozenset({"--api-key"}),
frozenset({"--api-key-file"}),
frozenset({"--ssl-key-file"}),
frozenset({"--ssl-cert-file"}),
# Single-model server -- Studio runs one model per llama-server
# process and serves its own UI. Enabling multi-model loading or
# llama-server's built-in web UI changes the surface clients see.
# ``--webui``/``--no-webui`` are the legacy spelling; current
# upstream uses ``--ui``/``--no-ui`` + ``--ui-*`` companions.
# Keep both so the denylist matches old and new llama-server
# binaries (Studio's prebuilt vs system-llama.cpp).
# Built-in web UI. --webui/--no-webui is the legacy spelling; upstream
# renamed to --ui/--no-ui + --ui-*. Keep both so prebuilt and system
# llama.cpp binaries match.
frozenset({"--webui", "--no-webui"}),
frozenset({"--ui", "--no-ui"}),
frozenset({"--ui-config"}),
@ -82,32 +57,44 @@ _DENYLIST_GROUPS: tuple[frozenset[str], ...] = (
frozenset({"--models-preset"}),
frozenset({"--models-max"}),
frozenset({"--models-autoload", "--no-models-autoload"}),
# Server-mode flips: --embedding / --rerank restrict llama-server to
# those endpoints, breaking Studio's /v1/chat/completions hop.
frozenset({"--embedding", "--embeddings"}),
frozenset({"--rerank", "--reranking"}),
# llama-server's own built-in tools flag would silently stack on top of
# Studio's --enable-tools / --disable-tools policy resolver.
frozenset({"--tools"}),
)
_DENYLIST: frozenset[str] = frozenset().union(*_DENYLIST_GROUPS)
def _flag_name(token: str) -> Optional[str]:
"""Return the flag name for a token, or None if it isn't a flag.
"""Flag name for ``token``, or None if it isn't a flag.
Peels ``--key=value`` to the bare ``--key``. Plain numeric values
like ``-1`` or ``-0.5`` (e.g. ``--seed -1``) are values, not flags;
llama-server short-form flags always start with a letter.
Peels `--key=value` to `--key`, treats `-1`/`-0.5` as values (shorts
always start with a letter), and normalises attached `-np8` / `-np-1` /
`-np8x` to `-np`. Mirrors the CLI's `_expand_attached_np_short`.
"""
token = token.strip()
if not token.startswith("-") or token in {"-", "--"}:
return None
if len(token) >= 2 and (token[1].isdigit() or token[1] == "."):
return None
return token.split("=", 1)[0]
name = token.split("=", 1)[0]
if len(name) > 3 and name.startswith("-np"):
suffix = name[3:]
if suffix[0].isdigit() or (
len(suffix) > 1 and suffix[0] in {"-", "+"} and suffix[1].isdigit()
):
return "-np"
return name
def validate_extra_args(args: Optional[Iterable[str]]) -> list[str]:
"""Validate user-supplied llama-server args.
Returns the args as a flat list ready to extend the llama-server
command. Raises ``ValueError`` (with the offending flag in the
message) the moment a token resolves to a Studio-managed flag.
"""
"""Validate user-supplied llama-server args. Returns a flat list ready to
extend the llama-server command; raises ``ValueError`` naming the
offending flag on the first managed token."""
if not args:
return []
out: list[str] = []
@ -120,23 +107,24 @@ def validate_extra_args(args: Optional[Iterable[str]]) -> list[str]:
f"and cannot be passed as an extra arg"
)
out.append(token)
parse_ctx_override(out)
parse_cache_override(out)
parse_split_mode_override(out)
return out
def is_managed_flag(flag: str) -> bool:
"""True if ``flag`` is a Studio-managed llama-server flag."""
return flag in _DENYLIST
"""True if ``flag`` is Studio-managed. Normalises via ``_flag_name`` so
`-np8` / `--parallel=8` classify like the canonical tokens."""
normalised = _flag_name(flag)
return normalised is not None and normalised in _DENYLIST
# Pass-through flags that shadow first-class ``LoadRequest`` fields
# (max_seq_length, cache_type_kv, speculative_type,
# chat_template_override). Stripped from inherited extras so they
# can't last-wins-override an Apply that re-sets the same first-class
# field.
# Pass-through flags that shadow first-class LoadRequest fields; stripped
# from inherited extras so they can't last-wins-override an Apply that
# re-sets the same field.
_CONTEXT_FLAGS: frozenset[str] = frozenset({"-c", "--ctx-size"})
_CACHE_FLAGS: frozenset[str] = frozenset(
{"-ctk", "--cache-type-k", "-ctv", "--cache-type-v"}
)
_CACHE_FLAGS: frozenset[str] = frozenset({"-ctk", "--cache-type-k", "-ctv", "--cache-type-v"})
_SPEC_FLAGS: frozenset[str] = frozenset(
{
"--spec-default",
@ -145,7 +133,13 @@ _SPEC_FLAGS: frozenset[str] = frozenset(
"--spec-ngram-size",
"--draft-min",
"--draft-max",
# MTP path (llama.cpp #22673).
# MTP path (llama.cpp #22673). --model-draft and aliases are
# Studio-managed since the separate-drafter support (Gemma 4): an
# inherited copy must not last-wins-override the auto-detected
# drafter. Explicit extras for the current load are never stripped.
"--model-draft",
"-md",
"--spec-draft-model",
"--spec-draft-n-max",
"--spec-draft-n-min",
"--spec-draft-p-min",
@ -164,17 +158,177 @@ _TEMPLATE_FLAGS: frozenset[str] = frozenset(
"--no-jinja",
}
)
# Multi-GPU split mode shadows the Tensor Parallelism toggle
# (--split-mode tensor). Pass-through stays allowed so users keep the
# row/none/layer modes the toggle doesn't expose, but it's stripped on
# inherit and reconciled into the round-tripped tensor_parallel state.
# --tensor-split is coupled to the split mode and is stripped with it: Studio
# owns the tensor-mode split ratios, so an inherited/stale --tensor-split must
# not last-wins-override Studio's computed asymmetric split.
_SPLIT_MODE_FLAGS: frozenset[str] = frozenset({"-sm", "--split-mode"})
_TENSOR_SPLIT_FLAGS: frozenset[str] = frozenset({"-ts", "--tensor-split"})
_SPLIT_SHADOWING_FLAGS: frozenset[str] = _SPLIT_MODE_FLAGS | _TENSOR_SPLIT_FLAGS
_SHADOWING_FLAGS: frozenset[str] = (
_CONTEXT_FLAGS | _CACHE_FLAGS | _SPEC_FLAGS | _TEMPLATE_FLAGS
_CONTEXT_FLAGS | _CACHE_FLAGS | _SPEC_FLAGS | _TEMPLATE_FLAGS | _SPLIT_SHADOWING_FLAGS
)
# Boolean flags inside _SHADOWING_FLAGS that take no value. The
# value-consuming heuristic in strip_shadowing_flags must skip just the
# flag for these, never the following token.
_BOOLEAN_SHADOWING_FLAGS: frozenset[str] = frozenset(
{"--spec-default", "--jinja", "--no-jinja"}
)
# Shadowing flags that take no value -- strip the flag only, not the next token.
_BOOLEAN_SHADOWING_FLAGS: frozenset[str] = frozenset({"--spec-default", "--jinja", "--no-jinja"})
def parse_ctx_override(args: Optional[Iterable[str]]) -> Optional[int]:
"""Return the last user-supplied ``-c`` / ``--ctx-size`` value.
Mirrors llama.cpp's last-wins parsing for the one numeric knob Studio's
load-time fit logic needs.
"""
if not args:
return None
tokens = [str(a) for a in args]
override: Optional[int] = None
i, n = 0, len(tokens)
while i < n:
tok = tokens[i]
flag = _flag_name(tok)
if flag is None or flag not in _CONTEXT_FLAGS:
i += 1
continue
if "=" in tok:
raw_value = tok.split("=", 1)[1]
i += 1
else:
if i + 1 >= n or _flag_name(tokens[i + 1]) is not None:
raise ValueError(f"llama-server flag '{flag}' requires an integer value")
raw_value = tokens[i + 1]
i += 2
try:
value = int(str(raw_value).strip())
except ValueError as exc:
raise ValueError(f"llama-server flag '{flag}' requires an integer value") from exc
if value < 0:
raise ValueError(f"llama-server flag '{flag}' requires a non-negative integer value")
override = value
return override
def resolve_requested_ctx(args: Optional[Iterable[str]], fallback_n_ctx: int) -> int:
"""Return the context size load_model should treat as requested.
Single source of truth for load_model's ctx-override conditional so
tests don't reimplement and assert against their own logic.
"""
override = parse_ctx_override(args)
return override if override is not None else fallback_n_ctx
def _last_flag_value(args: Optional[Iterable[str]], flags: frozenset[str]) -> Optional[str]:
"""Return the last-wins string value among ``flags`` in extras, or None.
Handles both ``--flag=value`` and ``--flag value`` forms and raises if a
matched flag has no (or an empty) value. Shared by the single-knob
last-wins parsers (cache type, split mode).
"""
if not args:
return None
tokens = [str(a) for a in args]
override: Optional[str] = None
i, n = 0, len(tokens)
while i < n:
tok = tokens[i]
flag = _flag_name(tok)
if flag is None or flag not in flags:
i += 1
continue
if "=" in tok:
raw_value = tok.split("=", 1)[1]
i += 1
else:
if i + 1 >= n or _flag_name(tokens[i + 1]) is not None:
raise ValueError(f"llama-server flag '{flag}' requires a value")
raw_value = tokens[i + 1]
i += 2
value = str(raw_value).strip()
if not value:
raise ValueError(f"llama-server flag '{flag}' requires a non-empty value")
override = value
return override
def parse_cache_override(args: Optional[Iterable[str]]) -> Optional[str]:
"""Return the last-wins cache type if extras pass cache flags.
Mirrors parse_ctx_override but for cache type. Recognises both -ctk
(key) and -ctv (value). When both flags appear, returns the last-wins
value, treating key and value cache flags as the same setting because
Studio's KV estimate has a single cache_type_kv knob.
"""
return _last_flag_value(args, _CACHE_FLAGS)
def resolve_cache_type_kv(
args: Optional[Iterable[str]], fallback_cache_type_kv: Optional[str]
) -> Optional[str]:
"""Return the cache type load_model should treat as requested.
Single source of truth for ``load_model``'s cache override conditional.
"""
override = parse_cache_override(args)
return override if override is not None else fallback_cache_type_kv
def parse_split_mode_override(args: Optional[Iterable[str]]) -> Optional[str]:
"""Return the last-wins ``--split-mode`` / ``-sm`` value from extras.
Mirrors parse_cache_override for the multi-GPU split mode. Returns the
raw mode string (e.g. ``tensor`` / ``row`` / ``none`` / ``layer``), or
None when extras don't set it.
"""
return _last_flag_value(args, _SPLIT_MODE_FLAGS)
def resolve_tensor_parallel(args: Optional[Iterable[str]], fallback_tensor_parallel: bool) -> bool:
"""Return the tensor-parallel state load_model should treat as requested.
A user-supplied ``--split-mode`` in extras last-wins-overrides the
toggle, so reconcile it back into the boolean: any explicit split mode
means tensor-parallel is on iff that mode is ``tensor``. Falls back to
the toggle value when extras don't set it.
"""
override = parse_split_mode_override(args)
if override is None:
return fallback_tensor_parallel
return override.strip().lower() == "tensor"
_MMPROJ_DISABLE_FLAGS: frozenset[str] = frozenset({"--no-mmproj", "--no-mmproj-auto"})
_MMPROJ_ENABLE_FLAGS: frozenset[str] = frozenset({"--mmproj-auto"})
def extra_args_disable_mmproj(args: Optional[Iterable[str]]) -> bool:
"""True when pass-through args opt out of vision mmproj loading.
llama-server parses --mmproj-auto / --no-mmproj / --no-mmproj-auto as one
boolean with last-wins semantics; mirror that here.
"""
if not args:
return False
disabled = False
for raw in args:
flag = _flag_name(str(raw))
if flag in _MMPROJ_DISABLE_FLAGS:
disabled = True
elif flag in _MMPROJ_ENABLE_FLAGS:
disabled = False
return disabled
def strip_shadowing_flags(
@ -184,17 +338,15 @@ def strip_shadowing_flags(
strip_cache: bool = True,
strip_spec: bool = True,
strip_template: bool = True,
strip_split_mode: bool = True,
) -> list[str]:
"""Strip flags that shadow first-class Studio settings.
Used when the route inherits a previous load's ``llama_extra_args``
so that an inherited ``-c 4096`` cannot override the current
request's ``max_seq_length`` (and equivalents for cache /
speculative / chat template). Each ``strip_*`` flag controls one
group; the route only strips groups whose corresponding first-class
field was actually supplied by the caller, so an inherited
``--chat-template-file`` survives an Apply that omits both
``llama_extra_args`` and ``chat_template_override``.
Used when inheriting a previous load's ``llama_extra_args`` so an
inherited `-c 4096` can't override the current `max_seq_length`
(same for cache / spec / template / split-mode). Each ``strip_*``
toggle controls one group; the route only strips groups whose
first-class field the caller actually supplied.
"""
shadowing: set[str] = set()
if strip_context:
@ -205,6 +357,8 @@ def strip_shadowing_flags(
shadowing |= _SPEC_FLAGS
if strip_template:
shadowing |= _TEMPLATE_FLAGS
if strip_split_mode:
shadowing |= _SPLIT_SHADOWING_FLAGS
tokens = [str(a) for a in (args or [])]
out: list[str] = []
@ -216,9 +370,8 @@ def strip_shadowing_flags(
out.append(tok)
i += 1
continue
# Drop this token. Boolean shadowing flags never carry a value;
# other shadowing flags consume the next token when it isn't a
# flag and the value isn't already packed as ``--key=value``.
# Drop the flag; also consume the next token unless it's boolean,
# already inline (`-c=4096`), or another flag.
if flag in _BOOLEAN_SHADOWING_FLAGS or "=" in tok:
i += 1
elif i + 1 < n and _flag_name(tokens[i + 1]) is None:
@ -226,3 +379,20 @@ def strip_shadowing_flags(
else:
i += 1
return out
def strip_split_mode_only(args: Optional[Iterable[str]]) -> Optional[list[str]]:
"""Remove the split-mode group (``--split-mode`` / ``-sm`` and the coupled
``--tensor-split`` / ``-ts``) from ``args``, keeping every other shadow flag.
Preserves a None/empty input so the inherit-vs-explicit-empty distinction
survives. Used where tensor mode is being forced off (downgrade / fallback)."""
if not args:
return args
return strip_shadowing_flags(
args,
strip_context = False,
strip_cache = False,
strip_spec = False,
strip_template = False,
strip_split_mode = True,
)

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

@ -0,0 +1,383 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
from __future__ import annotations
import asyncio
import json
import os
import shlex
import sys
import time
from typing import Any, Optional
from loggers import get_logger
logger = get_logger(__name__)
MCP_TOOL_PREFIX = "mcp__"
# A failed probe isn't cached (a recovered server must come back), but it's
# recorded so a down server isn't re-probed -- and the chat send re-hung for
# the full timeout -- on every message. Cool off for this long after a failure;
# much longer for OAuth, whose probe can hang up to _OAUTH_PROBE_TIMEOUT,
# so that hang doesn't recur every minute.
FAILED_PROBE_COOLOFF_SECONDS = 60.0
OAUTH_FAILED_PROBE_COOLOFF_SECONDS = 300.0
_oauth_token_store = None
def is_stdio(address: str) -> bool:
"""A non-HTTP address is a local stdio command, e.g.
'npx -y @modelcontextprotocol/server-filesystem /path'."""
return not address.strip().lower().startswith(("http://", "https://"))
def _split_windows_command_line(address: str) -> list[str]:
"""Parse a Windows command line using the same backslash/quote rules that
subprocess.list2cmdline() writes. This keeps trailing backslashes before a
closing quote from being doubled in the resulting argv."""
parts: list[str] = []
current: list[str] = []
in_quotes = False
backslashes = 0
arg_started = False
i = 0
while i < len(address):
ch = address[i]
if ch == "\\":
backslashes += 1
i += 1
continue
if ch == '"':
current.extend("\\" * (backslashes // 2))
if backslashes % 2:
current.append('"')
else:
in_quotes = not in_quotes
arg_started = True
backslashes = 0
i += 1
continue
if ch.isspace() and not in_quotes:
if backslashes:
current.extend("\\" * backslashes)
arg_started = True
backslashes = 0
if arg_started or current:
parts.append("".join(current))
current = []
arg_started = False
i += 1
while i < len(address) and address[i].isspace():
i += 1
continue
if backslashes:
current.extend("\\" * backslashes)
arg_started = True
backslashes = 0
current.append(ch)
arg_started = True
i += 1
if backslashes:
current.extend("\\" * backslashes)
arg_started = True
if in_quotes:
raise ValueError("No closing quotation")
if arg_started or current:
parts.append("".join(current))
return parts
def parse_stdio_command(address: str) -> list[str]:
"""Split a stdio command line into argv. Shared by route validation and the
transport so both agree on quoting (notably Windows backslash paths)."""
posix = sys.platform != "win32"
if posix:
return shlex.split(address, posix = posix)
if address.lstrip().startswith("'"):
raise ValueError("Single-quoted executables are not supported on Windows")
return _split_windows_command_line(address)
def join_stdio_command(parts: list[str]) -> str:
"""Inverse of parse_stdio_command: join argv into a single command string
that parse_stdio_command() splits back into ``parts`` on this platform.
Config files (issue #5936) carry structured command + args; storage holds
one string in the url field. Windows uses list2cmdline so spaced/backslash
paths round-trip through the posix=False quote-strip; posix uses shlex."""
if sys.platform == "win32":
import subprocess
return subprocess.list2cmdline(parts)
return shlex.join(parts)
def stdio_mcp_enabled() -> bool:
"""stdio MCP servers spawn local processes as the backend user (bypassing the
sandbox), so allowed only when the host is the user's own machine. On startup
a loopback bind defaults UNSLOTH_STUDIO_ALLOW_STDIO_MCP=1 (see
utils.host_policy.apply_stdio_mcp_loopback_default, called from run.py); the
Tauri app does the same. Off for Colab and any network (0.0.0.0) bind unless
an operator sets the var out-of-band; set it to 0 to force-disable.
When stdio is on only because of that loopback auto-default, an explicit
`unsloth studio run --disable-tools` turns it back off (a local stdio command
is server-side code execution). An explicit operator opt-in via the env var
still wins -- including the documented `=1` network opt-in, where the process
tool policy is False merely by the external-host default, not by choice."""
if os.environ.get("UNSLOTH_STUDIO_ALLOW_STDIO_MCP") != "1":
return False
from state.tool_policy import get_tool_policy
from utils.host_policy import loopback_default_active
if loopback_default_active() and get_tool_policy() is False:
return False
return True
# Probe timeouts for discovering a server's tool list. OAuth needs minutes for
# first-connect/expired-token browser sign-in; stdio allows for first-run
# package download (e.g. `npx -y ...`); HTTP fails fast.
_HTTP_PROBE_TIMEOUT = 8.0
_OAUTH_PROBE_TIMEOUT = 305.0
_STDIO_PROBE_TIMEOUT = 60.0
def probe_timeout(address: str, use_oauth: bool) -> float:
if use_oauth:
return _OAUTH_PROBE_TIMEOUT
return _STDIO_PROBE_TIMEOUT if is_stdio(address) else _HTTP_PROBE_TIMEOUT
def parse_server_headers(server: dict) -> Optional[dict]:
"""Parsed headers_json. For stdio servers this dict is the process env
instead of HTTP headers (see _client)."""
raw = server.get("headers_json")
if not raw:
return None
try:
parsed = json.loads(raw)
except (json.JSONDecodeError, ValueError):
return None
return parsed if isinstance(parsed, dict) else None
def _oauth_store():
global _oauth_token_store
if _oauth_token_store is None:
from key_value.aio._utils.sanitization import AlwaysHashStrategy
from key_value.aio.stores.filetree import FileTreeStore
from utils.paths.storage_roots import ensure_dir, studio_root
# Hash keys/collections — fastmcp uses raw URLs as keys, and FileTreeStore
# would treat the "://" as nested directories.
_oauth_token_store = FileTreeStore(
data_directory = ensure_dir(studio_root() / "mcp-oauth-tokens"),
key_sanitization_strategy = AlwaysHashStrategy(),
collection_sanitization_strategy = AlwaysHashStrategy(),
)
return _oauth_token_store
async def clear_oauth_tokens_async(url: str) -> None:
"""Drop any persisted OAuth tokens for ``url``. fastmcp keys tokens by MCP
URL, so on server delete / URL change / OAuth disable we must clear them, else
re-registering the same URL reuses the old account's token. Best-effort: store
/ OAuth failures must not 500 the delete / update route."""
try:
from fastmcp.client.auth import OAuth
auth = OAuth(mcp_url = url, token_storage = _oauth_store())
await auth.token_storage_adapter.clear()
except Exception as exc: # noqa: BLE001
# Cleanup is best-effort; the row delete still wins.
logger.warning("Failed to clear OAuth tokens for %s: %s", url, exc)
def _client(
url: str,
headers: Optional[dict],
use_oauth: bool = False,
):
from fastmcp import Client
if is_stdio(url):
# Belt-and-suspenders: never spawn unless stdio is enabled on this host.
if not stdio_mcp_enabled():
raise PermissionError("stdio MCP servers are disabled on this host")
from fastmcp.client.transports import StdioTransport
parts = parse_stdio_command(url)
if not parts:
raise ValueError(f"Empty stdio command: {url!r}")
# env vars ride the headers field (merged over the SDK default env).
# keep_alive=False tears the subprocess down so a one-shot call leaves no orphan.
return Client(
StdioTransport(
command = parts[0],
args = parts[1:],
env = headers or None,
keep_alive = False,
)
)
from fastmcp.client.transports import SSETransport, StreamableHttpTransport
from fastmcp.mcp_config import infer_transport_type_from_url
auth = None
if use_oauth:
from fastmcp.client.auth import OAuth
auth = OAuth(mcp_url = url, token_storage = _oauth_store())
transport_cls = (
SSETransport if infer_transport_type_from_url(url) == "sse" else StreamableHttpTransport
)
return Client(transport_cls(url = url, headers = headers or None, auth = auth))
async def list_tools_async(
url: str,
headers: Optional[dict] = None,
timeout: float = 5.0,
use_oauth: bool = False,
) -> list[dict]:
async def _fetch() -> list[dict]:
async with _client(url, headers, use_oauth) as client:
tools = await client.list_tools()
return [t.model_dump(exclude_none = True) for t in tools]
return await asyncio.wait_for(_fetch(), timeout = timeout)
# Discovered-tool cache, keyed by MCP server id. get_enabled_mcp_tools()
# probes a server only on a cache miss, keeping MCP discovery off the chat
# send's critical path -- tool schemas are stable within a session. The
# /refresh route warms it; a URL/header/OAuth change or a delete evicts it.
# Successful probes are cached indefinitely.
_tool_cache: dict[str, list[dict]] = {}
# server_id -> monotonic time before which a failed server must not be
# re-probed (see record_probe_failure). Cleared on a successful probe or
# eviction.
_probe_cooloff_until: dict[str, float] = {}
# MCP server fields whose change invalidates a server's discovered tools: the
# endpoint/auth used to probe it (url, headers, oauth) or whether it's used at
# all (is_enabled). A rename does not. The update route's eviction and
# get_enabled_mcp_tools' mid-probe guard both key off this so they can't drift.
TOOL_CACHE_INVALIDATING_FIELDS = frozenset({"url", "headers_json", "use_oauth", "is_enabled"})
def get_cached_tools(server_id: str) -> Optional[list[dict]]:
return _tool_cache.get(server_id)
def cache_tools(server_id: str, tools: list[dict]) -> None:
_tool_cache[server_id] = tools
_probe_cooloff_until.pop(server_id, None)
def record_probe_failure(server_id: str, use_oauth: bool = False) -> None:
cooloff = OAUTH_FAILED_PROBE_COOLOFF_SECONDS if use_oauth else FAILED_PROBE_COOLOFF_SECONDS
_probe_cooloff_until[server_id] = time.monotonic() + cooloff
def in_failure_cooloff(server_id: str) -> bool:
return _probe_cooloff_until.get(server_id, 0.0) > time.monotonic()
def invalidate_tool_cache(server_id: Optional[str] = None) -> None:
"""Evict one server's cached tools, or every entry when server_id is None."""
if server_id is None:
_tool_cache.clear()
_probe_cooloff_until.clear()
else:
_tool_cache.pop(server_id, None)
_probe_cooloff_until.pop(server_id, None)
def _flatten_result(result: Any) -> str:
parts = []
for block in getattr(result, "content", None) or []:
text = getattr(block, "text", None)
if text:
parts.append(str(text))
body = "\n".join(parts)
if not body:
structured = getattr(result, "structured_content", None)
body = str(structured) if structured is not None else ""
if getattr(result, "is_error", False):
# "Error: " prefix triggers tool_call_parser's TOOL_ERROR_PREFIXES nudge.
return f"Error: {body}" if body else "Error: tool returned no content"
return body
def call_tool_sync(
url: str,
headers: Optional[dict],
name: str,
args: dict,
timeout: Optional[float] = 300.0,
use_oauth: bool = False,
cancel_event = None,
) -> str:
"""Synchronously call an MCP tool.
``cancel_event``: optional ``threading.Event``. When set, the in-flight call is
cancelled and a cancellation Error returned. Polled alongside the tool call via
``asyncio.wait`` so a /cancel POST interrupts even mid-network-read.
"""
async def _call() -> Any:
async with _client(url, headers, use_oauth) as client:
return await client.call_tool(name, args)
async def _watch_cancel() -> None:
# 50 ms cadence keeps cancellation responsive without busy-looping;
# matches routes/inference.py's cancel watcher cadence.
while cancel_event is not None and not cancel_event.is_set():
await asyncio.sleep(0.05)
async def _race() -> Any:
# Check cancellation before spawning the call task so a pre-set event
# short-circuits before opening the transport / HTTP connection.
if cancel_event is not None and cancel_event.is_set():
raise _MCPCancelled
call_task = asyncio.create_task(_call())
if cancel_event is None:
return await asyncio.wait_for(call_task, timeout = timeout)
watch_task = asyncio.create_task(_watch_cancel())
try:
done, pending = await asyncio.wait(
{call_task, watch_task},
timeout = timeout,
return_when = asyncio.FIRST_COMPLETED,
)
finally:
for t in (call_task, watch_task):
if not t.done():
t.cancel()
if not done:
raise asyncio.TimeoutError
if call_task in done:
return call_task.result()
raise _MCPCancelled
try:
result = asyncio.run(_race())
except _MCPCancelled:
return f"Error: MCP tool '{name}' cancelled"
except asyncio.TimeoutError:
return f"Error: MCP tool '{name}' timed out after {timeout:g}s"
except Exception as exc:
logger.exception("MCP call_tool failed for %s: %s", name, exc)
return f"Error: MCP tool '{name}' failed: {exc}"
return _flatten_result(result)
class _MCPCancelled(Exception):
"""Internal sentinel raised when cancel_event fires before the tool returns."""

169
studio/backend/core/inference/mcp_config_import.py Normal file
View file

@ -0,0 +1,169 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Parse a standard ``mcpServers`` JSON config (Claude Desktop / Cursor / Cline
/ VS Code) into entries the existing MCP storage understands. See issue #5936.
A stdio entry (``command`` + ``args`` + ``env``) is joined into the single
command string the ``url`` field already stores; a remote entry (``url`` +
``headers``) maps straight through. Parsing never raises on a single bad entry:
it returns ``(entries, errors)`` so one malformed server can't sink the import.
"""
from __future__ import annotations
from dataclasses import dataclass
from typing import Optional
from core.inference.mcp_client import join_stdio_command
_SCALAR = (str, int, float, bool)
_UNSUPPORTED_STDIO_FIELDS = ("cwd", "envFile")
_UNSUPPORTED_TIMEOUT_FIELDS = ("timeout", "timeoutMs", "timeoutSeconds")
_HTTP_REMOTE_TYPES = ("http", "streamableHttp")
@dataclass
class ParsedMcpEntry:
display_name: str
url: str # joined command (stdio) or http(s) url (remote)
headers: Optional[dict[str, str]] # env vars (stdio) or http headers (remote)
is_stdio: bool
is_enabled: bool = True
use_oauth: bool = False
def _coerce_str_dict(value: dict) -> dict[str, str]:
return {str(k): str(v) for k, v in value.items()}
def _has_variable_reference(value: object) -> bool:
if isinstance(value, str):
return "${" in value
if isinstance(value, list):
return any(_has_variable_reference(item) for item in value)
if isinstance(value, dict):
return any(_has_variable_reference(item) for item in value.values())
return False
def _has_null_value(value: object) -> bool:
return isinstance(value, dict) and any(item is None for item in value.values())
def _enabled_from_spec(label: str, spec: dict) -> tuple[Optional[bool], Optional[str]]:
disabled = spec.get("disabled")
if disabled is None:
return True, None
if not isinstance(disabled, bool):
return None, f"{label}: 'disabled' must be true or false."
return not disabled, None
def _parse_entry(name: str, spec: object) -> tuple[Optional[ParsedMcpEntry], Optional[str]]:
label = str(name).strip()
if not label:
return None, "Server entry has an empty name."
if not isinstance(spec, dict):
return None, f"{label}: entry must be an object."
if _has_variable_reference(spec):
return None, f"{label}: VS Code variable references are not supported by import."
is_enabled, error = _enabled_from_spec(label, spec)
if error:
return None, error
has_command = bool(spec.get("command"))
has_url = bool(spec.get("url"))
if has_command and has_url:
return None, f"{label}: entry has both 'command' and 'url'; use one."
if not has_command and not has_url:
return None, f"{label}: entry needs a 'command' (stdio) or 'url' (remote)."
if has_command:
command = spec["command"]
if not isinstance(command, str):
return None, f"{label}: 'command' must be a string."
entry_type = spec.get("type")
if entry_type is not None and entry_type != "stdio":
return None, f"{label}: stdio entry has unsupported type {entry_type!r}."
sandbox_enabled = spec.get("sandboxEnabled")
if sandbox_enabled is not None and not isinstance(sandbox_enabled, bool):
return None, f"{label}: 'sandboxEnabled' must be true or false."
if sandbox_enabled:
return None, f"{label}: sandboxed stdio servers cannot be preserved by import."
unsupported = [field for field in _UNSUPPORTED_STDIO_FIELDS if spec.get(field) is not None]
if unsupported:
return None, f"{label}: import cannot preserve {', '.join(unsupported)}."
if spec.get("oauth") is not None:
return None, f"{label}: 'oauth' is only supported for remote servers."
args = spec.get("args") or []
if not isinstance(args, list) or not all(isinstance(a, _SCALAR) for a in args):
return None, f"{label}: 'args' must be a list of strings."
env = spec.get("env")
if env is not None and not isinstance(env, dict):
return None, f"{label}: 'env' must be an object."
if _has_null_value(env):
return None, f"{label}: null environment values are not supported by import."
url = join_stdio_command([command, *(str(a) for a in args)])
headers = _coerce_str_dict(env) if env else None
return ParsedMcpEntry(label, url, headers, True, is_enabled = is_enabled), None
url = spec["url"]
if not isinstance(url, str):
return None, f"{label}: 'url' must be a string."
url = url.strip()
entry_type = spec.get("type")
if entry_type is not None and entry_type not in (*_HTTP_REMOTE_TYPES, "sse"):
return None, f"{label}: remote entry has unsupported type {entry_type!r}."
unsupported_timeout = [
field for field in _UNSUPPORTED_TIMEOUT_FIELDS if spec.get(field) is not None
]
if unsupported_timeout:
return None, f"{label}: import cannot preserve {', '.join(unsupported_timeout)}."
url_infers_sse = url.rstrip("/").endswith("/sse")
if entry_type == "sse" and not url_infers_sse:
return None, f"{label}: explicit SSE transport cannot be preserved for this URL."
if entry_type in _HTTP_REMOTE_TYPES and url_infers_sse:
return None, f"{label}: explicit HTTP transport cannot be preserved for this URL."
oauth_raw = spec.get("oauth")
if oauth_raw is not None and not isinstance(oauth_raw, dict):
return None, f"{label}: 'oauth' must be an object."
headers_raw = spec.get("headers")
if headers_raw is not None and not isinstance(headers_raw, dict):
return None, f"{label}: 'headers' must be an object."
if _has_null_value(headers_raw):
return None, f"{label}: null header values are not supported by import."
headers = _coerce_str_dict(headers_raw) if headers_raw else None
return ParsedMcpEntry(
label,
url,
headers,
False,
is_enabled = is_enabled,
use_oauth = oauth_raw is not None,
), None
def parse_mcp_config(config: object) -> tuple[list[ParsedMcpEntry], list[str]]:
"""Parse a Claude-Desktop/Cursor/Cline/VS Code config. Accepts the
``mcpServers`` key (primary) or ``servers`` (VS Code alias). Returns
``(entries, errors)``; a bad entry adds an error rather than raising."""
if not isinstance(config, dict):
return [], ["Config must be a JSON object."]
servers_key = "mcpServers" if "mcpServers" in config else "servers"
servers = config.get(servers_key)
if servers is None:
return [], ["Config has no 'mcpServers' (or 'servers') object."]
if not isinstance(servers, dict):
return [], [f"'{servers_key}' must be an object mapping name -> server."]
entries: list[ParsedMcpEntry] = []
errors: list[str] = []
for name, spec in servers.items():
entry, error = _parse_entry(name, spec)
if error:
errors.append(error)
elif entry:
entries.append(entry)
return entries, errors

165
studio/backend/core/inference/mlx_inference.py
View file

@ -7,11 +7,40 @@ instead of torch/transformers for model loading and generation.
import threading
from typing import Optional, Generator
from core.inference.runtime_context import runtime_context_length
from loggers import get_logger
logger = get_logger(__name__)
def _build_generation_stats(prompt_n, prompt_tps, gen_n, gen_tps):
"""Map mlx stream stats onto the usage/timings shape llama-server emits."""
prompt_n = int(prompt_n or 0)
gen_n = int(gen_n or 0)
prompt_tps = float(prompt_tps or 0.0)
gen_tps = float(gen_tps or 0.0)
prompt_ms = (prompt_n / prompt_tps * 1000.0) if prompt_tps > 0 else 0.0
predicted_ms = (gen_n / gen_tps * 1000.0) if gen_tps > 0 else 0.0
return {
"usage": {
"prompt_tokens": prompt_n,
"completion_tokens": gen_n,
"total_tokens": prompt_n + gen_n,
},
"timings": {
"prompt_n": prompt_n,
"prompt_ms": prompt_ms,
"prompt_per_token_ms": (prompt_ms / prompt_n) if prompt_n > 0 else 0.0,
"prompt_per_second": prompt_tps,
"predicted_n": gen_n,
"predicted_ms": predicted_ms,
"predicted_per_token_ms": (predicted_ms / gen_n) if gen_n > 0 else 0.0,
"predicted_per_second": gen_tps,
"cache_n": 0,
},
}
class MLXInferenceBackend:
def __init__(self):
self.models = {}
@ -20,8 +49,9 @@ class MLXInferenceBackend:
self.loaded_local_models = []
self.device = "mlx"
self._generation_lock = threading.Lock()
# usage/timings of the latest generation; shipped on gen_done.
self.last_generation_stats = None
# MLX state
self._model = None
self._tokenizer = None
self._processor = None
@ -34,10 +64,9 @@ class MLXInferenceBackend:
def _configure_memory_limits(self):
"""Apply Metal memory caps before loading a model.
Mirrors MLXTrainer._configure_memory_limits's defaults:
memory_limit = 85% of recommended working-set,
wired_limit = min(recommended, memory_limit). Recorded so unload
can lower wired_limit back to release pinned RAM.
memory_limit = 85% of recommended working-set;
wired_limit = min(recommended, memory_limit). Recorded so unload can
lower wired_limit back to release pinned RAM.
"""
import mlx.core as mx
@ -78,18 +107,10 @@ class MLXInferenceBackend:
model_name = config.identifier if hasattr(config, "identifier") else str(config)
is_vision = getattr(config, "is_vision", False)
# GGUF guard. GGUF models are served via llama-server in the
# parent process, NOT via mlx-lm in this MLX subprocess. The
# route at studio/backend/routes/inference.py:592 (`if config.
# is_gguf:`) is responsible for sending GGUF traffic to the
# llama-server backend before reaching the MLX orchestrator.
# If we end up here with is_gguf=True, the route's
# `detect_gguf_model_remote` returned None on its first call
# (transient HF Hub flake) but the subprocess re-detection
# succeeded. The subprocess cannot reach into the parent's
# llama-server, so all we can do is raise loudly so the caller
# gets a clear error instead of a cryptic
# "config.json does not exist" from mlx_lm.utils.load_model.
# GGUF guard. GGUF models are served by llama-server in the parent
# process, not mlx-lm here. Reaching this with is_gguf=True means the
# route's first detection flaked (transient HF Hub) but the subprocess
# re-detected GGUF; raise loudly instead of a cryptic mlx_lm error.
if getattr(config, "is_gguf", False):
raise RuntimeError(
f"MLXInferenceBackend cannot load GGUF model '{model_name}': "
@ -102,7 +123,6 @@ class MLXInferenceBackend:
if hf_token:
import os
os.environ["HF_TOKEN"] = hf_token
self._configure_memory_limits()
@ -156,10 +176,10 @@ class MLXInferenceBackend:
"is_audio": False,
"audio_type": None,
"has_audio_input": False,
"context_length": runtime_context_length(self._model, max_seq_length),
}
# Capture chat_template_info so the worker IPC reply can ship
# it back to the parent and the route layer classifies
# capabilities the same way as the transformers / GGUF paths.
# Capture chat_template_info so the worker IPC reply ships it back and
# the route layer classifies capabilities like the other paths.
self._populate_chat_template_info(model_name)
logger.info("Model %s loaded successfully", model_name)
@ -168,10 +188,8 @@ class MLXInferenceBackend:
def _populate_chat_template_info(self, model_name: str) -> None:
"""Mirror InferenceBackend._load_chat_template_info for MLX.
Stores ``chat_template_info`` on ``self.models[model_name]``
with the resolved ``tokenizer.chat_template`` so
``_detect_safetensors_features`` (route layer) sees the same
template the model actually uses."""
Stores ``chat_template_info`` on ``self.models[model_name]`` with the
resolved ``tokenizer.chat_template``."""
entry = self.models.get(model_name)
if not entry:
return
@ -246,10 +264,8 @@ class MLXInferenceBackend:
max_new_tokens = 256,
repetition_penalty = 1.0,
cancel_event = None,
# Reasoning / tool kwargs forwarded by the route + worker -- the
# MLX path renders the template via apply_chat_template_for_
# generation so these are honoured the same way as the
# transformers path.
# Reasoning / tool kwargs forwarded by the route + worker; rendered via
# apply_chat_template_for_generation like the transformers path.
tools = None,
enable_thinking = None,
reasoning_effort = None,
@ -258,6 +274,9 @@ class MLXInferenceBackend:
if self._model is None:
raise RuntimeError("No model loaded")
# Reset so a failed run cannot surface stale stats.
self.last_generation_stats = None
# Build messages with system prompt
full_messages = []
if system_prompt:
@ -275,11 +294,9 @@ class MLXInferenceBackend:
{"type": "text", "text": content},
]
elif isinstance(content, list):
# Prepend image if not already there
# Prepend image if not already present
has_image = any(
p.get("type") == "image"
for p in content
if isinstance(p, dict)
p.get("type") == "image" for p in content if isinstance(p, dict)
)
if not has_image:
content.insert(0, {"type": "image"})
@ -349,9 +366,7 @@ class MLXInferenceBackend:
preserve_thinking = preserve_thinking,
)
if prompt is None:
raise RuntimeError(
"apply_chat_template returned None — tokenizer may be incompatible"
)
raise RuntimeError("apply_chat_template returned None — tokenizer may be incompatible")
sampler = make_sampler(
temp = temperature,
@ -360,8 +375,7 @@ class MLXInferenceBackend:
min_p = float(min_p or 0.0),
min_tokens_to_keep = 1,
)
# Only build a logits processor when we actually have a non-trivial
# repetition penalty (1.0 is the no-op value).
# Only build a logits processor for a non-trivial repetition penalty.
logits_processors = None
if repetition_penalty is not None and float(repetition_penalty) not in (
0.0,
@ -380,6 +394,7 @@ class MLXInferenceBackend:
type(self._tokenizer).__name__,
)
with self._generation_lock:
final_response = None
try:
gen_kwargs = dict(
prompt = prompt,
@ -393,8 +408,9 @@ class MLXInferenceBackend:
self._tokenizer,
**gen_kwargs,
):
final_response = response
token_ids.append(response.token)
# Decode full sequence with skip_special_tokens — same as GPU
# Decode full sequence with skip_special_tokens
cumulative = self._tokenizer.decode(
token_ids,
skip_special_tokens = True,
@ -405,9 +421,17 @@ class MLXInferenceBackend:
break
except Exception as e:
import traceback
logger.error("stream_generate failed:\n%s", traceback.format_exc())
raise
finally:
# Latch final cumulative stats for the usage/timings chunk.
if final_response is not None:
self.last_generation_stats = _build_generation_stats(
getattr(final_response, "prompt_tokens", 0),
getattr(final_response, "prompt_tps", 0.0),
getattr(final_response, "generation_tokens", 0),
getattr(final_response, "generation_tps", 0.0),
)
def _generate_vlm(
self,
@ -432,10 +456,9 @@ class MLXInferenceBackend:
apply_chat_template_for_generation,
)
# Pick the chat-template-aware caller: processors that expose
# their own apply_chat_template + chat_template attr (e.g.
# Qwen2.5-VL) use it directly; otherwise fall back to the
# nested tokenizer.
# Pick the chat-template-aware caller: processors with their own
# apply_chat_template + chat_template (e.g. Qwen2.5-VL) use it
# directly; else fall back to the nested tokenizer.
chat_target = self._processor
if (
getattr(self._processor, "apply_chat_template", None) is None
@ -453,8 +476,7 @@ class MLXInferenceBackend:
preserve_thinking = preserve_thinking,
)
# For VLM: always use mlx_vlm's stream_generate which handles
# pixel_values properly (passes None for text-only, image for VLM)
# mlx_vlm's stream_generate handles pixel_values (None for text-only)
images = [image] if image is not None else None
cumulative = ""
@ -464,11 +486,9 @@ class MLXInferenceBackend:
image is not None,
)
# mlx_vlm.stream_generate forwards **kwargs into generate_step, which
# accepts temp/top_p/top_k/repetition_penalty (and builds the sampler
# + logits_processors internally). Pass them through.
# NOTE: mlx_vlm.generate_step expects ``temperature=`` (long form) —
# passing ``temp=`` silently falls into **kwargs and is ignored,
# leaving generation stuck at the default 0.0 (greedy).
# builds the sampler + logits_processors internally.
# GOTCHA: generate_step expects ``temperature=`` (long form); ``temp=``
# silently falls into **kwargs and is ignored, stuck at greedy 0.0.
vlm_kwargs = dict(
max_tokens = max_new_tokens,
temperature = temperature,
@ -483,23 +503,36 @@ class MLXInferenceBackend:
vlm_kwargs["repetition_penalty"] = float(repetition_penalty)
with self._generation_lock:
for response in vlm_stream(
self._model,
self._processor,
prompt,
images,
**vlm_kwargs,
):
token_text = (
response.text if hasattr(response, "text") else str(response)
)
cumulative += token_text
yield cumulative
if cancel_event and cancel_event.is_set():
break
final_response = None
try:
for response in vlm_stream(
self._model,
self._processor,
prompt,
images,
**vlm_kwargs,
):
final_response = response
token_text = response.text if hasattr(response, "text") else str(response)
cumulative += token_text
yield cumulative
if cancel_event and cancel_event.is_set():
break
finally:
# mlx_vlm exposes the same stats fields as mlx_lm.
if final_response is not None:
self.last_generation_stats = _build_generation_stats(
getattr(final_response, "prompt_tokens", 0),
getattr(final_response, "prompt_tps", 0.0),
getattr(final_response, "generation_tokens", 0),
getattr(final_response, "generation_tps", 0.0),
)
def generate_with_adapter_control(
self, use_adapter = None, cancel_event = None, **gen_kwargs
self,
use_adapter = None,
cancel_event = None,
**gen_kwargs,
) -> Generator[str, None, None]:
# MLX LoRA adapter toggling not yet supported — generate normally
yield from self.generate_chat_response(cancel_event = cancel_event, **gen_kwargs)

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

@ -4,13 +4,12 @@
"""
Inference orchestrator — subprocess-based.
Provides the same API as InferenceBackend, but delegates all ML work
to a persistent subprocess. The subprocess is spawned on first model load
and stays alive for subsequent requests.
Same API as InferenceBackend, but delegates all ML work to a persistent
subprocess spawned on first model load and reused for later requests.
When switching between models that need different transformers versions
(e.g. GLM-4.7-Flash needs 5.x, Qwen needs 4.57.x), the old subprocess
is killed and a new one is spawned with the correct version.
When switching between models needing different transformers versions
(e.g. GLM-4.7-Flash needs 5.x, Qwen needs 4.57.x), the old subprocess is
killed and a new one spawned with the correct version.
Pattern follows core/training/training.py.
"""
@ -18,6 +17,7 @@ Pattern follows core/training/training.py.
import atexit
import base64
import os
import signal
import structlog
from loggers import get_logger
import multiprocessing as mp
@ -51,9 +51,8 @@ class InferenceOrchestrator:
"""
Inference backend orchestrator — subprocess-based.
Exposes the same API surface as InferenceBackend so routes/inference.py
needs minimal changes. Internally, all heavy ML operations happen in
a persistent subprocess.
Same API surface as InferenceBackend (so routes/inference.py needs
minimal changes); all heavy ML work happens in a persistent subprocess.
"""
def __init__(self):
@ -61,18 +60,15 @@ class InferenceOrchestrator:
self._proc: Optional[mp.Process] = None
self._cmd_queue: Any = None
self._resp_queue: Any = None
self._cancel_event: Any = None # mp.Event — set to cancel generation instantly
self._cancel_event: Any = None # mp.Event — set to cancel generation
self._lock = threading.Lock()
self._gen_lock = (
threading.Lock()
) # Serializes generation — one request at a time
self._gen_lock = threading.Lock() # Serializes generation
# Dispatcher state — for compare mode (adapter-controlled requests).
# Instead of serializing via _gen_lock, adapter-controlled requests
# send commands directly to the subprocess and read from per-request
# mailboxes. A dispatcher thread routes resp_queue events by request_id.
# Dispatcher state for compare mode (adapter-controlled requests):
# bypass _gen_lock, send commands directly, read from per-request
# mailboxes routed by a dispatcher thread on request_id.
self._mailboxes: dict[str, queue.Queue] = {}
self._mailbox_lock = threading.Lock() # Protects _mailboxes dict
self._mailbox_lock = threading.Lock()
self._dispatcher_thread: Optional[threading.Thread] = None
self._dispatcher_stop = threading.Event()
@ -88,16 +84,12 @@ class InferenceOrchestrator:
self._top_hub_cache: Optional[list[str]] = None
self._top_models_ready = threading.Event()
# Version tracking for subprocess reuse
self._current_transformers_major: Optional[str] = None # "4" or "5"
atexit.register(self._cleanup)
logger.info("InferenceOrchestrator initialized (subprocess mode)")
# Kick off background fetch of top models from HF
threading.Thread(
target = self._fetch_top_models, daemon = True, name = "top-models"
).start()
threading.Thread(target = self._fetch_top_models, daemon = True, name = "top-models").start()
# ------------------------------------------------------------------
# Default models (top GGUFs fetched dynamically from HF)
@ -105,14 +97,13 @@ class InferenceOrchestrator:
@property
def default_models(self) -> list[str]:
# Wait up to 5s for background HF fetch to finish
# Wait up to 5s for background HF fetch
self._top_models_ready.wait(timeout = 5)
top_gguf = self._top_gguf_cache or []
top_hub = self._top_hub_cache or []
# Curated static defaults first (editorial picks like new models),
# then HF download-ranked models to backfill.
# Send extras so the frontend still has 4 per category
# after removing already-downloaded models.
# Curated static defaults first, then HF download-ranked to backfill.
# Send extras so the frontend keeps 4 per category after removing
# downloaded ones.
result: list[str] = []
seen: set[str] = set()
for m in self._static_models + top_gguf + top_hub:
@ -125,7 +116,6 @@ class InferenceOrchestrator:
"""Fetch top GGUF and non-GGUF repos from unsloth by downloads."""
try:
import httpx
resp = httpx.get(
"https://huggingface.co/api/models",
params = {
@ -138,16 +128,13 @@ class InferenceOrchestrator:
)
if resp.status_code == 200:
models = resp.json()
# Top 40 GGUFs - frontend pages through them on-demand via
# infinite scroll, so we send a deep pool.
gguf_ids = [
m["id"] for m in models if m.get("id", "").upper().endswith("-GGUF")
][:40]
# Top 40 GGUFs (deep pool for frontend infinite scroll)
gguf_ids = [m["id"] for m in models if m.get("id", "").upper().endswith("-GGUF")][
:40
]
# Top 40 non-GGUF hub models
hub_ids = [
m["id"]
for m in models
if not m.get("id", "").upper().endswith("-GGUF")
m["id"] for m in models if not m.get("id", "").upper().endswith("-GGUF")
][:40]
if gguf_ids:
self._top_gguf_cache = gguf_ids
@ -199,16 +186,16 @@ class InferenceOrchestrator:
def _shutdown_subprocess(self, timeout: float = 10.0) -> None:
"""Gracefully shut down the inference subprocess."""
self._stop_dispatcher() # Stop dispatcher before killing subprocess
self._stop_dispatcher() # before killing subprocess
if self._proc is None or not self._proc.is_alive():
self._proc = None
return
# 1. Cancel any ongoing generation first (instant via mp.Event)
self._cancel_generation()
time.sleep(0.5) # Brief wait for generation to stop
time.sleep(0.5)
# 2. Drain stale responses from queue
# 2. Drain stale responses
self._drain_queue()
# 3. Send shutdown command
@ -250,9 +237,46 @@ class InferenceOrchestrator:
self._shutdown_subprocess(timeout = 5.0)
def _ensure_subprocess_alive(self) -> bool:
"""Check if subprocess is alive."""
"""True if the subprocess is alive."""
return self._proc is not None and self._proc.is_alive()
def _subprocess_crash_message(self, context: str) -> str:
"""Return a user-facing crash message with the worker exit status."""
context_label = {
"wait": "loading the model",
"generation": "generating a response",
"audio generation": "generating audio",
"audio input generation": "processing audio input",
}.get(context, context)
message = f"The inference worker stopped unexpectedly while {context_label}."
if self._proc is None:
return f"{message} Details: process missing."
exitcode = self._proc.exitcode
pid = self._proc.pid
if exitcode is None:
return f"{message} Details: pid={pid}."
if exitcode < 0:
signum = -exitcode
try:
sig_name = signal.Signals(signum).name
except ValueError:
sig_name = f"SIG{signum}"
suffix = ""
if sig_name == "SIGKILL":
suffix = (
" This usually means the system killed it under memory pressure. "
"Try a smaller model, lower context length, or close other GPU-heavy apps."
)
return (
f"{message}{suffix} " f"Details: pid={pid}, signal={sig_name}, exitcode={exitcode}."
)
return f"{message} Details: pid={pid}, exitcode={exitcode}."
# ------------------------------------------------------------------
# Queue helpers
# ------------------------------------------------------------------
@ -277,17 +301,19 @@ class InferenceOrchestrator:
except (EOFError, OSError, ValueError):
return None
def _wait_response(self, expected_type: str, timeout: float = 300.0) -> dict:
def _wait_response(
self,
expected_type: str,
timeout: float = 300.0,
) -> dict:
"""Block until a response of the expected type arrives.
Also handles 'status' and 'error' events during the wait.
Returns the matching response dict.
Raises RuntimeError on timeout or subprocess crash.
Also handles 'status' and 'error' events during the wait. Returns the
matching response dict; raises RuntimeError on timeout or crash.
The *timeout* is an **inactivity** timeout: it resets whenever the
subprocess sends a status message, so long-running operations (large
downloads, slow model loads) won't be killed as long as the subprocess
keeps reporting progress.
*timeout* is an **inactivity** timeout: it resets on each status
message, so long-running operations (large downloads, slow loads)
survive as long as the subprocess keeps reporting progress.
"""
deadline = time.monotonic() + timeout
@ -298,7 +324,7 @@ class InferenceOrchestrator:
if resp is None:
# Check subprocess health
if not self._ensure_subprocess_alive():
raise RuntimeError("Inference subprocess crashed during wait")
raise RuntimeError(self._subprocess_crash_message("wait"))
continue
rtype = resp.get("type", "")
@ -329,8 +355,7 @@ class InferenceOrchestrator:
)
raise RuntimeError(
f"Timeout waiting for '{expected_type}' response "
f"(no activity for {timeout}s)"
f"Timeout waiting for '{expected_type}' response " f"(no activity for {timeout}s)"
)
def _drain_queue(self) -> list:
@ -349,8 +374,8 @@ class InferenceOrchestrator:
def _drain_until_gen_done(self, timeout: float = 5.0) -> None:
"""Consume resp_queue events until gen_done/gen_error, discarding them.
Called after cancel to ensure stale tokens from the cancelled
generation don't leak into the next request.
Called after cancel so stale tokens from the cancelled generation
don't leak into the next request.
"""
deadline = time.monotonic() + timeout
while time.monotonic() < deadline:
@ -371,10 +396,9 @@ class InferenceOrchestrator:
def _start_dispatcher(self) -> None:
"""Start the dispatcher thread if not already running.
The dispatcher reads from the shared resp_queue and routes
responses to per-request mailbox queues. This allows multiple
adapter-controlled (compare) requests to be in-flight without
holding _gen_lock.
The dispatcher reads the shared resp_queue and routes responses to
per-request mailbox queues, letting multiple adapter-controlled
(compare) requests be in-flight without holding _gen_lock.
"""
if self._dispatcher_thread is not None and self._dispatcher_thread.is_alive():
return
@ -426,9 +450,8 @@ class InferenceOrchestrator:
mbox.put(resp)
continue
# No matching mailbox — might be for a _gen_lock reader or orphaned
# Push it back so _read_resp can pick it up. But we can't un-get
# from mp.Queue, so log a warning.
# No matching mailbox (a _gen_lock reader or orphaned). Can't
# un-get from mp.Queue, so just log.
if rtype not in ("status",):
logger.debug(
"Dispatcher: no mailbox for request_id=%s type=%s, dropping",
@ -453,16 +476,13 @@ class InferenceOrchestrator:
enable_thinking: Optional[bool] = None,
reasoning_effort: Optional[str] = None,
preserve_thinking: Optional[bool] = None,
stats_holder: Optional[dict] = None,
) -> Generator[str, None, None]:
"""Dispatched generation — sends command without holding _gen_lock.
Uses a per-request mailbox to receive tokens. This allows two
compare-mode requests to be queued in the subprocess simultaneously,
eliminating the inter-generation round-trip overhead.
The subprocess processes commands sequentially from its cmd_queue,
so generation is still serialized at the GPU level — we just avoid
the orchestrator-level lock contention.
Uses a per-request mailbox for tokens so two compare-mode requests can
be queued at once. The subprocess still runs commands sequentially, so
GPU work stays serialized; this only avoids orchestrator lock contention.
"""
if not self._ensure_subprocess_alive():
yield "Error: Inference subprocess is not running"
@ -528,22 +548,23 @@ class InferenceOrchestrator:
except queue.Empty:
# Timeout — check subprocess health
if not self._ensure_subprocess_alive():
yield "Error: Inference subprocess crashed during generation"
yield f"Error: {self._subprocess_crash_message('generation')}"
return
continue
rtype = resp.get("type", "")
if rtype == "token":
# Check cancel from route (e.g. SSE connection closed)
# Cancel from route (e.g. SSE connection closed)
if cancel_event is not None and cancel_event.is_set():
self._cancel_generation()
# Drain remaining events for this request
self._drain_mailbox(mailbox, timeout = 5.0)
return
yield resp.get("text", "")
elif rtype == "gen_done":
if stats_holder is not None:
stats_holder["stats"] = resp.get("stats")
return
elif rtype == "gen_error":
@ -553,7 +574,11 @@ class InferenceOrchestrator:
with self._mailbox_lock:
self._mailboxes.pop(request_id, None)
def _drain_mailbox(self, mailbox: queue.Queue, timeout: float = 5.0) -> None:
def _drain_mailbox(
self,
mailbox: queue.Queue,
timeout: float = 5.0,
) -> None:
"""Drain a mailbox until gen_done/gen_error, discarding tokens."""
deadline = time.monotonic() + timeout
while time.monotonic() < deadline:
@ -571,8 +596,8 @@ class InferenceOrchestrator:
def _wait_dispatcher_idle(self) -> None:
"""Wait for all dispatched requests to complete, then stop dispatcher.
Called by _generate_inner before using the _gen_lock path, to ensure
the dispatcher thread isn't competing for resp_queue reads.
Called by _generate_inner before the _gen_lock path so the dispatcher
thread isn't competing for resp_queue reads.
"""
if self._dispatcher_thread is None or not self._dispatcher_thread.is_alive():
return
@ -585,9 +610,9 @@ class InferenceOrchestrator:
break
time.sleep(0.1)
# Only stop dispatcher if all mailboxes drained. If compare
# requests are still active, leave the dispatcher running so
# their token routing isn't killed mid-stream.
# Only stop dispatcher if all mailboxes drained. If compare requests
# are still active, leave it running so their token routing isn't
# killed mid-stream.
with self._mailbox_lock:
still_active = bool(self._mailboxes)
if still_active:
@ -615,9 +640,8 @@ class InferenceOrchestrator:
) -> bool:
"""Load a model for inference.
Always spawns a fresh subprocess for each model load. This ensures
a clean Python interpreter — no stale unsloth patches, torch.compile
caches, or inspect.getsource() failures from a previous model.
Always spawns a fresh subprocess per load for a clean interpreter (no
stale unsloth patches, torch.compile caches, or getsource failures).
"""
from utils.transformers_version import needs_transformers_5
@ -646,16 +670,14 @@ class InferenceOrchestrator:
sub_config["resolved_gpu_ids"] = resolved_gpu_ids
sub_config["gpu_selection"] = gpu_selection
# Always kill existing subprocess and spawn fresh.
# Reusing a subprocess after unsloth patches torch internals
# causes inspect.getsource() failures on the next model load.
# Always kill the existing subprocess and spawn fresh: reusing one
# after unsloth patches torch internals breaks getsource on reload.
if self._ensure_subprocess_alive():
self._cancel_generation()
time.sleep(0.3)
self._shutdown_subprocess()
elif self._proc is not None:
# Dead subprocess — clean up
self._shutdown_subprocess(timeout = 2)
disable_xet = sub_config.get("disable_xet", False) or (
@ -677,24 +699,22 @@ class InferenceOrchestrator:
try:
resp = self._wait_response("loaded")
except DownloadStallError:
# First stall and Xet was enabled -> retry with Xet disabled
# First stall with Xet on -> retry with Xet disabled
if attempt == 0 and not disable_xet:
logger.warning(
"Download stalled for '%s' -- retrying with "
"HF_HUB_DISABLE_XET=1",
"Download stalled for '%s' -- retrying with HF_HUB_DISABLE_XET=1",
model_name,
)
self._shutdown_subprocess(timeout = 5)
disable_xet = True
continue
# Second stall (or already had xet disabled) -> give up
# Second stall (or xet already off) -> give up
self._shutdown_subprocess(timeout = 5)
raise RuntimeError(
f"Download stalled for '{model_name}' even with "
f"HF_HUB_DISABLE_XET=1 -- check your network connection"
)
# Got a response — check success
if resp.get("success"):
self._current_transformers_major = needed_major
model_info = resp.get("model_info", {})
@ -702,22 +722,20 @@ class InferenceOrchestrator:
self.models[self.active_model_name] = {
"is_vision": model_info.get("is_vision", False),
"is_lora": model_info.get("is_lora", False),
"is_mlx": model_info.get("is_mlx", False),
"display_name": model_info.get("display_name", model_name),
"is_audio": model_info.get("is_audio", False),
"audio_type": model_info.get("audio_type"),
"has_audio_input": model_info.get("has_audio_input", False),
"context_length": model_info.get("context_length"),
}
# Mirror chat_template_info so routes can classify
# capabilities without re-entering the subprocess.
# Mirror chat_template_info so routes can classify caps
# without re-entering the subprocess.
_tpl_info = model_info.get("chat_template_info")
if isinstance(_tpl_info, dict):
self.models[self.active_model_name]["chat_template_info"] = (
_tpl_info
)
self.models[self.active_model_name]["chat_template_info"] = _tpl_info
self.loading_models.discard(model_name)
logger.info(
"Model '%s' loaded successfully in subprocess", model_name
)
logger.info("Model '%s' loaded successfully in subprocess", model_name)
return True
else:
error = resp.get("error", "Failed to load model")
@ -746,7 +764,7 @@ class InferenceOrchestrator:
return True
if not self._ensure_subprocess_alive():
# No subprocess — just clear local state
# No subprocess — clear local state
self.models.pop(model_name, None)
if self.active_model_name == model_name:
self.active_model_name = None
@ -793,13 +811,16 @@ class InferenceOrchestrator:
enable_thinking: Optional[bool] = None,
reasoning_effort: Optional[str] = None,
preserve_thinking: Optional[bool] = None,
stats_holder: Optional[dict] = None,
) -> Generator[str, None, None]:
"""Generate response, streaming tokens from subprocess.
Optional ``tools`` / ``enable_thinking`` / ``reasoning_effort`` /
``preserve_thinking`` kwargs are forwarded into the worker so
``tokenizer.apply_chat_template`` can render tool schemas and
reasoning controls when the template understands them.
``tools`` / ``enable_thinking`` / ``reasoning_effort`` /
``preserve_thinking`` are forwarded so the template can render tool
schemas and reasoning controls.
``stats_holder``: caller-owned dict; on gen_done its "stats" key gets
the worker's usage/timings. Request-scoped to avoid cross-stream reads.
"""
yield from self._generate_inner(
messages = messages,
@ -817,6 +838,7 @@ class InferenceOrchestrator:
enable_thinking = enable_thinking,
reasoning_effort = reasoning_effort,
preserve_thinking = preserve_thinking,
stats_holder = stats_holder,
)
def generate_chat_completion_with_tools(
@ -838,25 +860,29 @@ class InferenceOrchestrator:
auto_heal_tool_calls: bool = True,
tool_call_timeout: int = 300,
session_id: Optional[str] = None,
rag_scope: Optional[dict] = None,
confirm_tool_calls: bool = False,
bypass_permissions: bool = False,
use_adapter: Optional[Union[bool, str]] = None,
stats_holder: Optional[dict] = None,
**_unused,
):
"""Run the safetensors agentic tool loop in this (parent)
process, calling the worker for each generation turn.
"""Run the safetensors agentic tool loop in the parent process,
calling the worker for each turn.
Yields the same event dicts as the GGUF tool loop so the route
layer can stream both backends through one helper. See
``safetensors_agentic.run_safetensors_tool_loop`` for the
event protocol.
Yields the same event dicts as the GGUF tool loop so the route layer
can stream both backends through one helper.
"""
from core.inference.safetensors_agentic import run_safetensors_tool_loop
from core.inference.tools import execute_tool
max_new_tokens = max_tokens if max_tokens and max_tokens > 0 else 2048
def _single_turn(conv: list):
# ``conv`` already carries any system message because the
# loop appends to a list seeded with system+user above.
def _single_turn(conv: list, *, active_tools: Optional[list[dict]] = None):
# ``conv`` already carries any system message. ``active_tools`` lets
# run_safetensors_tool_loop drop one-shot tools (e.g. render_html) from
# later same-response prompts.
turn_tools = active_tools if active_tools is not None else tools
common_kwargs = dict(
messages = conv,
system_prompt = "",
@ -868,10 +894,12 @@ class InferenceOrchestrator:
max_new_tokens = max_new_tokens,
repetition_penalty = repetition_penalty,
cancel_event = cancel_event,
tools = tools,
tools = turn_tools,
enable_thinking = enable_thinking,
reasoning_effort = reasoning_effort,
preserve_thinking = preserve_thinking,
# last turn wins, like the GGUF tool loop
stats_holder = stats_holder,
)
if use_adapter is not None:
yield from self.generate_with_adapter_control(
@ -895,23 +923,28 @@ class InferenceOrchestrator:
max_tool_iterations = max_tool_iterations,
tool_call_timeout = tool_call_timeout,
session_id = session_id,
rag_scope = rag_scope,
confirm_tool_calls = confirm_tool_calls,
bypass_permissions = bypass_permissions,
)
def generate_with_adapter_control(
self,
use_adapter: Optional[Union[bool, str]] = None,
cancel_event = None,
stats_holder: Optional[dict] = None,
**gen_kwargs,
) -> Generator[str, None, None]:
"""Generate with adapter control, streaming tokens from subprocess.
Uses the dispatcher path (no _gen_lock) so that compare-mode
requests don't block each other. The subprocess naturally
serializes them via its sequential command loop.
Uses the dispatcher path (no _gen_lock) so compare-mode requests
don't block each other; the subprocess serializes them via its
sequential command loop.
"""
yield from self._generate_dispatched(
use_adapter = use_adapter,
cancel_event = cancel_event,
stats_holder = stats_holder,
**gen_kwargs,
)
@ -932,12 +965,12 @@ class InferenceOrchestrator:
enable_thinking: Optional[bool] = None,
reasoning_effort: Optional[str] = None,
preserve_thinking: Optional[bool] = None,
stats_holder: Optional[dict] = None,
) -> Generator[str, None, None]:
"""Inner generation logic — sends command to subprocess, yields tokens.
Serialized by _gen_lock: only one generation runs at a time.
This prevents concurrent readers from consuming each other's
tokens off the shared resp_queue.
Serialized by _gen_lock (one generation at a time) so concurrent
readers don't consume each other's tokens off the shared resp_queue.
"""
if not self._ensure_subprocess_alive():
yield "Error: Inference subprocess is not running"
@ -947,14 +980,11 @@ class InferenceOrchestrator:
yield "Error: No active model"
return
# If the dispatcher is running (from a previous compare-mode request),
# wait for all dispatched requests to finish, then stop the dispatcher
# so we can safely read from resp_queue directly.
# Drain any prior compare-mode dispatcher so we can read resp_queue.
self._wait_dispatcher_idle()
# Serialize generation — single GPU, one generation at a time.
# Without this lock, two concurrent readers on the same resp_queue
# can consume and drop each other's token events.
# Serialize generation: two concurrent readers on resp_queue would
# consume and drop each other's token events.
with self._gen_lock:
yield from self._generate_locked(
messages = messages,
@ -972,6 +1002,7 @@ class InferenceOrchestrator:
enable_thinking = enable_thinking,
reasoning_effort = reasoning_effort,
preserve_thinking = preserve_thinking,
stats_holder = stats_holder,
)
def _generate_locked(
@ -991,6 +1022,7 @@ class InferenceOrchestrator:
enable_thinking: Optional[bool] = None,
reasoning_effort: Optional[str] = None,
preserve_thinking: Optional[bool] = None,
stats_holder: Optional[dict] = None,
) -> Generator[str, None, None]:
"""Actual generation logic — must be called under _gen_lock."""
request_id = str(uuid.uuid4())
@ -1016,8 +1048,7 @@ class InferenceOrchestrator:
if use_adapter is not None:
cmd["use_adapter"] = use_adapter
# Only forward template kwargs the caller actually set so older
# workers that ignore unknown keys still work.
# Only forward template kwargs the caller set, for older worker compat.
if tools is not None:
cmd["tools"] = tools
if enable_thinking is not None:
@ -1033,15 +1064,14 @@ class InferenceOrchestrator:
yield f"Error: {exc}"
return
# Yield tokens from response queue — we are the only reader
# because _gen_lock is held.
# We are the only resp_queue reader (under _gen_lock).
while True:
resp = self._read_resp(timeout = 30.0)
if resp is None:
# Check subprocess health
if not self._ensure_subprocess_alive():
yield "Error: Inference subprocess crashed during generation"
yield f"Error: {self._subprocess_crash_message('generation')}"
return
continue
@ -1058,17 +1088,18 @@ class InferenceOrchestrator:
return
if rtype == "token":
# Check cancel from route (e.g. SSE connection closed)
# Cancel from route (e.g. SSE connection closed)
if cancel_event is not None and cancel_event.is_set():
self._cancel_generation()
# Wait for the subprocess to acknowledge cancellation
# (gen_done/gen_error) so stale events don't leak into
# the next generation request.
# Wait for the cancel ack so stale events don't leak into
# the next request.
self._drain_until_gen_done(timeout = 5.0)
return
yield resp.get("text", "")
elif rtype == "gen_done":
if stats_holder is not None:
stats_holder["stats"] = resp.get("stats")
return
elif rtype == "gen_error":
@ -1102,7 +1133,7 @@ class InferenceOrchestrator:
) -> Tuple[bytes, int]:
"""Generate TTS audio. Returns (wav_bytes, sample_rate).
Blocking — sends command and waits for the complete audio response.
Blocking — sends command and waits for the full audio response.
"""
if not self._ensure_subprocess_alive():
raise RuntimeError("Inference subprocess is not running")
@ -1137,9 +1168,7 @@ class InferenceOrchestrator:
if resp is None:
if not self._ensure_subprocess_alive():
raise RuntimeError(
"Inference subprocess crashed during audio generation"
)
raise RuntimeError(self._subprocess_crash_message("audio generation"))
continue
rtype = resp.get("type", "")
@ -1229,11 +1258,9 @@ class InferenceOrchestrator:
request_id = str(uuid.uuid4())
# Convert numpy array to list for mp.Queue serialization
# numpy array -> list for mp.Queue serialization
audio_data = (
audio_array.tolist()
if hasattr(audio_array, "tolist")
else list(audio_array)
audio_array.tolist() if hasattr(audio_array, "tolist") else list(audio_array)
)
cmd = {
@ -1263,7 +1290,7 @@ class InferenceOrchestrator:
if resp is None:
if not self._ensure_subprocess_alive():
yield "Error: Inference subprocess crashed during audio input generation"
yield ("Error: " + self._subprocess_crash_message("audio input generation"))
return
continue
@ -1294,10 +1321,12 @@ class InferenceOrchestrator:
# Local helpers (no subprocess needed)
# ------------------------------------------------------------------
def resize_image(self, img, max_size: int = 800):
"""Resize image while maintaining aspect ratio.
No ML imports needed — runs locally in parent process.
"""
def resize_image(
self,
img,
max_size: int = 800,
):
"""Resize image preserving aspect ratio (runs locally, no ML imports)."""
if img is None:
return None
if img.size[0] > max_size or img.size[1] > max_size:
@ -1316,28 +1345,26 @@ class InferenceOrchestrator:
return base64.b64encode(buf.getvalue()).decode("ascii")
def get_current_model(self) -> Optional[str]:
"""Get currently active model name."""
"""Currently active model name."""
return self.active_model_name
def is_model_loading(self) -> bool:
"""Check if any model is currently loading."""
"""True if any model is loading."""
return len(self.loading_models) > 0
def get_loading_model(self) -> Optional[str]:
"""Get name of currently loading model."""
"""Name of the currently loading model."""
return next(iter(self.loading_models)) if self.loading_models else None
def check_vision_model_compatibility(self) -> bool:
"""Check if current model supports vision."""
"""True if the current model supports vision."""
if self.active_model_name and self.active_model_name in self.models:
return self.models[self.active_model_name].get("is_vision", False)
return False
def _is_gpt_oss_model(self, model_name: str = None) -> bool:
"""Parent-side gpt-oss detection so the safetensors route can run
the same guard without an IPC round-trip to the subprocess."""
"""Parent-side gpt-oss detection so the route avoids an IPC round-trip."""
from utils.datasets import is_gpt_oss_model_name
return is_gpt_oss_model_name(model_name or self.active_model_name or "")
@ -1346,7 +1373,7 @@ _inference_backend = None
def get_inference_backend() -> InferenceOrchestrator:
"""Get global inference backend instance (orchestrator)."""
"""Global inference backend instance (orchestrator)."""
global _inference_backend
if _inference_backend is None:
_inference_backend = InferenceOrchestrator()

255
studio/backend/core/inference/pricing.py
View file

@ -1,50 +1,23 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Static per-MTok pricing tables for external providers, plus a
``calculate_cost`` helper that turns an upstream ``usage`` block into
a USD figure for surfacing in the chat UI.
"""Per-MTok pricing tables and ``calculate_cost`` (usage block -> USD).
Neither the Anthropic Messages API nor the OpenAI Responses API
reports a ``cost`` field on the response. Both expose detailed token
counts (input, output, cache hits, server-tool invocations); pricing
multipliers live in the provider docs. We fold the docs into a static
table here, multiply by the usage block, and emit a per-turn cost +
running session total client-side.
Sources (verified live 2026-05-22):
- Anthropic models overview:
https://platform.claude.com/docs/en/about-claude/models/overview
- Anthropic prompt-caching multipliers (5m write 1.25x, 1h write 2x,
read 0.1x):
https://platform.claude.com/docs/en/build-with-claude/prompt-caching
- Anthropic web search ($10 / 1000 searches, code execution
free-with-paid when paired with the newer web tools):
https://platform.claude.com/docs/en/agents-and-tools/tool-use/web-search-tool
https://platform.claude.com/docs/en/agents-and-tools/tool-use/code-execution-tool
- OpenAI pricing page (input / output per MTok per model family):
https://platform.openai.com/docs/pricing
Sources: Anthropic prompt-caching docs (5m write 1.25x, 1h write 2x,
read 0.1x), web search ($10/1000), code execution; OpenAI pricing page.
"""
from __future__ import annotations
from typing import Any, Optional
# Per-million-token base pricing. `cache_5m_write_mult`, `cache_1h_write_mult`,
# `cache_read_mult` are multipliers ON `input_per_mtok` -- not absolute prices --
# matching how Anthropic publishes them (5m write = 1.25x base, etc.).
#
# `input_per_mtok` and `output_per_mtok` are USD per 1,000,000 tokens.
# Per-MTok base USD. Cache multipliers apply to `input_per_mtok`
# (not absolute prices), per Anthropic docs.
ANTHROPIC_PRICING: dict[str, dict[str, float]] = {
"claude-opus-4-7": {"input_per_mtok": 5.0, "output_per_mtok": 25.0},
"claude-opus-4-6": {"input_per_mtok": 5.0, "output_per_mtok": 25.0},
# Canonical 4.5 ids are referenced from backend defaults (e.g.
# PROVIDER_REGISTRY['anthropic'].default_models) without the date
# suffix. The dated ids ARE the canonical names per Anthropic's
# models overview, but lookups for the bare id ("claude-opus-4-5")
# don't prefix-match the dated key the other way around, so we
# alias both forms here. Otherwise calculate_cost returns
# priced=False + zero cost for the common ids.
# Alias bare + dated id: backend defaults use the bare form, which
# won't prefix-match the dated key.
"claude-opus-4-5": {"input_per_mtok": 5.0, "output_per_mtok": 25.0},
"claude-opus-4-5-20251101": {"input_per_mtok": 5.0, "output_per_mtok": 25.0},
"claude-opus-4-1": {"input_per_mtok": 15.0, "output_per_mtok": 75.0},
@ -59,19 +32,9 @@ ANTHROPIC_PRICING: dict[str, dict[str, float]] = {
}
OPENAI_PRICING: dict[str, dict[str, float]] = {
# All values verified against developers.openai.com/api/docs/pricing
# 2026-05-22. Update against the live pricing page on every model launch.
# Initial commit underbilled every gpt-5.x family 2-6x -- fixed here
# after PR review caught it via doc cross-check.
#
# `long_context_input_per_mtok` / `long_context_output_per_mtok` /
# `long_context_threshold` are populated when OpenAI publishes a
# second pricing tier for prompts above N input tokens. gpt-5.5 and
# gpt-5.4 cross over at 272k input tokens; the long-context rates
# are double the headline input price (and ~1.5x on output). Other
# families currently ship with a single rate (no `long_context_*`
# keys = no tier crossover). Reference:
# https://developers.openai.com/api/docs/pricing
# Verified against developers.openai.com/api/docs/pricing.
# `long_context_*` keys apply past the threshold (gpt-5.5/5.4: 272k);
# families without them ship a single rate.
"gpt-5.5": {
"input_per_mtok": 5.0,
"output_per_mtok": 30.0,
@ -91,43 +54,33 @@ OPENAI_PRICING: dict[str, dict[str, float]] = {
"gpt-5.4-mini": {"input_per_mtok": 0.75, "output_per_mtok": 4.5},
"gpt-5.4-nano": {"input_per_mtok": 0.20, "output_per_mtok": 1.25},
"gpt-5.3-codex": {"input_per_mtok": 1.75, "output_per_mtok": 14.0},
# chat-latest / gpt-5.3-chat-latest is an alias for the current
# ChatGPT model; same price as gpt-5.5.
# chat-latest aliases gpt-5.5.
"gpt-5.3-chat-latest": {"input_per_mtok": 5.0, "output_per_mtok": 30.0},
"chat-latest": {"input_per_mtok": 5.0, "output_per_mtok": 30.0},
# o-series and gpt-4.5: NOT currently listed on the pricing page.
# Removed to avoid silent-underbilling drift. Returning priced=False
# is honest; the UI can still render token counts. Restore with
# verified per-MTok rates if/when the page lists them again.
# o-series / gpt-4.5 left off the pricing page: omit so calculate_cost
# returns priced=False instead of silently $0.
}
# Shared multipliers (same across every Anthropic model).
# Shared multipliers (all Anthropic models).
ANTHROPIC_CACHE_5M_WRITE_MULT = 1.25
ANTHROPIC_CACHE_1H_WRITE_MULT = 2.0
ANTHROPIC_CACHE_READ_MULT = 0.1
# Anthropic fast-mode (Opus 4.6/4.7 only): 6x on input + output.
# https://platform.claude.com/docs/en/build-with-claude/fast-mode#pricing
ANTHROPIC_FAST_MODE_MULT = 6.0
# OpenAI: cache reads are 0.1x base input, cache writes are not billed
# separately (the first prefix-write request just pays normal input).
# OpenAI: cache reads 0.1x; cache writes pay input price.
OPENAI_CACHE_READ_MULT = 0.1
# Server-tool surcharges.
# Anthropic: $10 / 1000 web searches; code_execution is $0.05/hr after
# 50 free hours/day per org (no per-org visibility here, so the
# calculator reports the marginal rate).
# Server-tool surcharges. Anthropic code_exec: $0.05/hr marginal
# (50 free hours/day per org, not shown here).
ANTHROPIC_WEB_SEARCH_USD_PER_1K = 10.0
ANTHROPIC_CODE_EXEC_USD_PER_HOUR = 0.05
# OpenAI: web_search is billed at $10/1000 calls plus the model's
# token rate for the returned search content (already captured under
# input/output_tokens). The hosted shell tool bills per 20-minute
# session per container memory tier (1g/4g/16g/64g at
# $0.03/$0.12/$0.48/$1.92). Since Studio doesn't surface the memory
# tier in the cost ledger and most users land on the default 1g, we
# bill the 1g rate ($0.09/hour) and let the user inspect the OpenAI
# dashboard for the exact figure on heavier configs.
# Source: developers.openai.com/api/docs/pricing 2026-05-22.
# OpenAI container bills per memory tier; report the 1g default
# ($0.09/hr) since the tier isn't surfaced to the ledger.
OPENAI_WEB_SEARCH_USD_PER_1K = 10.0
OPENAI_CONTAINER_USD_PER_HOUR = 0.09 # 1g default tier; 3 x $0.03 / 60min
OPENAI_CONTAINER_USD_PER_HOUR = 0.09 # 1g default tier
def _lookup(provider: str, model: str) -> Optional[dict[str, float]]:
@ -142,41 +95,18 @@ def _lookup(provider: str, model: str) -> Optional[dict[str, float]]:
return None
if model in table:
return table[model]
# Fall back to a prefix match so date-suffixed snapshots
# ("gpt-5.5-2026-04-23") inherit the canonical-id prices.
for key, val in table.items():
if model.startswith(key):
return val
# Longest-prefix match on a dash boundary: dated snapshots inherit
# canonical prices, but "claude-opus-4-15" won't match "claude-opus-4-1".
for key in sorted(table, key = len, reverse = True):
if model.startswith(key) and (len(model) == len(key) or model[len(key)] == "-"):
return table[key]
return None
def calculate_cost(
provider: str,
model: str,
usage: dict[str, Any],
) -> dict[str, float]:
"""Return a per-turn USD cost breakdown.
def calculate_cost(provider: str, model: str, usage: dict[str, Any]) -> dict[str, float]:
"""Return a per-turn USD cost breakdown (per-bucket + total).
Returns a dict with the per-bucket cost AND the totals so the
frontend can render either a single number or a "where did the
money go" tooltip without re-doing the math:
{
"input_usd": 0.0042,
"output_usd": 0.012,
"cache_write_usd": 0.0001,
"cache_read_usd": 0.0008,
"server_tools_usd": 0.01,
"total_usd": 0.0271,
"billable_input_tokens": 5023, # input + cache_create + cache_read
"billable_output_tokens": 480,
"model_priced": "claude-opus-4-7",
"priced": true,
}
When the model isn't in the static table (new family, custom base
URL), `priced` is False and every USD field is 0.0; the frontend
can still show the token counts.
Unknown model -> ``priced`` False and USD fields 0.0 (token counts still report).
"""
prices = _lookup(provider, model)
out: dict[str, float] = {
@ -192,34 +122,61 @@ def calculate_cost(
"priced": bool(prices),
}
input_tokens = int(usage.get("input_tokens") or 0)
output_tokens = int(usage.get("output_tokens") or 0)
cache_creation = int(usage.get("cache_creation_input_tokens") or 0)
cache_read = int(usage.get("cache_read_input_tokens") or 0)
# OpenAI Responses reports cached tokens under input_tokens_details
# but ALSO folds them into the top-level input_tokens, so we don't
# add cache_read into the billable total again below (Anthropic
# excludes cache buckets from input_tokens, OpenAI includes them --
# the two providers differ here and the calculator must match).
if provider == "openai":
details = usage.get("input_tokens_details") or {}
# Accept raw (input_tokens/output_tokens) and Studio chat-style
# (prompt_tokens/completion_tokens) envelopes. Cache buckets differ:
# raw Anthropic: input_tokens EXCLUDES cache buckets
# raw OpenAI: input_tokens INCLUDES cache_read
# Studio Anthropic: prompt_tokens INCLUDES cache_creation + cache_read
# Studio OpenAI: prompt_tokens == raw input_tokens
# Clamp >=0 so corrupted payloads can't produce a negative bill.
cache_creation = max(0, int(usage.get("cache_creation_input_tokens") or 0))
cache_read_native_present = (
"cache_read_input_tokens" in usage and usage.get("cache_read_input_tokens") is not None
)
cache_read = max(0, int(usage.get("cache_read_input_tokens") or 0))
# Fall back to mirrored prompt_tokens_details only when native
# cache_read_input_tokens is absent; an explicit native 0 is
# authoritative, so a stale proxy mirror can't inflate cache_read.
if not cache_read_native_present:
details = usage.get("prompt_tokens_details") or {}
if isinstance(details, dict):
cache_read = max(cache_read, int(details.get("cached_tokens") or 0))
# OpenAI: cache_read already counted inside input_tokens.
cache_read = max(0, int(details.get("cached_tokens") or 0))
has_input_tokens = "input_tokens" in usage and usage.get("input_tokens") is not None
if has_input_tokens:
input_tokens = max(0, int(usage.get("input_tokens") or 0))
else:
# Chat-style: peel cache buckets back out for Anthropic to get
# the raw uncached prompt count.
prompt_tokens = max(0, int(usage.get("prompt_tokens") or 0))
if provider == "anthropic":
input_tokens = max(0, prompt_tokens - cache_creation - cache_read)
else:
input_tokens = prompt_tokens
# Prefer raw output_tokens even when 0 (an `or` would pick a stale
# completion_tokens).
if "output_tokens" in usage and usage.get("output_tokens") is not None:
output_tokens = max(0, int(usage.get("output_tokens") or 0))
else:
output_tokens = max(0, int(usage.get("completion_tokens") or 0))
if provider == "openai":
# Cached tokens land on input_tokens_details (raw Responses) or
# prompt_tokens_details (Studio chat-style).
for key in ("input_tokens_details", "prompt_tokens_details"):
details = usage.get(key) or {}
if isinstance(details, dict):
cache_read = max(cache_read, int(details.get("cached_tokens") or 0))
# OpenAI input_tokens already counts cache_read.
out["billable_input_tokens"] = input_tokens + cache_creation
else:
# Anthropic: input_tokens excludes cache_* buckets, add them all.
# Anthropic input_tokens excludes cache buckets; add them back.
out["billable_input_tokens"] = input_tokens + cache_creation + cache_read
out["billable_output_tokens"] = output_tokens
if not prices:
return out
# Long-context tier crossover (gpt-5.5 / gpt-5.4 today). OpenAI
# bills the whole turn at the long-context rate once the prompt
# crosses the threshold, NOT a per-token blend, so we pick a
# single (base, out_per) pair for this turn based on
# billable_input_tokens.
# Long-context tier: whole-turn flip (not per-token blend) once
# billable_input_tokens crosses the threshold.
lc_thresh = prices.get("long_context_threshold")
in_long_context_tier = (
lc_thresh is not None
@ -235,26 +192,32 @@ def calculate_cost(
base = prices["input_per_mtok"]
out_per = prices["output_per_mtok"]
# Anthropic fast-mode: 6x on input + output. Cache multipliers stack
# on top, so applying once to (base, out_per) flows into the
# cache_*_usd buckets below.
if provider == "anthropic" and usage.get("speed") == "fast":
base *= ANTHROPIC_FAST_MODE_MULT
out_per *= ANTHROPIC_FAST_MODE_MULT
if out["model_priced"]:
out["model_priced"] = f"{out['model_priced']} (fast)"
out["input_usd"] = (input_tokens / 1_000_000.0) * base
out["output_usd"] = (output_tokens / 1_000_000.0) * out_per
if provider == "anthropic":
# Split cache_creation across 5m / 1h buckets when the
# response surfaces the breakdown.
cc_breakdown = usage.get("cache_creation") or {}
cc_5m = int(cc_breakdown.get("ephemeral_5m_input_tokens") or 0)
cc_1h = int(cc_breakdown.get("ephemeral_1h_input_tokens") or 0)
# Split cache_creation into 5m / 1h buckets when surfaced.
# Tolerate non-dict (some proxies fold to an int total).
cc_raw = usage.get("cache_creation")
cc_breakdown = cc_raw if isinstance(cc_raw, dict) else {}
cc_5m = max(0, int(cc_breakdown.get("ephemeral_5m_input_tokens") or 0))
cc_1h = max(0, int(cc_breakdown.get("ephemeral_1h_input_tokens") or 0))
if cc_5m + cc_1h == 0 and cache_creation > 0:
# Fall back: assume default 5m pool when no breakdown is given.
# No breakdown -- assume default 5m pool.
cc_5m = cache_creation
out["cache_write_usd"] = (
cc_5m / 1_000_000.0
) * base * ANTHROPIC_CACHE_5M_WRITE_MULT + (
out["cache_write_usd"] = (cc_5m / 1_000_000.0) * base * ANTHROPIC_CACHE_5M_WRITE_MULT + (
cc_1h / 1_000_000.0
) * base * ANTHROPIC_CACHE_1H_WRITE_MULT
out["cache_read_usd"] = (
(cache_read / 1_000_000.0) * base * ANTHROPIC_CACHE_READ_MULT
)
out["cache_read_usd"] = (cache_read / 1_000_000.0) * base * ANTHROPIC_CACHE_READ_MULT
# Server-tool surcharges.
srv = usage.get("server_tool_use") or {}
if isinstance(srv, dict):
@ -265,24 +228,15 @@ def calculate_cost(
+ code_exec_hours * ANTHROPIC_CODE_EXEC_USD_PER_HOUR
)
else:
# OpenAI: cache writes share the base input price (no premium).
# Only cache reads get the 0.1x multiplier; subtract those from
# the input_usd we already counted so we don't double-bill.
# Anthropic excludes cache buckets from input_tokens, but
# OpenAI folds them in, so the math differs.
# OpenAI: cache writes pay base input, only reads get 0.1x.
# Subtract cached from already-counted input_usd to avoid
# double-billing (OpenAI folds cache into input_tokens).
if cache_read > 0:
non_cached_input = max(0, input_tokens - cache_read)
out["input_usd"] = (non_cached_input / 1_000_000.0) * base
out["cache_read_usd"] = (
(cache_read / 1_000_000.0) * base * OPENAI_CACHE_READ_MULT
)
# Server-tool surcharges. OpenAI doesn't include these on its
# `usage` object directly -- web_search invocations are counted
# from `ResponseFunctionWebSearch` items in the output array,
# and container hours come from the SSE translator's shell-tool
# accounting. Studio surfaces both under a normalised
# `openai_tool_use` key on the usage dict the SSE finaliser
# hands to this calculator.
out["cache_read_usd"] = (cache_read / 1_000_000.0) * base * OPENAI_CACHE_READ_MULT
# OpenAI server-tool surcharges arrive under `openai_tool_use`
# (normalised by the SSE finaliser from output items).
srv = usage.get("openai_tool_use") or {}
if isinstance(srv, dict):
web_searches = int(srv.get("web_search_requests") or 0)
@ -304,17 +258,14 @@ def calculate_cost(
def pricing_snapshot() -> dict[str, Any]:
"""Whole pricing table, for the /api/providers/pricing endpoint.
Returns a flat structure the frontend can hand to its cost
formatter without re-implementing the multipliers.
"""
"""Whole pricing table for the /api/providers/pricing endpoint."""
return {
"anthropic": {
"models": dict(ANTHROPIC_PRICING),
"cache_5m_write_mult": ANTHROPIC_CACHE_5M_WRITE_MULT,
"cache_1h_write_mult": ANTHROPIC_CACHE_1H_WRITE_MULT,
"cache_read_mult": ANTHROPIC_CACHE_READ_MULT,
"fast_mode_mult": ANTHROPIC_FAST_MODE_MULT,
"web_search_usd_per_1k": ANTHROPIC_WEB_SEARCH_USD_PER_1K,
"code_execution_usd_per_hour": ANTHROPIC_CODE_EXEC_USD_PER_HOUR,
},

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

@ -5,7 +5,7 @@
Static registry of supported external LLM providers.
All providers expose OpenAI-compatible /v1/chat/completions endpoints
with Bearer token authentication and SSE streaming support.
with Bearer token auth and SSE streaming.
"""
import re
@ -26,11 +26,8 @@ PROVIDER_REGISTRY: dict[str, dict[str, Any]] = {
"supports_tool_calling": True,
"auth_header": "Authorization",
"auth_prefix": "Bearer ",
# Keep the model picker scoped to the current generation. The remote
# /v1/models listing returns dozens of historical snapshots, fine-tunes
# and non-chat models (embeddings, TTS, image, moderation) that we
# never want to surface in the chat UI. Filtering here so backend
# is the single source of truth.
# Scope the picker to the current generation. /v1/models returns many
# historical snapshots, fine-tunes, and non-chat models we don't want.
"model_id_allowlist": re.compile(r"^(gpt-5\.[345]|gpt-4\.5|o3)(?:[-.]|$)"),
# Hide dated snapshots and the retired plain gpt-5.3 id.
"model_id_denylist": re.compile(r"^(gpt-5\.3)$|-\d{4}-\d{2}-\d{2}$"),
@ -46,11 +43,7 @@ PROVIDER_REGISTRY: dict[str, dict[str, Any]] = {
"claude-sonnet-4-5",
"claude-haiku-4-5",
],
# Anthropic /v1/models returns dated snapshot ids alongside the
# canonical names (e.g. claude-3-5-sonnet-20241022). Hide the
# YYYYMMDD-suffixed variants from the picker — same intent as the
# OpenAI denylist, just a different date format (no dashes between
# year/month/day).
# Hide YYYYMMDD-suffixed snapshot ids (e.g. claude-3-5-sonnet-20241022).
"model_id_denylist": re.compile(r"-\d{8}$"),
"supports_streaming": True,
"supports_vision": True,
@ -65,28 +58,61 @@ PROVIDER_REGISTRY: dict[str, dict[str, Any]] = {
},
"gemini": {
"display_name": "Google Gemini",
"base_url": "https://generativelanguage.googleapis.com/v1beta/openai",
# Curated lineup — Google's /v1beta/openai/models returns dozens
# of historical / experimental / embedding ids. Cap to the current
# 3.x family plus the rolling `*-latest` aliases.
# Native Gemini REST endpoint -- does NOT speak OpenAI Chat Completions;
# translated in `_stream_gemini` (external_provider.py).
# https://ai.google.dev/gemini-api/docs
"base_url": "https://generativelanguage.googleapis.com/v1beta",
# Curated lineup (ListModels returns many historical/experimental ids).
# Excluded on purpose:
# - `gemini-2.0-flash*` (retired 2026-06-01; 404 on use)
# - `gemini-3-pro-preview` (shut down 2026-03-09; auto-redirects to
# `gemini-3.1-pro-preview`, so we surface 3.1 directly).
"default_models": [
"gemini-3.1-pro-preview",
"gemini-3.5-flash",
"gemini-3.1-flash-lite",
"gemini-3-flash-preview",
"gemini-pro-latest",
"gemini-flash-latest",
"gemini-flash-lite-latest",
"gemini-2.5-pro",
"gemini-2.5-flash",
"gemini-2.5-flash-lite",
"gemini-3-pro-image-preview",
"gemini-3.1-flash-image-preview",
"gemini-2.5-flash-image",
],
"supports_streaming": True,
"supports_vision": True,
"supports_tool_calling": True,
"auth_header": "Authorization",
"auth_prefix": "Bearer ",
"notes": "OpenAI-compatible endpoint. API key from https://aistudio.google.com/apikey.",
# Native API takes the bare key on `x-goog-api-key`.
"auth_header": "x-goog-api-key",
"auth_prefix": "",
"openai_compatible": False,
"notes": (
"Native Gemini API. Translation lives in _stream_gemini. "
"API key from https://aistudio.google.com/apikey. "
"See https://ai.google.dev/gemini-api/docs for endpoint shapes."
),
# gemini-3-pro-preview was shut down 2026-03-09 and auto-aliased to
# gemini-3.1-pro-preview; drop it so users see one canonical card.
"model_id_deny_exact": ("gemini-3-pro-preview",),
# Chat-capable 3.5 / 3.1 / 3 / 2.5 families plus rolling *-latest
# aliases. Image-tier ids flow through the Nano Banana
# `responseModalities` path in `_stream_gemini`. Retired 2.0 ids
# excluded (they 404 on use).
"model_id_allowlist": re.compile(
r"^(gemini-3\.1-flash-lite|gemini-3-flash-preview|"
r"gemini-3\.1-pro-preview|gemini-pro-latest|"
r"gemini-flash-latest|gemini-flash-lite-latest)$"
r"^("
r"gemini-3\.5-(?:flash|pro)(?:-preview)?|"
r"gemini-3\.1-(?:flash|pro|flash-lite)(?:-preview)?(?:-customtools)?|"
r"gemini-3\.1-flash-image-preview|"
r"gemini-3-(?:flash|pro)(?:-preview)?|"
r"gemini-3-pro-image-preview|"
r"nano-banana-pro-preview|"
r"gemini-2\.5-pro|gemini-2\.5-flash|gemini-2\.5-flash-lite|"
r"gemini-2\.5-flash-image|"
r"gemini-pro-latest|gemini-flash-latest|gemini-flash-lite-latest"
r")$"
),
},
"deepseek": {
@ -135,14 +161,11 @@ PROVIDER_REGISTRY: dict[str, dict[str, Any]] = {
"kimi": {
"display_name": "Kimi",
"base_url": "https://api.moonshot.ai/v1",
# Current Kimi model lineup per the official docs:
# https://platform.kimi.ai/docs/models
# Listing/overview endpoints used to enumerate them:
# https://platform.kimi.ai/docs/api/list-models
# https://platform.kimi.ai/docs/api/overview
# kimi-k2.6 and kimi-k2.5 are the two SoTA multimodal models we
# surface in the picker; everything else (moonshot-v1-*, dated
# k2 previews) is filtered out by model_id_allowlist below.
# Surface only the two SoTA multimodal models (kimi-k2.6/k2.5);
# moonshot-v1-* and dated k2 previews are filtered by the allowlist.
# Docs: https://platform.kimi.ai/docs/models
# Listing/overview: https://platform.kimi.ai/docs/api/list-models
# https://platform.kimi.ai/docs/api/overview
"default_models": [
"kimi-k2.6",
"kimi-k2.5",
@ -154,10 +177,8 @@ PROVIDER_REGISTRY: dict[str, dict[str, Any]] = {
"auth_prefix": "Bearer ",
"notes": "Moonshot API key. China: use base URL https://api.moonshot.cn/v1",
"model_id_allowlist": re.compile(r"^kimi-k2\.[56]$"),
# Both k2.6 and k2.5 are reasoning-class. The API rejects custom
# sampling: "invalid temperature: only 1 is allowed for this model"
# (and the same shape for top_p). Strip both fields from the
# outbound body so the server falls back to its required defaults.
# Reasoning-class: the API rejects custom temperature/top_p ("only 1
# is allowed"). Strip both so the server uses its required defaults.
"body_omit": ("temperature", "top_p"),
},
"qwen": {
@ -179,9 +200,8 @@ PROVIDER_REGISTRY: dict[str, dict[str, Any]] = {
"huggingface": {
"display_name": "Hugging Face",
"base_url": "https://router.huggingface.co/v1",
# Seed the picker with a few popular ids so something is selectable
# before the live /v1/models call resolves. The remote listing is
# the source of truth — see model_list_mode below.
# Seed the picker before the live /v1/models call resolves; the remote
# listing (see model_list_mode) is the source of truth.
"default_models": [
"openai/gpt-oss-120b",
"deepseek-ai/DeepSeek-V3",
@ -199,29 +219,22 @@ PROVIDER_REGISTRY: dict[str, dict[str, Any]] = {
"returns the cross-provider chat catalog. See "
"https://huggingface.co/docs/inference-providers/index."
),
# /v1/models works on the HF router and returns the full chat-model
# catalog (state.org/model[:policy] ids). Switch to remote so users
# see live availability — the picker has a search box, and
# loadModels() merges defaults so default_models entries remain
# visible if the remote call fails.
# Remote so users see live availability; loadModels() merges defaults
# so they stay visible if the remote call fails.
"model_list_mode": "remote",
# Scope the catalog to first-party org repos we trust as primary
# sources. The HF /v1/models response is otherwise hundreds of
# ids long (community fine-tunes, mirrors, fp8 variants, etc.).
# Scope to trusted first-party org repos (the response is otherwise
# hundreds of community fine-tunes, mirrors, fp8 variants).
"model_id_allowlist": re.compile(
r"^(openai|deepseek-ai|google|meta-llama|Qwen|moonshotai|"
r"mistralai|zai-org)/"
r"^(openai|deepseek-ai|google|meta-llama|Qwen|moonshotai|mistralai|zai-org)/"
),
# Cap the post-filter list. /v1/models has no server-side limit
# or popularity sort, so this is just "first N matches" — pair it
# with the default_models seed so the most useful flagship ids
# are always among the top regardless of the API's order.
# Cap the post-filter list to first N matches (no server-side sort);
# default_models keeps flagship ids near the top.
"model_id_limit": 15,
},
"vllm": {
"display_name": "vLLM",
# User-supplied via provider_base_url; the route layer already falls
# back to the payload's base_url when the registry entry has none.
# User-supplied via provider_base_url; the route falls back to the
# payload's base_url when the registry entry has none.
"base_url": "",
"default_models": [],
"supports_streaming": True,
@ -229,15 +242,28 @@ PROVIDER_REGISTRY: dict[str, dict[str, Any]] = {
"supports_tool_calling": True,
"auth_header": "Authorization",
"auth_prefix": "Bearer ",
# Force /v1/chat/completions in stream_chat_completion — vLLM's
# /v1/responses rebuilds messages and runs them through the loaded
# model's chat template, which 400s on strict-alternation templates
# (Gemma 3 raises "Conversation roles must alternate user/assistant
# /user/assistant/..."). The chat-completions path takes messages
# verbatim and avoids that template gauntlet.
# Force /v1/chat/completions -- vLLM's /v1/responses rebuilds messages
# through the chat template, 400ing on strict-alternation templates
# (Gemma 3). The chat-completions path takes messages verbatim.
"notes": "Self-hosted vLLM server. Always routed to /v1/chat/completions.",
# Surfaced through the frontend's CUSTOM_PROVIDER_PRESETS, not the
# /api/providers/registry dropdown — see list_available_providers.
# Surfaced via the frontend's CUSTOM_PROVIDER_PRESETS, not the dropdown.
"hidden": True,
},
"custom": {
"display_name": "Custom",
# User-supplied via provider_base_url.
"base_url": "",
"default_models": [],
"supports_streaming": True,
"supports_vision": True,
"supports_tool_calling": True,
"auth_header": "Authorization",
"auth_prefix": "Bearer ",
"notes": (
"User-supplied OpenAI-compatible server. Routed to "
"/v1/chat/completions; /models is optional."
),
# Surfaced by the frontend's generic Custom option, not the dropdown.
"hidden": True,
},
"ollama": {
@ -273,7 +299,7 @@ PROVIDER_REGISTRY: dict[str, dict[str, Any]] = {
"openrouter": {
"display_name": "OpenRouter",
"base_url": "https://openrouter.ai/api/v1",
# Curated list for Studio's picker (explicitly locked, not live /models).
# Curated picker list (locked, not live /models).
"default_models": [
"openrouter/free",
"openai/gpt-4o",
@ -323,10 +349,8 @@ def get_base_url(provider_type: str) -> str | None:
def list_available_providers() -> list[dict[str, Any]]:
"""Return all registered providers (for the /registry endpoint).
Hidden entries (``"hidden": True``) are filtered out — they exist in the
registry only for backend lookups (e.g. ``supports_vision`` for vLLM) and
are surfaced in the frontend via ``CUSTOM_PROVIDER_PRESETS`` instead of
the cloud-provider dropdown.
Hidden entries are filtered out: they exist only for backend lookups and
are surfaced via ``CUSTOM_PROVIDER_PRESETS`` instead of the dropdown.
"""
result = []
for provider_type, info in PROVIDER_REGISTRY.items():

22
studio/backend/core/inference/runtime_context.py Normal file
View file

@ -0,0 +1,22 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Runtime context length helpers shared by inference backends."""
from __future__ import annotations
from typing import Any, Optional
def runtime_context_length(model: Any, fallback: Optional[int] = None) -> Optional[int]:
"""Return the effective context length Unsloth attached to a loaded model."""
for value in (getattr(model, "max_seq_length", None), fallback):
if isinstance(value, bool):
continue
try:
value_int = int(value)
except (TypeError, ValueError):
continue
if value_int > 0:
return value_int
return None

567
studio/backend/core/inference/safetensors_agentic.py
View file

@ -4,94 +4,141 @@
"""
Safetensors/transformers agentic tool loop.
Wraps a single-turn cumulative-text generator (the existing
``InferenceOrchestrator.generate_chat_response`` pipeline that streams
from a worker subprocess) with the tool-calling, thinking-block,
status, and metadata event protocol used by the GGUF path. Keeps the
front-end SSE shape identical across backends so the chat UI does not
care which engine actually ran the model.
Wraps a single-turn cumulative-text generator with the same tool-calling,
thinking-block, status, and metadata event protocol the GGUF path uses, so
the front-end SSE shape is identical across backends.
The GGUF path lives in ``llama_cpp.py`` and talks to llama-server's
structured ``delta.tool_calls`` directly. Native transformers has no
such structured channel, so this loop parses tool calls from the
cumulative text and dispatches them via ``core.inference.tools``.
Unlike the GGUF path (``llama_cpp.py``), which uses llama-server's structured
``delta.tool_calls``, native transformers has no such channel, so this loop
parses tool calls from the cumulative text and dispatches via
``core.inference.tools``.
"""
import json
import re
import threading
from typing import Callable, Generator, Optional
from urllib.parse import urlparse
from loggers import get_logger
from core.inference.tool_call_parser import (
_TOOL_ALL_PATS,
BUDGET_EXHAUSTED_NUDGE,
DUPLICATE_CALL_NUDGE,
TOOL_ERROR_NUDGE,
TOOL_ERROR_PREFIXES,
RAG_MAX_SEARCHES_PER_TURN,
RAG_SEARCH_CAP_NUDGE,
TOOL_XML_SIGNALS,
has_tool_signal,
parse_tool_calls_from_text,
strip_tool_markup,
)
from core.inference.tool_loop_controller import (
ToolLoopController,
coerce_tool_arguments,
status_for_tool,
tool_event_provenance,
)
from state.tool_approvals import (
TOOL_REJECTED_MESSAGE,
abort_tool_decision,
begin_tool_decision,
new_approval_id,
wait_tool_decision,
)
logger = get_logger(__name__)
# Buffer cap while waiting to disambiguate a possible tool-call prefix.
# Buffer cap while disambiguating a possible tool-call prefix.
_MAX_BUFFER_CHARS = 32
def strip_tool_markup_streaming(
text: str,
*,
auto_heal_tool_calls: bool = True,
tool_protocol_active: bool = False,
) -> str:
"""Strip open-ended tool XML from display text without trimming whitespace."""
if not (auto_heal_tool_calls or tool_protocol_active):
return text
for pat in _TOOL_ALL_PATS:
text = pat.sub("", text)
return text
def _strip_tool_markup_final(
text: str,
*,
auto_heal_tool_calls: bool,
tool_protocol_active: bool = False,
) -> str:
if not (auto_heal_tool_calls or tool_protocol_active):
return text
return strip_tool_markup(text, final = True)
def _status_for_tool(tool_name: str, arguments: dict) -> str:
"""Return a human-readable status line matching the GGUF path."""
if tool_name == "web_search":
url = (arguments.get("url") or "").strip()
if url:
parsed = urlparse(url)
if parsed.scheme in ("http", "https") and parsed.hostname:
host = parsed.hostname
if host.startswith("www."):
host = host[4:]
return f"Reading: {host}"
return "Reading page..."
query = arguments.get("query", "")
return f"Searching: {query}"
if tool_name == "python":
preview = (arguments.get("code") or "").strip().split("\n")[0][:60]
return f"Running Python: {preview}" if preview else "Running Python..."
if tool_name == "terminal":
preview = (arguments.get("command") or "")[:60]
return f"Running: {preview}" if preview else "Running command..."
return f"Calling: {tool_name}"
return status_for_tool(tool_name, arguments)
_CANONICAL_HEAL_ARG = {"python": "code", "terminal": "command"}
_FUNCTION_SIGNAL_RE = re.compile(r"<function=([\w-]+)>")
_TOOL_CALL_NAME_RE = re.compile(r'"name"\s*:\s*"([\w-]+)"')
def _coerce_arguments(raw_args, *, heal: bool, tool_name: str = "") -> dict:
"""Normalise tool ``arguments`` to a dict.
def _detect_render_html_tool_start(content: str) -> bool:
"""Return True when the first drained tool call is clearly render_html."""
function_match = _FUNCTION_SIGNAL_RE.search(content)
tool_call_index = content.find("<tool_call>")
if not function_match and tool_call_index < 0:
return False
Some templates emit a JSON string, others a bare query string. With
``heal=True`` we accept a bare string as ``{<canonical_key>: ...}``
so a Hermes-style call without proper JSON still runs the tool. The
canonical key is picked per tool: ``code`` for python, ``command``
for terminal, ``query`` for everything else (e.g. web_search).
"""
if isinstance(raw_args, dict):
return raw_args
if isinstance(raw_args, str):
try:
parsed = json.loads(raw_args)
if isinstance(parsed, dict):
return parsed
except (json.JSONDecodeError, ValueError):
pass
if heal:
key = _CANONICAL_HEAL_ARG.get(tool_name, "query")
return {key: raw_args}
return {"raw": raw_args}
return {}
if function_match and (tool_call_index < 0 or function_match.start() < tool_call_index):
return function_match.group(1) == "render_html"
if tool_call_index >= 0:
name_match = _TOOL_CALL_NAME_RE.search(content[tool_call_index:])
return bool(name_match and name_match.group(1) == "render_html")
return False
def _coerce_arguments_with_provenance(
raw_args,
*,
heal: bool,
tool_name: str = "",
):
"""Normalise tool ``arguments`` and report whether healing was applied."""
coerced = coerce_tool_arguments(raw_args, heal = heal, tool_name = tool_name)
return coerced.arguments, coerced.healed
def _coerce_arguments(
raw_args,
*,
heal: bool,
tool_name: str = "",
) -> dict:
arguments, _ = _coerce_arguments_with_provenance(
raw_args,
heal = heal,
tool_name = tool_name,
)
return arguments
def _tool_event_provenance(**flags: object) -> dict[str, object]:
return tool_event_provenance(**flags)
def _call_single_turn(single_turn, conversation: list, active_tools: list[dict]):
"""Call a single-turn generator with active tool schemas when supported."""
try:
return single_turn(conversation, active_tools = active_tools)
except TypeError as exc:
if "active_tools" not in str(exc):
raise
return single_turn(conversation)
def run_safetensors_tool_loop(
@ -105,44 +152,60 @@ def run_safetensors_tool_loop(
max_tool_iterations: int = 25,
tool_call_timeout: int = 300,
session_id: Optional[str] = None,
rag_scope: Optional[dict] = None,
confirm_tool_calls: bool = False,
bypass_permissions: bool = False,
) -> Generator[dict, None, None]:
"""Drive an agentic tool loop on top of a cumulative-text generator.
``single_turn(messages)`` must yield cumulative assistant text
(each yield is a snapshot including all previously emitted tokens).
The loop:
``single_turn(messages)`` must yield cumulative assistant text (each
yield is a snapshot of all tokens so far). The loop:
* Buffers the leading characters of every turn so it can decide
whether the model is about to emit a tool call. Plain content
starts streaming as soon as the buffer rules it out.
* On detecting ``<tool_call>`` or ``<function=`` in the cumulative
text, drains the rest of the turn silently and parses tool calls
out of the full content.
* Buffers each turn's leading chars to decide whether a tool call is
coming. Plain content streams once the buffer rules it out.
* On ``<tool_call>`` or ``<function=`` in the cumulative text, drains
the rest of the turn silently and parses tool calls from the content.
* Executes each tool via ``execute_tool``, appends the assistant
tool-call message and the tool result to the conversation, and
re-enters ``single_turn`` for the next iteration.
* After ``max_tool_iterations`` turns without a final answer, asks
the model once more to produce a final answer with no tools.
tool-call message and tool result, and re-enters ``single_turn``.
* After ``max_tool_iterations`` turns without a final answer, asks once
more for a final answer with no tools.
Yields event dicts matching the GGUF path:
* ``{"type": "status", "text": ...}`` -- empty string clears the badge.
* ``{"type": "content", "text": ...}`` -- cumulative cleaned text for
the current assistant turn (the consumer should diff against its
own ``prev_text`` cursor).
the current turn (consumer diffs against its own ``prev_text`` cursor).
* ``{"type": "tool_start", "tool_name", "tool_call_id", "arguments"}``
* ``{"type": "tool_end", "tool_name", "tool_call_id", "result"}``
"""
conversation = list(messages)
tool_call_history: list[tuple[str, bool]] = []
# Forced first-pass RAG (mirrors the GGUF loop) so doc Qs don't lose to web_search.
from core.inference.tools import build_rag_autoinject
_auto = None if confirm_tool_calls else build_rag_autoinject(conversation, rag_scope)
if _auto:
for _ev in _auto["events"]:
yield _ev
conversation.extend(_auto["messages"])
unrestricted_tools = not tools
tool_controller = ToolLoopController(
tools = None if unrestricted_tools else tools,
auto_heal_tool_calls = auto_heal_tool_calls,
)
# RAG: cap knowledge-base searches per assistant turn (controller-agnostic).
kb_search_count = 0
final_attempt_done = False
allowed_tool_names = {
(tool.get("function") or {}).get("name")
for tool in (tools or [])
if (tool.get("function") or {}).get("name")
}
next_call_id = 0
def _tool_succeeded(tool_name: str) -> bool:
key_prefix = f"{tool_name}:"
return any(
record.executed and not record.is_error and record.key.startswith(key_prefix)
for record in tool_controller.history
)
if max_tool_iterations <= 0:
# 0 = disabled (same contract as the GGUF loop).
yield {"type": "status", "text": ""}
@ -156,13 +219,26 @@ def run_safetensors_tool_loop(
if cancel_event is not None and cancel_event.is_set():
return
if final_attempt_done:
active_tools: list[dict] = []
else:
active_tools = tool_controller.active_tools()
if not active_tools and not unrestricted_tools:
final_attempt_done = True
active_tools = []
tool_protocol_active = not final_attempt_done and (unrestricted_tools or bool(active_tools))
tool_xml_signals = TOOL_XML_SIGNALS if tool_protocol_active else ()
detect_state = _state_buffering
content_buffer = ""
content_accum = ""
cumulative_display = ""
last_emitted = ""
provisional_render_html_started = False
provisional_render_html_id = f"call_{next_call_id}"
gen = single_turn(conversation)
gen = _call_single_turn(single_turn, conversation, active_tools)
prev_cumulative = ""
for cumulative in gen:
@ -170,7 +246,7 @@ def run_safetensors_tool_loop(
return
if not isinstance(cumulative, str):
continue # defensive: pipeline only yields strings
continue # defensive: pipeline yields only strings
delta = cumulative[len(prev_cumulative) :]
prev_cumulative = cumulative
@ -179,26 +255,68 @@ def run_safetensors_tool_loop(
content_accum += delta
if detect_state == _state_draining:
if (
not _tool_succeeded("render_html")
and any(
((tool.get("function") or {}).get("name") == "render_html")
for tool in active_tools
)
and not provisional_render_html_started
and _detect_render_html_tool_start(content_accum)
):
provisional_render_html_started = True
yield {
"type": "tool_start",
"tool_name": "render_html",
"tool_call_id": provisional_render_html_id,
"arguments": {},
"provenance": _tool_event_provenance(provisional = True),
}
continue
if detect_state == _state_streaming:
candidate = cumulative_display + delta
signal_pos = -1
for sig in TOOL_XML_SIGNALS:
for sig in tool_xml_signals:
p = candidate.find(sig)
if p >= 0 and (signal_pos < 0 or p < signal_pos):
signal_pos = p
if signal_pos >= 0:
before_tool = candidate[:signal_pos]
cleaned_before = strip_tool_markup(before_tool)
cleaned_before = strip_tool_markup_streaming(
before_tool,
auto_heal_tool_calls = auto_heal_tool_calls,
tool_protocol_active = tool_protocol_active,
)
if len(cleaned_before) > len(last_emitted):
last_emitted = cleaned_before
yield {"type": "content", "text": cleaned_before}
cumulative_display = candidate
detect_state = _state_draining
if (
not _tool_succeeded("render_html")
and any(
((tool.get("function") or {}).get("name") == "render_html")
for tool in active_tools
)
and not provisional_render_html_started
and _detect_render_html_tool_start(content_accum)
):
provisional_render_html_started = True
yield {
"type": "tool_start",
"tool_name": "render_html",
"tool_call_id": provisional_render_html_id,
"arguments": {},
"provenance": _tool_event_provenance(provisional = True),
}
continue
cumulative_display = candidate
cleaned = strip_tool_markup(cumulative_display)
cleaned = strip_tool_markup_streaming(
cumulative_display,
auto_heal_tool_calls = auto_heal_tool_calls,
tool_protocol_active = tool_protocol_active,
)
if len(cleaned) > len(last_emitted):
last_emitted = cleaned
yield {"type": "content", "text": cleaned}
@ -212,7 +330,7 @@ def run_safetensors_tool_loop(
is_match = False
is_prefix = False
for sig in TOOL_XML_SIGNALS:
for sig in tool_xml_signals:
if stripped.startswith(sig):
is_match = True
break
@ -221,13 +339,45 @@ def run_safetensors_tool_loop(
break
if is_match:
# Tool signal -- flush any visible prefix before DRAINING
# so the route sends it before tool_start.
cumulative_display += content_buffer
cleaned = strip_tool_markup_streaming(
cumulative_display,
auto_heal_tool_calls = auto_heal_tool_calls,
tool_protocol_active = tool_protocol_active,
)
if len(cleaned) > len(last_emitted):
last_emitted = cleaned
yield {"type": "content", "text": cleaned}
detect_state = _state_draining
if (
not _tool_succeeded("render_html")
and any(
((tool.get("function") or {}).get("name") == "render_html")
for tool in active_tools
)
and not provisional_render_html_started
and _detect_render_html_tool_start(content_accum)
):
provisional_render_html_started = True
yield {
"type": "tool_start",
"tool_name": "render_html",
"tool_call_id": provisional_render_html_id,
"arguments": {},
"provenance": _tool_event_provenance(provisional = True),
}
elif is_prefix and len(stripped) < _MAX_BUFFER_CHARS:
continue
else:
detect_state = _state_streaming
cumulative_display += content_buffer
cleaned = strip_tool_markup(cumulative_display)
cleaned = strip_tool_markup_streaming(
cumulative_display,
auto_heal_tool_calls = auto_heal_tool_calls,
tool_protocol_active = tool_protocol_active,
)
if len(cleaned) > len(last_emitted):
last_emitted = cleaned
yield {"type": "content", "text": cleaned}
@ -237,16 +387,24 @@ def run_safetensors_tool_loop(
return
if detect_state == _state_buffering:
# Buffer never resolved -- tool XML or plain content.
# Buffer never resolved -- tool XML or plain content?
stripped = content_buffer.lstrip()
if stripped and has_tool_signal(stripped):
if (
stripped
and tool_protocol_active
and any(sig in stripped for sig in tool_xml_signals)
):
detect_state = _state_draining
else:
if content_buffer:
cumulative_display += content_buffer
yield {
"type": "content",
"text": strip_tool_markup(cumulative_display, final = True),
"text": _strip_tool_markup_final(
cumulative_display,
auto_heal_tool_calls = auto_heal_tool_calls,
tool_protocol_active = False,
),
}
yield {"type": "status", "text": ""}
return
@ -254,19 +412,30 @@ def run_safetensors_tool_loop(
if detect_state == _state_streaming:
# No tool detected mid-stream -- check for late tool XML.
safety_tc = None
if has_tool_signal(content_accum):
saw_tool_signal = tool_protocol_active and any(
sig in content_accum for sig in tool_xml_signals
)
if saw_tool_signal:
safety_tc = parse_tool_calls_from_text(
content_accum,
id_offset = next_call_id,
allow_incomplete = auto_heal_tool_calls,
)
if not safety_tc:
# Final answer: streaming already emitted content.
# Skip a final=True re-strip so literal "<tool_call>"
# in prose survives when no real tool call parsed.
# Final answer: if a literal tool marker in prose was stripped
# during streaming but did not parse as a real call, restore the
# raw cumulative text for core callers. Route-level cleanup can
# still apply the Auto-Heal display policy.
if saw_tool_signal and content_accum:
yield {"type": "content", "text": content_accum}
yield {"type": "status", "text": ""}
return
tool_calls = safety_tc
content_text = strip_tool_markup(content_accum, final = True)
content_text = _strip_tool_markup_final(
content_accum,
auto_heal_tool_calls = auto_heal_tool_calls,
tool_protocol_active = True,
)
logger.info(
"Safetensors safety net: parsed %d tool call(s) from streamed content",
len(tool_calls),
@ -276,15 +445,39 @@ def run_safetensors_tool_loop(
tool_calls = parse_tool_calls_from_text(
content_accum,
id_offset = next_call_id,
allow_incomplete = auto_heal_tool_calls,
)
if not tool_calls and auto_heal_tool_calls:
# Parser found nothing -- surface raw content so any
# literal "<tool_call>" prose is preserved.
if not tool_calls:
# Parser found nothing. Auto-Heal-enabled display cleanup
# strips unparseable tool XML; disabled Auto-Heal preserves
# the raw text so literal/malformed markup stays visible.
if content_accum:
yield {"type": "content", "text": content_accum}
yield {
"type": "content",
"text": _strip_tool_markup_final(
content_accum,
auto_heal_tool_calls = auto_heal_tool_calls,
tool_protocol_active = False,
),
}
if provisional_render_html_started:
yield {
"type": "tool_end",
"tool_name": "render_html",
"tool_call_id": provisional_render_html_id,
"result": "Error: render_html tool call could not be parsed.",
"provenance": _tool_event_provenance(provisional = True),
}
yield {"type": "status", "text": ""}
return
content_text = strip_tool_markup(content_accum, final = True)
content_text = _strip_tool_markup_final(
content_accum,
auto_heal_tool_calls = auto_heal_tool_calls,
tool_protocol_active = True,
)
if tool_calls:
next_call_id += len(tool_calls)
if final_attempt_done:
# Final-answer turn re-called a tool -- stop the loop.
@ -294,99 +487,121 @@ def run_safetensors_tool_loop(
return
assistant_msg: dict = {"role": "assistant", "content": content_text}
if tool_calls:
assistant_msg["tool_calls"] = tool_calls
next_call_id += len(tool_calls)
conversation.append(assistant_msg)
assistant_appended = False
for tc in tool_calls or []:
func = tc.get("function", {}) or {}
tool_name = func.get("name", "") or ""
arguments = _coerce_arguments(
func.get("arguments", {}),
heal = auto_heal_tool_calls,
tool_name = tool_name,
provisional_match = (
provisional_render_html_started
and tool_name == "render_html"
and tc.get("id", "") == provisional_render_html_id
)
decision = tool_controller.prepare_call(tc, provisional = provisional_match)
yield {"type": "status", "text": _status_for_tool(tool_name, arguments)}
yield {
"type": "tool_start",
"tool_name": tool_name,
"tool_call_id": tc.get("id", ""),
"arguments": arguments,
}
tc_key = tool_name + str(arguments)
if allowed_tool_names and tool_name not in allowed_tool_names:
result = (
f"Error: tool '{tool_name}' is not enabled for this "
"request. Use one of the enabled tools or provide a "
"final answer."
if not decision.should_execute:
if content_text and not assistant_appended:
conversation.append(assistant_msg)
assistant_appended = True
completion = tool_controller.record_noop(decision)
conversation.append(completion.model_message())
logger.info(
"Suppressed local safetensors tool call as internal no-op: "
f"action={decision.action} tool={decision.tool_name}"
)
break
if not assistant_appended:
assistant_msg["tool_calls"] = [decision.as_assistant_tool_call()]
conversation.append(assistant_msg)
assistant_appended = True
else:
already_ran_ok = any(
k == tc_key and not err for k, err in tool_call_history
)
if already_ran_ok:
result = DUPLICATE_CALL_NUDGE
else:
eff_timeout = (
None if tool_call_timeout >= 9999 else tool_call_timeout
assistant_msg.setdefault("tool_calls", []).append(decision.as_assistant_tool_call())
# Bypass wins over the confirm gate at the loop level too, so a
# direct internal caller passing both flags never prompts.
needs_confirm = bool(confirm_tool_calls) and not bypass_permissions
approval_id = new_approval_id() if needs_confirm else ""
decision_slot = begin_tool_decision(session_id, approval_id) if needs_confirm else None
start_event = decision.tool_start_event()
start_event["approval_id"] = approval_id
start_event["awaiting_confirmation"] = needs_confirm
try:
yield {"type": "status", "text": decision.status_text}
yield start_event
if (
decision_slot is not None
and wait_tool_decision(
decision_slot,
approval_id,
cancel_event = cancel_event,
)
try:
result = execute_tool(
tool_name,
arguments,
cancel_event = cancel_event,
timeout = eff_timeout,
session_id = session_id,
)
except Exception as exc:
logger.exception("Tool %s raised: %s", tool_name, exc)
result = f"Error: tool raised an exception: {exc}"
== "deny"
):
decision_slot = None
yield {
"type": "tool_end",
"tool_name": decision.tool_name,
"tool_call_id": decision.tool_call_id,
"result": TOOL_REJECTED_MESSAGE,
"provenance": decision.provenance,
}
denied_message = {
"role": "tool",
"name": decision.tool_name,
"content": TOOL_REJECTED_MESSAGE,
}
if decision.tool_call_id:
denied_message["tool_call_id"] = decision.tool_call_id
conversation.append(denied_message)
continue
decision_slot = None
finally:
if decision_slot is not None:
abort_tool_decision(decision_slot, approval_id)
yield {
"type": "tool_end",
"tool_name": tool_name,
"tool_call_id": tc.get("id", ""),
"result": result,
}
eff_timeout = None if tool_call_timeout >= 9999 else tool_call_timeout
# RAG: cap paraphrased KB re-searches that slip past the dup guard.
if (
decision.tool_name == "search_knowledge_base"
and kb_search_count >= RAG_MAX_SEARCHES_PER_TURN
):
result = RAG_SEARCH_CAP_NUDGE
else:
try:
result = execute_tool(
decision.tool_name,
decision.arguments,
cancel_event = cancel_event,
timeout = eff_timeout,
session_id = session_id,
rag_scope = rag_scope,
disable_sandbox = bypass_permissions,
)
except Exception as exc:
logger.exception("Tool %s raised: %s", decision.tool_name, exc)
result = f"Error: tool raised an exception: {exc}"
if decision.tool_name == "search_knowledge_base":
kb_search_count += 1
is_error = isinstance(result, str) and result.lstrip().startswith(
TOOL_ERROR_PREFIXES
)
tool_call_history.append((tc_key, is_error))
# Strip frontend image sentinel from the model's view.
# Cut at the first occurrence so leading and consecutive
# sentinels are both removed.
result_for_model = result
if isinstance(result_for_model, str) and "__IMAGES__:" in result_for_model:
result_for_model = result_for_model.split("__IMAGES__:", 1)[0].rstrip()
if is_error:
result_for_model = result_for_model + TOOL_ERROR_NUDGE
tool_msg: dict = {
"role": "tool",
"name": tool_name,
"content": result_for_model,
}
tool_call_id = tc.get("id")
if tool_call_id:
tool_msg["tool_call_id"] = tool_call_id
conversation.append(tool_msg)
completion = tool_controller.record_result(decision, result)
yield completion.tool_end_event()
conversation.append(completion.tool_message())
# Clear the status badge before the next turn.
yield {"type": "status", "text": ""}
if tool_controller.force_final_answer:
final_attempt_done = True
continue
if not unrestricted_tools and not tool_controller.active_tools():
final_attempt_done = True
continue
if iteration + 1 >= max_tool_iterations and not final_attempt_done:
# Budget exhausted; nudge a final plain answer.
final_attempt_done = True
conversation.append(
{
"role": "user",
"content": BUDGET_EXHAUSTED_NUDGE,
}
)
conversation.append({"role": "user", "content": BUDGET_EXHAUSTED_NUDGE})
yield {"type": "status", "text": ""}

70
studio/backend/core/inference/tensor_fallback.py Normal file
View file

@ -0,0 +1,70 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Tensor-parallel -> layer-split auto-fallback for GGUF loads.
Kept in its own module (no FastAPI / httpx deps) so the orchestration can be
unit-tested with a fake loader, without a GPU or a running llama-server.
"""
from __future__ import annotations
import logging
from typing import Awaitable, Callable, Optional
from core.inference.llama_server_args import (
resolve_tensor_parallel,
strip_split_mode_only,
)
logger = logging.getLogger(__name__)
async def load_with_tensor_fallback(
attempt_load: Callable[[bool, Optional[list[str]]], Awaitable[bool]],
*,
requested_tensor: bool,
extra_args: Optional[list[str]],
label: str = "",
cancelled: Optional[Callable[[], bool]] = None,
) -> bool:
"""Run a GGUF load with the tensor-parallel -> layer-split auto-fallback.
``attempt_load(tensor_parallel, extra_args)`` performs one load and returns
True on success; it *raises* on a hard crash (llama-server aborts on some
archs / older builds), which is treated the same as a False return.
Tensor mode can be requested by the toggle or by a ``--split-mode tensor``
in ``extra_args`` (an allowed shadow flag), so the retry is keyed on whether
tensor mode is actually engaged, and it strips ``--split-mode`` from the
extras so the layer retry can't relaunch the same failing tensor load. A
non-tensor load keeps its original contract and propagates exceptions.
``cancelled()`` distinguishes a real tensor-start failure from a user
cancellation: ``attempt_load`` also returns False when the load was
cancelled, so without this the helper would restart a load the user just
cancelled.
"""
tensor_requested = resolve_tensor_parallel(extra_args, requested_tensor)
try:
success = await attempt_load(requested_tensor, extra_args)
except Exception as exc:
if not tensor_requested:
raise
logger.warning("Tensor-parallel load raised for '%s': %s", label, exc)
success = False
if success or not tensor_requested:
return success
# The first attempt returned False because the user cancelled, not because
# tensor mode is unsupported -- do not relaunch the cancelled load.
if cancelled is not None and cancelled():
return success
logger.warning(
"Tensor-parallel load failed for '%s'; retrying with layer split "
"(this model may not support tensor parallelism)",
label,
)
return await attempt_load(False, strip_split_mode_only(extra_args))

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

@ -11,15 +11,16 @@ import json
import re
# _TOOL_CLOSED_PATS: closed pairs only. _TOOL_ALL_PATS: also trailing
# unclosed runs so truncated tails don't leak markup.
# _TOOL_CLOSED_PATS: closed pairs only. _TOOL_ALL_PATS: also trailing unclosed
# runs so truncated tails don't leak markup. The [\w-] name set matches OpenAI's
# so hyphenated MCP tool names (mcp__srv__list-issues) parse like built-ins.
_TOOL_CLOSED_PATS = [
re.compile(r"<tool_call>.*?</tool_call>", re.DOTALL),
re.compile(r"<function=\w+>.*?</function>", re.DOTALL),
re.compile(r"<function=[\w-]+>.*?</function>", re.DOTALL),
]
_TOOL_ALL_PATS = _TOOL_CLOSED_PATS + [
re.compile(r"<tool_call>.*$", re.DOTALL),
re.compile(r"<function=\w+>.*$", re.DOTALL),
re.compile(r"<function=[\w-]+>.*$", re.DOTALL),
]
@ -46,6 +47,12 @@ DUPLICATE_CALL_NUDGE = (
"provide your final answer now."
)
RENDER_HTML_REPEAT_NUDGE = (
"Error: render_html was already called for this response. Do not call "
"render_html again in this response unless the user asks for changes. "
"Provide the final answer now."
)
TOOL_ERROR_NUDGE = (
"\n\nThe tool call encountered an issue. Please try a different "
"approach or rephrase your request."
@ -57,14 +64,38 @@ BUDGET_EXHAUSTED_NUDGE = (
"any more tools."
)
# The exact-args dup guard misses paraphrased re-searches, so also cap executed
# KB searches per turn, then nudge.
RAG_MAX_SEARCHES_PER_TURN = 3
RAG_SEARCH_CAP_NUDGE = (
"You have already searched the knowledge base several times this turn. "
"Do not search again. Answer the question using the passages already "
"retrieved above; if they do not contain the answer, say so plainly."
)
# Pre-compiled patterns reused by ``parse_tool_calls_from_text``.
_TC_JSON_START_RE = re.compile(r"<tool_call>\s*\{")
_TC_FUNC_START_RE = re.compile(r"<function=(\w+)>\s*")
_TC_FUNC_START_RE = re.compile(r"<function=([\w-]+)>\s*")
_TC_END_TAG_RE = re.compile(r"</tool_call>")
_TC_FUNC_CLOSE_RE = re.compile(r"\s*</function>\s*$")
_TC_PARAM_START_RE = re.compile(r"<parameter=(\w+)>\s*")
# [\w-] so hyphenated MCP param names (issue-number) aren't dropped.
_TC_PARAM_START_RE = re.compile(r"<parameter=([\w-]+)>\s*")
_TC_PARAM_CLOSE_RE = re.compile(r"\s*</parameter>\s*$")
_PARAM_CLOSE_TAG = "</parameter>"
_FUNC_CLOSE_TAG = "</function>"
def _inside_open_parameter(content: str, pos: int) -> bool:
"""Return True when ``pos`` falls inside an unclosed parameter value."""
last_param_start = -1
for match in _TC_PARAM_START_RE.finditer(content, 0, pos):
last_param_start = match.start()
if last_param_start < 0:
return False
last_param_close = content.rfind(_PARAM_CLOSE_TAG, 0, pos)
last_func_close = content.rfind(_FUNC_CLOSE_TAG, 0, pos)
return last_param_start > max(last_param_close, last_func_close)
def strip_tool_markup(text: str, *, final: bool = False) -> str:
@ -80,7 +111,12 @@ def strip_tool_markup(text: str, *, final: bool = False) -> str:
return text.strip() if final else text
def parse_tool_calls_from_text(content: str, *, id_offset: int = 0) -> list[dict]:
def parse_tool_calls_from_text(
content: str,
*,
id_offset: int = 0,
allow_incomplete: bool = True,
) -> list[dict]:
"""Parse OpenAI-format ``tool_calls`` from model text.
Returns a list of ``{"id", "type", "function": {"name", "arguments"}}``
@ -94,15 +130,17 @@ def parse_tool_calls_from_text(content: str, *, id_offset: int = 0) -> list[dict
- XML-style function blocks:
``<function=name><parameter=k>v</parameter></function>``
Closing tags (``</tool_call>``, ``</function>``, ``</parameter>``)
are all optional since models frequently omit them.
``allow_incomplete=True`` keeps the historical healing behavior for
missing closing tags. ``allow_incomplete=False`` accepts only
well-formed wrappers so disabled Auto-Heal can still parse valid
local tool protocol without repairing truncated output.
"""
tool_calls: list[dict] = []
# Pattern 1: <tool_call>{json}. Balanced-brace scan that skips
# braces inside JSON strings.
# Pattern 1: <tool_call>{json}. Balanced-brace scan, skipping braces in
# JSON strings.
for m in _TC_JSON_START_RE.finditer(content):
brace_start = m.end() - 1 # position of the opening {
brace_start = m.end() - 1 # opening {
depth, i = 0, brace_start
in_string = False
while i < len(content):
@ -122,39 +160,41 @@ def parse_tool_calls_from_text(content: str, *, id_offset: int = 0) -> list[dict
if depth == 0:
break
i += 1
if depth == 0:
json_str = content[brace_start : i + 1]
try:
obj = json.loads(json_str)
tc = {
"id": f"call_{id_offset + len(tool_calls)}",
"type": "function",
"function": {
"name": obj.get("name", ""),
"arguments": obj.get("arguments", {}),
},
}
if isinstance(tc["function"]["arguments"], dict):
tc["function"]["arguments"] = json.dumps(
tc["function"]["arguments"]
)
tool_calls.append(tc)
except (json.JSONDecodeError, ValueError):
pass
if depth != 0:
continue
if not allow_incomplete:
tail_after_json = content[i + 1 :].lstrip()
if _TC_END_TAG_RE.match(tail_after_json) is None:
continue
json_str = content[brace_start : i + 1]
try:
obj = json.loads(json_str)
tc = {
"id": f"call_{id_offset + len(tool_calls)}",
"type": "function",
"function": {
"name": obj.get("name", ""),
"arguments": obj.get("arguments", {}),
},
}
if isinstance(tc["function"]["arguments"], dict):
tc["function"]["arguments"] = json.dumps(tc["function"]["arguments"])
tool_calls.append(tc)
except (json.JSONDecodeError, ValueError):
pass
# Pattern 2: <function=name><parameter=k>v... -- closing tags
# optional; don't use </function> as body boundary because code
# values can contain that literal.
# Pattern 2: <function=name><parameter=k>v... -- closing tags optional;
# </function> isn't a body boundary since code values can contain it.
if not tool_calls:
func_starts = list(_TC_FUNC_START_RE.finditer(content))
func_starts = [
fm
for fm in _TC_FUNC_START_RE.finditer(content)
if not _inside_open_parameter(content, fm.start())
]
for idx, fm in enumerate(func_starts):
func_name = fm.group(1)
body_start = fm.end()
next_func = (
func_starts[idx + 1].start()
if idx + 1 < len(func_starts)
else len(content)
)
next_func = func_starts[idx + 1].start() if idx + 1 < len(func_starts) else len(content)
end_tag = _TC_END_TAG_RE.search(content[body_start:])
if end_tag:
body_end = body_start + end_tag.start()
@ -162,18 +202,37 @@ def parse_tool_calls_from_text(content: str, *, id_offset: int = 0) -> list[dict
body_end = len(content)
body_end = min(body_end, next_func)
body = content[body_start:body_end]
body = _TC_FUNC_CLOSE_RE.sub("", body)
if not allow_incomplete:
# Bound the body at the closing </function> tag rather than
# the end of the response, so a complete call followed by
# trailing prose is still accepted (matching the JSON-style
# <tool_call> path, which already tolerates trailing text).
# rfind picks the last </function>, so a literal </function>
# inside a code parameter value stays in the body.
close_idx = body.rfind(_FUNC_CLOSE_TAG)
if close_idx < 0:
continue
body = body[:close_idx]
else:
body = _TC_FUNC_CLOSE_RE.sub("", body)
arguments: dict = {}
param_starts = list(_TC_PARAM_START_RE.finditer(body))
if len(param_starts) == 1:
# Single param: take everything to body end so
# embedded </parameter> in code strings is preserved.
# Single param: take everything to body end so an embedded
# </parameter> in code strings is preserved.
pm = param_starts[0]
val = body[pm.end() :]
val = _TC_PARAM_CLOSE_RE.sub("", val)
if not allow_incomplete:
stripped_val = val.rstrip()
if not stripped_val.endswith(_PARAM_CLOSE_TAG):
continue
val = stripped_val[: -len(_PARAM_CLOSE_TAG)]
else:
val = _TC_PARAM_CLOSE_RE.sub("", val)
arguments[pm.group(1)] = val.strip()
else:
valid_params = True
for pidx, pm in enumerate(param_starts):
param_name = pm.group(1)
val_start = pm.end()
@ -183,8 +242,17 @@ def parse_tool_calls_from_text(content: str, *, id_offset: int = 0) -> list[dict
else len(body)
)
val = body[val_start:next_param]
val = _TC_PARAM_CLOSE_RE.sub("", val)
if not allow_incomplete:
stripped_val = val.rstrip()
if not stripped_val.endswith(_PARAM_CLOSE_TAG):
valid_params = False
break
val = stripped_val[: -len(_PARAM_CLOSE_TAG)]
else:
val = _TC_PARAM_CLOSE_RE.sub("", val)
arguments[param_name] = val.strip()
if not valid_params:
continue
tc = {
"id": f"call_{id_offset + len(tool_calls)}",

412
studio/backend/core/inference/tool_loop_controller.py Normal file
View file

@ -0,0 +1,412 @@
# SPDX-License-Identifier: AGPL-3.0-only
# Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
"""Shared controller state for Studio local agentic tool loops.
This module is intentionally dependency-light: it owns only per-response
ledger state and value objects used by the GGUF and safetensors loops.
Route/SSE conversion, tool execution, and model streaming stay in the
backend-specific modules.
"""
from __future__ import annotations
import copy
import json
from dataclasses import dataclass, field
from typing import Any, Literal, Mapping, Sequence
from urllib.parse import urlparse
from core.inference.tool_call_parser import TOOL_ERROR_NUDGE, TOOL_ERROR_PREFIXES
_CANONICAL_HEAL_ARG = {
"python": "code",
"terminal": "command",
"render_html": "code",
}
_ONE_SHOT_TOOLS = frozenset({"render_html"})
NoopReason = Literal["duplicate", "disabled", "render_html_repeat"]
ToolAction = Literal["execute", "duplicate", "disabled", "render_html_repeat"]
@dataclass(frozen = True)
class CoercedArguments:
"""Normalized tool arguments plus whether healing changed the shape."""
arguments: dict[str, Any]
healed: bool = False
@dataclass(frozen = True)
class ToolCallDecision:
"""Decision made before any visible tool event is emitted."""
action: ToolAction
tool_name: str
arguments: dict[str, Any]
tool_call_id: str = ""
key: str = ""
provenance: dict[str, Any] = field(default_factory = dict)
status_text: str = ""
noop_result: str = ""
@property
def should_execute(self) -> bool:
return self.action == "execute"
@property
def emit_visible_events(self) -> bool:
"""Only real executions should become frontend-visible tool cards."""
return self.should_execute
@property
def noop_reason(self) -> NoopReason | None:
if self.action == "execute":
return None
return self.action
def tool_start_payload(self) -> dict[str, Any]:
"""Build the payload fields for a real tool_start event."""
return {
"tool_name": self.tool_name,
"tool_call_id": self.tool_call_id,
"arguments": self.arguments,
"provenance": self.provenance,
}
def tool_start_event(self) -> dict[str, Any]:
"""Build the existing backend event shape for a real execution."""
return {"type": "tool_start", **self.tool_start_payload()}
def as_assistant_tool_call(self) -> dict[str, Any]:
"""Return an OpenAI-style tool_call with normalized arguments."""
tool_call: dict[str, Any] = {
"type": "function",
"function": {
"name": self.tool_name,
"arguments": json.dumps(
self.arguments,
ensure_ascii = False,
sort_keys = True,
separators = (",", ":"),
),
},
}
if self.tool_call_id:
tool_call["id"] = self.tool_call_id
return tool_call
@dataclass(frozen = True)
class ToolCallCompletion:
"""Result/nudge that should be fed back to the next model turn."""
decision: ToolCallDecision
result: str
is_error: bool = False
executed: bool = False
def tool_end_payload(self) -> dict[str, Any]:
"""Build the payload fields for a real tool_end event."""
return {
"tool_name": self.decision.tool_name,
"tool_call_id": self.decision.tool_call_id,
"result": self.result,
"provenance": self.decision.provenance,
}
def tool_end_event(self) -> dict[str, Any]:
"""Build the existing backend event shape for a real execution result."""
return {"type": "tool_end", **self.tool_end_payload()}
def tool_message(self) -> dict[str, Any]:
"""Return the OpenAI-compatible tool message for a real execution."""
if not self.executed:
raise ValueError("No-op completions are internal nudges, not tool messages")
return self.model_message()
def model_message(self) -> dict[str, Any]:
"""Return the internal message appended before the next generation.
Executed calls keep the existing OpenAI-compatible ``role=tool``
continuation. No-op controller decisions are not real tool output, so
they are fed back as a hidden user nudge rather than a normal tool
result.
"""
if not self.executed:
return {"role": "user", "content": self.result}
content = strip_result_for_model(self.result)
if self.is_error:
content = content + TOOL_ERROR_NUDGE
message: dict[str, Any] = {
"role": "tool",
"name": self.decision.tool_name,
"content": content,
}
if self.decision.tool_call_id:
message["tool_call_id"] = self.decision.tool_call_id
return message
@dataclass(frozen = True)
class _ToolCallRecord:
key: str
is_error: bool
executed: bool
action: ToolAction
def _json_default(value: Any) -> str:
return str(value)
def canonical_tool_call_key(tool_name: str, arguments: Mapping[str, Any]) -> str:
"""Return a stable key for duplicate detection."""
canonical_args = json.dumps(
dict(arguments),
ensure_ascii = False,
sort_keys = True,
separators = (",", ":"),
default = _json_default,
)
return f"{tool_name}:{canonical_args}"
def coerce_tool_arguments(
raw_args: Any,
*,
heal: bool,
tool_name: str = "",
) -> CoercedArguments:
"""Normalize model-emitted ``function.arguments`` to a dictionary."""
if isinstance(raw_args, Mapping):
return CoercedArguments(dict(raw_args), False)
if isinstance(raw_args, str):
try:
parsed = json.loads(raw_args)
if isinstance(parsed, Mapping):
return CoercedArguments(dict(parsed), False)
except (json.JSONDecodeError, ValueError):
pass
if heal:
key = _CANONICAL_HEAL_ARG.get(tool_name, "query")
return CoercedArguments({key: raw_args}, True)
return CoercedArguments({"raw": raw_args}, False)
return CoercedArguments({}, False)
def tool_event_provenance(**flags: object) -> dict[str, object]:
"""Return provenance metadata with falsey flags omitted."""
provenance: dict[str, object] = {"source": "local"}
for key, value in flags.items():
if value is not None and value is not False:
provenance[key] = value
return provenance
def status_for_tool(tool_name: str, arguments: Mapping[str, Any]) -> str:
"""Return the status text already used by local tool streams."""
if tool_name == "web_search":
url = str(arguments.get("url") or "").strip()
if url:
parsed = urlparse(url)
if parsed.scheme in ("http", "https") and parsed.hostname:
host = parsed.hostname
if host.startswith("www."):
host = host[4:]
return f"Reading: {host}"
return "Reading page..."
return f"Searching: {arguments.get('query', '')}"
if tool_name == "python":
preview = str(arguments.get("code") or "").strip().split("\n")[0][:60]
return f"Running Python: {preview}" if preview else "Running Python..."
if tool_name == "terminal":
preview = str(arguments.get("command") or "")[:60]
return f"Running: {preview}" if preview else "Running command..."
return f"Calling: {tool_name}"
def is_tool_error(result: str) -> bool:
return isinstance(result, str) and result.lstrip().startswith(TOOL_ERROR_PREFIXES)
def strip_result_for_model(result: str) -> str:
"""Remove frontend-only sentinels (image paths, RAG source map) before
feeding the result back to the model."""
for sentinel in ("__IMAGES__:", "__RAG_SOURCES__:"):
if sentinel in result:
result = result.split(sentinel, 1)[0].rstrip()
return result
def _tool_name_from_schema(tool: Mapping[str, Any]) -> str:
function = tool.get("function")
if not isinstance(function, Mapping):
return ""
name = function.get("name")
return str(name or "")
def _noop_result(reason: NoopReason, tool_name: str) -> str:
if reason == "duplicate":
return (
"The previous tool request was not executed because this exact "
"tool call already completed successfully. Do not repeat the same "
"tool call. Continue with a different enabled tool if that would "
"materially help, or provide the final answer if you have enough "
"information."
)
if reason == "render_html_repeat":
return (
"render_html completed successfully earlier in this assistant "
"response. Do not call render_html again unless the user asks for "
"changes. Do not mention this internal instruction. Provide only "
"the requested final note or answer."
)
return (
f"The previous tool request was not executed because tool "
f"'{tool_name}' is not enabled for this request. Provide the "
"final answer now without calling more tools."
)
class ToolLoopController:
"""Per-response ledger for local agentic tool loops."""
def __init__(
self,
*,
tools: Sequence[Mapping[str, Any]] | None,
auto_heal_tool_calls: bool = True,
one_shot_tools: frozenset[str] = _ONE_SHOT_TOOLS,
duplicate_noop_limit: int = 2,
) -> None:
self._restrict_to_allowed = tools is not None
self._tools = [copy.deepcopy(dict(tool)) for tool in (tools or [])]
self._allowed_tool_names = {
name for name in (_tool_name_from_schema(tool) for tool in self._tools) if name
}
self._auto_heal_tool_calls = auto_heal_tool_calls
self._one_shot_tools = one_shot_tools
self._completed_one_shot_tools: set[str] = set()
self._successful_keys: set[str] = set()
self._duplicate_noop_counts: dict[str, int] = {}
self._duplicate_noop_limit = max(1, duplicate_noop_limit)
self._history: list[_ToolCallRecord] = []
self._force_final_answer = False
@property
def history(self) -> tuple[_ToolCallRecord, ...]:
return tuple(self._history)
@property
def force_final_answer(self) -> bool:
"""True once a terminal no-op should transition to a no-tools pass."""
return self._force_final_answer
def active_tools(self) -> list[dict[str, Any]]:
"""Return tools still worth advertising to the next model call."""
if self._force_final_answer:
return []
active: list[dict[str, Any]] = []
for tool in self._tools:
name = _tool_name_from_schema(tool)
if name in self._completed_one_shot_tools:
continue
active.append(copy.deepcopy(tool))
return active
def prepare_call(
self,
tool_call: Mapping[str, Any],
*,
forced: bool = False,
provisional: bool = False,
) -> ToolCallDecision:
"""Classify a parsed tool call before any visible event is yielded."""
function = tool_call.get("function")
function = function if isinstance(function, Mapping) else {}
tool_name = str(function.get("name") or "").strip()
coerced = coerce_tool_arguments(
function.get("arguments", {}),
heal = self._auto_heal_tool_calls,
tool_name = tool_name,
)
key = canonical_tool_call_key(tool_name, coerced.arguments)
provenance = tool_event_provenance(
healed = coerced.healed,
forced = forced,
provisional = provisional,
)
action: ToolAction = "execute"
noop = ""
if tool_name in self._completed_one_shot_tools:
action = "render_html_repeat"
noop = _noop_result("render_html_repeat", tool_name)
elif self._restrict_to_allowed and tool_name not in self._allowed_tool_names:
action = "disabled"
noop = _noop_result("disabled", tool_name)
elif key in self._successful_keys:
action = "duplicate"
noop = _noop_result("duplicate", tool_name)
return ToolCallDecision(
action = action,
tool_name = tool_name,
arguments = coerced.arguments,
tool_call_id = str(tool_call.get("id") or ""),
key = key,
provenance = provenance,
status_text = status_for_tool(tool_name, coerced.arguments),
noop_result = noop,
)
def record_result(self, decision: ToolCallDecision, result: Any) -> ToolCallCompletion:
"""Record a real tool execution and return model/frontend payload helpers."""
result_text = result if isinstance(result, str) else str(result)
failed = is_tool_error(result_text)
self._history.append(
_ToolCallRecord(
key = decision.key,
is_error = failed,
executed = True,
action = decision.action,
)
)
if not failed:
self._successful_keys.add(decision.key)
if decision.tool_name in self._one_shot_tools:
self._completed_one_shot_tools.add(decision.tool_name)
return ToolCallCompletion(
decision = decision,
result = result_text,
is_error = failed,
executed = True,
)
def record_noop(self, decision: ToolCallDecision) -> ToolCallCompletion:
"""Record a controller no-op without creating visible tool output."""
self._history.append(
_ToolCallRecord(
key = decision.key,
is_error = False,
executed = False,
action = decision.action,
)
)
if decision.action == "duplicate":
duplicate_count = self._duplicate_noop_counts.get(decision.key, 0) + 1
self._duplicate_noop_counts[decision.key] = duplicate_count
if duplicate_count >= self._duplicate_noop_limit:
self._force_final_answer = True
elif decision.action in ("disabled", "render_html_repeat"):
self._force_final_answer = True
return ToolCallCompletion(
decision = decision,
result = decision.noop_result,
is_error = False,
executed = False,
)

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

File diff suppressed because it is too large Load diff

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